[TRTLLM-14541][fix] VisualGen: deterministic autotuner tactics across runs and ranks

[luyiyun1021]

Description

Multi-GPU VisualGen (LTX-2, Wan 2.2) DiT output was nondeterministic: a fresh autotuning run produced different frames than a persistent-cache load of the same run, and different ranks silently diverged. This PR fixes the three root causes so that a tune-run is byte-identical to a cache-load, and every rank runs the identical kernels.

1. CUDA graphs were captured while the autotuner was still tuning. Warmup ran torch.compile + autotune + CUDA-graph capture in one pass, so each rank baked its own in-flight tactic choices into the graphs before tactics were finalized. Rank-level nsys --cuda-graph-trace=node confirmed this: a rank whose local tactic for a shape was cutlass had no cuBLASLt (nvjet_*) kernel in its graph at that shape, whereas every cache-load rank uniformly baked the cache's tactic.

2. NVFP4 GEMM tactic selection flips between numerically-distinct backends on tiny perf gaps. The GEMM autotuner picks among cutlass / cuBLASLt / cuda_core by measured time. On several LTX/Wan shapes cutlass and cuBLASLt are within timing noise of each other, so the winner flips per-launch and per-rank; the two backends are numerically distinct (~1e-4/element), so a flip changes the output. (Intra-backend flips are bit-exact, which is how the cross-backend flip was isolated: forcing a single backend is fully deterministic.)

3. The autotuner's cross-rank tactic sync never engaged for VisualGen (mapping/mesh gap). The autotuner unifies per-rank tactics through its Distributed backend keyed off a Mapping. VisualGen's to_llm_mapping() reports tp_size == 1 (its parallelism is on cfg/ulysses), so the autotuner saw a world of 1 and never synced. Even forcing a world-sized mapping is not enough: the autotuner's tp_cp_allgather gathers over the mapping's TP process group, and VisualGen's DeviceMeshTopologyImpl.device_mesh is a build-once class singleton whose TP axis spans a single rank — so the gather returned only the local cache and the merge was a silent no-op (verified: gathered=1 on 8 ranks).

Fix — two-phase warmup + a one-shot world-domain post-tune tactic merge, generic to all VisualGen pipelines (mirrors the LLM _run_autotuner_warmup, which also disables graph capture while tuning and syncs tactics through the autotuner's distributed backend):

• CUDAGraphRunner.capture_enabled flag + BasePipeline.no_cuda_graph_capture() context manager: warmup forwards run eagerly (no capture) while tuning.
• AutoTuner.post_tune_merge_tactics() + autotune(post_tune_merge_dist=dist): a new one-shot collective that all-gathers every rank's whole profiling cache and keeps the fastest tactic per key, run once at the autotune() context exit before the cache is saved. Unlike the per-op distributed_tuning_strategy sync it runs once and does not require the tuned ops to be invoked symmetrically across ranks, so it is safe for VisualGen's non-SPMD warmup (rank-0-only text encoder, parallel-VAE subsets).
• VisualGenMapping.to_autotuner_mapping() + _VisualGenAutotuneDist: give the autotuner a world-sized mapping so it treats all ranks as one tuning group, and a torch.distributed communicator whose gather spans the default world group (bypassing the VisualGen device mesh's single-rank TP axis). autotune() attaches that mapping up front (for a correct rank) but only wires the communicator at exit, so _is_distributed() stays False and per-op sync is skipped while phase-1 tuning runs; the merge then runs once at exit.
• PipelineLoader: phase 1 tunes eagerly with capture disabled and passes post_tune_merge_dist to autotune(), which merges tactics across all ranks and saves the cache at context exit → phase 2 re-runs warmup to capture graphs from the finalized, merged tactics.
• autotuner._serialize_cache_data: serialize enum tactics by .value so Fp4QuantTactic round-trips (otherwise ast.literal_eval crashes when loading any cache that contains it).

Trade-off: warmup now runs the pipeline forward twice (tune, then capture), which increases load-time warmup. Scope: this makes a tune-run byte-identical to a persistent-cache load and merges tactics across ranks within a run. Full cross-launch determinism (two independent no-cache tuning runs) still requires the persistent autotuner cache to freeze the tactic table, since problem 2's near-tie selection is timing-dependent per launch.

Deterministic cache workflow

Determinism is gated by the TLLM_AUTOTUNER_CACHE_PATH env var. autotune() loads the cache when that file already exists (otherwise it tunes), and always re-saves the cross-rank-merged table on context exit — so a build-once-then-load flow reproduces byte-identical output on every run:

# Build once: the file does not exist yet -> tune, merge across ranks, write the cache.
TLLM_AUTOTUNER_CACHE_PATH=/path/to/vg_cache.json <run VisualGen>

# Every later run: the file exists -> load it, capture from it, reproduce byte-for-byte.
TLLM_AUTOTUNER_CACHE_PATH=/path/to/vg_cache.json <run VisualGen>

With this PR the build run itself is already byte-identical to the loads (tune-run == cache-load), so the first, cache-building run is usable output rather than a throwaway. Cross-launch determinism across two independent no-cache runs still requires this persistent cache, since the tactic table depends on per-launch near-tie timing.

Test Coverage

Verified byte-identical tune-run vs cache-load output (mp4 md5) on B200 x8:

Model	Config	tune-run mp4 md5	cache-load mp4 md5	graphs captured
LTX-2	768x1280x121, 40 steps, cfg2 x ulysses4, mixed backends	5a8bfb2f	5a8bfb2f (equal)	16/16
Wan 2.2 A14B	480x832x81, static NVFP4, two-expert, cfg2 x ulysses4	b06b734f	b06b734f (equal)	32/32

Pre-fix baseline on the same LTX config: tune-run vs cache-load mp4 md5 differed on every launch, and three fresh tuning runs produced three different outputs. AutoTuner.post_tune_merge_tactics gets unit coverage in tests/unittest/_torch/misc/test_autotuner.py (min-time winner + rank-subset keys kept + single-rank no-op); the end-to-end determinism check is multi-GPU + real-checkpoint, so its scripts are attached to TRTLLM-14541 rather than added to CI.

PR Checklist

• PR description clearly explains what and why.

• PR Follows TRT-LLM CODING GUIDELINES.

• Test cases are provided for new code paths.

• Any new dependencies have been scanned for license and vulnerabilities.

• CODEOWNERS updated if ownership changes.

• Documentation updated as needed.

• Update tava architecture diagram if significant design change.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

• Ensured nondeterministic multi-GPU VisualGen output is addressed by preventing premature CUDA graph capture during tuning and enabling capture only after cross-rank tactic unification with finalized tactics.
• Implemented two-phase warmup in visual_gen/pipeline_loader.py:
  • Phase 1: run warmup under pipeline.no_cuda_graph_capture() with tuning graph capture disabled (skip_dynamic_tuning_buckets=True) and an optional distributed post-merge hook.
  • Phase 2: rerun warmup after leaving the tuning context so kernels/tactics are consistent across ranks and CUDA graph capture can occur deterministically.
• Added instance-level CUDA graph gating in visual_gen/cuda_graph_runner.py (capture_enabled) and a pipeline-scoped context manager BasePipeline.no_cuda_graph_capture() that snapshots/restores runner state via finally to avoid leaking capture settings.
• Made autotune results byte-identical between tune-runs and persistent-cache loads by:
  • Adding autotune(..., post_tune_merge_dist=...) and AutoTuner.post_tune_merge_tactics() to all-gather per-rank in-memory profiling caches and merge by cache key while keeping the minimum recorded time.
  • Fixing enum-valued tactic cache serialization: AutoTunerProfilingCache._serialize_cache_data now serializes repr(tactic.value) for enum.Enum tactics, matching the existing deserialization path and preventing enum literal round-trip failures.
• Improved distributed tactic synchronization by introducing VisualGenMapping.to_autotuner_mapping() (world-sized autotuning group) and _VisualGenAutotuneDist (object all-gather via dist.all_gather_object) to unify autotuner state across ranks.
• Added state hygiene: the autotune context temporarily re-attaches a distributed-state attachment for independent tuning, then restores the autotuner singleton’s previous distributed mapping/state after post-merge to avoid cross-context contamination.

QA Engineer Review

Test changes

tests/unittest/_torch/misc/test_autotuner.py

• Added:
  • test_post_tune_merge_tactics_min_time_and_subset_kept
  • test_post_tune_merge_tactics_single_rank_noop
  • test_profiling_cache_enum_tactic_roundtrip
  • test_autotune_post_tune_merge_before_save

Coverage in tests/integration/test_lists/

• None of the newly added test functions are referenced under tests/integration/test_lists/ (no matches found for the test names).

Verdict

insufficient

[luyiyun1021]

/bot run --disable-fail-fast

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The changes add enum-compatible autotuner cache serialization, cross-rank profiling-cache merging, scoped CUDA graph capture suppression, full-world visual generation mappings, and a two-phase distributed warmup flow.

Changes

Visual generation autotuning

Layer / File(s)	Summary
Autotuner cache serialization and merging tensorrt_llm/_torch/autotuner.py	The autotune context temporarily configures distributed state, merges rank-local profiling caches by minimum recorded time before saving, restores prior state, and serializes enum tactics using their underlying values.
Scoped CUDA graph capture control tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py, tensorrt_llm/_torch/visual_gen/pipeline.py	Uncaptured graph keys execute eagerly when capture is disabled, and the pipeline context manager restores each runner’s prior capture state.
Distributed visual generation warmup tensorrt_llm/_torch/visual_gen/mapping.py, tensorrt_llm/_torch/visual_gen/pipeline_loader.py	Warmup uses a full-world autotuning mapping and communicator, performs eager autotuned warmup, merges caches, and runs a final warmup.
Autotuner validation tests/unittest/_torch/misc/test_autotuner.py	Tests cover distributed merging, single-rank no-op behavior, enum persistence, save ordering, and distributed-state restoration.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PipelineLoader
  participant BasePipeline
  participant CUDAGraphRunner
  participant AutoTuner
  participant DistributedWorld
  PipelineLoader->>BasePipeline: disable CUDA graph capture
  PipelineLoader->>AutoTuner: run autotuned warmup
  AutoTuner->>CUDAGraphRunner: execute uncaptured keys eagerly
  AutoTuner->>DistributedWorld: gather profiling caches
  DistributedWorld-->>AutoTuner: return rank cache dictionaries
  AutoTuner->>AutoTuner: apply fastest entries and save cache
  PipelineLoader->>BasePipeline: run final warmup
  BasePipeline->>CUDAGraphRunner: restore capture states

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16794: Related autotuner profiling-cache enum serialization and deserialization handling.

Suggested reviewers: chzblych, junyixu-nv, nvshreyas

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the PR template and clearly summarizes the main change to VisualGen autotuner determinism.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections and is detailed and relevant.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tensorrt_llm/_torch/visual_gen/pipeline.py (1)

210-225: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add the required return annotation.

no_cuda_graph_capture() is a new function but has no return annotation. Add Iterator[None] (or the repository’s equivalent) for the @contextlib.contextmanager generator.

As per coding guidelines, every function must be annotated.

Proposed fix

+from collections.abc import Iterator
+
 `@contextlib.contextmanager`
-def no_cuda_graph_capture(self):
+def no_cuda_graph_capture(self) -> Iterator[None]:

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 210 - 225, Update
no_cuda_graph_capture with the repository’s standard generator return
annotation, using Iterator[None] (or the equivalent typing symbol) while
preserving its existing context-manager behavior.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline_loader.py`:
- Around line 40-55: Update _VisualGenAutotuneDist.__init__ to type mapping as
Mapping, and annotate tp_cp_allgather’s obj parameter as object with a
list[object] return type. Add or reuse the appropriate Mapping import, keeping
the override free of Any.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 210-225: Update no_cuda_graph_capture with the repository’s
standard generator return annotation, using Iterator[None] (or the equivalent
typing symbol) while preserving its existing context-manager behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7d7e8fea-e77c-48b7-a459-77c158593170

📥 Commits

Reviewing files that changed from the base of the PR and between 3d86c72 and 1f00262.

📒 Files selected for processing (6)

• tensorrt_llm/_torch/autotuner.py
• tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
• tensorrt_llm/_torch/visual_gen/mapping.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py
• tests/unittest/_torch/misc/test_autotuner.py

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61235 [ run ] triggered by Bot. Commit: 10270a2 Link to invocation

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

tensorrt_llm/_torch/autotuner.py (1)

777-786: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Narrow the cache-serialization exception handler.

This changed serialization path catches every exception, although the expected round-trip failures are parse or assertion failures. Restrict it to the expected exception types so unrelated defects are not downgraded to a warning.

As per coding guidelines, “Avoid broad exception handling; catch specific exceptions instead of using bare except:.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/autotuner.py` around lines 777 - 786, In the tactic
round-trip validation within the autotuner cache-serialization path, narrow the
handler around ast.literal_eval and the assertion to catch only the expected
parsing and assertion exception types. Preserve the existing warning_once
behavior for those failures while allowing unrelated exceptions to propagate.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 311-315: Update the autotune flow around setup_distributed_state
and cache I/O so rank is derived only after
autotuner.setup_distributed_state(*unify_tactics) installs the distributed
mapping. Use the resulting rank, or unify_tactics[0].rank, for subsequent cache
load/save operations while preserving cross-rank tactic unification.

In `@tests/unittest/_torch/misc/test_autotuner.py`:
- Around line 1270-1305: Extend the autotuner tests beyond direct cache merging:
add a profiling-cache save/load round-trip using an enum tactic and a focused
autotune context test that passes unify_tactics and verifies cross-rank
unification occurs before persistence. Anchor these additions to the existing
AutoTuner and profiling_cache test setup, and include the required changed-test
coverage and coverage verdict.

---

Nitpick comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 777-786: In the tactic round-trip validation within the autotuner
cache-serialization path, narrow the handler around ast.literal_eval and the
assertion to catch only the expected parsing and assertion exception types.
Preserve the existing warning_once behavior for those failures while allowing
unrelated exceptions to propagate.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a558b706-3094-49a8-9579-13c86f592b84

📥 Commits

Reviewing files that changed from the base of the PR and between 1f00262 and 10270a2.

📒 Files selected for processing (6)

• tensorrt_llm/_torch/autotuner.py
• tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
• tensorrt_llm/_torch/visual_gen/mapping.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py
• tests/unittest/_torch/misc/test_autotuner.py

🚧 Files skipped from review as they are similar to previous changes (3)

• tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
• tensorrt_llm/_torch/visual_gen/mapping.py
• tensorrt_llm/_torch/visual_gen/pipeline.py

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61241 [ run ] triggered by Bot. Commit: f526572 Link to invocation

[tensorrt-cicd]

PR_Github #61235 [ run ] completed with state ABORTED. Commit: 10270a2

Link to invocation

[luyiyun1021]

/bot run --disable-fail-fast

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 282-286: Update the autotuning context around post_tune_merge_dist
and the corresponding cleanup near the later exit block to save the singleton
autotuner’s prior mapping and _dist before setup_distributed_state, then restore
both after merge/cache persistence completes. Use a nested finally so
restoration occurs on success and failure, preventing later sessions from
retaining the temporary full-world distributed state.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 00511959-cbe7-4485-964a-5522bdab91f4

📥 Commits

Reviewing files that changed from the base of the PR and between f526572 and 4c9d830.

📒 Files selected for processing (6)

• tensorrt_llm/_torch/autotuner.py
• tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
• tensorrt_llm/_torch/visual_gen/mapping.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py
• tests/unittest/_torch/misc/test_autotuner.py

🚧 Files skipped from review as they are similar to previous changes (4)

• tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tensorrt_llm/_torch/visual_gen/mapping.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py

[tensorrt-cicd]

PR_Github #61255 [ run ] triggered by Bot. Commit: 4c9d830 Link to invocation

[tensorrt-cicd]

PR_Github #61241 [ run ] completed with state ABORTED. Commit: f526572

Link to invocation

[luyiyun1021]

/bot run --disable-fail-fast

[JunyiXu-nv]

LGTM

[zhenhuaw-me]

@hyukn can you help to take a look at this PR? Thanks!

[zhenhuaw-me]

LGTM. Thanks!

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61877 [ run ] triggered by Bot. Commit: c8c85d3 Link to invocation

[tensorrt-cicd]

PR_Github #61877 [ run ] completed with state FAILURE. Commit: c8c85d3
/LLM/main/L0_MergeRequest_PR pipeline #50070 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[luyiyun1021]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62097 [ run ] triggered by Bot. Commit: c8c85d3 Link to invocation

[github-actions]

Removed the "ci: post-merge approved" label because @zhenhuaw-me could not be verified as an active member of NVIDIA/trt-llm-ci-approvers. Ask a member of that team to apply it.

[tensorrt-cicd]

PR_Github #62097 [ run ] completed with state SUCCESS. Commit: c8c85d3
/LLM/main/L0_MergeRequest_PR pipeline #50280 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation