From 9a34ccf50230332ce0982874b63b18bacdbe8093 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Fri, 21 Nov 2025 02:48:09 +0000 Subject: [PATCH 1/3] [TRTLLM-7963][fix] Remove the duplicated shape profile generating to avoid the host overhead. This PR skips the shape profile generating process if the profile has already been found in the cache under tuning mode. This is a prerequisite for nested autotuning. Because host overhead might be included during the profiling process of the high-level op. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tensorrt_llm/_torch/autotuner.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 4121be04c552..9c2ff7cd54c9 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -698,16 +698,14 @@ 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) + # Early return if it's not tuning, use cache found one or fallback one if not self.is_tuning_mode: - is_cache_hit, best_runner_id, best_tactic, min_time = self.profiling_cache.search_cache( - custom_op, runners, input_shapes, tuning_config) best_runner = runners[best_runner_id] # TODO: check the stored runner and tactic can implement this shape here - # Should not directly try (runner, tactic) here, or it will hurt a lot of inference perf. - - # Record the cache miss config. - # Expect no cache miss in inference. Thus, any cache miss should be recorded. + # Log the cache miss. Expect no cache miss in inference. if not is_cache_hit: logger.warning_once( f"[AutoTunner] Using the fallback tactic, due to cache miss on input shapes={input_shapes}", @@ -715,6 +713,10 @@ def choose_one( return (best_runner, best_tactic) + # 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: + return (runners[best_runner_id], best_tactic) + assert len(runners) > 0, "At least one runner is required" assert all([isinstance(r, TunableRunner) for r in runners]), \ "All Given runners must be subclass of TunableRunner" From 580e9d4f57ad000a5ba2ecd496c6e397b44cdba6 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Fri, 21 Nov 2025 03:26:30 +0000 Subject: [PATCH 2/3] Change CUDA graph based profiling to default on. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tensorrt_llm/_torch/autotuner.py | 2 +- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 9c2ff7cd54c9..b94830b30c03 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -99,7 +99,7 @@ class TuningConfig: constraint_specs: Tuple[ConstraintSpec, ...] = () tune_max_num_tokens: int = None inputs_pre_hook: Callable = None - use_cuda_graph: bool = False + use_cuda_graph: bool = True @dataclass(unsafe_hash=True) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 319eaba62663..5a660a179ca4 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -37,7 +37,6 @@ class CuteDSLNVFP4BlackwellLinear(TunableRunner): 0, 0, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2), ), constraint_specs=(ConstraintSpec(2, 0, fp4_scale_infer_shape), ), - use_cuda_graph=True, ) def __init__(self, alpha: float, output_dtype: torch.dtype): From d2de14703db8069f963fde1082a6265a473db589 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Fri, 21 Nov 2025 08:21:38 +0000 Subject: [PATCH 3/3] Implement heuristic tuning repeats. Tune fewer times when kernel takes long enough time to complete. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tensorrt_llm/_torch/autotuner.py | 81 ++++++++++++++++++++------------ 1 file changed, 50 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index b94830b30c03..29e2ec0b4e95 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -526,7 +526,7 @@ class AutoTuner: _CUDA_GRAPH_DELAY_MICRO_SECS = 100 _instance = None - def __init__(self, warmup=3, repeat=10, stream_delay_micro_secs=1000): + def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000): self.repeat = repeat self.warmup = warmup self.stream_delay_micro_secs = stream_delay_micro_secs @@ -883,43 +883,62 @@ def _profile_single_kernel( are used to ensure accurate timing. """ stream = torch.cuda.current_stream() - graph = torch.cuda.CUDAGraph() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - with torch.cuda.stream(stream): - # warm up, no timing - for _ in range(self.warmup): - runner(inputs, tactic=tactic, **kwargs) - - if use_cuda_graph: - with torch.cuda.graph(graph): - for _ in range(self.repeat): - runner(inputs, tactic=tactic, **kwargs) + # If the warm up time is longer than 0.5ms, we will profile the kernel with fewer repeats. + profile_fewer_repeat = 2 + short_profile_threshold_ms = 1 + + avg_time = float('inf') + + def pure_profile(stream: torch.cuda.Stream, repeat: int): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + graph = torch.cuda.CUDAGraph() + + with torch.cuda.stream(stream): + if use_cuda_graph: + with torch.cuda.graph(graph): + for _ in range(repeat): + runner(inputs, tactic=tactic, **kwargs) + + stream.synchronize() + + # Delay the profiled kernel launch to eliminate affects of host time overhead in profiling. + # TODO: This is build time sensitive, O(tactic_num * impl_num * num_profile * tunable_ops) + # Consider apply a preprofiling to estimate the kernel execution time, then decide the necessity. + if use_cuda_graph: + delay_kernel(self._CUDA_GRAPH_DELAY_MICRO_SECS, stream) + else: + delay_kernel(self.stream_delay_micro_secs, stream) - stream.synchronize() + start.record() - # Delay the profiled kernel launch to eliminate affects of host time overhead in profiling. - # TODO: This is build time sensitive, O(tactic_num * impl_num * num_profile * tunable_ops) - # Consider apply a preprofiling to estimate the kernel execution time, then decide the necessity. - if use_cuda_graph: - delay_kernel(self._CUDA_GRAPH_DELAY_MICRO_SECS, stream) - else: - delay_kernel(self.stream_delay_micro_secs, stream) + if use_cuda_graph: + graph.replay() + else: + for _ in range(repeat): + runner(inputs, tactic=tactic, **kwargs) - start.record() + end.record() + stream.synchronize() - if use_cuda_graph: - graph.replay() - else: - for _ in range(self.repeat): - runner(inputs, tactic=tactic, **kwargs) + return start.elapsed_time(end) / repeat - end.record() + for _ in range(self.warmup): + runner(inputs, tactic=tactic, **kwargs) - stream.synchronize() + fewer_repeat_avg_time = pure_profile(stream, profile_fewer_repeat) - avg_time = start.elapsed_time(end) / self.repeat + disable_short_profile = os.environ.get( + "TLLM_AUTOTUNER_DISABLE_SHORT_PROFILE", "0") == "1" + if fewer_repeat_avg_time > short_profile_threshold_ms and not disable_short_profile: + print( + f"[Autotuner] Few repeat estimated time is longer than {short_profile_threshold_ms}ms, directly use the few repeat estimated time to avoid redundant profiling." + ) + # directly use the few repeat estimated time to avoid redundant profiling + avg_time = fewer_repeat_avg_time + else: + # profile the kernel with the full repeat to get precise time + avg_time = pure_profile(stream, self.repeat) shapes = self._get_input_sizes(inputs) logger.debug(