diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index eaee69e54ebb..99fbafaf8797 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/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, diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 3bacfa74071e..4a25e93856d9 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/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(