diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index b613d351db6a..5fd9d3bc4a1a 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -203,6 +203,30 @@ curl http://localhost:8000/server_info } ``` +### Kernel builds ahead of first use + +An op that compiles its kernels on first use lists them when it runs, for every shape it +sweeps down from that call's and has not listed yet, as the autotuner sweeps its tuning +buckets; worker threads build them in parallel. An op that the autotuner tunes has three +states: the autotuner warmup tunes it, a build state builds the kernels of every shape, +runner and tactic its tuning will time on the same worker threads without timing them, and +serving runs the tuned tactic. + +For a model with such an op, the engine first loads dummy weights and runs the autotuner +warmup's two forwards on them (a context batch at the maximum token count and one +generation request) in the build state, then loads the checkpoint while the kernels build. +The startup log shows these as the `rehearsal`, `checkpoint` and `autotuner_builds` phases, +which `model_engine_creation_seconds` then does not include; `autotuner_builds` waits for +the autotuned ops' builds. Supported ops: MiniMax-M3's MSA attention (`fmha_sm100`), and in +the build state, DeepGEMM's FP8 block-scale GEMMs on Blackwell and FlashInfer's MXFP8 GEMM. + +| Environment variable | Default | Effect | +|----------------------|---------|--------| +| `TLLM_KERNEL_BUILDER` | `1` | `0` compiles each kernel on its first use instead. | +| `TLLM_KERNEL_BUILDER_REHEARSAL` | `1` | `0` loads the checkpoint first; the kernels then build from the warmup on. | +| `TLLM_KERNEL_BUILDER_AUTOTUNER` | `1` | `0` leaves the autotuned ops' kernels to the autotuner warmup. | +| `TLLM_KERNEL_BUILDER_THREADS` | this rank's share of the CPUs, up to 32 | Worker threads per process. | + ## Unified Cache for Compiled Artifacts Set `TRTLLM_CACHE_DIR` before importing TensorRT-LLM to place its compiled diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py index bc1301e57c1b..604f3d9f5389 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py @@ -21,6 +21,7 @@ from __future__ import annotations +import contextlib from typing import TYPE_CHECKING, Optional import torch @@ -145,7 +146,10 @@ def _aligned_nvfp4_dequant_scales( if cache.device != kv_scale_quant_orig.device: raise ValueError("MiniMax-M3 NVFP4 scale reload must preserve the scale buffer device") if refresh or getattr(attn, "_msa_nvfp4_dequant_scale_source_ptr", None) != source_ptr: - cache[:, 0].copy_(kv_scale_quant_orig[1:3]) + # A buffer a forward created is an inference tensor, which only inference mode + # may update in place. + with torch.inference_mode() if cache.is_inference() else contextlib.nullcontext(): + cache[:, 0].copy_(kv_scale_quant_orig[1:3]) attn._msa_nvfp4_dequant_scale_source_ptr = source_ptr k_global_scale = cache[0, :1] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_kernel_builds.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_kernel_builds.py new file mode 100644 index 000000000000..33a0696afa2c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_kernel_builds.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fmha_sm100's kernels, built ahead through ``KernelBuilder``. + +fmha_sm100 compiles one kernel per variant (dtype, query tile, single warp group, sparse +mode, page size, split KV, pack factor) on first use, one variant at a time per process. +Its plan picks the variant from the longest query segment and the split count; the call +site fixes the rest. These two are its build dimensions (``_CONFIG``, ``_PROXY_CONFIG``). + +``install`` wraps the package's ``fmha_sm100_plan`` and ``fmha_sm100``. On a run in a +build pass, the plan and the run are replayed for one query segment of each shape the +dimensions sweep from that call's and no earlier run of its call signature swept, with +the package's variant lookup recording the variant it is asked for instead of compiling +it. Each variant found is submitted to ``KernelBuilder``, whose workers compile it in the +directory and under the file lock the package's own lookup uses. The plan, split-KV +reduction and sparse top-k modules need nothing from a call: ``reaches`` submits them, or +else the first run in a build pass. + +In a rehearsal every run is such a replay: it compiles and launches nothing, and its +outputs hold zeros and an empty maximum score. +""" + +from __future__ import annotations + +import functools +import threading +import time +import types +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.kernel_builder import BuildConfig, DynamicSpec, KernelBuilder +from tensorrt_llm.logger import logger + +_OP = "fmha_sm100" + + +def _query_len_buckets(longest: int) -> Tuple[int, ...]: + """Query segment lengths up to the longest a call has: each length up to 129 crosses a + boundary of the pack factor, the single warp group or the 128-row query tile; 256 takes + the 256-row tile.""" + return tuple(range(1, min(longest, 129) + 1)) + ((256,) if longest >= 256 else ()) + + +_QUERY_LEN = DynamicSpec("qo_len", _query_len_buckets) +# A call's query segment lengths. +_CONFIG = BuildConfig((_QUERY_LEN,)) +# The MSA indexer's maximum-score pass also chooses whether to split KV, per batch +# (msa_indexer._proxy_kv_splits); a variant depends only on whether it splits. +_PROXY_CONFIG = BuildConfig((_QUERY_LEN, DynamicSpec("num_kv_splits", (1, 2)))) +# Plan options a replay does not pass through: it plans its own lengths and offsets, and a +# swept split count. +_PLAN_SKIP = ("qo_offset", "device", "stream", "num_kv_splits") +# Run options that do not choose the variant. +_RUN_SKIP = ("sm_scale", "q_scale", "k_scale", "v_scale", "o_scale", "check_input_valid") +# Run options a replay replaces: its query is its own, and its outputs are discarded. +_RUN_REPLACED = ("q_offset_override", "out", "max_score") +_PLAN_ARGS = "_tllm_kernel_builder_plan_args" +_INSTALLED = "_tllm_kernel_builder_installed" +# The modules every MSA model loads: (name, library, the package's compile function). +_MODULES = ( + ("plan", "fmha_sm100_plan.so", "_do_compile_plan"), + ("reduction", "fmha_sm100_reduction.so", "_do_compile_reduction"), + ("sparse_topk", "sparse_topk_select.so", "_do_compile_sparse_topk"), +) + +_replay = threading.local() # .found: the variants a replay on this thread asked for +_package: Optional["_Package"] = None # set by install +_install_lock = threading.Lock() + + +class _Stub: + """A variant or the split-KV reduction module whose calls do nothing, for a replay.""" + + def __getattr__(self, name: str) -> Callable[..., None]: + return lambda *args, **kwargs: None + + +def _tag(value: Any) -> Any: + """What an argument contributes to a call signature: a tensor its dtype and rank.""" + if isinstance(value, torch.Tensor): + return ("tensor", value.dtype, value.dim()) + if isinstance(value, (list, tuple)): + return tuple(_tag(item) for item in value) + try: + hash(value) + except TypeError: + return ("unhashable", type(value).__name__) + return value + + +class _Package: + """The fmha_sm100 entry points ``install`` wraps, and the originals they call.""" + + def __init__(self, module: types.ModuleType, api: types.ModuleType, jit: types.ModuleType): + self.jit = jit + self.manager = jit._variant_manager + self.variant_key = jit._variant_key_from_runtime + self.sparse_prefill_min_len = api._prefill_qlen_threshold(True) + # Checked here so that a package without them is left alone. + for name in ("_compile_only", "_load_templates", "_loaded", "_ensure_all_module"): + getattr(self.manager, name) + for name in ("_built_dir", "_acquire_file_lock", "_release_file_lock"): + getattr(jit, name) + for _, _, compile_fn in _MODULES: + getattr(jit, compile_fn) + self.plan_fn, self.run_fn = module.fmha_sm100_plan, module.fmha_sm100 + self.variant_fn, self.reduction_fn = api.get_fmha_variant, api.get_reduction_module + self.looked_up: set = set() + self.modules_submitted = False + + def plan(self, qo_segment_lens, kv_segment_lens, *args, **kwargs): + plan_info = self.plan_fn(qo_segment_lens, kv_segment_lens, *args, **kwargs) + options = {key: value for key, value in kwargs.items() if key not in _PLAN_SKIP} + num_kv_splits = kwargs.get("num_kv_splits", 1) + for sub_plan in plan_info[3:]: + if isinstance(sub_plan, dict): + sub_plan[_PLAN_ARGS] = (args, options, qo_segment_lens, num_kv_splits) + return plan_info + + def submit_modules(self) -> None: + if self.modules_submitted: + return + self.modules_submitted = True + builder = KernelBuilder.get() + for name, library, compile_fn in _MODULES: + builder.submit( + _OP, name, functools.partial(self.build_module, name, library, compile_fn) + ) + + def run(self, q, k, v, plan_info, *args, **kwargs): + builder = KernelBuilder.get() + if ( + builder.active() + and not args + and isinstance(plan_info, tuple) + and not torch.cuda.is_current_stream_capturing() + ): + self.submit_modules() + planned = next( + (p[_PLAN_ARGS] for p in plan_info[3:] if isinstance(p, dict) and _PLAN_ARGS in p), + None, + ) + if planned is not None: + plan_args, plan_options, qo_segment_lens, num_kv_splits = planned + run_options = {key: v_ for key, v_ in kwargs.items() if key not in _RUN_SKIP} + proxy = run_options.get("output_o") is False and plan_options.get("output_maxscore") + config = _PROXY_CONFIG if proxy else _CONFIG + signature = ( + tuple(_tag(arg) for arg in plan_args), + tuple(sorted((key, _tag(value)) for key, value in plan_options.items())), + q.dtype, + tuple(q.shape[1:]), + k.dtype, + tuple(k.shape[1:]), + tuple(sorted((key, _tag(value)) for key, value in run_options.items())), + None if proxy else num_kv_splits, + ) + lens = torch.as_tensor(qo_segment_lens) + observed = {"qo_len": max(int(lens.max()), 1) if lens.numel() else 1} + if proxy: + observed["num_kv_splits"] = num_kv_splits + builder.prebuild( + _OP, + signature, + config, + observed, + functools.partial( + self._list, q, k, v, plan_args, plan_options, num_kv_splits, run_options + ), + ) + if builder.rehearsing(): + return self._rehearse(q, k, v, plan_info, *args, **kwargs) + return self.run_fn(q, k, v, plan_info, *args, **kwargs) + + def _rehearse(self, q, k, v, plan_info, *args, **kwargs): + """The package's run with its variant and split-KV reduction stubbed; its outputs + are set to what a run over no keys gives: zeros and a maximum score of -inf.""" + _replay.found = {} + try: + result = self.run_fn(q, k, v, plan_info, *args, **kwargs) + finally: + _replay.found = None + out, max_score = result if isinstance(result, tuple) else (result, None) + for tensor, value in ( + (out, 0.0), + (kwargs.get("out"), 0.0), + (max_score, float("-inf")), + (kwargs.get("max_score"), float("-inf")), + ): + if isinstance(tensor, torch.Tensor): + tensor.fill_(value) + return result + + def _list( + self, q, k, v, plan_args, plan_options, num_kv_splits, run_options, shapes + ) -> List[Tuple[str, Callable]]: + """Replay the plan and the run for one query segment of each shape; return the + variants asked for. + + A shape the package rejects is skipped: a call of that shape would fail too. + """ + options = {key: value for key, value in run_options.items() if key not in _RUN_REPLACED} + sparse = plan_options.get("kv_block_num", -1) > 0 + sparse_mode = plan_options.get("sparse_kernel_mode", "auto") + found: Dict[str, dict] = {} + _replay.found = found + try: + for shape in shapes: + qo_len = shape["qo_len"] + if sparse and ( + sparse_mode == "prefill" + or (sparse_mode == "auto" and qo_len > self.sparse_prefill_min_len) + ): + # The sparse prefill path, whose CuTe DSL kernels are not variants. + continue + lens = torch.tensor([qo_len], dtype=torch.int32) + q_replay = torch.empty((qo_len, *q.shape[1:]), dtype=q.dtype, device=q.device) + splits = shape.get("num_kv_splits", num_kv_splits) + try: + plan_info = self.plan_fn( + lens, lens.clone(), *plan_args, num_kv_splits=splits, **plan_options + ) + self.run_fn(q_replay, k, v, plan_info, **options) + except (ValueError, RuntimeError, AssertionError) as e: + logger.debug(f"[KernelBuilder] {_OP}: no variant for {shape}: {e}") + finally: + _replay.found = None + # The package loads its two source templates on the first compile, unlocked: load + # them here, before any worker compiles. + self.manager._load_templates() + self.manager._ensure_all_module() + bundled = self.manager._all_module + return [ + (name, functools.partial(self._build, name, params)) + for name, params in found.items() + if bundled is None or not hasattr(bundled, f"run_{name}") + ] + + def _build(self, name: str, params: dict) -> bool: + cache_dir = self.jit._built_dir(name, f"{name}.so") + fd = self.jit._acquire_file_lock(cache_dir / ".compile.lock") + try: + if (cache_dir / f"{name}.so").exists(): + return False + self.manager._compile_only(name, params, cache_dir) + return True + finally: + self.jit._release_file_lock(fd) + + def build_module(self, name: str, library: str, compile_fn: str) -> bool: + cache_dir = self.jit._built_dir(name, library) + fd = self.jit._acquire_file_lock(cache_dir.parent / f"{name}.lock") + try: + if (cache_dir / library).exists(): + return False + getattr(self.jit, compile_fn)(cache_dir) + return True + finally: + self.jit._release_file_lock(fd) + + def get_variant(self, *args): + found: Optional[Dict[str, dict]] = getattr(_replay, "found", None) + if found is not None: + name, params = self.variant_key(*args) + found.setdefault(name, params) + return _Stub() + if args in self.looked_up: + return self.variant_fn(*args) + try: + name = self.variant_key(*args)[0] + except ValueError: + return self.variant_fn(*args) + missing = ( + name not in self.manager._loaded + and not (self.jit._built_dir(name, f"{name}.so") / f"{name}.so").exists() + ) + t0 = time.perf_counter() + variant = self.variant_fn(*args) + seconds = time.perf_counter() - t0 + self.looked_up.add(args) + # A variant loaded from the package's bundle or a finished build takes milliseconds. + if missing and seconds >= 1.0: + KernelBuilder.get().op_compiled(_OP, name, seconds) + return variant + + def get_reduction_module(self): + if getattr(_replay, "found", None) is not None: + return _Stub() + return self.reduction_fn() + + +def reaches(model_config: Any, model: Any = None) -> bool: + """True when the model runs MiniMax-M3's MSA attention on fmha_sm100, whose modules + are then submitted.""" + from tensorrt_llm._utils import is_sm_100f + from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + from .msa_utils import msa_package_available, require_msa_module + + config = getattr(model_config, "sparse_attention_config", None) + if not isinstance(config, MiniMaxM3SparseAttentionConfig) or config.implementation != "msa": + return False + if not is_sm_100f() or not msa_package_available(): + return False + require_msa_module() + if _package is None: + return False + _package.submit_modules() + return True + + +def install(module: types.ModuleType) -> None: + """Route ``module``'s plan, run and variant lookup through ``KernelBuilder``; once.""" + global _package + if getattr(module, _INSTALLED, False): + return + with _install_lock: + if getattr(module, _INSTALLED, False): + return + if KernelBuilder.get().enabled: + try: + from fmha_sm100 import api, jit + + package = _Package(module, api, jit) + except Exception as e: # noqa: BLE001 - the kernels stay on the package's own path + logger.warning( + f"[KernelBuilder] fmha_sm100's kernels are left to their first use: " + f"{type(e).__name__}: {e}" + ) + else: + _package = package + module.fmha_sm100_plan = package.plan + module.fmha_sm100 = package.run + api.get_fmha_variant = package.get_variant + api.get_reduction_module = package.get_reduction_module + setattr(module, _INSTALLED, True) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py index 774c21d88126..51fb4d7da414 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py @@ -96,6 +96,9 @@ def require_msa_module() -> types.ModuleType: "MiniMax-M3 MSA attention requires the fmha_sm100 kernels packaged " "with TensorRT-LLM. Reinstall TensorRT-LLM from a complete build." ) from exc + from .msa_kernel_builds import install + + install(fmha_sm100) return fmha_sm100 diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index e5a6cbfa94d9..12b32cc036e2 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -22,19 +22,22 @@ import itertools import json import os +import threading import time import zlib from abc import ABC, abstractmethod from dataclasses import dataclass, field -from functools import lru_cache +from functools import lru_cache, partial from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import (Any, Callable, Dict, Hashable, List, Optional, Set, Tuple, + Union) import torch from cuda.bindings import driver import tensorrt_llm from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.kernel_builder import KernelBuilder from tensorrt_llm._torch.nccl_window_graph import ( abandon_nccl_window_graph_owner, nccl_window_allocation_suppressed, nccl_window_graph_capture, release_nccl_window_graph_owner) @@ -49,6 +52,10 @@ # Unique tag to avoid collisions with other comms PP_COMM_TAG_AUTOTUNING = 30000 +# TLLM_KERNEL_BUILDER_AUTOTUNER=0 keeps the build state from building anything. +_BUILD_ENV = "TLLM_KERNEL_BUILDER_AUTOTUNER" +# The KernelBuilder op name of the build state's builds is this plus the custom op. +BUILD_OP_PREFIX = "autotuner:" class DistributedTuningStrategy(enum.Enum): @@ -149,11 +156,13 @@ 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 + parallel_jit_warmup (int): Before profiling, build every not-yet-tuned profile + (``TunableRunner.build``) 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. Nonzero also has the build state (``build_kernels``) build + the op's kernels on KernelBuilder's worker threads. Opted-in runners' ``build`` + must be thread-safe and must not use NCCL windows. 0 disables; -1 uses 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. @@ -289,6 +298,48 @@ def unique_id(self): """ return tuple(self.__dict__.values()) + def build(self, + inputs: List[torch.Tensor], + *, + tactic: Any = -1, + **kwargs) -> None: + """Compile the kernels ``tactic`` runs on ``inputs``; the result is discarded. + + The build state and the parallel JIT warmup call it on worker threads + (``TuningConfig.parallel_jit_warmup``). This default runs the tactic once, + which compiles what a first call compiles. A runner that can compile + without launching overrides it. + """ + self(inputs, tactic=tactic, **kwargs) + + def build_key(self, profile: OptimizationProfile, tactic: Any) -> Hashable: + """Builds with equal keys compile the same kernels, so only one of them runs. + + This default distinguishes every profile and tactic. + """ + return (profile.get_opt_shapes(), str(tactic)) + + +@contextlib.contextmanager +def build_kernels(): + """The autotuner's build state, for forwards whose outputs are discarded. + + An op that tunes has three states: it tunes inside ``autotune()``, it builds + here, and otherwise it runs its tuned tactic. Here ``choose_one`` queues the + builds of every profile, runner and tactic the op's tuning would time on + KernelBuilder's worker threads (``AutoTuner.prebuild``), without timing them, + and returns the tactic the op would run. An op that builds this way returns + zeros of its output shape instead of running it. A nested ``autotune()`` leaves + the build state on. + """ + autotuner = AutoTuner.get() + old_mode = autotuner.is_build_mode + autotuner.is_build_mode = True + try: + yield + finally: + autotuner.is_build_mode = old_mode + @contextlib.contextmanager def autotune(tune_mode: bool = True, @@ -944,6 +995,28 @@ def _spec_in_bounds(spec, shapes_list) -> bool: return True +def _interleaved(profiles: List[OptimizationProfile], + num_threads: int) -> List[OptimizationProfile]: + """``profiles`` reordered for ``num_threads`` workers: one from each of as many + contiguous blocks of the sweep in turn, so the builds in flight come from + far-apart shapes and mostly select different kernels rather than several + threads compiling the same one.""" + size = -(-len(profiles) // max(1, num_threads)) + blocks = [profiles[i * size:(i + 1) * size] for i in range(num_threads)] + return [b[j] for j in range(size) for b in blocks if j < len(b)] + + +_build_streams = threading.local() + + +def _thread_build_stream() -> torch.cuda.Stream: + """The calling thread's own stream for builds, made on first use.""" + stream = getattr(_build_streams, "stream", None) + if stream is None: + stream = _build_streams.stream = torch.cuda.Stream() + return stream + + def _default_parallel_jit_warmup_threads() -> int: """Threads for ``parallel_jit_warmup=-1``: this rank's CPUs, up to 32. @@ -972,6 +1045,11 @@ class AutoTuner: This class handles automatic performance tuning of tensor operations by profiling different implementations and caching the best performing configurations. + ``choose_one`` behaves by state: inside ``autotune()`` it profiles and caches the + best (runner, tactic); inside ``build_kernels()`` it builds the kernels of every + one without timing them (``prebuild``); otherwise it returns the cached choice, + or the fallback. + Args: warmup (int): Number of warmup iterations before profiling (default: 3) repeat (int): Number of profiling iterations for averaging (default: 10) @@ -995,7 +1073,13 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000): # 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.is_build_mode = False self.skip_dynamic_tuning_buckets = False + # The build state's bookkeeping, shared with the worker threads: the + # choose_one calls that queued their builds, and the build keys built. + self._build_lock = threading.Lock() + self._build_calls: set[Tuple] = set() + self._built_keys: set[Tuple] = set() # Timing backend: globaltimer kernel vs cuda events. # TLLM_PROFILING_TIMER env var overrides auto-detection: @@ -1201,6 +1285,18 @@ def choose_one( input_shapes = tuple(self._get_input_sizes(inputs)) + if self.is_build_mode: + # Build what the tuning would time, and return what the op runs now. + # Nothing is profiled or cached, and no rank waits for another. + self.prebuild(custom_op, runners, tuning_config, inputs, **kwargs) + _, best_runner_id, best_tactic, _ = self.profiling_cache.search_cache( + custom_op, + runners, + input_shapes, + tuning_config, + apply_map_to_tuning_buckets=True) + return (runners[best_runner_id], best_tactic) + # 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 @@ -1624,11 +1720,104 @@ def _profile_runners_impl( return best_runner_id, best_tactic, min_time, has_tuning_failure_occurred + def prebuild(self, custom_op: str, runners: List[TunableRunner], + tuning_config: TuningConfig, inputs: List[torch.Tensor], + **kwargs) -> None: + """In the build state, queue the builds of every profile the op's tuning + would profile for ``inputs`` on KernelBuilder's worker threads. + + Each build builds every runner's every valid tactic for its profile + (``_build_profile``). The profiles are swept down from the inputs' shapes, + whatever ``skip_dynamic_tuning_buckets`` is, since the tuning that follows + may tune each of them. Nothing is timed or cached. ``choose_one`` calls it in + the build state; an op calls it for a backend its forward does not run. + + Only an op that opts in with ``TuningConfig.parallel_jit_warmup`` builds; a + MERGE-strategy op, whose runners are collective, does not. A call with the + runners and input shapes of an earlier one returns at once. + """ + if (not self.is_build_mode or tuning_config.parallel_jit_warmup == 0 + or tuning_config.distributed_tuning_strategy + == DistributedTuningStrategy.MERGE + or os.environ.get(_BUILD_ENV, "1") == "0"): + return + builder = KernelBuilder.get() + if not builder.enabled: + return + runner_ids = tuple((type(runner).__name__, str(runner.unique_id())) + for runner in runners) + call = (custom_op, runner_ids, tuple(self._get_input_sizes(inputs)), + tuple(getattr(t, "dtype", None) for t in inputs)) + with self._build_lock: + if call in self._build_calls: + return + self._build_calls.add(call) + profiles = self._optimization_profiles( + tuning_config, inputs, skip_dynamic_tuning_buckets=False) + # The builds read the inputs on streams of their own, after the work + # queued so far, which produced them. + ready = torch.cuda.Event() + ready.record() + queued = sum( + builder.submit( + BUILD_OP_PREFIX + custom_op, (runner_ids, p.get_opt_shapes()), + partial(self._build_profile, + custom_op, + runners, + p, + inputs, + tuning_config, + None, + kwargs, + ready=ready)) + for p in _interleaved(profiles, builder.threads)) + logger.info(f"[Autotuner] {custom_op}: {queued} of {len(profiles)} " + "profiles queued to build") + + def _build_profile(self, + custom_op: str, + runners: List[TunableRunner], + profile: OptimizationProfile, + inputs: List[torch.Tensor], + tuning_config: TuningConfig, + stream: Optional[torch.cuda.Stream], + kwargs: Dict[str, Any], + ready: Optional[torch.cuda.Event] = None) -> bool: + """Build every runner's every valid tactic for ``profile``, on its inputs. + + A tactic whose ``TunableRunner.build_key`` this process built already is + skipped. Runs under inference mode on ``stream``, or else on a stream of + the calling thread's own, after ``ready``, and returns whether anything + was built. + """ + stream = stream if stream is not None else _thread_build_stream() + if ready is not None: + stream.wait_event(ready) + built = False + with torch.inference_mode(), 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: + runner_id = (type(runner).__name__, str(runner.unique_id())) + for tactic in runner.get_valid_tactics(tensors, profile, + **kwargs): + key = (custom_op, runner_id, + runner.build_key(profile, tactic)) + with self._build_lock: + if key in self._built_keys: + continue + self._built_keys.add(key) + runner.build(tensors, tactic=tactic, **kwargs) + built = True + stream.synchronize() + return built + 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. + """Build each not-yet-tuned profile 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 @@ -1653,12 +1842,7 @@ def _parallel_jit_warmup(self, custom_op: str, runners: List[TunableRunner], 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)] + order = _interleaved(pending, num_threads) 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. @@ -1668,15 +1852,8 @@ def _parallel_jit_warmup(self, custom_op: str, runners: List[TunableRunner], 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() + self._build_profile(custom_op, runners, profile, inputs, + tuning_config, stream, kwargs) except Exception as e: # noqa: BLE001 - the profiling loop reports it failures.append(e) @@ -1896,13 +2073,18 @@ def elapsed_time(): return avg_time def _optimization_profiles( - self, tuning_config: TuningConfig, - inputs: List[torch.Tensor]) -> List[OptimizationProfile]: + self, + tuning_config: TuningConfig, + inputs: List[torch.Tensor], + skip_dynamic_tuning_buckets: Optional[bool] = None + ) -> List[OptimizationProfile]: """Generate optimization profiles for autotuning. Args: tuning_config (TuningConfig): Tuning configuration inputs (List[torch.Tensor]): List of input tensors + skip_dynamic_tuning_buckets (Optional[bool]): Overrides the tuner's + ``skip_dynamic_tuning_buckets`` when given. Returns: List of OptimizationProfile objects representing different configurations @@ -1923,13 +2105,15 @@ def _optimization_profiles( generated_profiles: List[OptimizationProfile] = [] dynamic_dims = [] + if skip_dynamic_tuning_buckets is None: + skip_dynamic_tuning_buckets = self.skip_dynamic_tuning_buckets for spec in tuning_config.dynamic_tensor_specs: assert callable(spec.gen_tuning_buckets) or isinstance(spec.gen_tuning_buckets, (list, tuple)), \ "The given dynamic dimension must provide a opt value generation function or a list of opt values" if not _spec_in_bounds(spec, base_profile.shapes): continue - if self.skip_dynamic_tuning_buckets: + if skip_dynamic_tuning_buckets: if spec.map_to_tuning_buckets is not None: # Still include the bucketed value of the actual shape so the # cache key used during profiling (raw) aligns with the key @@ -2149,6 +2333,9 @@ def clear_cache(self) -> None: """Clear the profiling cache.""" self.profiling_cache.clear() self._keys_profiled_in_process.clear() + with self._build_lock: + self._build_calls.clear() + self._built_keys.clear() if hasattr(torch.ops.trtllm, "clear_allreduce_tactic_cache"): torch.ops.trtllm.clear_allreduce_tactic_cache() diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 715b12dd1dca..feb79e216c57 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -786,11 +786,13 @@ class FlashInferMXFP8GemmRunner(TunableRunner): CUTLASS = -1 # -1 is the AutoTuner fallback tactic. CUTE_DSL = 0 + # build() only loads FlashInfer's CUTLASS module, so it is safe on threads. tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, get_power_of_2_num_tokens_buckets, next_positive_power_of_2), ), constraint_specs=(ConstraintSpec(1, 0, _mxfp8_scale_infer_shape), ), + parallel_jit_warmup=-1, ) def __init__(self, output_dtype: torch.dtype) -> None: @@ -803,6 +805,23 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: return [self.CUTLASS, self.CUTE_DSL] + def build(self, + inputs: List[torch.Tensor], + *, + tactic: int = CUTLASS, + **kwargs) -> None: + """CUTLASS: build and load FlashInfer's MXFP8 GEMM module, without running it. + + CuTe DSL: nothing; its kernels compile when the tactic is profiled. + """ + if tactic == self.CUTLASS: + from flashinfer.gemm.gemm_base import get_cutlass_mxfp8_gemm_module + get_cutlass_mxfp8_gemm_module(torch.cuda.get_device_capability()[0]) + + def build_key(self, profile: OptimizationProfile, tactic: int) -> str: + # One module per tactic serves every shape. + return str(tactic) + def forward( self, inputs: List[torch.Tensor], @@ -898,6 +917,20 @@ def mxfp8_gemm_autotuned( gemm_inputs, tune)) +def prebuild_flashinfer_mxfp8_gemm(act: torch.Tensor, act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype) -> None: + """In the autotuner's build state, queue the builds of the FlashInfer MXFP8 GEMM + that ``mxfp8_gemm_autotuned`` and FlashInfer's own MXFP8 tuning time.""" + tuner = AutoTuner.get() + if not tuner.is_build_mode: + return + runner = FlashInferMXFP8GemmRunner(output_dtype) + tuner.prebuild(_FLASHINFER_MXFP8_GEMM_AUTOTUNED_OP, [runner], + runner.tuning_config, [act, act_scale, weight, weight_scale]) + + def flashinfer_mxfp8_gemm_cute_dsl( act: torch.Tensor, act_scale: torch.Tensor, @@ -2374,12 +2407,17 @@ def fp8_prequantized_swap_ab_gemm( disable_ue8m0_cast: bool = False, ) -> torch.Tensor: runner = Fp8PrequantizedSwapABGemmRunner(output_dtype, disable_ue8m0_cast) - _, best_tactic = AutoTuner.get().choose_one( + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( "trtllm::fp8_prequantized_swap_ab_gemm", [runner], Fp8PrequantizedSwapABGemmRunner.tuning_config, [activation, activation_scale, weight, weight_scale], ) + if tuner.is_build_mode: + # The build state queued this GEMM's builds; nothing runs. + return activation.new_zeros((activation.size(0), weight.size(0)), + dtype=output_dtype) return runner( inputs=[activation, activation_scale, weight, weight_scale], tactic=best_tactic, @@ -2433,6 +2471,10 @@ def fp8_swap_ab_gemm( fp8SwapABGemmRunner.tuning_config, [input, weight, weight_scale], ) + if tuner.is_build_mode: + # The build state queued this GEMM's builds; nothing runs. + return input.new_zeros((input.size(0), weight.size(0)), + dtype=output_dtype) return gemm_runner( inputs=[input, weight, weight_scale], tactic=best_tactic, diff --git a/tensorrt_llm/_torch/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py new file mode 100644 index 000000000000..d3fa89e71c10 --- /dev/null +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``KernelBuilder``: kernels built on worker threads before the calls that need them. + +An op whose kernels compile on first use, and that does not go through the autotuner, +registers them here, the way an op that does registers its tactics with the autotuner. +The forward supplies the largest shape; the op sweeps its own shape space from it. Inside +a build pass, the op's ``BuildConfig`` turns the shape a call has into the shapes to build +for, as a ``TuningConfig`` turns it into tuning buckets, and the op lists the kernels of +those its call signature has not listed yet. Worker threads then build them in parallel. +The op's own later call finds a kernel built, or waits on the op's own lock for a build +that is still running. + +An op that goes through the autotuner builds through it instead: in the autotuner's +build state (``autotuner.build_kernels``), ``choose_one`` queues here the builds of every +shape, runner and tactic the op's tuning will time. + +A rehearsal is a build pass over the autotuner warmup's forwards on dummy weights, run +before the checkpoint loads, in the autotuner's build state: the ops list their kernels +without building or launching them, and the pass does not wait, so the builds run while +the weights load. It runs for a model that one of the op families in ``_FAMILIES`` +reaches. + +``TLLM_KERNEL_BUILDER=0`` turns it off: no pass is active and nothing is submitted. +``TLLM_KERNEL_BUILDER_REHEARSAL=0`` keeps the builds but not the rehearsal: they start +at the warmup. ``TLLM_KERNEL_BUILDER_THREADS`` sets the number of worker threads +(default: this rank's share of the CPUs, up to 32, as for the autotuner's parallel JIT +warmup). +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import importlib +import itertools +import os +import queue +import threading +import time +from typing import ( + Any, + Callable, + Dict, + Hashable, + Iterable, + Iterator, + List, + Mapping, + Optional, + Tuple, + Union, +) + +import torch + +from tensorrt_llm.logger import logger + +__all__ = ["BuildConfig", "DynamicSpec", "KernelBuilder"] + +_ENV = "TLLM_KERNEL_BUILDER" +_REHEARSAL_ENV = "TLLM_KERNEL_BUILDER_REHEARSAL" +_THREADS_ENV = "TLLM_KERNEL_BUILDER_THREADS" +_PROGRESS_EVERY_S = 60.0 + +# The op families a rehearsal can reach: name -> "module:function". The function takes the +# model config and the model as created, returns True when the model reaches the op, and +# then queues the builds the configuration alone determines. A module is imported only +# when asked. +_FAMILIES = { + "fmha_sm100": ( + "tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_kernel_builds:reaches" + ), + "deep_gemm": "tensorrt_llm._torch.modules.linear:deep_gemm_reaches", +} + +# A build returns True when it compiled the kernel, False when the kernel was already there. +Build = Callable[[], bool] + +_QUEUED, _DONE, _FAILED = "queued", "done", "failed" + + +@dataclasses.dataclass(frozen=True) +class DynamicSpec: + """A dimension an op's kernels depend on, swept as the autotuner sweeps a + ``DynamicTensorSpec``. + + Args: + name: The dimension's key in the shape a call has and in the shapes swept. + gen_buckets: The values to build for: a tuple, or a function of the value the + call has (the largest a forward supplies) returning them. The call's own + value is always one of them. + """ + + name: str + gen_buckets: Union[Tuple[int, ...], Callable[[int], Iterable[int]]] + + +@dataclasses.dataclass(frozen=True) +class BuildConfig: + """The dimensions an op's kernels depend on, as a ``TuningConfig`` declares its own.""" + + dynamic_specs: Tuple[DynamicSpec, ...] = () + + def shapes(self, observed: Mapping[str, int]) -> List[Dict[str, int]]: + """Every combination of the dimensions' buckets for a call that has ``observed``.""" + axes = [] + for spec in self.dynamic_specs: + value = observed[spec.name] + buckets = spec.gen_buckets(value) if callable(spec.gen_buckets) else spec.gen_buckets + axes.append(sorted(set(buckets) | {value})) + names = [spec.name for spec in self.dynamic_specs] + return [dict(zip(names, point)) for point in itertools.product(*axes)] + + +@dataclasses.dataclass +class _OpStats: + signatures: int = 0 + list_s: float = 0.0 + submitted: int = 0 + built: int = 0 + present: int = 0 + failed: int = 0 + build_s: float = 0.0 + late: int = 0 + late_s: float = 0.0 + waited: int = 0 + waited_s: float = 0.0 + + def summary(self) -> str: + parts = [ + f"{self.signatures} call signatures listed in {self.list_s:.1f}s", + f"{self.submitted} kernels submitted: {self.built} built in {self.build_s:.1f}s, " + f"{self.present} already built, {self.failed} failed", + ] + if self.waited: + parts.append(f"the op waited {self.waited_s:.1f}s for {self.waited} of them") + if self.late: + parts.append(f"{self.late} compiled by the op itself in {self.late_s:.1f}s") + return "; ".join(parts) + + +class KernelBuilder: + """Process-wide queue of kernel builds, run on worker threads. + + ``prebuild`` is called at an op's entry point; outside a build pass it returns after + one check. ``submit`` queues one build directly, for builds that need only the + configuration. ``op_compiled`` records a kernel the op compiled itself. + """ + + _instance: Optional["KernelBuilder"] = None + _instance_lock = threading.Lock() + + def __init__( + self, + threads: Optional[int] = None, + enabled: Optional[bool] = None, + rehearse: Optional[bool] = None, + ): + self.enabled = os.environ.get(_ENV, "1") != "0" if enabled is None else enabled + if rehearse is None: + rehearse = os.environ.get(_REHEARSAL_ENV, "1") != "0" + self.rehearse = self.enabled and rehearse + if threads is None: + if _THREADS_ENV in os.environ: + threads = int(os.environ[_THREADS_ENV]) + else: + from .autotuner import _default_parallel_jit_warmup_threads + + threads = _default_parallel_jit_warmup_threads() + self.threads = max(1, threads) + self._lock = threading.Lock() + self._idle = threading.Condition(self._lock) + self._queue: "queue.Queue[Tuple[str, Hashable, Build]]" = queue.Queue() + self._workers: List[threading.Thread] = [] + self._device: Optional[int] = None + self._passes = 0 + self._rehearsals = 0 + self._pending = 0 + # Queued and running builds per op. + self._pending_by_op: Dict[str, int] = {} + # The shapes listed so far for each (op, call signature). + self._listed: Dict[Tuple[str, Hashable], set] = {} + self._states: Dict[Tuple[str, Hashable], str] = {} + self._stats: Dict[str, _OpStats] = {} + + @classmethod + def get(cls) -> "KernelBuilder": + if cls._instance is None: + with cls._instance_lock: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + # Build passes. + + @contextlib.contextmanager + def build_pass(self, name: str, wait: bool = True) -> Iterator[None]: + """Inside, ops list the kernels of the shapes they sweep. The outermost pass logs + what was built at its end, after waiting for the builds unless ``wait`` is False.""" + if not self.enabled: + yield + return + with self._lock: + self._passes += 1 + try: + yield + except BaseException: + with self._lock: + self._passes -= 1 + raise + with self._lock: + self._passes -= 1 + last = self._passes == 0 + if last: + waited = self.join() if wait else 0.0 + for line in self.report(name, waited): + logger.info(line) + + @contextlib.contextmanager + def rehearsal(self) -> Iterator[None]: + """A build pass over forwards on dummy weights. Ops list their kernels without + building or launching them, and the pass does not wait for the builds.""" + if not self.enabled: + yield + return + with self._lock: + self._rehearsals += 1 + try: + with self.build_pass("rehearsal", wait=False): + yield + finally: + with self._lock: + self._rehearsals -= 1 + + def active(self) -> bool: + return self.enabled and self._passes > 0 + + def rehearsing(self) -> bool: + """True inside a rehearsal: an op lists its kernels, then skips its own launch.""" + return self.enabled and self._rehearsals > 0 + + def wants_rehearsal(self, model_config: Any, model: Any = None) -> bool: + """True when the rehearsal is on and an op family in ``_FAMILIES`` reaches + ``model``, as created from ``model_config``. Each family that does queues the + builds its configuration determines.""" + if not self.rehearse: + return False + wanted = False + for name, target in _FAMILIES.items(): + module_name, function = target.split(":") + try: + reaches = getattr(importlib.import_module(module_name), function) + wanted = bool(reaches(model_config, model)) or wanted + except Exception as e: # noqa: BLE001 - the family's kernels build at the warmup + logger.warning( + f"[KernelBuilder] {name}: could not tell whether the model reaches it; " + f"its kernels build at the warmup. {type(e).__name__}: {e}" + ) + return wanted + + # Registration. + + def prebuild( + self, + op: str, + signature: Hashable, + config: BuildConfig, + observed: Mapping[str, int], + list_builds: Callable[[List[Dict[str, int]]], Iterable[Tuple[Hashable, Build]]], + ) -> None: + """In a build pass, submit every build ``list_builds(shapes)`` returns as + ``(key, build)``, where ``shapes`` are those ``config`` sweeps for a call that has + ``observed`` and that no earlier call of ``(op, signature)`` swept. + + A failure to list leaves the kernels to the op. + """ + if not self.active(): + return + swept = config.shapes(observed) + with self._lock: + listed = self._listed.setdefault((op, signature), set()) + new_signature = not listed + shapes = [shape for shape in swept if tuple(sorted(shape.items())) not in listed] + if not shapes: + return + listed.update(tuple(sorted(shape.items())) for shape in shapes) + t0 = time.perf_counter() + try: + builds = list(dict(list_builds(shapes)).items()) + except Exception as e: # noqa: BLE001 - the op still builds what it needs on first use + logger.warning( + f"[KernelBuilder] {op}: listing the kernels of a call signature failed; the op " + f"builds them on first use. {type(e).__name__}: {e}" + ) + return + list_s = time.perf_counter() - t0 + queued = sum(self.submit(op, key, build) for key, build in self._rotated(builds)) + with self._lock: + stats = self._op_stats(op) + stats.signatures += new_signature + stats.list_s += list_s + logger.info( + f"[KernelBuilder] {op}: {len(builds)} kernels over {len(shapes)} new shapes of " + f"{'a new' if new_signature else 'a'} call signature, {queued} queued " + f"({list_s:.1f}s to list)" + ) + + def submit(self, op: str, key: Hashable, build: Build) -> bool: + """Queue one build; False when ``(op, key)`` was already queued or built.""" + if not self.enabled: + return False + with self._lock: + if (op, key) in self._states: + return False + self._states[(op, key)] = _QUEUED + self._op_stats(op).submitted += 1 + self._pending += 1 + self._pending_by_op[op] = self._pending_by_op.get(op, 0) + 1 + self._start_workers() + self._queue.put((op, key, build)) + return True + + def state(self, op: str, key: Hashable) -> Optional[str]: + """``"queued"``, ``"done"``, ``"failed"``, or None for a kernel never submitted.""" + with self._lock: + return self._states.get((op, key)) + + def stats(self, op: str) -> _OpStats: + """A copy of ``op``'s counts so far.""" + with self._lock: + return dataclasses.replace(self._op_stats(op)) + + def op_compiled(self, op: str, key: Hashable, seconds: float) -> None: + """Record a kernel the op's own call had to build or wait for.""" + if not self.enabled: + return + with self._lock: + stats = self._op_stats(op) + waited = self._states.get((op, key)) in (_QUEUED, _DONE) + if waited: + stats.waited += 1 + stats.waited_s += seconds + else: + stats.late += 1 + stats.late_s += seconds + if waited: + logger.info(f"[KernelBuilder] {op} {key}: the op waited {seconds:.1f}s for its build") + else: + logger.info( + f"[KernelBuilder] {op} {key}: not built ahead, the op compiled it in {seconds:.1f}s" + ) + + # Waiting and reporting. + + def join(self, timeout: Optional[float] = None, prefix: Optional[str] = None) -> float: + """Wait for every queued build, or for those of the ops whose name starts with + ``prefix``; returns the seconds waited.""" + t0 = last = time.perf_counter() + with self._idle: + while (pending := self._pending_for(prefix)) > 0: + now = time.perf_counter() + if timeout is not None and now - t0 >= timeout: + break + if now - last >= _PROGRESS_EVERY_S: + logger.info(f"[KernelBuilder] waiting {now - t0:.0f}s for {pending} builds") + last = now + self._idle.wait(timeout=1.0) + return time.perf_counter() - t0 + + def _pending_for(self, prefix: Optional[str]) -> int: + if prefix is None: + return self._pending + return sum(n for op, n in self._pending_by_op.items() if op.startswith(prefix)) + + def report(self, name: str, waited: float = 0.0) -> List[str]: + with self._lock: + lines = [ + f"[KernelBuilder] {name}: {op}: {stats.summary()}" + for op, stats in sorted(self._stats.items()) + ] + pending = self._pending + if lines and waited >= 0.1: + lines.append(f"[KernelBuilder] {name}: waited {waited:.1f}s for the builds at its end") + if pending: + lines.append(f"[KernelBuilder] {name}: {pending} builds continue after it") + return lines + + # Workers. + + def _op_stats(self, op: str) -> _OpStats: + return self._stats.setdefault(op, _OpStats()) + + def _rotated(self, builds: List[Tuple[Hashable, Build]]) -> List[Tuple[Hashable, Build]]: + """Start each GPU's list at a different place: ranks sharing a cache directory then + build different kernels first instead of waiting on each other's locks.""" + if not builds or not torch.cuda.is_available(): + return builds + offset = torch.cuda.current_device() % len(builds) + return builds[offset:] + builds[:offset] + + def _start_workers(self) -> None: + if self._device is None and torch.cuda.is_available(): + self._device = torch.cuda.current_device() + while len(self._workers) < min(self.threads, self._pending): + worker = threading.Thread( + target=self._work, name=f"tllm-kernel-builder-{len(self._workers)}", daemon=True + ) + worker.start() + self._workers.append(worker) + + def _work(self) -> None: + device_set = False + while True: + op, key, build = self._queue.get() + t0 = time.perf_counter() + state = _DONE + # Every item is settled, whatever it raises: an unsettled one would hang join(). + try: + if not device_set and self._device is not None: + torch.cuda.set_device(self._device) + device_set = True + built = build() + except BaseException as e: # noqa: BLE001 - the op builds a failed kernel on first use + state, built = _FAILED, False + logger.warning( + f"[KernelBuilder] {op} {key}: build failed; the op builds it on first use. " + f"{type(e).__name__}: {e}" + ) + # A build can hold the tensors it was made from: drop it before the next get(). + del build + seconds = time.perf_counter() - t0 + with self._idle: + self._states[(op, key)] = state + stats = self._op_stats(op) + if state == _FAILED: + stats.failed += 1 + elif built: + stats.built += 1 + stats.build_s += seconds + else: + stats.present += 1 + self._pending -= 1 + self._pending_by_op[op] -= 1 + # Waits by op prefix end as an op's builds finish, not only at idle. + self._idle.notify_all() diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 0879a830fb20..f0bbeea9e598 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -20,7 +20,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( - BufferKind, mxfp8_quantize_gemm_autotuned) + BufferKind, mxfp8_quantize_gemm_autotuned, prebuild_flashinfer_mxfp8_gemm) from tensorrt_llm._torch.peft.lora.layer import LoraLayer from tensorrt_llm._utils import is_device_integrated, mpi_disabled from tensorrt_llm.bindings import ipc_nvls_supported @@ -1402,6 +1402,20 @@ def load_weights_fused_gate_up_linear( shard_offset, shard_size) +def deep_gemm_reaches(model_config, model) -> bool: + """Whether ``model`` runs an FP8 block-scale linear through DeepGEMM, whose + kernels the autotuner builds in its build state (a KernelBuilder op family).""" + del model_config + if model is None or not is_sm_100f(): + return False + return any( + isinstance(getattr(module, "quant_method", None), + FP8BlockScalesLinearMethod) + and not (getattr(module, "use_cute_dsl_blockscaling_mm", False) + or getattr(module, "disable_deep_gemm", False)) + for module in model.modules()) + + class FP8BlockScalesLinearMethod(UnquantizedLinearMethod): # fp8_block_scaling_gemm does not support writing into an NCCL window buffer. @@ -3799,6 +3813,13 @@ def apply(self, module: Linear, input: torch.Tensor, # Dynamic MXFP8 activation quantization (swizzled SF layout), # then the CUTLASS block-scaled e4m3xe4m3 GEMM. act_e4m3, act_sf = torch.ops.trtllm.mxfp8_quantize(input, True) + if self.needs_flashinfer_autotune and not is_torch_compiling(): + # The warmup will tune FlashInfer's GEMM: in the autotuner's + # build state, build it now. + prebuild_flashinfer_mxfp8_gemm(act_e4m3, act_sf, + module.weight, + module.weight_scale, + module.dtype) use_flashinfer = self.backend == "flashinfer" or ( self.backend == "auto" and not is_torch_compiling() and (_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.get() or diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 7419e578f422..e8cb6df93419 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -435,6 +435,17 @@ def warmup(self, resource_manager: ResourceManager) -> None: self._eager_workspace_reclaimer = None self._eager_workspace_reclaimer = self._warmup.run(resource_manager) + def rehearse(self, resource_manager: ResourceManager) -> None: + """The warmup's rehearsal forwards. The state they bind to ``resource_manager`` + (attention and speculative metadata, cached input copies) goes with it.""" + try: + self._warmup.rehearse(resource_manager) + finally: + self.attn_metadata = None + self.spec_metadata = None + if self._h2d_cache is not None: + self._h2d_cache.clear() + def release_graphs(self, *, release_nccl_window_owners: bool = True) -> None: self._graphs.release(release_nccl_window_owners=release_nccl_window_owners) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 9998624cffd8..b82f2ad1f187 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -15,12 +15,14 @@ import torch from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.autotuner import AutoTuner, autotune, build_kernels from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.distributed.ops import warmup_named_allreduce_topologies +from tensorrt_llm._torch.kernel_builder import KernelBuilder from tensorrt_llm._torch.memory_buffer_utils import clear_memory_buffers from tensorrt_llm._torch.models.modeling_utils import timing_metric from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.nccl_window_graph import nccl_window_allocation_suppressed from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import ( @@ -297,10 +299,89 @@ def run(self, resource_manager: ResourceManager) -> EagerWorkspaceReclaimer | No ) # Warm here, not in AllReduce.__init__, to scope it to the LLM executor's graph capture. warmup_named_allreduce_topologies(self.model) - # Retain a partial summary when a phase raises. - with self._warmup_timer: + # Retain a partial summary when a phase raises. Ops list their kernels for + # KernelBuilder as calls reach new shapes; the pass waits for the builds at its end. + with self._warmup_timer, KernelBuilder.get().build_pass("warmup"): return self._warmup_scheduled(resource_manager, kv_cache_manager) + def rehearse(self, resource_manager: ResourceManager) -> None: + """The autotuner warmup's forwards, run eagerly in the autotuner's build state + on whatever weights the model holds, inside the caller's KernelBuilder rehearsal. + + They supply the largest shapes: each op the forwards reach sweeps its own shapes + down from them, and an op that tunes queues the builds its tuning will time. + Every rank calls it. Its requests are released and its workspaces freed. + """ + kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.KV_CACHE_MANAGER + ) + if kv_cache_manager is None or not self._context_warmups_supported: + return + if self.mapping.cp_size > 1 and self.mapping.cp_config.get("cp_type") != CpType.HELIX: + return + AutoTuner.get() + if self._config.enable_autotuner: + # The MXFP8 methods dispatch as the autotuner warmup will tune them. + _select_mxfp8_methods(self) + warmup_configs, _ = _autotuner_warmup_configs( + self, resource_manager, generation_only=self.mapping.has_cp_helix() + ) + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner + + try: + # force_eager: no torch.compile traces, whole-model or op-level. Window + # requests come from plain memory, so the NCCL window pool grows only with + # the warmup that sizes it, not with these eager forwards' buffers too. + with ( + torch.inference_mode(), + torch.compiler.set_stance("force_eager"), + self._torch_compile_bypass(), + nccl_window_allocation_suppressed(), + self._graphs.no_cuda_graph(), + build_kernels(), + ): + for num_tokens, num_gen_requests in warmup_configs: + self._rehearse_batch(resource_manager, num_tokens, num_gen_requests) + except BaseException: + self._reset_moe_alltoall_state() + raise + finally: + torch.cuda.synchronize() + MoERunner.clear_all_workspaces() + clear_memory_buffers() + gc.collect() + torch.cuda.empty_cache() + + def _rehearse_batch( + self, resource_manager: ResourceManager, num_tokens: int, num_gen_requests: int + ) -> None: + """One forward of the rehearsal, on a batch as the autotuner warmup builds it.""" + warmup_request = self._dummy_requests.create_warmup_request( + resource_manager, num_tokens, num_gen_requests + ) + with self._dummy_requests.release_batch_context(warmup_request, resource_manager) as batch: + if not self._agreement.should_run_warmup_batch( + batch, + num_tokens, + f"rehearsal, num_tokens={num_tokens}, num_gen_requests={num_gen_requests}", + ): + return + with self._warmup_timer.phase( + f"rehearsal shape num_tokens={num_tokens}, num_gen_requests={num_gen_requests}", + record=False, + log_start=True, + ): + with discard_nccl_window_tensor_outputs( + (getattr(self, "_warmup_inputs", None), batch) + ): + self._forward_warmup( + batch, + resource_manager, + enable_spec_decode=self._config.is_spec_decode, + runtime_draft_len=get_static_draft_len(self._config.spec_config), + ) + torch.cuda.synchronize() + def _get_max_shape_warmup_requests( self, resource_manager: ResourceManager ) -> list[tuple[int, int]]: @@ -1037,70 +1118,14 @@ def _run_autotuner_warmup( ``generation_only`` tunes on one batch of ``max_batch_size`` generation requests instead of a context batch, for Helix CP generation workers. """ - from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( - IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, - MXFP8GemmRunner, - ) - from tensorrt_llm._torch.modules.linear import ( - MXFP8LinearMethod, - flashinfer_mxfp8_autotune, - mxfp8_gemm_backend_override, - ) + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner + from tensorrt_llm._torch.modules.linear import flashinfer_mxfp8_autotune enable_trtllm_autotuner = self._config.enable_autotuner if not enable_trtllm_autotuner: return - mxfp8_methods = [] - for module in self.model.modules(): - quant_method = getattr(module, "quant_method", None) - if isinstance(quant_method, MXFP8LinearMethod): - mxfp8_methods.append(quant_method) - - # This engine owns startup warmup, so it explicitly opts its MXFP8 - # methods into native tuning. Standalone modules and engine paths that - # skip this warmup remain on the direct native op. - for method in mxfp8_methods: - method.enable_native_autotune() - - # Native and FlashInfer tuning are independent. Capture native - # eligibility before enabling graph-only FlashInfer dispatch. - native_mxfp8_methods = [method for method in mxfp8_methods if method.needs_native_autotune] - compile_all_batches = ( - self._config.torch_compile_enabled and not self._config.torch_compile_prefill_only - ) - if compile_all_batches and mxfp8_gemm_backend_override() is None: - # Compiled auto dispatch uses native; do not tune unused backends. - # Prefill-only compile retains the eager generation-graph policy. - for method in mxfp8_methods: - method.disable_flashinfer_auto() - use_mxfp8_flashinfer_graph_default = ( - self._graphs.enabled - and not compile_all_batches - and mxfp8_gemm_backend_override() is None - and any( - getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) - for module in self.model.modules() - ) - ) - if use_mxfp8_flashinfer_graph_default: - # CuTeDSL is the alternative backend; PP has no graph-pass handoff. - tune_with_cute_dsl = ( - IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE and not self.mapping.has_pp() - ) - for quant_method in mxfp8_methods: - quant_method.enable_flashinfer_auto() - quant_method.tune_decode_graph_backends = ( - tune_with_cute_dsl and quant_method.uses_flashinfer - ) - # An explicit auto setting stays intact, but compiled auto dispatch - # uses native GEMM. Do not run or mark an unused FlashInfer tuning pass. - flashinfer_mxfp8_methods = [ - method - for method in mxfp8_methods - if method.needs_flashinfer_autotune - and (not compile_all_batches or method.backend == "flashinfer") - ] + mxfp8_methods, native_mxfp8_methods, flashinfer_mxfp8_methods = _select_mxfp8_methods(self) # Every TP and PP rank must make the same backend decision before any # rank returns or enters a tuning forward with model collectives. @@ -1133,25 +1158,9 @@ def _run_autotuner_warmup( enable_flashinfer_mxfp8_autotuner = bool(flashinfer_mxfp8_methods) enable_native_mxfp8_autotuner = bool(native_mxfp8_methods) - kv_cache_manager = resource_manager.get_resource_manager( - ResourceManagerType.KV_CACHE_MANAGER + warmup_configs, curr_max_num_tokens = _autotuner_warmup_configs( + self, resource_manager, generation_only ) - if generation_only: - tokens_per_gen_request = 1 + self._config.max_total_draft_tokens - num_gen_requests = min( - self._config.max_batch_size, - self._config.max_num_tokens // tokens_per_gen_request, - ) - curr_max_num_tokens = num_gen_requests * tokens_per_gen_request - else: - token_num_upper_bound = min( - self._config.max_num_tokens, - self._config.max_batch_size * (self._config.max_seq_len - 1), - ) - curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( - token_num_upper_bound=token_num_upper_bound, - max_num_draft_tokens=self._config.original_max_draft_len, - ) # The warmup runs again after the KV cache is rebuilt; when the token # shape is unchanged and no MXFP8 tuning is pending, the profiling @@ -1181,18 +1190,6 @@ def _run_autotuner_warmup( f"FlashInfer MXFP8={enable_flashinfer_mxfp8_autotuner})..." ) - if generation_only: - warmup_configs = [(curr_max_num_tokens, num_gen_requests)] - else: - warmup_configs = [(curr_max_num_tokens, 0)] - if ( - self.guided_decoder is None - and not self.mapping.has_pp() - and not is_context_only_kv_cache(kv_cache_manager) - ): - # Add generation request to warmup the autotuner cache. - warmup_configs.append((1 + self._config.max_total_draft_tokens, 1)) - def run_autotuner_pass(autotune_context: Any, synchronize_trtllm_cache: bool) -> bool: """Run one isolated tuning pass with fresh synthetic batches.""" ran_forward = False @@ -1608,3 +1605,107 @@ def _prewarm_minimax_m3_msa_indexer( if self.mapping.tp_size > 1 and self.dist is not None: self.dist.tp_allgather(1) logger.info("MiniMax-M3 MSA indexer prewarm complete") + + +def _autotuner_warmup_configs( + warmup: DecoderWarmup, resource_manager: ResourceManager, generation_only: bool = False +) -> tuple[list[tuple[int, int]], int]: + """The autotuner warmup's batches as ``(num_tokens, num_gen_requests)``, and its + largest token count. A function of the warmup, not a method: tests run the warmup's + methods on stand-in objects. + + A context batch of the most tokens this rank's KV cache holds, and one generation + request with its draft tokens; with ``generation_only``, one batch of + ``max_batch_size`` generation requests instead. + """ + kv_cache_manager = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) + if generation_only: + tokens_per_gen_request = 1 + warmup._config.max_total_draft_tokens + num_gen_requests = min( + warmup._config.max_batch_size, + warmup._config.max_num_tokens // tokens_per_gen_request, + ) + curr_max_num_tokens = num_gen_requests * tokens_per_gen_request + return [(curr_max_num_tokens, num_gen_requests)], curr_max_num_tokens + token_num_upper_bound = min( + warmup._config.max_num_tokens, + warmup._config.max_batch_size * (warmup._config.max_seq_len - 1), + ) + curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( + token_num_upper_bound=token_num_upper_bound, + max_num_draft_tokens=warmup._config.original_max_draft_len, + ) + warmup_configs = [(curr_max_num_tokens, 0)] + if ( + warmup.guided_decoder is None + and not warmup.mapping.has_pp() + and not is_context_only_kv_cache(kv_cache_manager) + ): + # Add generation request to warmup the autotuner cache. + warmup_configs.append((1 + warmup._config.max_total_draft_tokens, 1)) + return warmup_configs, curr_max_num_tokens + + +def _select_mxfp8_methods(warmup: DecoderWarmup) -> tuple[list, list, list]: + """This rank's MXFP8 linear methods, and those the autotuner warmup tunes. A + function of the warmup, not a method: tests run the warmup's methods on + stand-in objects. + + Returns ``(all, native, flashinfer)``. It also sets each method's dispatch + for the warmup and serving, the same way on every call; the warmup then + agrees the FlashInfer choice across ranks. + """ + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE, + ) + from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod, mxfp8_gemm_backend_override + + mxfp8_methods = [] + for module in warmup.model.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, MXFP8LinearMethod): + mxfp8_methods.append(quant_method) + + # This engine owns startup warmup, so it explicitly opts its MXFP8 + # methods into native tuning. Standalone modules and engine paths that + # skip this warmup remain on the direct native op. + for method in mxfp8_methods: + method.enable_native_autotune() + + # Native and FlashInfer tuning are independent. Capture native + # eligibility before enabling graph-only FlashInfer dispatch. + native_mxfp8_methods = [method for method in mxfp8_methods if method.needs_native_autotune] + compile_all_batches = ( + warmup._config.torch_compile_enabled and not warmup._config.torch_compile_prefill_only + ) + if compile_all_batches and mxfp8_gemm_backend_override() is None: + # Compiled auto dispatch uses native; do not tune unused backends. + # Prefill-only compile retains the eager generation-graph policy. + for method in mxfp8_methods: + method.disable_flashinfer_auto() + use_mxfp8_flashinfer_graph_default = ( + warmup._graphs.enabled + and not compile_all_batches + and mxfp8_gemm_backend_override() is None + and any( + getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) + for module in warmup.model.modules() + ) + ) + if use_mxfp8_flashinfer_graph_default: + # CuTeDSL is the alternative backend; PP has no graph-pass handoff. + tune_with_cute_dsl = IS_FLASHINFER_MXFP8_CUTE_DSL_AVAILABLE and not warmup.mapping.has_pp() + for quant_method in mxfp8_methods: + quant_method.enable_flashinfer_auto() + quant_method.tune_decode_graph_backends = ( + tune_with_cute_dsl and quant_method.uses_flashinfer + ) + # An explicit auto setting stays intact, but compiled auto dispatch + # uses native GEMM. Do not run or mark an unused FlashInfer tuning pass. + flashinfer_mxfp8_methods = [ + method + for method in mxfp8_methods + if method.needs_flashinfer_autotune + and (not compile_all_batches or method.backend == "flashinfer") + ] + return mxfp8_methods, native_mxfp8_methods, flashinfer_mxfp8_methods diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 514e9efd7399..eaed5e8d996f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -128,7 +128,11 @@ def __init__( checkpoint_loader: BaseCheckpointLoader | None = None, model_weights_memory_tag: str | None = None, model_weights_restore_mode=None, + defer_checkpoint: bool = False, ): + """``defer_checkpoint``: for a model a KernelBuilder rehearsal reaches, build + the engine on dummy weights; ``has_deferred_checkpoint`` then tells, and the + caller runs ``rehearse`` and ``load_deferred_checkpoint``.""" # cleanup() also runs on an engine whose construction failed part-way. self._cleanup_done = False self.model_loader: ModelLoader | None = None @@ -148,7 +152,11 @@ def __init__( checkpoint_loader=checkpoint_loader, model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, + defer_checkpoint=defer_checkpoint, ) + if self.has_deferred_checkpoint and not self.model.model_config.is_generation: + # No KV cache to rehearse with. + self.model_loader.load_deferred_checkpoint(self.model) if not self.model.model_config.is_generation: self.llm_args = resolve_scheduled_encoder_args(self.llm_args) self._validate_breakable_cuda_graph_compatibility() @@ -262,6 +270,7 @@ def _load_model( checkpoint_loader: BaseCheckpointLoader | None, model_weights_memory_tag: str | None, model_weights_restore_mode, + defer_checkpoint: bool = False, ) -> None: """Load the model, or take the given one, at the requested capacity.""" self.moe_load_balancer: BaseMoeLoadBalancer | None = None @@ -278,7 +287,9 @@ def _load_model( ) # Open checkpoint and load the LLM module object. self.model, moe_load_balancer = self.model_loader.load( - checkpoint_dir=model_path, checkpoint_loader=checkpoint_loader + checkpoint_dir=model_path, + checkpoint_loader=checkpoint_loader, + defer_checkpoint=defer_checkpoint, ) if isinstance(moe_load_balancer, BaseMoeLoadBalancer): self.moe_load_balancer = moe_load_balancer @@ -1072,6 +1083,27 @@ def release_executor_resources(self, *, release_model_owned_graphs: bool) -> Non if isinstance(self._runner, DecoderRunner): self._runner.attn_metadata = None + @property + def has_deferred_checkpoint(self) -> bool: + """True while the model holds the dummy weights of a deferred checkpoint load.""" + return self.model_loader is not None and self.model_loader.has_deferred_checkpoint + + def rehearse(self, resource_manager: ResourceManager) -> None: + """Warmup forwards on the dummy weights, inside the caller's KernelBuilder + rehearsal, so the ops list their kernels; every rank calls it.""" + if not isinstance(self._runner, DecoderRunner): + return + self.set_warmup_purpose("rehearsal") + with self.warmup_scope(): + self._runner.rehearse(resource_manager) + + def load_deferred_checkpoint(self) -> None: + """Load the checkpoint the model loader deferred, as a weight update does.""" + self.unwrap_compiled_model_for_refit() + self.model_loader.load_deferred_checkpoint(self.model) + self.restore_compiled_model_after_refit(None) + self.invalidate_multimodal_encoder_cache() + def unwrap_compiled_model_for_refit(self) -> None: """Release compiled state before refit. diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 632f4b6ce007..84bda75b0701 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -8,9 +8,9 @@ import traceback import warnings from contextlib import AbstractContextManager, contextmanager -from dataclasses import replace +from dataclasses import dataclass, replace from enum import Enum -from typing import Any, Callable, Iterator, Optional, Tuple +from typing import Any, Callable, Iterable, Iterator, Optional import torch @@ -49,6 +49,7 @@ from ...llmapi.llm_args import CheckpointIOPolicy, LoadFormat from .._experimental.modeling_v2 import modeling_v2_resolve_on_all_ranks +from ..kernel_builder import KernelBuilder from ..locality_domain.policy import LocalityDomainPolicy from ..model_config import ModelConfig from ..models import AutoModelForCausalLM @@ -453,6 +454,47 @@ def validate_lm_head_output_dtype(model: torch.nn.Module, f"{type(model).__name__}: no LMHead was constructed.") +# Bounds for uniform dummy values of the 1-byte types torch cannot sample directly. They are +# not the widest possible, only small enough to avoid NaNs. FP4 bytes are 2 packed values: +# up to 64, the high 4 bits stay 0. +_DUMMY_QUANTIZED_BOUNDS = { + torch.float8_e4m3fn: (-3.0, 3.0), + torch.int8: (-3.0, 3.0), + float4_e2m1x2: (0, 64), +} + + +def _fill_dummy_quantized(tensor: torch.Tensor, + generator: torch.Generator) -> None: + """Uniform values of a type in `_DUMMY_QUANTIZED_BOUNDS`, sampled in FP16 and cast.""" + low, high = _DUMMY_QUANTIZED_BOUNDS[tensor.dtype] + values = torch.empty(tensor.shape, + dtype=torch.float16, + device=tensor.device) + tensor.copy_( + values.uniform_(low, high, generator=generator).to(tensor.dtype)) + + +def fill_rehearsal_values(tensors: Iterable[torch.Tensor], + low: float = -1e-3, + high: float = 1e-3, + seed: int = 0) -> None: + """Values a forward can run on, for tensors that have none of their own. + + Floating and quantized tensors get dummy values as `initialize_dummy_weights` gives + them; any other tensor gets zeros, so an index stays in range. + """ + for tensor in tensors: + generator = torch.Generator(device=tensor.device) + generator.manual_seed(seed) + if tensor.dtype in _DUMMY_QUANTIZED_BOUNDS: + _fill_dummy_quantized(tensor, generator) + elif torch.is_floating_point(tensor) and tensor.element_size() >= 2: + tensor.uniform_(low, high, generator=generator) + else: + tensor.zero_() + + def initialize_dummy_weights( model: torch.nn.Module, low: float = -1e-3, @@ -470,23 +512,6 @@ def initialize_dummy_weights( to not generate NaNs/inf for these. """ - def _get_random_min_max(dtype: torch.dtype) -> Tuple[int, int]: - """Return safe (min, max) bounds for uniform sampling of `dtype`.""" - # These values are not necessarily the largest possible min/max, - # they need to be small enough to avoid NaNs. - if dtype in (torch.float8_e4m3fn, torch.int8): - return (-3.0, 3.0) - - elif dtype == float4_e2m1x2: - # These correspond to bits of 2 packed FP4 values. - # Because we only go up to 64, the high 4 bits will - # always be 0. But this is fine - we just need values - # that won't generate NaNs. - return (0, 64) - - else: - raise NotImplementedError(f"Unknown quantized type: {dtype}.") - # Calibration scalars (input_scale / inv_input_scale / kv_scales / # inv_kv_scales / alpha / scalar_alpha) must keep their create_weights # default (typically 1.0). Randomizing them breaks FP8 attention output @@ -510,18 +535,9 @@ def _get_random_min_max(dtype: torch.dtype) -> Tuple[int, int]: dtype = param.data.dtype if param.data.element_size() < 2: - # We need to do a cast/round since torch doesn't have uniform_ - # support for these dtypes. - tmp_param = torch.empty(param.data.shape, - dtype=torch.float16, - device=param.data.device) - - quant_min, quant_max = _get_random_min_max(dtype) - tmp_param = tmp_param.uniform_(quant_min, - quant_max, - generator=generator) - - param.data.copy_(tmp_param.to(dtype)) + if dtype not in _DUMMY_QUANTIZED_BOUNDS: + raise NotImplementedError(f"Unknown quantized type: {dtype}.") + _fill_dummy_quantized(param.data, generator) # Note: no need to to mess with int32 params, these are probably # constants and not weights. @@ -871,6 +887,43 @@ class ModelLoaderMetricNames(Enum): WEIGHT_MANIFEST_SECONDS = "weight_manifest_seconds" +class _ModuleTree: + """Each module's submodules and attributes, to put the module tree back as it was. + + Putting back a submodule slot is not enough: assigning a module to an attribute + (e.g. a draft's `embed_tokens = None` tied to the target's) moves the name from + the attributes to the submodules. + """ + + def __init__(self, model: torch.nn.Module): + self._modules = [(module, dict(module._modules), { + name: value + for name, value in vars(module).items() + if name not in ("_parameters", "_buffers", "_modules") + }) for module in model.modules()] + + def restore(self) -> None: + for module, children, attributes in self._modules: + module._modules.clear() + module._modules.update(children) + for name, value in attributes.items(): + module.__dict__.setdefault(name, value) + + +@dataclass +class _DeferredCheckpoint: + """A checkpoint load `ModelLoader.load` left for `load_deferred_checkpoint`.""" + config: ModelConfig + checkpoint_dir: str + checkpoint_loader: BaseCheckpointLoader + loads_draft_weights: bool + post_transform_qualification: PostTransformQualificationDecision + # The module tree as created. + tree: _ModuleTree + # Each parameter's (module, name, shape, dtype) as created. + parameters: list[tuple[torch.nn.Module, str, torch.Size, torch.dtype]] + + _CheckpointStartupMetadata = dict[str, str | dict[str, str | bool | None] | None] @@ -1094,6 +1147,7 @@ def __init__(self, self._weight_pool_proxy = None self._gms_backend = None self._checkpoint_loader: Optional[BaseCheckpointLoader] = None + self._deferred: Optional[_DeferredCheckpoint] = None # Mostly weight loading and processing time metrics, updated when load() is called. self._metrics: dict[str, float] = {} self._startup_metadata: _CheckpointStartupMetadata = {} @@ -1263,6 +1317,7 @@ def load( self, checkpoint_dir: str, checkpoint_loader: BaseCheckpointLoader, + defer_checkpoint: bool = False, ) -> tuple[DecoderModelForCausalLM, BaseMoeLoadBalancer | None]: """ Loads the model, its weights, and applies necessary configurations. @@ -1270,11 +1325,15 @@ def load( Args: checkpoint_dir: The directory of the model checkpoint. checkpoint_loader: The loader object for model checkpoints. + defer_checkpoint: For a model a KernelBuilder rehearsal reaches, + return it with dummy weights instead; `has_deferred_checkpoint` + then tells, and `load_deferred_checkpoint` loads the checkpoint. Returns: The loaded and initialized PyTorch model. """ self._metrics = {} + self._deferred = None self._checkpoint_loader = checkpoint_loader config = self._load_and_validate_config(checkpoint_dir, checkpoint_loader) @@ -1366,6 +1425,7 @@ def load( ) memo: dict[torch.Tensor, torch.Tensor] = {} + materialized: list[torch.Tensor] = [] if (self.model_weights_memory_tag is not None and load_format != LoadFormat.GMS): @@ -1426,6 +1486,8 @@ def init_meta_tensor(t: torch.Tensor): return memo[t] model._apply(init_meta_tensor) + # The tensors that hold no value of their own. + materialized = list(memo.values()) # Ensure everything is at least on CUDA # No-op if worked as expected @@ -1446,41 +1508,31 @@ def init_meta_tensor(t: torch.Tensor): # prior to zero-copy mapping, then refreshes derived state after # real GMS tensors are bound. gms_post_load_handled = False + # Deferred only where the checkpoint can load later as it would now: + # the AUTO loader fills tensors materialized from meta, nothing + # else holds them (MoE load balancer, pinned host copies), and the + # module tree is the Modeling V1 one the refit hooks cover. + if (defer_checkpoint and materialized + and load_format == LoadFormat.AUTO and + checkpoint_loader.checkpoint_format != CheckpointFormat.MX + and not requires_standard_hf_loading + and moe_load_balancer is None and getattr( + model.model_config, "modeling_v2_target", None) is None + and KernelBuilder.get().wants_rehearsal(config, model)): + self._deferred = self._load_rehearsal_weights( + model, + config, + materialized, + checkpoint_dir=checkpoint_dir, + checkpoint_loader=checkpoint_loader, + loads_draft_weights=loads_draft_weights, + post_transform_qualification=post_transform_qualification) + return model, moe_load_balancer + del materialized if load_format == LoadFormat.AUTO: - # Pass model= so format-specific loaders (e.g. MX) can - # write weights directly into parameter buffers via P2P. - # Generic loaders ignore model=; loaders that can consume a - # live module reference (MX) use it for direct writes. - load_weights_kwargs: dict = { - "mapping": self.mapping, - "model": model, - # Generic loaders ignore it; MXCheckpointLoader pops it. - "source_identity": self._source_identity, - } - if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: - load_weights_kwargs["model_config"] = config - load_weights_kwargs["load_config"] = self.llm_args - load_weights_kwargs[ - "post_transform_protocol_version"] = self._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION - # If a separate draft model still needs a raw disk load, - # do not accept post-transform bytes for only the target - # model. Enable this only after target and draft subgraphs - # have an explicit mixed-layout policy. - load_weights_kwargs[ - "allow_post_transform_weights"] = post_transform_qualification.qualified - if post_transform_qualification.qualified: - load_weights_kwargs[ - "prepare_post_transform_receiver"] = self._setup_aliases - - weights_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( - model, 'llm_checkpoint_dir') else checkpoint_dir) - weights_preloaded = self._materialize_checkpoint_weights( - checkpoint_loader, weights_checkpoint_dir, model, config, - load_weights_kwargs) - - if loads_draft_weights: - self._materialize_draft_checkpoint_weights( - checkpoint_loader, model) + weights_preloaded = self._load_checkpoint_weights( + checkpoint_loader, checkpoint_dir, model, config, + loads_draft_weights, post_transform_qualification) elif load_format == LoadFormat.GMS: # GPU Memory Service path: weight tensors live in a @@ -1757,60 +1809,218 @@ def init_meta_tensor_in_pool(t: torch.Tensor): raise NotImplementedError( f"No load support for load format: {load_format}") - with timing_metric( - ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, - self._metrics): - if not gms_post_load_handled: - checkpoint_loader.post_load_apply( - model, weights_preloaded=weights_preloaded) - mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( - checkpoint_loader, - model, - weights_preloaded=weights_preloaded, - qualification=post_transform_qualification) - if mx_staged_receiver_path: - self._setup_aliases(model) - self._mark_weights_transformed(model) - self._walk_cache_state(model) - else: - self._walk_full_post_load(model) - self._post_load_publish( - checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=weights_preloaded, - qualification=post_transform_qualification) - - # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, - # `register_weight_slots_after_to_cuda` and `finalize_model` - # must run INSIDE `mem_pool_scope` and BEFORE `finalize_write` - # so MoE allocations become part of the committed layout that - # RO peers receive. Today they run outside the pool and after - # commit, which would silently produce a broken MoE routing - # state on RO peers — that combination is REJECTED at config - # time by `TorchLlmArgs.validate_gms_moe_compat` in - # `tensorrt_llm/llmapi/llm_args.py`. When implementing the - # follow-up, drop the validator gate AFTER moving these calls - # into the pool scope. - if isinstance(moe_load_balancer, BaseMoeLoadBalancer): - moe_load_balancer.register_weight_slots_after_to_cuda() - logger.info("moe_load_balancer finalizing model...") - moe_load_balancer.finalize_model() - logger.info("moe_load_balancer finalize model done") + self._post_load( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + post_transform_qualification=post_transform_qualification, + gms_post_load_handled=gms_post_load_handled, + moe_load_balancer=moe_load_balancer) - if torch.cuda.is_available(): # CPU guard - torch.cuda.synchronize() - # Reclaim segments freed during per-module post_load_weights (e.g. - # MegaMoE _transform_main_weights releases ~5-6 GiB of redundant - # weight Parameters via `.data = empty(0)` that PyTorch's caching - # allocator otherwise holds onto). Returning them to the driver - # gives downstream stages (KV cache estimation, attention workspace - # alloc, autotuner warmup symmetric-fabric setup) full visibility - # of free HBM. Single one-shot call after all modules are - # finalized; per-layer empty_cache here is unsafe because it can - # perturb NVLink barrier synchronization in multi-rank DG init. - torch.cuda.empty_cache() + self._finish_load(checkpoint_loader, model, weights_preloaded) + return model, moe_load_balancer + + @property + def has_deferred_checkpoint(self) -> bool: + """True after `load(defer_checkpoint=True)` returned a model with dummy weights.""" + return self._deferred is not None + def load_deferred_checkpoint(self, model: DecoderModelForCausalLM) -> None: + """Load the checkpoint a `load(defer_checkpoint=True)` left out. + + The model goes back to the state `load` had it in before it put dummy + weights in: the module tree as created, and each parameter in its + creation layout, through `pre_reload_weights` as for a weight update. + The rest of the load then runs as `load` runs it. + """ + deferred, self._deferred = self._deferred, None + assert deferred is not None, "no checkpoint load was deferred" + with timing_metric( + ModelLoaderMetricNames.TOTAL_MODEL_LOADING_SECONDS.value, + self._metrics): + torch.cuda.synchronize() + deferred.tree.restore() + for module in model.modules(): + if hasattr(module, "pre_reload_weights") and not getattr( + module, "_weights_removed", False): + module.pre_reload_weights() + for module, name, shape, dtype in deferred.parameters: + param = module._parameters.get(name) + if param is None or param.shape != shape or param.dtype != dtype: + raise RuntimeError( + f"{type(module).__name__}.{name} is " + f"{None if param is None else (tuple(param.shape), param.dtype)} " + f"after the KernelBuilder rehearsal and its post-load " + f"hooks, not {(tuple(shape), dtype)} as created, so the " + f"checkpoint cannot load into it. Set " + f"TLLM_KERNEL_BUILDER_REHEARSAL=0 to load it first.") + self._reset_weights_transformed(model) + weights_preloaded = self._load_checkpoint_weights( + deferred.checkpoint_loader, deferred.checkpoint_dir, model, + deferred.config, deferred.loads_draft_weights, + deferred.post_transform_qualification) + self._post_load(deferred.checkpoint_loader, + model, + checkpoint_dir=deferred.checkpoint_dir, + weights_preloaded=weights_preloaded, + post_transform_qualification=deferred. + post_transform_qualification, + gms_post_load_handled=False, + moe_load_balancer=None) + self._finish_load(deferred.checkpoint_loader, model, weights_preloaded) + + def _load_rehearsal_weights( + self, model: DecoderModelForCausalLM, config: ModelConfig, + materialized: list[torch.Tensor], *, checkpoint_dir: str, + checkpoint_loader: BaseCheckpointLoader, loads_draft_weights: bool, + post_transform_qualification: PostTransformQualificationDecision + ) -> "_DeferredCheckpoint": + """What `LoadFormat.DUMMY` loads, into the tensors that hold no value. + + Returns what `load_deferred_checkpoint` needs to load the checkpoint: + the module tree and parameters as created, and the load's inputs. + """ + logger.info( + "Loading dummy weights for a KernelBuilder rehearsal; the checkpoint " + "loads after it.") + deferred = _DeferredCheckpoint( + config=config, + checkpoint_dir=checkpoint_dir, + checkpoint_loader=checkpoint_loader, + loads_draft_weights=loads_draft_weights, + post_transform_qualification=post_transform_qualification, + tree=_ModuleTree(model), + parameters=[(module, name, param.shape, param.dtype) + for module in model.modules() + for name, param in module._parameters.items() + if param is not None], + ) + with timing_metric( + ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, + self._metrics): + fill_rehearsal_values(materialized) + if loads_draft_weights: + model.draft_model.load_weights_from_target_model(model) + self._walk_full_post_load(model) + torch.cuda.synchronize() + torch.cuda.empty_cache() + if LOW_M_GEMM_ACTIVE: + prepare_low_m_gemm(model) + return deferred + + def _load_checkpoint_weights( + self, checkpoint_loader: BaseCheckpointLoader, checkpoint_dir: str, + model: DecoderModelForCausalLM, config: ModelConfig, + loads_draft_weights: bool, + post_transform_qualification: PostTransformQualificationDecision + ) -> bool: + """The `LoadFormat.AUTO` load of the target and draft checkpoints. + + Returns whether the checkpoint loader preloaded the weights. + """ + # Pass model= so format-specific loaders (e.g. MX) can + # write weights directly into parameter buffers via P2P. + # Generic loaders ignore model=; loaders that can consume a + # live module reference (MX) use it for direct writes. + load_weights_kwargs: dict = { + "mapping": self.mapping, + "model": model, + # Generic loaders ignore it; MXCheckpointLoader pops it. + "source_identity": self._source_identity, + } + if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: + load_weights_kwargs["model_config"] = config + load_weights_kwargs["load_config"] = self.llm_args + load_weights_kwargs[ + "post_transform_protocol_version"] = self._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION + # If a separate draft model still needs a raw disk load, + # do not accept post-transform bytes for only the target + # model. Enable this only after target and draft subgraphs + # have an explicit mixed-layout policy. + load_weights_kwargs[ + "allow_post_transform_weights"] = post_transform_qualification.qualified + if post_transform_qualification.qualified: + load_weights_kwargs[ + "prepare_post_transform_receiver"] = self._setup_aliases + + weights_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( + model, 'llm_checkpoint_dir') else checkpoint_dir) + weights_preloaded = self._materialize_checkpoint_weights( + checkpoint_loader, weights_checkpoint_dir, model, config, + load_weights_kwargs) + + if loads_draft_weights: + self._materialize_draft_checkpoint_weights(checkpoint_loader, model) + return weights_preloaded + + def _post_load( + self, checkpoint_loader: BaseCheckpointLoader, + model: DecoderModelForCausalLM, *, checkpoint_dir: str, + weights_preloaded: bool, + post_transform_qualification: PostTransformQualificationDecision, + gms_post_load_handled: bool, + moe_load_balancer: BaseMoeLoadBalancer | None) -> None: + """The post-load hooks, the MoE load balancer's finalization, and the + release of the memory they freed.""" + with timing_metric( + ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, + self._metrics): + if not gms_post_load_handled: + checkpoint_loader.post_load_apply( + model, weights_preloaded=weights_preloaded) + mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + qualification=post_transform_qualification) + if mx_staged_receiver_path: + self._setup_aliases(model) + self._mark_weights_transformed(model) + self._walk_cache_state(model) + else: + self._walk_full_post_load(model) + self._post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + qualification=post_transform_qualification) + + # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, + # `register_weight_slots_after_to_cuda` and `finalize_model` + # must run INSIDE `mem_pool_scope` and BEFORE `finalize_write` + # so MoE allocations become part of the committed layout that + # RO peers receive. Today they run outside the pool and after + # commit, which would silently produce a broken MoE routing + # state on RO peers — that combination is REJECTED at config + # time by `TorchLlmArgs.validate_gms_moe_compat` in + # `tensorrt_llm/llmapi/llm_args.py`. When implementing the + # follow-up, drop the validator gate AFTER moving these calls + # into the pool scope. + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): + moe_load_balancer.register_weight_slots_after_to_cuda() + logger.info("moe_load_balancer finalizing model...") + moe_load_balancer.finalize_model() + logger.info("moe_load_balancer finalize model done") + + if torch.cuda.is_available(): # CPU guard + torch.cuda.synchronize() + # Reclaim segments freed during per-module post_load_weights (e.g. + # MegaMoE _transform_main_weights releases ~5-6 GiB of redundant + # weight Parameters via `.data = empty(0)` that PyTorch's caching + # allocator otherwise holds onto). Returning them to the driver + # gives downstream stages (KV cache estimation, attention workspace + # alloc, autotuner warmup symmetric-fabric setup) full visibility + # of free HBM. Single one-shot call after all modules are + # finalized; per-layer empty_cache here is unsafe because it can + # perturb NVLink barrier synchronization in multi-rank DG init. + torch.cuda.empty_cache() + + def _finish_load(self, checkpoint_loader: BaseCheckpointLoader, + model: DecoderModelForCausalLM, + weights_preloaded: bool) -> None: + """The weight manifest, the loading metrics and the low-M GEMM dispatch.""" manifest_result = self._dump_final_weight_manifest( checkpoint_loader, model, weights_preloaded=weights_preloaded) if manifest_result is not None: @@ -1824,7 +2034,6 @@ def init_meta_tensor_in_pool(t: torch.Tensor): f"Model loading metrics for rank {self.mapping.rank}: {metrics}") if LOW_M_GEMM_ACTIVE: prepare_low_m_gemm(model) - return model, moe_load_balancer def _check_gms_source_identity(self, gms_backend) -> None: """Pre-materialize SourceIdentity gate for the GMS read-only path. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index c5dcbbbfc672..27b8ac4e5442 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -67,6 +67,7 @@ from .model_engine import ModelEngine from .model_loader import ModelLoader, construct_checkpoint_loader from .py_executor import PyExecutor +from .rehearsal import join_autotuner_builds, rehearse from .scheduler import gen_preemption _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS = (90, 100, 103, 107, 120, 121) @@ -932,6 +933,11 @@ def allocation_scope(current_stage: ExecutorMemoryType): model_weights_restore_mode = sleep_config.restore_modes[ ExecutorMemoryType.MODEL_WEIGHTS_MAIN] + # A model a KernelBuilder rehearsal reaches is built on dummy weights; its + # checkpoint loads after the rehearsal below. Not with what takes the + # weights over as they load (DWDP) or captures graphs on them (the + # multimodal encoders' CUDA graphs). + mm_config = llm_args.multimodal_config model_engine = ModelEngine( model_path=checkpoint_dir, llm_args=llm_args, @@ -942,6 +948,10 @@ def allocation_scope(current_stage: ExecutorMemoryType): checkpoint_loader=checkpoint_loader, model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, + defer_checkpoint=(dwdp_manager is None + and dense_dwdp_manager is None + and (mm_config is None + or not mm_config.encoder_cuda_graph)), ) is_generation = model_engine.model.model_config.is_generation @@ -1329,6 +1339,14 @@ def allocation_scope(current_stage: ExecutorMemoryType): is_disagg=is_disagg, ) + if model_engine.has_deferred_checkpoint: + with _startup_timer.phase("rehearsal"): + rehearse(model_engine, kv_cache_creator) + with _startup_timer.phase("checkpoint"): + model_engine.load_deferred_checkpoint() + with _startup_timer.phase("autotuner_builds"): + join_autotuner_builds() + if not skip_est: log_memory_usage("after loading weights") diff --git a/tensorrt_llm/_torch/pyexecutor/rehearsal.py b/tensorrt_llm/_torch/pyexecutor/rehearsal.py new file mode 100644 index 000000000000..bfe9f4109372 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/rehearsal.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KernelBuilder rehearsal: warmup forwards on dummy weights, before the checkpoint loads. + +For a model a KernelBuilder op family reaches, the model engine is built on dummy weights +(``ModelEngine(defer_checkpoint=True)``). ``rehearse`` then runs the autotuner warmup's +forwards on throwaway KV cache managers inside a KernelBuilder rehearsal, in the +autotuner's build state: the ops list their kernels, and the kernels build on worker +threads while the checkpoint loads. ``join_autotuner_builds`` then waits for the builds +the autotuned ops queued. +""" + +from __future__ import annotations + +import copy +import gc +from typing import TYPE_CHECKING + +import torch + +from tensorrt_llm.logger import logger + +from ..autotuner import BUILD_OP_PREFIX +from ..kernel_builder import KernelBuilder +from ..speculative import get_spec_resource_manager +from .resource_manager import ResourceManager, ResourceManagerType + +if TYPE_CHECKING: + from ._util import KvCacheCreator + from .model_engine import ModelEngine + + +def rehearse(model_engine: ModelEngine, kv_cache_creator: KvCacheCreator) -> None: + """Run the rehearsal on KV cache managers sized as for the capacity estimation. + + Without the estimation there is no small KV cache to rehearse with, and nothing runs. + The managers are shut down and their memory freed before it returns, so the + checkpoint loads with the memory it would have had. ``kv_cache_creator`` is left + as it was. Every rank calls it. + """ + # The estimation changes the creator's KV cache config and rebinds the other + # fields it changes: a copy with a config of its own leaves the original as it was. + creator = copy.copy(kv_cache_creator) + creator._kv_cache_config = kv_cache_creator._kv_cache_config.model_copy(deep=True) + if not creator.try_prepare_estimation(): + return + resources: dict = {} + try: + creator.build_managers(resources, estimating_kv_cache=True) + spec_resource_manager = get_spec_resource_manager(model_engine) + if spec_resource_manager is not None: + resources[ResourceManagerType.SPEC_RESOURCE_MANAGER] = spec_resource_manager + with KernelBuilder.get().rehearsal(): + model_engine.rehearse(ResourceManager(resources)) + finally: + if ResourceManagerType.KV_CACHE_MANAGER in resources: + creator.teardown_managers(resources) + resources.clear() + gc.collect() + torch.cuda.empty_cache() + + +def join_autotuner_builds() -> None: + """Wait for the builds the rehearsal's autotuned ops queued, then free their memory. + + They run their kernels on inputs of their own, so they end before the KV cache + estimation measures free memory. Builds that only compile keep running. Every rank + calls it. + """ + waited = KernelBuilder.get().join(prefix=BUILD_OP_PREFIX) + torch.cuda.empty_cache() + if waited >= 0.1: + logger.info(f"[KernelBuilder] waited {waited:.1f}s for the autotuner's builds") diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 544c0067d240..70d8d165d38f 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -198,6 +198,7 @@ l0_cpu: - unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py - unittest/_torch/test_custom_config_registration.py - unittest/_torch/test_hf_quant_config.py + - unittest/_torch/test_kernel_builder.py - unittest/_torch/test_mmap_utils.py - unittest/_torch/test_connector.py - unittest/_torch/test_connector_adp.py diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py new file mode 100644 index 000000000000..ad7941237f10 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fmha_sm100 through KernelBuilder: a call signature's variants listed by replaying the +package's plan and run, and built where and under the lock the package's own lookup uses.""" + +import pytest +import torch + +from tensorrt_llm._torch import kernel_builder +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import msa_kernel_builds +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + msa_package_available, + require_msa_module, +) + + +@pytest.fixture +def fmha_sm100(monkeypatch): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 (Blackwell) required") + if not msa_package_available(): + pytest.skip("fmha_sm100 (MSA) not importable") + monkeypatch.setattr(kernel_builder.KernelBuilder, "_instance", kernel_builder.KernelBuilder()) + module = require_msa_module() + if msa_kernel_builds._package is None: + pytest.skip("TLLM_KERNEL_BUILDER=0") + return module + + +@pytest.fixture +def no_compile(monkeypatch): + from fmha_sm100 import jit + + def compile_variant(name, params): + raise AssertionError(f"listing compiled {name}") + + monkeypatch.setattr(jit._variant_manager, "_compile_only", compile_variant) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) +def test_listing_replays_the_query_lengths_up_to_the_call_s(fmha_sm100, no_compile, dtype): + from fmha_sm100 import jit + + page_size = head_dim = 128 + num_heads, num_kv_heads, qo_len = 8, 1, 100 + q = torch.zeros(qo_len, num_heads, head_dim, device="cuda").to(dtype) + k = torch.zeros(4, num_kv_heads, page_size, head_dim, device="cuda").to(dtype) + kv_indices = torch.arange(4, dtype=torch.int32, device="cuda") + options = dict(num_kv_heads=num_kv_heads, page_size=page_size, causal=True) + run_options = {"kv_indices": kv_indices, "output_maxscore": False} + + def variants(longest): + shapes = msa_kernel_builds._CONFIG.shapes({"qo_len": longest}) + builds = msa_kernel_builds._package._list( + q, k, k, (num_heads,), options, 1, run_options, shapes + ) + names = [name for name, _ in builds] + assert len(names) == len(set(names)) + return set(names) + + # A call up to 256 rows reaches the 128- and 256-row tiles; one of 8 rows reaches fewer. + assert len(variants(256)) >= 2 + assert variants(8) < variants(256) + # Workers compile in parallel: the package's lazily loaded templates are loaded already. + assert jit._variant_manager._inst_template is not None + assert jit._variant_manager._run_template is not None + + +def test_the_indexer_maximum_score_pass_also_lists_its_split_variants(fmha_sm100, no_compile): + page_size = head_dim = 128 + num_index_heads, qo_len = 4, 128 + q = torch.zeros(qo_len, num_index_heads, head_dim, device="cuda", dtype=torch.bfloat16) + k = torch.zeros(1, 1, page_size, head_dim, device="cuda", dtype=torch.bfloat16) + kv_indices = torch.arange(1, dtype=torch.int32, device="cuda") + options = dict( + page_size=page_size, + causal=True, + num_qo_heads=num_index_heads, + num_kv_heads=1, + output_maxscore=True, + ) + run_options = {"kv_indices": kv_indices, "output_o": False, "output_maxscore": True} + + def variants(config, observed): + builds = msa_kernel_builds._package._list( + q, k, k, (), options, 1, run_options, config.shapes(observed) + ) + return {name for name, _ in builds} + + unsplit = variants(msa_kernel_builds._CONFIG, {"qo_len": qo_len}) + swept = variants(msa_kernel_builds._PROXY_CONFIG, {"qo_len": qo_len, "num_kv_splits": 1}) + assert unsplit < swept + + +def test_a_rehearsal_run_lists_and_launches_nothing(fmha_sm100, no_compile, monkeypatch): + builder = kernel_builder.KernelBuilder.get() + submitted = [] + monkeypatch.setattr(builder, "submit", lambda op, key, build: submitted.append(key) or True) + page_size = head_dim = 128 + num_heads, num_kv_heads, qo_len = 8, 1, 100 + q = torch.ones(qo_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + k = torch.ones(4, num_kv_heads, page_size, head_dim, device="cuda", dtype=torch.bfloat16) + kv_indices = torch.arange(4, dtype=torch.int32, device="cuda") + lens = torch.tensor([qo_len], dtype=torch.int32) + out = torch.full_like(q, float("nan")) + with builder.rehearsal(): + plan = fmha_sm100.fmha_sm100_plan( + lens, + lens.clone(), + num_heads, + num_kv_heads=num_kv_heads, + page_size=page_size, + num_kv_splits=1, + causal=True, + ) + fmha_sm100.fmha_sm100(q, k, k, plan, kv_indices=kv_indices, out=out, output_maxscore=False) + # The run listed its variants, compiled none, and left zeros where its output goes. + modules = {name for name, _, _ in msa_kernel_builds._MODULES} + assert len([key for key in submitted if key not in modules]) >= 2 + assert torch.equal(out, torch.zeros_like(out)) + + +def test_kernels_are_built_once_where_the_package_loads_them(fmha_sm100, monkeypatch, tmp_path): + from fmha_sm100 import jit + + compiled = [] + + def compile_variant(name, params, cache_dir): + compiled.append(name) + (cache_dir / f"{name}.so").write_bytes(b"") + + def compile_plan(cache_dir): + compiled.append("plan") + cache_dir.mkdir(parents=True, exist_ok=True) + (cache_dir / "fmha_sm100_plan.so").write_bytes(b"") + + monkeypatch.setattr(jit, "_built_dir", lambda name, library: tmp_path / name) + monkeypatch.setattr(jit._variant_manager, "_compile_only", compile_variant) + monkeypatch.setattr(jit, "_do_compile_plan", compile_plan) + package = msa_kernel_builds._package + assert package._build("v", {}) is True + assert package._build("v", {}) is False + assert package.build_module("plan", "fmha_sm100_plan.so", "_do_compile_plan") is True + assert package.build_module("plan", "fmha_sm100_plan.so", "_do_compile_plan") is False + assert compiled == ["v", "plan"] diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index b920f61adb37..33e61d2d3909 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -2463,3 +2463,201 @@ def node_ranks(): else: monkeypatch.setenv("LOCAL_WORLD_SIZE", env_local) assert autotuner._default_parallel_jit_warmup_threads() == expected + + +@pytest.fixture +def kernel_builder(monkeypatch): + """A KernelBuilder of the test's own, for the autotuner's build state.""" + from tensorrt_llm._torch import kernel_builder as kb + builder = kb.KernelBuilder(threads=4, enabled=True) + monkeypatch.setattr(kb.KernelBuilder, "_instance", builder) + monkeypatch.delenv("TLLM_KERNEL_BUILDER_AUTOTUNER", raising=False) + monkeypatch.delenv("TLLM_AUTOTUNER_PARALLEL_JIT_WARMUP", raising=False) + AutoTuner.get().clear_cache() + yield builder + builder.join(timeout=600) + AutoTuner.get().clear_cache() + + +_BUILD_BUCKETS = (4, 8, 16, 32, 64, 128) + + +def _build_config(**kwargs) -> TuningConfig: + return TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, _BUILD_BUCKETS), ), + **kwargs) + + +def _build_inputs() -> List[torch.Tensor]: + return [ + torch.randn(128, 64, device="cuda"), + torch.randn(64, 128, device="cuda") + ] + + +def test_the_build_state_builds_every_profile_without_tuning(kernel_builder): + """In the build state, choose_one builds every profile its tuning would profile + on KernelBuilder's threads and returns the fallback without caching anything. + The tuning that follows still profiles every profile, and builds nothing again.""" + tuner = AutoTuner.get() + ThreadRecordingRunner.calls.clear() + config = _build_config(parallel_jit_warmup=4) + inputs = _build_inputs() + op = "autotuner_test::build_state" + caller = threading.get_ident() + with autotuner.build_kernels(): + assert tuner.is_build_mode + _, tactic = tuner.choose_one(op, [ThreadRecordingRunner()], config, + inputs) + # Another layer's call queues nothing more. + tuner.choose_one(op, [ThreadRecordingRunner()], config, inputs) + assert not tuner.is_build_mode + assert tactic == -1 + kernel_builder.join(timeout=60) + assert sorted( + m for m, _ in ThreadRecordingRunner.calls) == sorted(_BUILD_BUCKETS) + assert all(t != caller for _, t in ThreadRecordingRunner.calls) + assert len(tuner.profiling_cache) == 0 + stats = kernel_builder.stats(autotuner.BUILD_OP_PREFIX + op) + assert (stats.submitted, stats.failed) == (len(_BUILD_BUCKETS), 0) + + ThreadRecordingRunner.calls.clear() + with autotune(): + tuner.choose_one(op, [ThreadRecordingRunner()], config, inputs) + assert sorted(m for m, t in ThreadRecordingRunner.calls + if t == caller) == sorted(_BUILD_BUCKETS) + assert all(t == caller for _, t in ThreadRecordingRunner.calls) + + +@pytest.mark.parametrize("threads, strategy, env", [ + (0, DistributedTuningStrategy.INDEPENDENT, None), + (4, DistributedTuningStrategy.MERGE, None), + (4, DistributedTuningStrategy.INDEPENDENT, "0"), +]) +def test_the_build_state_builds_only_ops_that_opt_in(kernel_builder, + monkeypatch, threads, + strategy, env): + """An op that does not opt in with parallel_jit_warmup, a collective (MERGE) + op, and TLLM_KERNEL_BUILDER_AUTOTUNER=0 build nothing.""" + if env is not None: + monkeypatch.setenv("TLLM_KERNEL_BUILDER_AUTOTUNER", env) + tuner = AutoTuner.get() + ThreadRecordingRunner.calls.clear() + config = _build_config(parallel_jit_warmup=threads, + distributed_tuning_strategy=strategy) + op = f"autotuner_test::build_state_opt_in_{threads}_{strategy.value}_{env}" + with autotuner.build_kernels(): + _, tactic = tuner.choose_one(op, [ThreadRecordingRunner()], config, + _build_inputs()) + kernel_builder.join(timeout=60) + assert tactic == -1 + assert kernel_builder.stats(autotuner.BUILD_OP_PREFIX + op).submitted == 0 + assert ThreadRecordingRunner.calls == [] + + +def test_a_runner_builds_through_its_own_build_once_per_build_key( + kernel_builder): + """A runner that compiles without launching overrides build(); builds with + equal build_keys build once. A nested autotune(tune_mode=False) keeps the build + state, whose sweep ignores skip_dynamic_tuning_buckets.""" + built = [] + + class ModuleRunner(TunableRunner): + + def get_valid_tactics(self, inputs, profile, **kwargs): + return [-1, 0] + + def forward(self, /, inputs, *, tactic=-1, **kwargs): + raise AssertionError("the build state launched the runner") + + def build(self, inputs, *, tactic=-1, **kwargs): + built.append(tactic) + + def build_key(self, profile, tactic): + return str(tactic) + + tuner = AutoTuner.get() + op = "autotuner_test::build_state_build_key" + with autotuner.build_kernels(), autotune(tune_mode=False, + skip_dynamic_tuning_buckets=True): + assert tuner.is_build_mode + tuner.choose_one(op, [ModuleRunner()], + _build_config(parallel_jit_warmup=4), _build_inputs()) + kernel_builder.join(timeout=60) + assert sorted(built) == [-1, 0] + assert kernel_builder.stats(autotuner.BUILD_OP_PREFIX + + op).submitted == len(_BUILD_BUCKETS) + + +def test_a_failed_build_is_left_to_tuning(kernel_builder): + """A build that fails leaves its kernels to the tuning, which still profiles + every profile on the calling thread.""" + + class FailingBuildRunner(ThreadRecordingRunner): + calls: list[tuple[int, int]] = [] + + def build(self, inputs, *, tactic=-1, **kwargs): + raise RuntimeError("nvcc failed") + + tuner = AutoTuner.get() + config = _build_config(parallel_jit_warmup=4) + op = "autotuner_test::build_state_failure" + with autotuner.build_kernels(): + tuner.choose_one(op, [FailingBuildRunner()], config, _build_inputs()) + kernel_builder.join(timeout=60) + assert kernel_builder.stats(autotuner.BUILD_OP_PREFIX + + op).failed == len(_BUILD_BUCKETS) + with autotune(): + tuner.choose_one(op, [FailingBuildRunner()], config, _build_inputs()) + assert {m for m, _ in FailingBuildRunner.calls} == set(_BUILD_BUCKETS) + assert not tuner.stats.failed_profiling_count.get(op) + + +@pytest.mark.skipif(not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] != 10, + reason="DeepGEMM's FP8 block-scale GEMMs run on SM100f") +def test_fp8_swap_ab_gemm_returns_zeros_and_builds_in_the_build_state( + kernel_builder): + """In the build state the DeepGEMM GEMM launches nothing: it returns zeros of + its output shape, and its sweep builds on KernelBuilder's threads.""" + from _torch.helpers import per_block_cast_to_fp8_e8m0 + + import tensorrt_llm._torch.custom_ops.torch_custom_ops # noqa: F401 + + m, n, k = 64, 1024, 1024 + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + weight, weight_scale = per_block_cast_to_fp8_e8m0( + torch.randn((n, k), device="cuda", dtype=torch.bfloat16)) + with autotuner.build_kernels(): + output = torch.ops.trtllm.fp8_swap_ab_gemm(a, weight, weight_scale) + assert output.shape == (m, n) and output.dtype == torch.bfloat16 + assert torch.count_nonzero(output) == 0 + kernel_builder.join(timeout=600) + stats = kernel_builder.stats(autotuner.BUILD_OP_PREFIX + + "trtllm::fp8_swap_ab_gemm") + assert stats.submitted == len(set(deep_gemm_jit_warmup_buckets(m)) | {m}) + assert stats.failed == 0 + + +def test_flashinfer_mxfp8_gemm_builds_its_module_without_a_launch(monkeypatch): + """FlashInfer's MXFP8 GEMM runner builds its CUTLASS tactic by loading the + module every shape uses, without running it; its CuTe DSL tactic builds + nothing.""" + gemm_base = pytest.importorskip("flashinfer.gemm.gemm_base") + from tensorrt_llm._torch.custom_ops.torch_custom_ops import \ + FlashInferMXFP8GemmRunner + + loaded = [] + monkeypatch.setattr(gemm_base, "get_cutlass_mxfp8_gemm_module", + lambda sm_major: loaded.append(sm_major)) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + monkeypatch.setattr( + runner, "forward", + lambda *args, **kwargs: pytest.fail("the build launched the GEMM")) + for tactic in runner.get_valid_tactics([], OptimizationProfile()): + runner.build([], tactic=tactic) + assert loaded == [torch.cuda.get_device_capability()[0]] + small = OptimizationProfile([[DynamicDim(8, 8, 16)]]) + large = OptimizationProfile([[DynamicDim(4096, 4096, 8192)]]) + assert runner.build_key(small, runner.CUTLASS) == runner.build_key( + large, runner.CUTLASS) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 3ac45b512476..d6be4eb8dcd3 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -166,8 +166,10 @@ def test_nvfp4_cache_fused_layout_validation(layout: str, expected: bool | None) ), ], ) +# A warmup forward creates the scale buffer inside inference mode. +@pytest.mark.parametrize("inference_mode", [False, True], ids=["eager", "inference_mode"]) def test_nvfp4_reload_preserves_attention_scale_storage( - monkeypatch: pytest.MonkeyPatch, device: str + monkeypatch: pytest.MonkeyPatch, device: str, inference_mode: bool ) -> None: """The model reload hook updates eager and captured readers at stable addresses.""" target = MiniMaxM3ForCausalLM.__new__(MiniMaxM3ForCausalLM) @@ -190,7 +192,10 @@ def load_scales(self, *, weights, **kwargs): monkeypatch.setattr(SpecDecOneEngineForCausalLM, "load_weights", load_scales) monkeypatch.setenv("TRTLLM_LOAD_KV_SCALES", "1") - k_scale, v_scale = _aligned_nvfp4_dequant_scales(decoder.self_attn.attn, projection.kv_scales) + with torch.inference_mode(inference_mode): + k_scale, v_scale = _aligned_nvfp4_dequant_scales( + decoder.self_attn.attn, projection.kv_scales + ) pointers = [ x.data_ptr() for x in (projection.kv_scales, projection.inv_kv_scales, k_scale, v_scale) ] diff --git a/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py b/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py new file mode 100644 index 000000000000..998d1f844222 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A checkpoint load deferred past a KernelBuilder rehearsal ends where a direct load does.""" + +import pytest +import torch + +from tensorrt_llm._torch.kernel_builder import KernelBuilder +from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader +from tensorrt_llm._torch.pyexecutor.model_loader import ( + ModelLoader, + _ModuleTree, + fill_rehearsal_values, +) +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.mapping import Mapping + +cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory): + transformers = pytest.importorskip("transformers") + config = transformers.LlamaConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=128, + torch_dtype="bfloat16", + ) + torch.manual_seed(0) + path = tmp_path_factory.mktemp("llama") + transformers.LlamaForCausalLM(config).to(torch.bfloat16).save_pretrained(path) + return str(path) + + +def _load(checkpoint, defer_checkpoint): + loader = ModelLoader( + llm_args=TorchLlmArgs(model=checkpoint), + mapping=Mapping(), + spec_config=None, + sparse_attention_config=None, + max_num_tokens=128, + max_seq_len=128, + ) + model, _ = loader.load(checkpoint, HfCheckpointLoader(), defer_checkpoint=defer_checkpoint) + return loader, model + + +def _state(model): + return {name: tensor.detach().clone() for name, tensor in model.state_dict().items()} + + +@cuda +def test_a_deferred_load_matches_a_direct_one(checkpoint, monkeypatch): + _, direct = _load(checkpoint, defer_checkpoint=False) + expected = _state(direct) + del direct + + monkeypatch.setattr(KernelBuilder, "wants_rehearsal", lambda self, model_config, model: True) + loader, model = _load(checkpoint, defer_checkpoint=True) + assert loader.has_deferred_checkpoint + rehearsal = _state(model) + # The model holds dummy weights until the checkpoint loads. + assert any(not torch.equal(rehearsal[name], expected[name]) for name in expected) + loader.load_deferred_checkpoint(model) + assert not loader.has_deferred_checkpoint + loaded = _state(model) + assert loaded.keys() == expected.keys() + for name, tensor in expected.items(): + assert torch.equal(loaded[name], tensor), name + + +@cuda +def test_no_deferral_for_a_model_no_op_family_reaches(checkpoint): + loader, _ = _load(checkpoint, defer_checkpoint=True) + assert not loader.has_deferred_checkpoint + + +@pytest.mark.cpu_only +def test_the_module_tree_goes_back_as_created(): + draft = torch.nn.Module() + draft.embed_tokens = None # an attribute, not a submodule + draft.lm_head = torch.nn.Linear(2, 4) + root = torch.nn.Module() + root.draft = draft + root.embed_tokens = torch.nn.Embedding(4, 2) + tree = _ModuleTree(root) + # Tying the draft to the target moves the name to the submodules. + draft.embed_tokens = root.embed_tokens + draft.lm_head = None + tree.restore() + assert draft.embed_tokens is None and "embed_tokens" not in draft._modules + assert isinstance(draft.lm_head, torch.nn.Linear) + + +@cuda +def test_rehearsal_values_are_bounded_and_indices_zero(): + floats = torch.full((64,), float("nan"), device="cuda", dtype=torch.bfloat16) + fp8 = torch.empty(64, device="cuda", dtype=torch.float8_e4m3fn) + index = torch.full((64,), -7, device="cuda", dtype=torch.int32) + fill_rehearsal_values([floats, fp8, index]) + assert floats.abs().max().item() <= 1e-3 + assert fp8.float().abs().max().item() <= 3.0 + assert torch.equal(index, torch.zeros_like(index)) diff --git a/tests/unittest/_torch/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py new file mode 100644 index 000000000000..50fb8634af7e --- /dev/null +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KernelBuilder: builds an op lists for the shapes it sweeps, run on worker threads.""" + +import sys +import threading +import types + +import pytest + +from tensorrt_llm._torch import kernel_builder +from tensorrt_llm._torch.kernel_builder import BuildConfig, DynamicSpec, KernelBuilder + +pytestmark = pytest.mark.cpu_only + + +def _prebuild(builder, signature, list_builds): + """``prebuild`` for an op without build dimensions.""" + builder.prebuild("op", signature, BuildConfig(), {}, lambda shapes: list_builds()) + + +def _recorder(built, key, compiled=True): + def build(): + built.append(key) + return compiled + + return build + + +def test_kernels_are_listed_once_per_signature_inside_a_pass(): + builder = KernelBuilder(threads=2, enabled=True) + listed, built = [], [] + + def list_builds(): + listed.append(True) + return [(name, _recorder(built, name)) for name in ("a", "b", "c")] + + _prebuild(builder, "sig", list_builds) + assert listed == [] # no pass, no listing + + with builder.build_pass("test"): + _prebuild(builder, "sig", list_builds) + _prebuild(builder, "sig", list_builds) + # A kernel another signature also reaches is built once. + _prebuild(builder, "sig2", lambda: [("a", _recorder(built, "a"))]) + # The pass waited for the builds. + assert listed == [True] + assert sorted(built) == ["a", "b", "c"] + assert builder.state("op", "b") == "done" + assert not builder.active() + stats = builder.stats("op") + assert (stats.signatures, stats.submitted, stats.built) == (2, 3, 3) + + +def test_a_failed_build_is_left_to_the_op(): + builder = KernelBuilder(threads=2, enabled=True) + built = [] + + def fail(): + raise RuntimeError("nvcc failed") + + def exit_(): + raise SystemExit(1) + + builds = [("bad", fail), ("exit", exit_), ("good", _recorder(built, "good"))] + with builder.build_pass("test"): + _prebuild(builder, "sig", lambda: builds) + # The pass ended: even an item raising outside Exception was settled. + assert built == ["good"] + assert builder.state("op", "bad") == builder.state("op", "exit") == "failed" + # The op compiling the failed kernel itself counts as a late build, not a wait. + builder.op_compiled("op", "bad", 80.0) + stats = builder.stats("op") + assert (stats.submitted, stats.built, stats.failed, stats.late, stats.waited) == (3, 1, 2, 1, 0) + + +def test_a_listing_failure_leaves_the_kernels_to_the_op(): + builder = KernelBuilder(threads=1, enabled=True) + + def list_builds(): + raise ValueError("unsupported shape") + + with builder.build_pass("test"): + _prebuild(builder, "sig", list_builds) + assert builder.state("op", "anything") is None + + +def test_the_op_waits_for_a_build_in_flight(): + builder = KernelBuilder(threads=1, enabled=True) + started, release = threading.Event(), threading.Event() + + def slow(): + started.set() + release.wait(timeout=30) + return True + + builder.submit("op", "k", slow) + assert started.wait(timeout=30) + assert builder.state("op", "k") == "queued" + builder.op_compiled("op", "k", 5.0) + release.set() + builder.join(timeout=30) + assert builder.state("op", "k") == "done" + stats = builder.stats("op") + assert (stats.built, stats.waited, stats.waited_s, stats.late) == (1, 1, 5.0, 0) + + +def test_an_op_builds_the_shapes_its_dimensions_sweep_from_the_call(): + config = BuildConfig( + ( + DynamicSpec("m", lambda m: [x for x in (1, 2, 4, 8) if x <= m]), + DynamicSpec("splits", (1, 2)), + ) + ) + # A function's buckets come from the call's value, a tuple's are fixed, and the call's + # own value is always one of them. + assert config.shapes({"m": 6, "splits": 1}) == [ + {"m": m, "splits": splits} for m in (1, 2, 4, 6) for splits in (1, 2) + ] + builder = KernelBuilder(threads=2, enabled=True) + swept, built = [], [] + + def list_builds(shapes): + swept.extend(shapes) + # Shapes that reach the same kernel list it once. + return [(shape["m"], _recorder(built, shape["m"])) for shape in shapes] + + with builder.build_pass("test"): + builder.prebuild("op", "sig", config, {"m": 6, "splits": 1}, list_builds) + # A smaller call of the signature sweeps nothing new; a larger one only its new shapes. + builder.prebuild("op", "sig", config, {"m": 4, "splits": 1}, list_builds) + assert len(swept) == 8 + builder.prebuild("op", "sig", config, {"m": 9, "splits": 1}, list_builds) + assert swept[8:] == [ + {"m": 8, "splits": 1}, + {"m": 8, "splits": 2}, + {"m": 9, "splits": 1}, + {"m": 9, "splits": 2}, + ] + assert sorted(built) == [1, 2, 4, 6, 8, 9] + assert builder.stats("op").signatures == 1 + + +def test_a_rehearsal_leaves_its_builds_running_for_the_next_pass(): + builder = KernelBuilder(threads=1, enabled=True) + release = threading.Event() + + def slow(): + release.wait(timeout=30) + return True + + assert not builder.rehearsing() + with builder.rehearsal(): + assert builder.rehearsing() and builder.active() + _prebuild(builder, "sig", lambda: [("k", slow)]) + # The rehearsal did not wait: the build is still running. + assert not builder.rehearsing() and not builder.active() + assert builder.state("op", "k") == "queued" + release.set() + with builder.build_pass("warmup"): + # The signature was listed in the rehearsal. + _prebuild(builder, "sig", lambda: pytest.fail("listed twice")) + # The warmup pass waited for it. + assert builder.state("op", "k") == "done" + + +def test_the_rehearsal_runs_for_a_model_an_op_family_reaches(monkeypatch): + reached = [] + family = types.ModuleType("family") + family.reaches = ( + lambda model_config, model: reached.append((model_config, model)) or model_config == "msa" + ) + monkeypatch.setitem(sys.modules, "family", family) + monkeypatch.setattr(kernel_builder, "_FAMILIES", {"op": "family:reaches"}) + builder = KernelBuilder(threads=1, enabled=True) + assert builder.wants_rehearsal("msa", "model") + assert not builder.wants_rehearsal("dense", "model") + assert reached == [("msa", "model"), ("dense", "model")] + assert not KernelBuilder(threads=1, enabled=True, rehearse=False).wants_rehearsal("msa") + assert not KernelBuilder(threads=1, enabled=False).wants_rehearsal("msa") + + +def test_join_waits_for_the_builds_of_the_ops_a_prefix_names(): + builder = KernelBuilder(threads=2, enabled=True) + release = threading.Event() + builder.submit("other", "k", lambda: release.wait(timeout=30)) + builder.submit("autotuner:op", "k", lambda: True) + assert builder.join(timeout=30, prefix="autotuner:") < 30 + assert builder.state("autotuner:op", "k") == "done" + assert builder.state("other", "k") == "queued" + release.set() + builder.join(timeout=30) + assert builder.state("other", "k") == "done" + + +def test_disabled(): + builder = KernelBuilder(threads=1, enabled=False) + listed = [] + with builder.build_pass("test"): + assert not builder.active() + _prebuild(builder, "sig", lambda: listed.append(True) or []) + with builder.rehearsal(): + assert not builder.rehearsing() + assert listed == [] + assert not builder.submit("op", "k", lambda: True)