[TRTLLM-16731][fix] Autotuner: skip redundant per-profile work for exclude_from_cache ops in tuning mode

tensorrt_llm/_torch/autotuner.py

@@ -321,8 +321,10 @@ def autotune(tune_mode: bool = True,
     # record the old tuning mode
     old_mode = autotuner.is_tuning_mode
     old_skip = autotuner.skip_dynamic_tuning_buckets
+    old_tuned_excluded_ops = autotuner._tuned_excluded_ops
     autotuner.is_tuning_mode = tune_required
     autotuner.skip_dynamic_tuning_buckets = skip_dynamic_tuning_buckets
+    autotuner._tuned_excluded_ops = set()
     # Avoid making compiled AllReduce paths initialize or inspect AutoTuner.
     from tensorrt_llm._torch.distributed.ops import \
         set_allreduce_autotuner_tuning_mode
@@ -342,6 +344,7 @@ def autotune(tune_mode: bool = True,
     finally:
         autotuner.is_tuning_mode = old_mode
         autotuner.skip_dynamic_tuning_buckets = old_skip
+        autotuner._tuned_excluded_ops = old_tuned_excluded_ops
         set_allreduce_autotuner_tuning_mode(old_mode)
         if autotune_enabled:
             torch.ops.trtllm.set_fine_grained_sync_disabled_override(False)
@@ -944,6 +947,10 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000):
         self.stream_delay_micro_secs = stream_delay_micro_secs
         self.profiling_cache = AutoTunerProfilingCache()
         self._primed_cached_tactics: set[tuple[str, str, str, str]] = set()
+        # exclude_from_cache ops whose profiles all have cache entries in the
+        # current autotune() context, keyed by op, runners, input shapes and
+        # profile settings.
+        self._tuned_excluded_ops: set[tuple] = set()
         self.is_tuning_mode = False
         self.skip_dynamic_tuning_buckets = False
 
@@ -1157,10 +1164,20 @@ def choose_one(
         # — in tuning mode — force the full profile path so the JIT
         # warmup actually runs even if a stale on-disk entry was loaded.
         # In inference mode we still consult the in-process cache, which
-        # was populated by this process's own warmup.
+        # was populated by this process's own warmup. Once this process has
+        # profiled every profile of the op for these shapes and profile
+        # settings in the current autotune() context, later calls use the
+        # in-process cache like any other op.
+        excluded_key = None
         if tuning_config.exclude_from_cache:
             self.profiling_cache.add_excluded_op(custom_op)
-        if tuning_config.exclude_from_cache and self.is_tuning_mode:
+            excluded_key = (custom_op,
+                            tuple(str(r.unique_id()) for r in runners),
+                            input_shapes, tuning_config.dynamic_tensor_specs,
+                            tuning_config.constraint_specs,
+                            tuning_config.tune_max_num_tokens)
+        if (excluded_key is not None and self.is_tuning_mode
+                and excluded_key not in self._tuned_excluded_ops):
             is_cache_hit = False
             best_runner_id, best_tactic, min_time = (
                 self.profiling_cache.fallback_entry())
@@ -1224,7 +1241,6 @@ def choose_one(
                 tuning_config.distributed_tuning_strategy):
             with self._shared_profile_graph_pool():
                 for p in profiles:
-                    tensors = self._prepare_input_tensors(p, inputs)
                     is_cache_hit, *_ = self.profiling_cache.search_cache(
                         custom_op,
                         runners,
@@ -1233,6 +1249,7 @@ def choose_one(
                         apply_map_to_tuning_buckets=False,
                     )
                     if not is_cache_hit:
+                        tensors = self._prepare_input_tensors(p, inputs)
                         # Initialize runner and tactic as None in case of no valid tactic or runners are found
                         with nvtx_range(
                                 f"{custom_op}, shape {p.get_opt_shapes()}"):
@@ -1243,6 +1260,15 @@ def choose_one(
 
         self._maybe_sync_cache_data(tuning_config.distributed_tuning_strategy,
                                     custom_op)
+        if excluded_key is not None and all(
+                self.profiling_cache.search_cache(
+                    custom_op,
+                    runners,
+                    p.get_opt_shapes(),
+                    tuning_config,
+                    apply_map_to_tuning_buckets=False,
+                )[0] for p in profiles):
+            self._tuned_excluded_ops.add(excluded_key)
         self._prime_cached_tactics(custom_op, runners, tuning_config, inputs,
                                    **kwargs)
 

tests/unittest/_torch/misc/test_autotuner.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import contextlib
+import dataclasses
 import enum
 import itertools
 import json
@@ -326,6 +327,106 @@ def fail_once(self, inputs, *, tactic=-1, **kwargs) -> torch.Tensor:
     assert JitGemmRunner.calls == [((3, 64), 0)]
 
 
+def test_exclude_from_cache_op_reuses_profiles_within_window(
+        monkeypatch) -> None:
+    """Profile an exclude_from_cache op once per autotune() window and prepare
+    inputs only for profiles that miss the cache."""
+    x = torch.randn(3, 64)
+    w = torch.randn(64, 128)
+    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
+        input_idx=0,
+        dim_idx=0,
+        gen_tuning_buckets=(3, 4, 5),
+        map_to_tuning_buckets=lambda value: value), ),
+                                 exclude_from_cache=True)
+    op = "test_exclude_from_cache_op_reuses_profiles_within_window"
+    tuner = AutoTuner.get()
+    tuner.clear_cache()
+    prepared, profiled = [], []
+    prepare = tuner._prepare_input_tensors
+    optimization_profiles = tuner._optimization_profiles
+    monkeypatch.setattr(
+        tuner, "_prepare_input_tensors",
+        lambda p, inputs: prepared.append(p.get_opt_shapes()) or prepare(
+            p, inputs))
+    monkeypatch.setattr(
+        tuner, "_optimization_profiles",
+        lambda config, inputs: profiled.append(op) or optimization_profiles(
+            config, inputs))
+
+    with autotune():
+        _, tactic = tuner.choose_one(op, [GemmRunner()], tuning_config, [x, w])
+        assert (len(profiled), len(prepared)) == (1, 3)
+        # Same op and shapes in the same window: served from the cache.
+        assert tuner.choose_one(op, [GemmRunner()], tuning_config,
+                                [x, w])[1] == tactic
+        assert (len(profiled), len(prepared)) == (1, 3)
+
+    # A new window walks the profiles again, but every profile is cached,
+    # so no input is prepared.
+    with autotune():
+        assert tuner.choose_one(op, [GemmRunner()], tuning_config,
+                                [x, w])[1] == tactic
+    assert (len(profiled), len(prepared)) == (2, 3)
+
+    # A nested context's record does not leak into the enclosing one.
+    with autotune():
+        with autotune(skip_dynamic_tuning_buckets=True):
+            tuner.choose_one(op, [GemmRunner()], tuning_config, [x, w])
+        tuner.choose_one(op, [GemmRunner()], tuning_config, [x, w])
+    assert (len(profiled), len(prepared)) == (4, 3)
+
+    # Different profile settings for the same op and shapes are profiled.
+    wider_config = dataclasses.replace(tuning_config, tune_max_num_tokens=8)
+    with autotune():
+        tuner.choose_one(op, [GemmRunner()], tuning_config, [x, w])
+        tuner.choose_one(op, [GemmRunner()], wider_config, [x, w])
+    assert len(profiled) == 6
+
+
+def test_exclude_from_cache_op_is_reprofiled_until_every_profile_is_cached(
+        monkeypatch) -> None:
+    """A profile that failed to produce a cache entry is retried on the next
+    call in the same autotune() context."""
+    x = torch.randn(3, 64)
+    w = torch.randn(64, 128)
+    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
+        input_idx=0,
+        dim_idx=0,
+        gen_tuning_buckets=(3, 4, 5),
+        map_to_tuning_buckets=lambda value: value), ),
+                                 exclude_from_cache=True)
+    op = "test_exclude_from_cache_op_is_reprofiled_until_every_profile_is_cached"
+    tuner = AutoTuner.get()
+    tuner.clear_cache()
+    profiled = []
+    profile_runners = tuner._profile_runners
+
+    def fail_first_profile(custom_op, runners, tensors, profile, config,
+                           **kwargs):
+        profiled.append(profile.get_opt_shapes()[0][0])
+        if len(profiled) == 1:
+            return None, -1, float("inf"), True
+        return profile_runners(custom_op, runners, tensors, profile, config,
+                               **kwargs)
+
+    monkeypatch.setattr(tuner, "_profile_runners", fail_first_profile)
+    generated = []
+    optimization_profiles = tuner._optimization_profiles
+    monkeypatch.setattr(
+        tuner, "_optimization_profiles",
+        lambda config, inputs: generated.append(op) or optimization_profiles(
+            config, inputs))
+    counts = []
+    with autotune():
+        for _ in range(3):
+            tuner.choose_one(op, [GemmRunner()], tuning_config, [x, w])
+            counts.append((len(generated), len(profiled)))
+    # Call 1 profiles all 3 buckets (the first fails), call 2 retries only the
+    # failed one, call 3 is served from the cache without generating profiles.
+    assert counts == [(1, 3), (2, 4), (2, 4)]
+
+
 @torch.library.custom_op("autotuner_test::get_best_gemm_tactic",
                          mutates_args=())
 def get_best_gemm_tactic(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: