diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 9ef81873b510..d6cf5f8baad8 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/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 "", + # 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, diff --git a/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py b/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py index 7b29b030e3d4..bab4c4a15d43 100644 --- a/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py +++ b/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: diff --git a/tensorrt_llm/_torch/visual_gen/mapping.py b/tensorrt_llm/_torch/visual_gen/mapping.py index 55f6a2d8b0ac..531d28a39657 100644 --- a/tensorrt_llm/_torch/visual_gen/mapping.py +++ b/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, + ) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 13cbe77bdf12..b44457c5ed25 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/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 diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index 3551371fa173..0ac542c2fd60 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -22,12 +22,14 @@ 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 @@ -35,6 +37,25 @@ 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() diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 898f01415de0..b7327c4ee088 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/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("") 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