[None][perf] Autotuner: distributed tuning for mhc + single-(runner,tactic) shortcut

tensorrt_llm/_torch/autotuner.py

@@ -124,6 +124,9 @@ class TuningConfig:
             Notice that not all tuning processes can benefit from this feature.
         use_cuda_graph (bool): Whether to use CUDA graph for the tuning process.
         distributed_tuning_strategy (DistributedTuningStrategy): Strategy for distributed tuning.
+        exclude_from_cache (bool): Skip disk persistence of tuning results for this op.
+            Use for JIT-driven ops where caching the tactic across processes would skip
+            the JIT warmup. In-process cache still works.
     """
     dynamic_tensor_specs: Tuple[DynamicTensorSpec, ...] = ()
     constraint_specs: Tuple[ConstraintSpec, ...] = ()
@@ -132,6 +135,7 @@ class TuningConfig:
     use_cold_l2_cache: bool = False
     use_cuda_graph: bool = True
     distributed_tuning_strategy: DistributedTuningStrategy = DistributedTuningStrategy.INDEPENDENT
+    exclude_from_cache: bool = False
 
 
 @dataclass(unsafe_hash=True)
@@ -409,6 +413,11 @@ def __init__(self):
         # Maps custom_op name -> DistributedTuningStrategy
         self.independent_op: Set[str] = set()
 
+        # Ops whose tuning results must NOT be persisted to disk (see the
+        # `exclude_from_cache` field on TuningConfig). Populated lazily
+        # from choose_one when an excluded op is first seen.
+        self.excluded_op: Set[str] = set()
+
         # Cache metadata for local storage and validation
         self.lib_version = tensorrt_llm.__version__
         self.creation_timestamp = time.time()
@@ -428,6 +437,7 @@ def __len__(self) -> int:
     def clear(self) -> None:
         self.cache.clear()
         self.independent_op.clear()
+        self.excluded_op.clear()
 
     def fallback_entry(self) -> Tuple:
         # runner_id = 0, tactic = -1
@@ -497,6 +507,19 @@ def add_independent_op(self, custom_op: str,
         if strategy != DistributedTuningStrategy.INDEPENDENT:
             self.independent_op.add(custom_op)
 
+    def add_excluded_op(self, custom_op: str) -> None:
+        """Mark `custom_op` as excluded from disk persistence.
+
+        Entries for excluded ops still live in the in-memory cache (so
+        within one process the lookup-and-reuse path works normally),
+        but `_partition_cache_by_strategy` will drop them at `save_cache`
+        time so they are never written to disk. Combined with the
+        force-retune branch in `AutoTuner.choose_one`, this guarantees
+        the tuning process (and any JIT warmup it triggers as a side
+        effect) re-runs on every process startup.
+        """
+        self.excluded_op.add(custom_op)
+
     def _partition_cache_by_strategy(
             self) -> Tuple[Dict[Tuple, Tuple], Dict[Tuple, Tuple]]:
         """Partition cache entries into shared and rank-specific caches.
@@ -505,12 +528,16 @@ def _partition_cache_by_strategy(
             A tuple of (shared_cache, rank_cache) where:
             - shared_cache: entries for non-INDEPENDENT ops (BROADCAST, MERGE, PARALLEL)
             - rank_cache: entries for INDEPENDENT ops
+            Entries belonging to `self.excluded_op` are dropped from both
+            partitions so they never reach disk.
         """
         shared_cache = {}
         rank_cache = {}
 
         for key, value in self.cache.items():
             custom_op = key[0]  # First element of cache key is custom_op name
+            if custom_op in self.excluded_op:
+                continue
             if custom_op not in self.independent_op:
                 rank_cache[key] = value
             else:
@@ -524,6 +551,7 @@ def save_cache(self, file_path: Union[str, Path], rank: int) -> None:
         Cache entries are organized based on distributed strategy:
         - INDEPENDENT ops are saved per-rank (rank_0, rank_1, ...)
         - Non-INDEPENDENT ops (BROADCAST, MERGE, PARALLEL) are saved in a shared dict
+        - Ops marked `exclude_from_cache=True` are dropped entirely.
 
         Args:
             file_path: Path where to save the cache
@@ -1014,12 +1042,28 @@ def choose_one(
             })
 
         input_shapes = tuple(self._get_input_sizes(inputs))
-        is_cache_hit, best_runner_id, best_tactic, min_time = self.profiling_cache.search_cache(
-            custom_op,
-            runners,
-            input_shapes,
-            tuning_config,
-            apply_map_to_tuning_buckets=True)
+
+        # For ops opted out of the persistent cache (e.g. JIT-driven
+        # kernels whose "tuning" triggers compilation as a side effect),
+        # mark them in the cache so save_cache filters them out, and
+        # — 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.
+        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:
+            is_cache_hit = False
+            best_runner_id, best_tactic, min_time = (
+                self.profiling_cache.fallback_entry())
+        else:
+            is_cache_hit, best_runner_id, best_tactic, min_time = (
+                self.profiling_cache.search_cache(
+                    custom_op,
+                    runners,
+                    input_shapes,
+                    tuning_config,
+                    apply_map_to_tuning_buckets=True))
 
         # Early return if it's not tuning, use cache found one or fallback one
         if not self.is_tuning_mode:
@@ -1126,6 +1170,15 @@ def _profile_runners(
         # If the inputs_pre_hook is provided, it will be called before profiling.
         if tuning_config.inputs_pre_hook is not None:
             input_tensors = tuning_config.inputs_pre_hook(input_tensors)
+
+        # Pre-walk the runners so we know the total candidate (runner, tactic)
+        # pair count before deciding whether to short-circuit profiling. The
+        # count is taken BEFORE distributed-strategy splitting; the PARALLEL
+        # strategy can leave a rank with one tactic locally even though the
+        # full set has several, and that case must still go through the
+        # timed profile path so the merge picks a real winner.
+        candidates = []
+        total_pairs = 0
         for runner_id, runner in enumerate(runners):
             # TODO: use FakeTensor here.
             runner_arg_names = {
@@ -1134,52 +1187,94 @@ def _profile_runners(
             }
             all_valid_tactics = runner.get_valid_tactics(
                 input_tensors, profile, **kwargs)
-
-            valid_tactics = self._maybe_parallelize_tactics(
-                all_valid_tactics, tuning_config.distributed_tuning_strategy)
-            if "do_preparation" in runner_arg_names and len(valid_tactics) > 0:
-                runner(
-                    input_tensors,
-                    tactic=-1,
-                    do_preparation=True,
-                    **kwargs,
+            total_pairs += len(all_valid_tactics)
+            candidates.append(
+                (runner_id, runner, runner_arg_names, all_valid_tactics))
+
+        if total_pairs == 1:
+            # Single-pair shortcut. There is nothing to compare against, so
+            # the timed warmup + profile-repeat loop is pure overhead. We
+            # fire one forward call on every rank to drive any JIT side
+            # effects (e.g. DeepGEMM cubin compile for fp8_swap_ab_gemm),
+            # then record the pair in the cache directly. `do_preparation`
+            # still runs for runners that opt into it.
+            runner_id, runner, runner_arg_names, valid_tactics = next(
+                c for c in candidates if len(c[3]) == 1)
+            tac = valid_tactics[0]
+            if "do_preparation" in runner_arg_names:
+                runner(input_tensors, tactic=-1, do_preparation=True, **kwargs)
+            try:
+                with nvtx_range(f"r{runner_id}, tactic {tac} (single-pair)"):
+                    runner(input_tensors, tactic=tac, **kwargs)
+                best_runner_id, best_tactic = runner_id, tac
+                # No comparable measurement exists; min_time stays 0.0 to
+                # distinguish "recorded without profiling" from a real
+                # timing.
+                min_time = 0.0
+            except Exception as e:
+                shapes = self._get_input_sizes(input_tensors)
+                logger.warning_once(
+                    f"[Autotuner] Single-pair run failed for custom_op={custom_op}, runner={runner}, tactic={tac}, shapes={shapes}. Error: {e}",
+                    key=(custom_op, "warning_autotuning_single_pair_failure"),
                 )
+                self.stats.failed_profiling_count[custom_op].add(
+                    self.profiling_cache.get_cache_key(
+                        custom_op,
+                        runner,
+                        profile.get_opt_shapes(),
+                        tuning_config,
+                        apply_map_to_tuning_buckets=False))
+                has_tuning_failure_occurred = True
+        else:
+            for runner_id, runner, runner_arg_names, all_valid_tactics in candidates:
+                valid_tactics = self._maybe_parallelize_tactics(
+                    all_valid_tactics,
+                    tuning_config.distributed_tuning_strategy)
+                if "do_preparation" in runner_arg_names and len(
+                        valid_tactics) > 0:
+                    runner(
+                        input_tensors,
+                        tactic=-1,
+                        do_preparation=True,
+                        **kwargs,
+                    )
 
-            for tac in valid_tactics:
-                try:
-                    with nvtx_range(f"r{runner_id}, tactic {tac}"):
-                        time_measured = self._profile_single_kernel(
-                            runner=runner,
-                            inputs=input_tensors,
-                            tactic=tac,
-                            tuning_config=tuning_config,
-                            use_cuda_graph=tuning_config.use_cuda_graph,
-                            **kwargs,
+                for tac in valid_tactics:
+                    try:
+                        with nvtx_range(f"r{runner_id}, tactic {tac}"):
+                            time_measured = self._profile_single_kernel(
+                                runner=runner,
+                                inputs=input_tensors,
+                                tactic=tac,
+                                tuning_config=tuning_config,
+                                use_cuda_graph=tuning_config.use_cuda_graph,
+                                **kwargs,
+                            )
+                    except Exception as e:
+                        # Handle None tensors for optional inputs
+                        shapes = self._get_input_sizes(input_tensors)
+                        logger.warning_once(
+                            f"[Autotuner] Failed when profiling runner={runner}, tactic={tac}, shapes={shapes}. Error: {e}",
+                            key=(custom_op,
+                                 "warning_autotuning_profile_failure"),
                         )
-                except Exception as e:
-                    # Handle None tensors for optional inputs
-                    shapes = self._get_input_sizes(input_tensors)
-                    logger.warning_once(
-                        f"[Autotuner] Failed when profiling runner={runner}, tactic={tac}, shapes={shapes}. Error: {e}",
-                        key=(custom_op, "warning_autotuning_profile_failure"),
-                    )
 
-                    # Record the failed profiling combinations
-                    self.stats.failed_profiling_count[custom_op].add(
-                        self.profiling_cache.get_cache_key(
-                            custom_op,
-                            runner,
-                            profile.get_opt_shapes(),
-                            tuning_config,
-                            apply_map_to_tuning_buckets=False))
-
-                    # Set time_measured to inf to notify the failure of the tactic. This can happen when `get_valid_tactics` mistakenly return wrong tactics
-                    # or some runtime error occurs during profiling.
-                    time_measured = float('inf')
-                    has_tuning_failure_occurred = True
-                if time_measured < min_time:
-                    min_time = time_measured
-                    best_runner_id, best_tactic = runner_id, tac
+                        # Record the failed profiling combinations
+                        self.stats.failed_profiling_count[custom_op].add(
+                            self.profiling_cache.get_cache_key(
+                                custom_op,
+                                runner,
+                                profile.get_opt_shapes(),
+                                tuning_config,
+                                apply_map_to_tuning_buckets=False))
+
+                        # Set time_measured to inf to notify the failure of the tactic. This can happen when `get_valid_tactics` mistakenly return wrong tactics
+                        # or some runtime error occurs during profiling.
+                        time_measured = float('inf')
+                        has_tuning_failure_occurred = True
+                    if time_measured < min_time:
+                        min_time = time_measured
+                        best_runner_id, best_tactic = runner_id, tac
 
         if best_runner_id is not None:
             # At least one valid (runner, tactic) pair is found

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -1593,8 +1593,19 @@ def forward(
 class fp8SwapABGemmRunner(TunableRunner):
     """Runs quantize + DeepGemm FP8 GEMM. Single tactic for JIT warmup."""
 
-    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
-        0, 0, deep_gemm_gen_tuning_buckets), ), )
+    # DeepGemm performs JIT compilation as a side effect of the tuning
+    # call. If we let the autotuner persist this op's tactic to disk, a
+    # subsequent process startup loads the cached tactic and skips the
+    # tuning path entirely — meaning the DeepGemm JIT warmup never runs
+    # and inference falls back to a slower uncompiled path.
+    # `exclude_from_cache=True` keeps the in-process cache working but
+    # ensures the tuning (and the JIT warmup it triggers) repeats on
+    # every process startup.
+    tuning_config = TuningConfig(
+        dynamic_tensor_specs=(DynamicTensorSpec(
+            0, 0, deep_gemm_gen_tuning_buckets), ),
+        exclude_from_cache=True,
+    )
 
     def __init__(self, output_dtype: torch.dtype, disable_ue8m0_cast: bool,
                  quant_tactic: int):

tensorrt_llm/_torch/modules/mhc/mhc_cuda.py

@@ -22,6 +22,7 @@
 from tensorrt_llm._torch.autotuner import (
     AutoTuner,
     ConstraintSpec,
+    DistributedTuningStrategy,
     DynamicTensorSpec,
     OptimizationProfile,
     TunableRunner,
@@ -271,6 +272,9 @@ class MhcPreMappingRunner(TunableRunner):
                 infer_shape=lambda shapes: shapes[0][0],
             ),
         ),
+        # PARALLEL fans tactics out across TP ranks (each rank times a
+        # disjoint subset, then merge picks the global best).
+        distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL,
     )
 
     def __init__(
@@ -844,6 +848,7 @@ class MhcFusedHcRunner(TunableRunner):
             # comb_mix_prev (input[3]) dim 0 = M
             ConstraintSpec(input_idx=3, dim_idx=0, infer_shape=lambda shapes: shapes[0][0]),
         ),
+        distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL,
     )
 
     def __init__(

tests/unittest/_torch/misc/test_autotuner.py

@@ -986,3 +986,79 @@ def forward(self,
               f"{event_ms:>16.4f} {gt_ms:>17.4f} "
               f"{abs_diff:>14.4f} {rel_diff * 100:>13.2f}")
     print("-" * 102)
+
+
+def test_single_pair_shortcut(monkeypatch):
+    """Single (runner, tactic) candidate must bypass the timed profile loop.
+
+    When ``_profile_runners`` sees exactly one (runner, tactic) pair, it
+    must (1) skip ``_profile_single_kernel`` entirely, (2) still fire the
+    ``do_preparation`` hook for runners that opt in, (3) fire exactly one
+    ``forward()`` to drive any JIT side effect, and (4) record the pair
+    in the profiling cache. Multi-tactic ops in the same fixture must
+    still use the timed path.
+    """
+
+    profile_calls: List[Any] = []
+
+    def _track(self, runner, inputs, tactic, tuning_config, **kwargs):
+        profile_calls.append(tactic)
+        return 1.0 + len(profile_calls) * 0.01
+
+    monkeypatch.setattr(AutoTuner, "_profile_single_kernel", _track)
+
+    forward_calls: List[tuple] = []
+
+    class PrepRunner(TunableRunner):
+
+        def unique_id(self):
+            return ()
+
+        def get_valid_tactics(self, inputs: List[FakeTensor],
+                              profile: OptimizationProfile,
+                              **kwargs) -> List[int]:
+            return [0]
+
+        def forward(self,
+                    /,
+                    inputs: List[torch.Tensor],
+                    *,
+                    tactic: int = -1,
+                    do_preparation: bool = False,
+                    **kwargs) -> torch.Tensor:
+            forward_calls.append((tactic, do_preparation))
+            if do_preparation:
+                return None
+            x, w = inputs
+            return x @ w
+
+    tuner = AutoTuner.get()
+    tuner.clear_cache()
+    x = torch.randn(M, 64, device="cuda")
+    w = torch.randn(64, 128, device="cuda")
+
+    # Single (runner, tactic): shortcut must fire.
+    op_single = "autotuner_test::single_pair_shortcut"
+    with autotune():
+        _, tactic = tuner.choose_one(op_single, [PrepRunner()], TuningConfig(),
+                                     [x, w])
+    assert tactic == 0
+    assert profile_calls == [], (
+        f"_profile_single_kernel must not be called for single-pair op; "
+        f"got {profile_calls}")
+    assert forward_calls == [
+        (-1, True), (0, False)
+    ], (f"Expected do_preparation then exactly one forward(tactic=0); "
+        f"got {forward_calls}")
+    assert len(tuner.profiling_cache.get_specific_custom_op(op_single)) == 1, (
+        "single-pair shortcut must still record the (runner, tactic) entry")
+
+    # Multi-tactic on the same fixture: timed profile path must still run.
+    forward_calls.clear()
+    op_multi = "autotuner_test::single_pair_shortcut_multi"
+    with autotune():
+        tuner.choose_one(op_multi, [GemmRunner()], TuningConfig(), [x, w])
+    # GemmRunner exposes 3 tactics -> 3 profile calls.
+    assert len(profile_calls) == 3, (
+        f"Multi-tactic op must hit _profile_single_kernel per tactic; "
+        f"got {len(profile_calls)} ({profile_calls})")