[TRTLLM-10288][perf] Reduce AutoTuner host overhead in inference hot path

tensorrt_llm/_torch/autotuner.py

@@ -253,6 +253,21 @@ def unique_id(self):
         """
         return tuple(self.__dict__.values())
 
+    @property
+    def _cache_key_prefix(self) -> Tuple[str, str]:
+        """Cached (class_name, unique_id_str) tuple — computed once per runner instance.
+
+        This avoids re-computing runner.__class__.__name__ and str(runner.unique_id())
+        on every get_cache_key() call during inference. These values are constant for a
+        given runner instance.
+        """
+        try:
+            return self.__cache_key_prefix
+        except AttributeError:
+            self.__cache_key_prefix = (self.__class__.__name__,
+                                       str(self.unique_id()))
+            return self.__cache_key_prefix
+
 
 @contextlib.contextmanager
 def autotune(tune_mode: bool = True,
@@ -298,6 +313,9 @@ def autotune(tune_mode: bool = True,
         autotuner.skip_dynamic_tuning_buckets = old_skip
         if autotune_enabled:
             logger.info("[Autotuner] Autotuning process ends")
+            # Invalidate the choose_one cache since the profiling cache
+            # has been updated with new entries.
+            autotuner._choose_one_cache.clear()
 
         # save cache
         if cache_path is not None:
@@ -447,11 +465,15 @@ def get_cache_key(
         tuning_config: TuningConfig,
         apply_map_to_tuning_buckets: bool = True,
     ) -> Tuple:
+        # Use cached (class_name, unique_id_str) prefix to avoid
+        # recomputing str(runner.unique_id()) on every call (~1-2us saving).
+        prefix = runner._cache_key_prefix
         return (
             custom_op,
-            runner.__class__.__name__,
-            str(runner.unique_id()),
-            AutoTuner.get()._find_nearest_profile(
+            prefix[0],
+            prefix[1],
+            # Call classmethod directly — avoids AutoTuner.get() singleton lookup.
+            AutoTuner._find_nearest_profile(
                 input_shapes,
                 tuning_config.dynamic_tensor_specs,
                 tuning_config.constraint_specs,
@@ -760,6 +782,14 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000):
         # Last captured choose_one() contexts
         self._last_capture: Optional['AutoTuner.TacticsCapture'] = None
 
+        # Inference fast-path cache for choose_one().
+        # Maps (custom_op, runner_ids, *bucketed_dims) -> (runner_id, tactic).
+        # On cache hit, resolves the runner via runners[runner_id] from the
+        # caller's list — never caches runner instances directly.
+        # Skips the full search_cache → _find_nearest_profile chain on
+        # repeated calls.
+        self._choose_one_cache: dict = {}
+
         # Dsitributed tuning state
         self._dist: Optional[Distributed] = None
         self._has_received_cache: bool = False
@@ -888,6 +918,31 @@ def choose_one(
             Runner authors are suggested to provide a fallback implementation for each runner to avoid potential issues.
         """
 
+        # ---- NON-TUNING PATH: fast-key cache ----
+        if not self.is_tuning_mode and self._active_capture is None:
+            fast_key = self._make_fast_key(custom_op, runners, tuning_config,
+                                           inputs)
+            cached = self._choose_one_cache.get(fast_key)
+            if cached is not None:
+                best_runner_id, best_tactic = cached
+                return (runners[best_runner_id], best_tactic)
+
+            # Cache miss — full resolution via search_cache
+            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)
+            if not is_cache_hit:
+                logger.warning_once(
+                    f"[AutoTuner] {custom_op} using the fallback tactic, "
+                    f"due to cache miss on input shapes={input_shapes}",
+                    key=(custom_op, "warning_autotuning_cache_miss_fallback"))
+            self._choose_one_cache[fast_key] = (best_runner_id, best_tactic)
+            return (runners[best_runner_id], best_tactic)
+
+        # ---- TUNING / CAPTURE / REPLAY ----
+
         # Check if we're in replay mode via active TacticsCapture
         if self._active_capture is not None and self._active_capture.is_replaying(
         ):
@@ -935,27 +990,13 @@ 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)
-
-        # Early return if it's not tuning, use cache found one or fallback one
-        if not self.is_tuning_mode:
-            best_runner = runners[best_runner_id]
-            # TODO: check the stored runner and tactic can implement this shape here
-            # Log the cache miss. Expect no cache miss in inference.
-            if not is_cache_hit:
-                logger.warning_once(
-                    f"[AutoTuner] {custom_op} using the fallback tactic, due to cache miss on input shapes={input_shapes}",
-                    key=(custom_op, "warning_autotuning_cache_miss_fallback"))
-
-            return (best_runner, best_tactic)
+        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)
 
-        # If it's tuning mode and cache hit, return the best runner and tactic to avoid redundant profiling.
-        if self.is_tuning_mode and is_cache_hit:
+        # Cache hit — skip profiling (avoids redundant work during tuning/capture).
+        if is_cache_hit:
             return (runners[best_runner_id], best_tactic)
 
         # PP rank does not have cache hit, so we try to receive the cache from the previous rank
@@ -1025,6 +1066,38 @@ def choose_one(
                 custom_op, 0) + tuning_end_time - tuning_start_time
         return (runners[runner_id], tactic)
 
+    # ------------------------------------------------------------------
+    # _make_fast_key — bucketed cache key for choose_one fast path
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _make_fast_key(custom_op, runners, tuning_config, inputs):
+        """Build a minimal cache key from runner identity + bucketed dynamic dims.
+
+        Includes each runner's ``_cache_key_prefix`` (class name + unique_id)
+        so that different runner configurations for the same ``custom_op``
+        do not collide.  For the dynamic dimensions we extract only the 1–2
+        integer values that actually vary at inference time (e.g. batch size)
+        and bucket them.  This gives an O(R + K) key where R = len(runners)
+        (typically 1–3) and K = len(dynamic_tensor_specs) (typically 1).
+        """
+        runner_ids = tuple(r._cache_key_prefix for r in runners)
+        dynamic_specs = tuning_config.dynamic_tensor_specs
+        if not dynamic_specs:
+            return (custom_op, runner_ids)
+        max_tokens = tuning_config.tune_max_num_tokens
+        dyn_vals = []
+        for spec in dynamic_specs:
+            inp = inputs[spec.input_idx]
+            if isinstance(inp, torch.Tensor):
+                v = spec.map_to_tuning_buckets(inp.size(spec.dim_idx))
+            else:
+                v = 0
+            if max_tokens is not None:
+                v = min(v, max_tokens)
+            dyn_vals.append(v)
+        return (custom_op, runner_ids, *dyn_vals)
+
     def _profile_runners(
         self,
         custom_op: str,

tests/unittest/_torch/misc/test_autotuner.py

@@ -846,6 +846,111 @@ def has_op_in_section(section_data: dict, op_name: str) -> bool:
     return True
 
 
+def test_choose_one_cache_hit():
+    """choose_one() should cache (runner_id, tactic) and reuse on subsequent calls."""
+    x, w = torch.randn(16, 64), torch.randn(64, 128)
+    runner = GemmRunner()
+    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
+        input_idx=0,
+        dim_idx=0,
+        gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
+        map_to_tuning_buckets=next_positive_power_of_2), ), )
+
+    tuner = AutoTuner.get()
+    tuner.clear_cache()
+    tuner._choose_one_cache.clear()
+
+    # Tune first
+    with autotune():
+        tuner.choose_one("test_choose_one_cache", [runner], tuning_config,
+                         [x, w])
+
+    # First choose_one — cache miss, populates _choose_one_cache
+    runner1, tactic1 = tuner.choose_one("test_choose_one_cache", [runner],
+                                        tuning_config, [x, w])
+    assert runner1 is runner
+    assert isinstance(tactic1, int)
+
+    fast_key = AutoTuner._make_fast_key("test_choose_one_cache", [runner],
+                                        tuning_config, [x, w])
+    assert fast_key in tuner._choose_one_cache, \
+        "_choose_one_cache should have entry after first call"
+
+    # Second choose_one — cache hit, same values
+    runner2, tactic2 = tuner.choose_one("test_choose_one_cache", [runner],
+                                        tuning_config, [x, w])
+    assert runner1 is runner2, "Cached choose_one should return same runner"
+    assert tactic1 == tactic2, "Cached choose_one should return same tactic"
+
+    # Re-tuning should clear _choose_one_cache
+    with autotune():
+        tuner.choose_one("test_choose_one_cache", [runner], tuning_config,
+                         [x, w])
+    assert len(tuner._choose_one_cache) == 0, \
+        "_choose_one_cache should be cleared after re-tuning"
+
+
+def test_choose_one_cache_different_runners_same_op():
+    """choose_one() cache must not collide when different runners share the same custom_op.
+
+    Reproduces the bug where _choose_one_cache key was (custom_op, *bucketed_dims)
+    without runner identity, causing a runner configured for one dtype to be
+    returned for a call with a different dtype (e.g. int8 runner used for int4 data).
+    """
+
+    class TypedRunner(TunableRunner):
+        """Runner whose output depends on a config parameter (simulates weight_dtype)."""
+
+        def __init__(self, scale: float):
+            super().__init__()
+            self.scale = scale
+
+        def unique_id(self):
+            return (self.scale, )
+
+        def get_valid_tactics(self, inputs, profile, **kwargs):
+            return [0]
+
+        def forward(self, /, inputs, *, tactic=0, **kwargs):
+            return inputs[0] @ inputs[1] * self.scale
+
+    x, w = torch.randn(16, 64), torch.randn(64, 128)
+    runner_a = TypedRunner(scale=1.0)
+    runner_b = TypedRunner(scale=2.0)
+    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
+        input_idx=0,
+        dim_idx=0,
+        gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
+        map_to_tuning_buckets=next_positive_power_of_2), ), )
+
+    tuner = AutoTuner.get()
+    tuner.clear_cache()
+    tuner._choose_one_cache.clear()
+
+    # Tune both runners under the SAME custom_op name
+    op_name = "test_cache_collision"
+    with autotune():
+        tuner.choose_one(op_name, [runner_a], tuning_config, [x, w])
+    with autotune():
+        tuner.choose_one(op_name, [runner_b], tuning_config, [x, w])
+
+    # First call — runner_a
+    r1, t1 = tuner.choose_one(op_name, [runner_a], tuning_config, [x, w])
+    result_a = r1(inputs=[x, w], tactic=t1)
+
+    # Second call — runner_b (MUST NOT return runner_a's cached result)
+    r2, t2 = tuner.choose_one(op_name, [runner_b], tuning_config, [x, w])
+    result_b = r2(inputs=[x, w], tactic=t2)
+
+    expected_a = x @ w * 1.0
+    expected_b = x @ w * 2.0
+    assert torch.allclose(result_a, expected_a, atol=1e-5), \
+        f"Runner A should produce scale=1.0 result"
+    assert torch.allclose(result_b, expected_b, atol=1e-5), \
+        f"Runner B should produce scale=2.0 result, got scale={result_b[0,0]/expected_a[0,0]:.1f}x (cache collision!)"
+    assert r1 is not r2, "Different runners must not be aliased by cache"
+
+
 @pytest.mark.skipif(torch.cuda.device_count() < 2,
                     reason="Requires at least 2 GPUs for this test")
 @pytest.mark.parametrize(