From ed45baf818dfb533d3a27f3a6ec24df2e0e18681 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Wed, 4 Nov 2026 06:37:23 +0000 Subject: [PATCH 1/4] [None][perf] autotuner: let excluded ops hit entries this process tuned exclude_from_cache ops (the DeepGEMM JIT-warmup runners) were forced to miss the cache on every tuning-mode call, so every layer re-ran the whole M sweep's bookkeeping (deep-copying all profiles and allocating random inputs for each) although this process had already tuned them. Only an entry this process profiled proves the JIT side effect ran here; entries loaded from disk or received from peers still force a re-tune. Input tensors are now built only for profiles that miss the cache. Kimi-K3 DEP8 cold start: about 230 s less autotuner/mamba warmup. Test: test_excluded_op_reuses_the_sweep_profiled_in_this_process. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- tensorrt_llm/_torch/autotuner.py | 32 +++++++++----- tests/unittest/_torch/misc/test_autotuner.py | 45 ++++++++++++++++++++ 2 files changed, 67 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 9ef87cc93d04..b19c9ddd696f 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -957,6 +957,9 @@ 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() + # Cache keys whose runner ran in this process (so any JIT it drives + # has happened here), as opposed to keys loaded from disk or peers. + self._keys_profiled_in_process: set[Tuple] = set() self.is_tuning_mode = False self.skip_dynamic_tuning_buckets = False @@ -1173,18 +1176,25 @@ def choose_one( # 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( + 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 (tuning_config.exclude_from_cache and self.is_tuning_mode + and is_cache_hit and self.profiling_cache.get_cache_key( custom_op, - runners, + runners[best_runner_id], input_shapes, tuning_config, - apply_map_to_tuning_buckets=True)) + apply_map_to_tuning_buckets=True) + not in self._keys_profiled_in_process): + # Only an entry this process profiled proves its JIT side + # effects already ran here; re-profile anything else. + is_cache_hit = False + best_runner_id, best_tactic, min_time = ( + self.profiling_cache.fallback_entry()) # Early return if it's not tuning, use cache found one or fallback one if not self.is_tuning_mode: @@ -1237,7 +1247,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, @@ -1246,6 +1255,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()}"): @@ -1558,6 +1568,7 @@ def _profile_runners_impl( # TODO: use named tuple to make it more readable self.profiling_cache[cache_key] = (best_runner_id, best_tactic, min_time) + self._keys_profiled_in_process.add(cache_key) self.stats.tuned_op_profiled_configs[custom_op] += 1 elif local_pairs == 0 and total_pairs > 0: @@ -2033,6 +2044,7 @@ def _prepare_input_tensors_with_batches( def clear_cache(self) -> None: """Clear the profiling cache.""" self.profiling_cache.clear() + self._keys_profiled_in_process.clear() if hasattr(torch.ops.trtllm, "clear_allreduce_tactic_cache"): torch.ops.trtllm.clear_allreduce_tactic_cache() diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index fa3986eff4ea..ef34a8b8fecd 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -2271,3 +2271,48 @@ def test_parallel_deal_start_is_stable_and_spread(): assert len(per_rank) == 1, (m, per_rank) starts.append(per_rank.pop()) assert set(starts) == set(range(group_size)), sorted(set(starts)) + + +def test_excluded_op_reuses_the_sweep_profiled_in_this_process(monkeypatch): + """Once this process has profiled an exclude_from_cache op's sweep, later + tuning-mode calls hit those entries instead of walking the sweep again, and + input tensors are only built for the profiles that missed the cache.""" + tuner = AutoTuner.get() + tuner.clear_cache() + counts = {"sweeps": 0, "inputs": 0} + real_profiles = tuner._optimization_profiles + real_inputs = tuner._prepare_input_tensors + + def profiles(*args, **kwargs): + counts["sweeps"] += 1 + return real_profiles(*args, **kwargs) + + def inputs(*args, **kwargs): + counts["inputs"] += 1 + return real_inputs(*args, **kwargs) + + monkeypatch.setattr(tuner, "_optimization_profiles", profiles) + monkeypatch.setattr(tuner, "_prepare_input_tensors", inputs) + config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, (4, 8, 16, 32)), ), + exclude_from_cache=True) + w = torch.randn(64, 128, device="cuda") + op = "autotuner_test::excluded_op_sweep_reuse" + with autotune(): + tuner.choose_one(op, [GemmRunner()], config, + [torch.randn(M, 64, device="cuda"), w]) + assert counts == {"sweeps": 1, "inputs": 4} + # The same op in later layers, at swept and unswept shapes. + for m in (M, M // 2, M): + tuner.choose_one(op, [GemmRunner()], config, + [torch.randn(m, 64, device="cuda"), w]) + assert counts == {"sweeps": 1, "inputs": 4} + + # An entry this process did not profile (e.g. received from a peer) does + # not prove the op's JIT ran here, so the call walks the sweep again. + tuner._keys_profiled_in_process.clear() + with autotune(): + tuner.choose_one(op, [GemmRunner()], config, + [torch.randn(M, 64, device="cuda"), w]) + assert counts["sweeps"] == 2 + tuner.clear_cache() From 6354444fbe303f558a000ef1abac17322e771239 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Wed, 4 Nov 2026 06:38:29 +0000 Subject: [PATCH 2/4] [None][perf] autotuner: drive the DeepGEMM JIT of a tuning sweep on a thread pool DeepGEMM's JIT (deep_jit) releases the GIL while nvcc runs, so the kernel compiles that a tuning sweep triggers can overlap within one process. TuningConfig.parallel_jit_warmup runs each not-yet-tuned profile once on a thread pool before the profiling loop, which then finds the kernels built. Calls in flight are interleaved across the sweep so that they mostly select different kernels. -1 uses this rank's share of the host's cores (up to 32); TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP overrides the thread count and 0 disables it. The three DeepGEMM JIT-warmup runners opt in with -1. Also correct comments that said DeepGEMM forks nvcc while holding the GIL. Kimi-K3 DEP8 on GB300, every rank with private empty JIT caches: DeepGEMM swap-AB JIT ~330 s -> ~69 s per rank, LLM() 494.7 -> 220.7 s. Test: test_parallel_jit_warmup_runs_pending_profiles_on_a_thread_pool. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- tensorrt_llm/_torch/autotuner.py | 81 +++++++++++++++++++ .../_torch/custom_ops/torch_custom_ops.py | 12 ++- tensorrt_llm/_torch/utils.py | 4 +- tests/unittest/_torch/misc/test_autotuner.py | 55 +++++++++++++ 4 files changed, 148 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index b19c9ddd696f..693e9743915b 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -13,6 +13,7 @@ # limitations under the License. import ast +import concurrent.futures import contextlib import copy import enum @@ -147,6 +148,14 @@ class TuningConfig: 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. + parallel_jit_warmup (int): Before profiling, run every not-yet-tuned profile once + on this many threads. For runners whose first call per shape JIT-compiles a + kernel and releases the GIL while compiling (DeepGEMM), those compiles then + overlap, and the profiling loop that follows finds them built. Opted-in + runners must be thread-safe and must not use NCCL windows. 0 disables; -1 uses + this rank's share of the host's cores (up to 32): neighbouring sweep points + often select the same kernel and compile it twice, so more threads keep + paying off. TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP overrides the thread count. """ dynamic_tensor_specs: Tuple[DynamicTensorSpec, ...] = () constraint_specs: Tuple[ConstraintSpec, ...] = () @@ -156,6 +165,7 @@ class TuningConfig: use_cuda_graph: bool = True distributed_tuning_strategy: DistributedTuningStrategy = DistributedTuningStrategy.INDEPENDENT exclude_from_cache: bool = False + parallel_jit_warmup: int = 0 @dataclass(unsafe_hash=True) @@ -1245,6 +1255,8 @@ def choose_one( # Synchronize ranks before profiling if self._should_current_rank_tune( tuning_config.distributed_tuning_strategy): + self._parallel_jit_warmup(custom_op, runners, profiles, inputs, + tuning_config, **kwargs) with self._shared_profile_graph_pool(): for p in profiles: is_cache_hit, *_ = self.profiling_cache.search_cache( @@ -1588,6 +1600,75 @@ def _profile_runners_impl( return best_runner_id, best_tactic, min_time, has_tuning_failure_occurred + def _parallel_jit_warmup(self, custom_op: str, runners: List[TunableRunner], + profiles: List[OptimizationProfile], + inputs: List[torch.Tensor], + tuning_config: TuningConfig, **kwargs) -> None: + """Run each not-yet-tuned profile once on a thread pool to drive its JIT. + + See ``TuningConfig.parallel_jit_warmup``. Purely a warmup: results are + discarded, and a failure here is left to the profiling loop, which + runs every profile again and handles errors as usual. + """ + num_threads = int( + os.environ.get("TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP", + tuning_config.parallel_jit_warmup)) + if num_threads < 0: + num_threads = min(32, (os.cpu_count() or 1) // + max(1, torch.cuda.device_count())) + if num_threads < 2: + return + pending = [ + p for p in profiles if not self.profiling_cache.search_cache( + custom_op, + runners, + p.get_opt_shapes(), + tuning_config, + apply_map_to_tuning_buckets=False, + )[0] + ] + num_threads = min(num_threads, len(pending)) + if num_threads < 2: + return + # Interleave contiguous blocks of the sweep, so the calls in flight + # come from far-apart shapes and mostly select different kernels + # rather than several threads compiling the same one. + size = -(-len(pending) // num_threads) + blocks = [pending[i * size:(i + 1) * size] for i in range(num_threads)] + order = [b[j] for j in range(size) for b in blocks if j < len(b)] + device = torch.cuda.current_device() + # One stream for all workers: the allocator caches their blocks where + # the profiling loop below allocates, and GPU time is not the cost here. + stream = torch.cuda.current_stream() + failures = [] + + def warm(profile: OptimizationProfile) -> None: + torch.cuda.set_device(device) + try: + with torch.cuda.stream(stream): + tensors = self._prepare_input_tensors(profile, inputs) + if tuning_config.inputs_pre_hook is not None: + tensors = tuning_config.inputs_pre_hook(tensors) + for runner in runners: + for tactic in runner.get_valid_tactics( + tensors, profile, **kwargs): + runner(tensors, tactic=tactic, **kwargs) + stream.synchronize() + except Exception as e: # noqa: BLE001 - the profiling loop reports it + failures.append(e) + + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor( + max_workers=num_threads, + thread_name_prefix="autotuner_jit_warmup") as pool: + list(pool.map(warm, order)) + logger.info(f"[Autotuner] {custom_op}: ran {len(order)} profiles on " + f"{num_threads} threads to drive their JIT in " + f"{time.perf_counter() - start:.1f}s" + + (f"; {len(failures)} failed and are left to profiling " + f"({type(failures[0]).__name__}: {failures[0]})" + if failures else "")) + @staticmethod def _run_profile_runner(runner: TunableRunner, inputs: List[torch.Tensor], tactic: Any, **kwargs) -> None: diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 672c9b9f20e8..4c4ec41debfd 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -2232,6 +2232,12 @@ def forward( return a +# DeepGemm compiles each layout with nvcc and releases the GIL meanwhile, so +# the autotuner drives the JIT of a sweep's distinct layouts on a thread pool +# (TuningConfig.parallel_jit_warmup; -1 sizes it from the host's cores). +_DEEP_GEMM_JIT_WARMUP_THREADS = -1 + + class fp8SwapABGemmRunner(TunableRunner): """Runs quantize + DeepGemm FP8 GEMM. Single tactic for JIT warmup.""" @@ -2247,6 +2253,7 @@ class fp8SwapABGemmRunner(TunableRunner): dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, deep_gemm_jit_warmup_buckets), ), exclude_from_cache=True, + parallel_jit_warmup=_DEEP_GEMM_JIT_WARMUP_THREADS, ) def __init__(self, output_dtype: torch.dtype, disable_ue8m0_cast: bool, @@ -2294,14 +2301,14 @@ class Fp8PrequantizedSwapABGemmRunner(TunableRunner): """Runs DeepGemm with pre-quantized FP8 activations and packed scales.""" # The same step-16 grid as the other two DeepGemm runners: a layout no - # bucket selects is compiled mid-inference instead, and DeepGemm forks - # nvcc while holding the GIL. + # bucket selects is compiled mid-inference instead. tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, deep_gemm_jit_warmup_buckets), ), constraint_specs=(ConstraintSpec( 1, 0, lambda input_shapes: input_shapes[0][0]), ), exclude_from_cache=True, + parallel_jit_warmup=_DEEP_GEMM_JIT_WARMUP_THREADS, ) def __init__(self, output_dtype: torch.dtype, @@ -2452,6 +2459,7 @@ class Fp8BlockScalingGemmRunner(TunableRunner): 0, 0, deep_gemm_jit_warmup_buckets), ), tune_max_num_tokens=4096, exclude_from_cache=True, + parallel_jit_warmup=_DEEP_GEMM_JIT_WARMUP_THREADS, ) def get_valid_tactics( diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 72abd57ee01c..5446c88e83a7 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -445,8 +445,8 @@ def deep_gemm_jit_warmup_buckets(max_m: int): DeepGemm picks its tile layout from a heuristic over M and compiles one kernel per selected layout. A layout no bucket selects gets compiled - mid-inference instead, and DeepGemm forks nvcc while holding the GIL, so - that compile stalls every rank of the attention-DP group. + mid-inference instead, inside that rank's forward, so that compile stalls + every rank of the attention-DP group. Step 16 is exactly the right spacing, and it is needed over the whole range. In ``deepgemm/csrc/jit_kernels/heuristics/sm100.hpp`` the selection diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index ef34a8b8fecd..313ecff659e4 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -23,6 +23,7 @@ import statistics import sys import tempfile +import threading from typing import Any, List import cloudpickle @@ -2316,3 +2317,57 @@ def inputs(*args, **kwargs): [torch.randn(M, 64, device="cuda"), w]) assert counts["sweeps"] == 2 tuner.clear_cache() + + +class ThreadRecordingRunner(TunableRunner): + """One tactic. Records the thread that ran each shape, as a runner whose + first call per shape JIT-compiles a kernel would compile it there.""" + calls: list[tuple[int, int]] = [] + + 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, + **kwargs) -> torch.Tensor: + type(self).calls.append((inputs[0].shape[0], threading.get_ident())) + return inputs[0] @ inputs[1] + + +@pytest.mark.parametrize("threads, env, expect_pool", [ + (0, None, False), + (4, None, True), + (4, "0", False), +]) +def test_parallel_jit_warmup_runs_pending_profiles_on_a_thread_pool( + monkeypatch, threads, env, expect_pool): + """With parallel_jit_warmup, every not-yet-tuned profile runs once on a + thread pool before the profiling loop, which still profiles every profile + on the calling thread. TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP overrides it.""" + if env is None: + monkeypatch.delenv("TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP", raising=False) + else: + monkeypatch.setenv("TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP", env) + tuner = AutoTuner.get() + tuner.clear_cache() + ThreadRecordingRunner.calls.clear() + buckets = (4, 8, 16, 32, 64, 128) + config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, buckets), ), + parallel_jit_warmup=threads) + x = torch.randn(128, 64, device="cuda") + w = torch.randn(64, 128, device="cuda") + caller = threading.get_ident() + with autotune(): + tuner.choose_one(f"autotuner_test::parallel_jit_warmup_{threads}_{env}", + [ThreadRecordingRunner()], config, [x, w]) + + on_caller = {m for m, t in ThreadRecordingRunner.calls if t == caller} + on_pool = sorted(m for m, t in ThreadRecordingRunner.calls if t != caller) + assert on_caller == set(buckets) + assert on_pool == (sorted(buckets) if expect_pool else []) + tuner.clear_cache() From 1132620b073df8708171d6cc5729812ddb91018b Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 5 Nov 2026 02:55:12 +0000 Subject: [PATCH 3/4] [None][perf] autotuner: size the default JIT warmup pool from the rank's CPU affinity Address review: parallel_jit_warmup=-1 divided os.cpu_count() by the GPUs this process sees, which overcommits when each rank sees one GPU and ignores cpuset limits. It now counts the CPUs in this process's affinity mask (which the compilers it starts inherit). Unless the rank is pinned to a subset of the host's CPUs, it divides them by the ranks on the node (local_mpi_size(), or LOCAL_WORLD_SIZE without MPI), up to 32. Also add a test that a profile failing on the warmup pool is left to the profiling loop, and a test of the default thread count. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- tensorrt_llm/_torch/autotuner.py | 35 ++++++++-- .../_torch/custom_ops/torch_custom_ops.py | 2 +- tests/unittest/_torch/misc/test_autotuner.py | 65 +++++++++++++++++++ 3 files changed, 95 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 693e9743915b..1674f45d6a5f 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -40,7 +40,8 @@ nccl_window_graph_capture, release_nccl_window_graph_owner) from tensorrt_llm._torch.nccl_window_tensor_scope import \ nccl_window_tensor_scope -from tensorrt_llm._utils import confidential_compute_enabled, nvtx_range +from tensorrt_llm._utils import (confidential_compute_enabled, local_mpi_size, + mpi_disabled, nvtx_range) from tensorrt_llm.bindings.internal.runtime import (delay_kernel, record_global_timer) from tensorrt_llm.logger import logger @@ -153,9 +154,10 @@ class TuningConfig: kernel and releases the GIL while compiling (DeepGEMM), those compiles then overlap, and the profiling loop that follows finds them built. Opted-in runners must be thread-safe and must not use NCCL windows. 0 disables; -1 uses - this rank's share of the host's cores (up to 32): neighbouring sweep points - often select the same kernel and compile it twice, so more threads keep - paying off. TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP overrides the thread count. + the CPUs this rank may run on, divided by the ranks on the node unless the + rank is pinned to a subset, up to 32: neighbouring sweep points often select + the same kernel and compile it twice, so more threads keep paying off. + TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP overrides the thread count. """ dynamic_tensor_specs: Tuple[DynamicTensorSpec, ...] = () constraint_specs: Tuple[ConstraintSpec, ...] = () @@ -942,6 +944,28 @@ def _spec_in_bounds(spec, shapes_list) -> bool: return True +def _default_parallel_jit_warmup_threads() -> int: + """Threads for ``parallel_jit_warmup=-1``: this rank's CPUs, up to 32. + + The compilers the warmup starts inherit this process's CPU affinity. A rank + pinned to a subset of the host's CPUs has that subset to itself; unpinned + ranks share all of them, so each takes its share by the ranks on the node. + """ + try: + cpus = len(os.sched_getaffinity(0)) + except AttributeError: # no affinity API (non-Linux) + cpus = os.cpu_count() or 1 + if cpus >= (os.cpu_count() or cpus): + local_ranks = int(os.environ.get("LOCAL_WORLD_SIZE", "1")) + if not mpi_disabled(): + try: + local_ranks = local_mpi_size() + except Exception: # noqa: BLE001 - MPI cannot split by node here + pass + cpus //= max(1, local_ranks) + return min(32, cpus) + + class AutoTuner: """AutoTuner for optimizing TensorRT LLM operations. @@ -1614,8 +1638,7 @@ def _parallel_jit_warmup(self, custom_op: str, runners: List[TunableRunner], os.environ.get("TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP", tuning_config.parallel_jit_warmup)) if num_threads < 0: - num_threads = min(32, (os.cpu_count() or 1) // - max(1, torch.cuda.device_count())) + num_threads = _default_parallel_jit_warmup_threads() if num_threads < 2: return pending = [ diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 4c4ec41debfd..6ed789beb864 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -2234,7 +2234,7 @@ def forward( # DeepGemm compiles each layout with nvcc and releases the GIL meanwhile, so # the autotuner drives the JIT of a sweep's distinct layouts on a thread pool -# (TuningConfig.parallel_jit_warmup; -1 sizes it from the host's cores). +# (TuningConfig.parallel_jit_warmup; -1 sizes it from this rank's CPUs). _DEEP_GEMM_JIT_WARMUP_THREADS = -1 diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 313ecff659e4..f25fe1a6180e 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -2371,3 +2371,68 @@ def test_parallel_jit_warmup_runs_pending_profiles_on_a_thread_pool( assert on_caller == set(buckets) assert on_pool == (sorted(buckets) if expect_pool else []) tuner.clear_cache() + + +def test_parallel_jit_warmup_worker_failure_is_left_to_profiling(monkeypatch): + """A profile that fails on the warmup pool does not fail choose_one: the + profiling loop still runs every profile on the calling thread.""" + monkeypatch.delenv("TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP", raising=False) + caller = threading.get_ident() + + class WorkerFailingRunner(ThreadRecordingRunner): + calls: list[tuple[int, int]] = [] + + def forward(self, /, inputs, *, tactic=-1, **kwargs): + if threading.get_ident() != caller: + raise RuntimeError("fails on warmup threads only") + return super().forward(inputs, tactic=tactic, **kwargs) + + tuner = AutoTuner.get() + tuner.clear_cache() + buckets = (4, 8, 16, 32, 64, 128) + config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, buckets), ), + parallel_jit_warmup=4) + x = torch.randn(128, 64, device="cuda") + w = torch.randn(64, 128, device="cuda") + with autotune(): + tuner.choose_one("autotuner_test::parallel_jit_warmup_worker_failure", + [WorkerFailingRunner()], config, [x, w]) + + assert {m for m, _ in WorkerFailingRunner.calls} == set(buckets) + assert all(t == caller for _, t in WorkerFailingRunner.calls) + assert not tuner.stats.failed_profiling_count.get( + "autotuner_test::parallel_jit_warmup_worker_failure") + tuner.clear_cache() + + +@pytest.mark.parametrize( + "host_cpus, affinity, local_ranks, env_local, mpi_off, expected", [ + (144, 144, 4, None, False, 32), + (64, 64, 8, None, False, 8), + (64, 8, 8, None, False, 8), + (64, 64, 1, "16", True, 4), + (64, 64, RuntimeError("no node split"), None, False, 32), + ]) +def test_default_parallel_jit_warmup_threads(monkeypatch, host_cpus, affinity, + local_ranks, env_local, mpi_off, + expected): + """parallel_jit_warmup=-1 uses the CPUs this rank may run on: unpinned ranks + share the host's CPUs with the other ranks on the node, a rank pinned to a + subset has it to itself, and the result is capped at 32.""" + + def node_ranks(): + if isinstance(local_ranks, Exception): + raise local_ranks + return local_ranks + + monkeypatch.setattr(os, "cpu_count", lambda: host_cpus) + monkeypatch.setattr(os, "sched_getaffinity", + lambda pid: set(range(affinity))) + monkeypatch.setattr(autotuner, "local_mpi_size", node_ranks) + monkeypatch.setattr(autotuner, "mpi_disabled", lambda: mpi_off) + if env_local is None: + monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False) + else: + monkeypatch.setenv("LOCAL_WORLD_SIZE", env_local) + assert autotuner._default_parallel_jit_warmup_threads() == expected From 1095b01d11bcbfb23d559d1afc6926f82ad2758b Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 5 Nov 2026 14:34:23 +0000 Subject: [PATCH 4/4] [None][fix] autotuner: keep the Hopper block-scale FP8 GEMM warmup serial Fp8BlockScalingGemmRunner opted in to parallel_jit_warmup, but on Hopper its kernels come from the C++ DeepGEMM JIT (cpp/include/tensorrt_llm/ deep_gemm), and torch ops run without the GIL. Its RuntimeCache is a plain unordered_map that Compiler::build reads and writes with no lock, so the warmup threads raced on it. On DGX_H100 this hung FP8 block-scale models during warmup (TestQwen3_8B::test_fp8_block_scales, TestLagunaXS:: test_fp8, test_disaggregated_simple_qwen3[Qwen3-8B-FP8]). On SM100 the op uses prebuilt cubins and compiles nothing, so the pool bought nothing there. Only the two runners built on the deep_gemm package, which touches its caches with the GIL held, keep the thread pool. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 6ed789beb864..564407cd4de1 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -2232,8 +2232,9 @@ def forward( return a -# DeepGemm compiles each layout with nvcc and releases the GIL meanwhile, so -# the autotuner drives the JIT of a sweep's distinct layouts on a thread pool +# The deep_gemm package compiles each layout with nvcc and releases the GIL +# meanwhile (its kernel caches are only touched with the GIL held), so the +# autotuner drives the JIT of a sweep's distinct layouts on a thread pool # (TuningConfig.parallel_jit_warmup; -1 sizes it from this rank's CPUs). _DEEP_GEMM_JIT_WARMUP_THREADS = -1 @@ -2453,13 +2454,15 @@ def _( # cubins and compiles nothing. class Fp8BlockScalingGemmRunner(TunableRunner): # Without exclude_from_cache, a warm disk cache short-circuits tuning and - # the JIT warmup never runs. + # the JIT warmup never runs. The warmup stays serial (no + # parallel_jit_warmup): on Hopper the kernels come from the C++ DeepGEMM + # JIT in cpp/include/tensorrt_llm/deep_gemm, whose runtime cache is not + # safe to use from several threads, and the op runs without the GIL. tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, deep_gemm_jit_warmup_buckets), ), tune_max_num_tokens=4096, exclude_from_cache=True, - parallel_jit_warmup=_DEEP_GEMM_JIT_WARMUP_THREADS, ) def get_valid_tactics(