[TRTLLM-14541][fix] VisualGen: byte-identical tune-run vs cache-load on multi-GPU

tensorrt_llm/_torch/autotuner.py

@@ -261,7 +261,8 @@ def unique_id(self):
 @contextlib.contextmanager
 def autotune(tune_mode: bool = True,
              cache_path: str = None,
-             skip_dynamic_tuning_buckets: bool = False):
+             skip_dynamic_tuning_buckets: bool = False,
+             post_tune_merge_dist: Optional[Distributed] = None):
     """Context manager for autotuning with distributed support.
 
     Args:
@@ -271,8 +272,21 @@ def autotune(tune_mode: bool = True,
             _optimization_profiles() so only actual input shapes from warmup
             are profiled. Useful for workloads (e.g. diffusion) where the
             LLM-oriented M-bucket sweep is unnecessary.
+        post_tune_merge_dist: Optional ``Distributed``. When set, tactics are
+            tuned per-rank independently and merged across ranks in one
+            collective at context exit (before the cache is saved) — for
+            pipelines whose warmup is not SPMD and so cannot use the per-op
+            distributed strategies.
     """
     autotuner = AutoTuner.get()
+    if post_tune_merge_dist is not None:
+        # Save the singleton's prior state: the full-world dist attached for the
+        # merge is temporary and must not leak into later sessions. Attach the
+        # real (world-sized) mapping now for a correct rank, but with no dist
+        # yet: _is_distributed() stays False so per-op sync is skipped while a
+        # non-SPMD pipeline tunes; the dist is attached at exit to merge.
+        prev_mapping, prev_dist = autotuner.mapping, autotuner._dist
+        autotuner.setup_distributed_state(post_tune_merge_dist.mapping, None)
     rank = autotuner.mapping.rank
 
     # if cache_path is provided, use the rank-specific file
@@ -303,10 +317,24 @@ def autotune(tune_mode: bool = True,
         if autotune_enabled:
             logger.info("[Autotuner] Autotuning process ends")
 
-        # save cache
-        if cache_path is not None:
-            logger.info(f"[Autotuner] Saving cache to {cache_path}")
-            autotuner.profiling_cache.save_cache(cache_path, rank)
+        try:
+            # Merge tactics across ranks in one collective before the cache is
+            # saved, so non-SPMD pipelines can tune independently and still agree.
+            if post_tune_merge_dist is not None:
+                autotuner.setup_distributed_state(post_tune_merge_dist.mapping,
+                                                  post_tune_merge_dist)
+                autotuner.post_tune_merge_tactics()
+
+            # save cache
+            if cache_path is not None:
+                logger.info(f"[Autotuner] Saving cache to {cache_path}")
+                autotuner.profiling_cache.save_cache(cache_path, rank)
+        finally:
+            # Restore the singleton's prior distributed state (see enter) so the
+            # temporary full-world state does not persist past this context.
+            if post_tune_merge_dist is not None:
+                autotuner.mapping = prev_mapping
+                autotuner._dist = prev_dist
 
 
 @dataclass
@@ -760,9 +788,17 @@ def _serialize_cache_data(self,
             # Convert any simple object to string for JSON compatibility
             key_str = str(key)
             runner_id, tactic, min_time = value
-            tactic_str = repr(tactic)
+            # Enum tactics (e.g. Fp4QuantTactic) repr as "<Fp4QuantTactic.TRTLLM: -1>",
+            # which ast.literal_eval can't parse back -> load_cache crash. Serialize
+            # the underlying value (reload yields that value; an IntEnum compares
+            # equal to it and kernels take the int).
+            is_enum = isinstance(tactic, enum.Enum)
+            tactic_check = tactic.value if is_enum else tactic
+            tactic_str = repr(tactic_check)
             try:
-                assert tactic == ast.literal_eval(
+                # Verify the serialized value round-trips (compare against the
+                # value we store, not the enum member — else non-IntEnum fails).
+                assert tactic_check == ast.literal_eval(
                     tactic_str
                 ), f"Tactic is not compatible with json.dumps/json.loads"
             except Exception as e:
@@ -1954,6 +1990,19 @@ def _merge_cache_data(self, custom_op: str):
 
         self.profiling_cache.merge_cache_data(merged_cache_data)
 
+    def post_tune_merge_tactics(self) -> None:
+        """Merge tactics across ranks after tuning: all-gather every rank's whole
+        profiling cache and keep the fastest tactic per key. One-shot
+        (session-level, not per-op). Collective — all ranks must call it."""
+        if not self._is_distributed():
+            return
+        merged = dict()
+        for data in self._dist.tp_cp_allgather(obj=self.profiling_cache.cache):
+            for key, value in data.items():
+                if value[-1] < merged.get(key, [float('inf')])[-1]:
+                    merged[key] = value
+        self.profiling_cache.merge_cache_data(merged)
+
     def _broadcast_cache_data(
         self,
         custom_op: str,

tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py

@@ -72,6 +72,8 @@ def __init__(
     ):
         self.config = config
         self.enabled = config.use_cuda_graph
+        # When False, uncaptured keys run eagerly instead of triggering capture.
+        self.capture_enabled = True
         self._shared_pool = shared_pool
         self._extra_key_fns: Dict[str, ExtraKeyFn] = {}
 
@@ -194,6 +196,8 @@ def wrapper(*args, **kwargs):
             key = self.get_graph_key(*args, **kwargs)
 
             if key not in self.graphs:
+                if not self.capture_enabled:
+                    return fn(*args, **kwargs)
                 self.capture(key, fn, args, kwargs)
                 return self.replay(key, args, kwargs)
             else:

tensorrt_llm/_torch/visual_gen/mapping.py

@@ -555,3 +555,12 @@ def to_llm_mapping(self) -> Mapping:
             rank=self.tp_rank,
             tp_size=self.tp_size,
         )
+
+    def to_autotuner_mapping(self) -> Mapping:
+        """Mapping that makes the autotuner treat all world ranks as one tuning
+        group (tp_size == world_size), so its post-tune cross-rank merge engages."""
+        return Mapping(
+            world_size=self.world_size,
+            rank=self._rank,
+            tp_size=self.world_size,
+        )

tensorrt_llm/_torch/visual_gen/pipeline.py

@@ -1,3 +1,4 @@
+import contextlib
 import itertools
 import os
 import time
@@ -206,6 +207,18 @@ def _setup_cuda_graphs(self):
             model.forward = runner.wrap(model.forward)
             self._cuda_graph_runners[name] = runner
 
+    @contextlib.contextmanager
+    def no_cuda_graph_capture(self):
+        """Run wrapped forwards eagerly instead of capturing new CUDA graphs."""
+        prev = {name: r.capture_enabled for name, r in self._cuda_graph_runners.items()}
+        for r in self._cuda_graph_runners.values():
+            r.capture_enabled = False
+        try:
+            yield
+        finally:
+            for name, r in self._cuda_graph_runners.items():
+                r.capture_enabled = prev[name]
+
     @property
     def rank(self):
         return dist.get_rank() if dist.is_initialized() else 0

tensorrt_llm/_torch/visual_gen/pipeline_loader.py

@@ -22,19 +22,40 @@
 import torch.distributed as dist
 
 from tensorrt_llm._torch.autotuner import autotune
+from tensorrt_llm._torch.distributed.communicator import Distributed, TorchDist
 from tensorrt_llm._torch.models.modeling_utils import MetaInitMode
 from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import (
     CUTE_AVAILABLE,
 )
 from tensorrt_llm.llmapi.utils import download_hf_model
 from tensorrt_llm.logger import logger
+from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.visual_gen.args import VisualGenArgs
 
 from .config import DiffusionPipelineConfig
 from .mapping import VisualGenMapping
 from .models import AutoPipeline
 from .pipeline_registry import PIPELINE_REGISTRY, PipelineComponent
 
+
+class _VisualGenAutotuneDist(TorchDist):
+    """Autotuner communicator whose collective spans the whole world.
+
+    VisualGen's device mesh has no TP axis over the tuned GEMMs (its parallelism
+    is on cfg/ulysses), so the mesh's tp group would gather a single rank. Gather
+    over the default world group instead, since every rank runs the transformer.
+    """
+
+    def __init__(self, mapping: Mapping):
+        Distributed.__init__(self, mapping)
+        assert dist.is_initialized()
+
+    def tp_cp_allgather(self, obj: object) -> list[object]:
+        gathered: list[object] = [None] * dist.get_world_size()
+        dist.all_gather_object(gathered, obj)
+        return gathered
+
+
 if TYPE_CHECKING:
     from .models import BasePipeline
 
@@ -305,10 +326,32 @@ def load(
 
         if not skip_warmup:
             if config.torch_compile.enable_autotune:
-                with autotune(
-                    cache_path=os.environ.get("TLLM_AUTOTUNER_CACHE_PATH"),
-                    skip_dynamic_tuning_buckets=True,
-                ):
+                cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH")
+                post_tune_merge_dist = None
+                if dist.is_initialized() and dist.get_world_size() > 1:
+                    amap = config.visual_gen_mapping.to_autotuner_mapping()
+                    post_tune_merge_dist = _VisualGenAutotuneDist(amap)
+
+                if post_tune_merge_dist is None:
+                    # Single rank: no cross-rank merge, so tune and capture in one
+                    # pass (the graph runner captures only after its own warmup
+                    # steps have finalized the tactic).
+                    with autotune(cache_path=cache_path, skip_dynamic_tuning_buckets=True):
+                        pipeline.warmup()
+                else:
+                    # Multi rank: capture graphs only AFTER tactics are merged
+                    # across ranks, otherwise each rank bakes its own pre-merge
+                    # tactic. Phase 1 tunes with capture off; the autotuner merges
+                    # at autotune() exit; phase 2 captures the merged tactics.
+                    with (
+                        pipeline.no_cuda_graph_capture(),
+                        autotune(
+                            cache_path=cache_path,
+                            skip_dynamic_tuning_buckets=True,
+                            post_tune_merge_dist=post_tune_merge_dist,
+                        ),
+                    ):
+                        pipeline.warmup()
                     pipeline.warmup()
             else:
                 pipeline.warmup()

tests/unittest/_torch/misc/test_autotuner.py

@@ -1,3 +1,4 @@
+import enum
 import itertools
 import json
 import math
@@ -1265,3 +1266,90 @@ def _cublas_call():
         f"CuteDSL heuristic kernel ({heuristic_us:.2f} us) is "
         f">{cublas_tolerance:.2f}x slower than cuBLAS NVFP4 "
         f"({cublas_us:.2f} us) for M={m}, N={n}, K={k}")
+
+
+def test_post_tune_merge_tactics_min_time_and_subset_kept():
+    tuner = autotuner.AutoTuner()
+    tuner.mapping = Mapping(world_size=2, rank=0, tp_size=2)
+
+    r0 = {
+        ("gemm", "X"): (0, ("cutlass", 10), 1.0),
+        ("q", "A"): (0, ("trtllm", ), 1.0)
+    }
+    r1 = {
+        ("gemm", "X"): (0, ("cublaslt", 0), 0.5),
+        ("vae", "B"): (0, ("cutlass", 2), 1.0)
+    }
+
+    class _FakeDist:
+
+        def tp_cp_allgather(self, obj):
+            return [r0, r1]
+
+    tuner._dist = _FakeDist()
+    tuner.profiling_cache.cache = dict(r0)
+    tuner.post_tune_merge_tactics()
+
+    cache = tuner.profiling_cache.cache
+    assert cache[("gemm", "X")] == (0, ("cublaslt", 0), 0.5)
+    assert cache[("q", "A")] == r0[("q", "A")]
+    assert cache[("vae", "B")] == r1[("vae", "B")]
+
+
+def test_post_tune_merge_tactics_single_rank_noop():
+    tuner = autotuner.AutoTuner()
+    tuner.mapping = Mapping(world_size=1, rank=0, tp_size=1)
+    tuner._dist = None
+    original = {("gemm", "X"): (0, ("cutlass", 10), 1.0)}
+    tuner.profiling_cache.cache = dict(original)
+    tuner.post_tune_merge_tactics()
+    assert tuner.profiling_cache.cache == original
+
+
+def test_profiling_cache_enum_tactic_roundtrip(tmp_path):
+    # An enum tactic must survive save -> load: repr("<Enum.X: v>") isn't
+    # ast.literal_eval-parsable, so the cache serializes tactic.value instead.
+    class _Tac(enum.IntEnum):
+        TRTLLM = -1
+
+    key = ("op::x", "Runner", "(1,)")
+    src = autotuner.AutoTuner().profiling_cache
+    src.cache[key] = (0, _Tac.TRTLLM, 1.0)
+    path = str(tmp_path / "cache.json")
+    src.save_cache(path, rank=0)  # must not raise
+
+    dst = autotuner.AutoTuner().profiling_cache
+    dst.load_cache(path, rank=0)
+    assert dst.cache[key][1] == _Tac.TRTLLM  # IntEnum compares equal to -1
+
+
+def test_autotune_post_tune_merge_before_save(tmp_path):
+    # autotune(post_tune_merge_dist=...) must merge across ranks and persist the
+    # merged winner, then restore the singleton's prior distributed state.
+    tuner = AutoTuner.get()
+    tuner.profiling_cache.cache.clear()
+    r0 = {("gemm", "X"): (0, ("cutlass", 10), 1.0)}
+    r1 = {("gemm", "X"): (0, ("cublaslt", 0), 0.5)}
+    tuner.profiling_cache.cache = dict(r0)
+
+    class _FakeDist:
+        mapping = Mapping(world_size=2, rank=0, tp_size=2)
+
+        def tp_cp_allgather(self, obj):
+            return [r0, r1]
+
+    prev_mapping, prev_dist = tuner.mapping, tuner._dist
+    path = str(tmp_path / "cache.json")
+    with autotune(cache_path=path, post_tune_merge_dist=_FakeDist()):
+        pass
+
+    # Merged in memory (fastest tactic won) and persisted before the context
+    # exits (the saved file carries the merged winner).
+    assert tuner.profiling_cache.cache[("gemm", "X")] == (0, ("cublaslt", 0),
+                                                          0.5)
+    persisted = autotuner.AutoTuner().profiling_cache
+    persisted.load_cache(path, rank=0)
+    assert persisted.cache[("gemm", "X")][1] == ("cublaslt", 0)
+    # The temporary full-world distributed state was restored.
+    assert tuner.mapping is prev_mapping
+    assert tuner._dist is prev_dist