From 919d773ff68480ef9aa4fc22e466dbccaa1f4831 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Mon, 9 Nov 2026 06:17:32 -0800 Subject: [PATCH 01/11] [TRTLLM-17679][feat] KernelBuilder: build fmha_sm100 kernel variants on worker threads during warmup KernelBuilder (tensorrt_llm/_torch/kernel_builder.py) queues the kernel builds of ops that compile on first use without going through the autotuner, and runs them on worker threads. The decoder warmup is a build pass: an op lists, on the first sight of each call signature, the kernels the signature reaches over its runtime shapes, and the pass waits for the builds at its end. MiniMax-M3's MSA attention (fmha_sm100) is the first op. Its plan and run are wrapped once, at import. The first run of each call signature in the warmup replays the package's plan and run for one query segment of every length up to 129 and of 256, with the variant lookup recording the variant instead of compiling it. Workers compile each variant in the directory and under the file lock the package's own lookup uses; the first run also queues the plan, split-KV reduction and sparse top-k modules. TLLM_KERNEL_BUILDER=0 turns it off. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 11 + .../minimax_m3/kernels/msa_kernel_builds.py | 256 +++++++++++++++ .../sparse/minimax_m3/kernels/msa_utils.py | 3 + tensorrt_llm/_torch/kernel_builder.py | 297 ++++++++++++++++++ .../engine/runners/decoder/warmup.py | 6 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../sparse/msa/test_msa_kernel_builds.py | 78 +++++ tests/unittest/_torch/test_kernel_builder.py | 107 +++++++ 8 files changed, 757 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_kernel_builds.py create mode 100644 tensorrt_llm/_torch/kernel_builder.py create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py create mode 100644 tests/unittest/_torch/test_kernel_builder.py diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index b613d351db6a..54aa136176bb 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -203,6 +203,17 @@ curl http://localhost:8000/server_info } ``` +### Kernel builds ahead of first use + +During the warmup, an op that compiles its kernels on first use lists them the first time +it runs, and worker threads build them in parallel. Supported ops: MiniMax-M3's MSA +attention (`fmha_sm100`). + +| Environment variable | Default | Effect | +|----------------------|---------|--------| +| `TLLM_KERNEL_BUILDER` | `1` | `0` compiles each kernel on its first use instead. | +| `TLLM_KERNEL_BUILDER_THREADS` | `8` | 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/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..a3ad9c646ce7 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_kernel_builds.py @@ -0,0 +1,256 @@ +# 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. + +``install`` wraps the package's ``fmha_sm100_plan`` and ``fmha_sm100``. On the first run +of each call signature in a build pass, the plan and the run are replayed for one query +segment of every length in ``_QUERY_LENS``, 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 first run in a build pass also submits the plan, split-KV +reduction and sparse top-k modules, which need nothing from a call. +""" + +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 KernelBuilder +from tensorrt_llm.logger import logger + +_OP = "fmha_sm100" +# One query segment of each length: up to 129 crosses every boundary of the pack factor, +# the single warp group and the 128-row query tile; 256 takes the 256-row tile. +_QUERY_LENS = tuple(range(1, 130)) + (256,) +# Plan options a replay does not pass through: it plans its own lengths and offsets. +_PLAN_SKIP = ("qo_offset", "device", "stream") +# 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} + for sub_plan in plan_info[3:]: + if isinstance(sub_plan, dict): + sub_plan[_PLAN_ARGS] = (args, options) + return plan_info + + 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() + ): + if not self.modules_submitted: + self.modules_submitted = True + for name, library, compile_fn in _MODULES: + builder.submit( + _OP, name, functools.partial(self.build_module, name, library, compile_fn) + ) + 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 = planned + run_options = {key: v_ for key, v_ in kwargs.items() if key not in _RUN_SKIP} + 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())), + ) + builder.prebuild( + _OP, + signature, + functools.partial(self._list, q, k, v, plan_args, plan_options, run_options), + ) + return self.run_fn(q, k, v, plan_info, *args, **kwargs) + + def _list(self, q, k, v, plan_args, plan_options, run_options) -> List[Tuple[str, Callable]]: + """Replay the plan and the run for each query length; return the variants asked for. + + A length 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 qo_len in _QUERY_LENS: + 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) + try: + plan_info = self.plan_fn(lens, lens.clone(), *plan_args, **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 {qo_len} query rows: {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 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/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py new file mode 100644 index 000000000000..05720fa83cf6 --- /dev/null +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -0,0 +1,297 @@ +# 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. On the first call of each call signature inside a build pass, the +op lists every kernel that signature reaches over its runtime shapes; 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. + +``TLLM_KERNEL_BUILDER=0`` turns it off: no pass is active and nothing is submitted. +``TLLM_KERNEL_BUILDER_THREADS`` sets the number of worker threads (default 8). +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import os +import queue +import threading +import time +from typing import Callable, Dict, Hashable, Iterable, Iterator, List, Optional, Tuple + +import torch + +from tensorrt_llm.logger import logger + +__all__ = ["KernelBuilder"] + +_ENV = "TLLM_KERNEL_BUILDER" +_THREADS_ENV = "TLLM_KERNEL_BUILDER_THREADS" +_DEFAULT_THREADS = 8 +_PROGRESS_EVERY_S = 60.0 + +# 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 +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): + self.enabled = os.environ.get(_ENV, "1") != "0" if enabled is None else enabled + if threads is None: + threads = int(os.environ.get(_THREADS_ENV, str(_DEFAULT_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._pending = 0 + self._signatures: set = 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) -> Iterator[None]: + """Ops list their kernels on first sight inside; the outermost pass waits for the + builds at its end and logs what was built.""" + 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() + for line in self.report(name, waited): + logger.info(line) + + def active(self) -> bool: + return self.enabled and self._passes > 0 + + # Registration. + + def prebuild( + self, + op: str, + signature: Hashable, + list_builds: Callable[[], Iterable[Tuple[Hashable, Build]]], + ) -> None: + """On the first sight of ``(op, signature)`` in a build pass, submit every build + ``list_builds()`` returns as ``(key, build)``. + + A failure to list leaves the kernels to the op. + """ + if not self.active(): + return + with self._lock: + if (op, signature) in self._signatures: + return + self._signatures.add((op, signature)) + t0 = time.perf_counter() + try: + builds = list(list_builds()) + 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 += 1 + stats.list_s += list_s + logger.info( + f"[KernelBuilder] {op}: {len(builds)} kernels for a new call signature, " + f"{queued} queued ({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._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) -> float: + """Wait for every queued build; returns the seconds waited.""" + t0 = last = time.perf_counter() + with self._idle: + while self._pending > 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 {self._pending} builds" + ) + last = now + self._idle.wait(timeout=1.0) + return time.perf_counter() - t0 + + 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()) + ] + if lines and waited >= 0.1: + lines.append(f"[KernelBuilder] {name}: waited {waited:.1f}s for the builds at its end") + 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}" + ) + 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 + if self._pending == 0: + self._idle.notify_all() diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index e4763d201dd1..21daf6504308 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -18,6 +18,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner, autotune 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 @@ -298,8 +299,9 @@ 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 on first sight; 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 _get_max_shape_warmup_requests( diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ff302874d5d6..39e7899b2aa6 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -212,6 +212,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..15f88ee201b1 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -0,0 +1,78 @@ +# 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.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) +def test_listing_replays_every_query_length_without_compiling(fmha_sm100, monkeypatch, dtype): + from fmha_sm100 import jit + + def no_compile(name, params): + raise AssertionError(f"listing compiled {name}") + + monkeypatch.setattr(jit._variant_manager, "_compile_only", no_compile) + 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, num_kv_splits=1, causal=True) + builds = msa_kernel_builds._package._list( + q, k, k, (num_heads,), options, {"kv_indices": kv_indices, "output_maxscore": False} + ) + names = [name for name, _ in builds] + assert len(names) == len(set(names)) + # A query tile of 128 rows and one of 256 cover at least two variants. + assert len(names) >= 2 + # 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_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/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py new file mode 100644 index 000000000000..9b014a84a049 --- /dev/null +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KernelBuilder: builds listed by an op on first sight, run on worker threads.""" + +import threading + +import pytest + +from tensorrt_llm._torch.kernel_builder import KernelBuilder + +pytestmark = pytest.mark.cpu_only + + +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")] + + builder.prebuild("op", "sig", list_builds) + assert listed == [] # no pass, no listing + + with builder.build_pass("test"): + builder.prebuild("op", "sig", list_builds) + builder.prebuild("op", "sig", list_builds) + # A kernel another signature also reaches is built once. + builder.prebuild("op", "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"): + builder.prebuild("op", "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"): + builder.prebuild("op", "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_disabled(): + builder = KernelBuilder(threads=1, enabled=False) + listed = [] + with builder.build_pass("test"): + assert not builder.active() + builder.prebuild("op", "sig", lambda: listed.append(True) or []) + assert listed == [] + assert not builder.submit("op", "k", lambda: True) From 7461616d13f84835c8edc9faf8dffcd8ea58262e Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Mon, 9 Nov 2026 22:52:10 -0800 Subject: [PATCH 02/11] [TRTLLM-17679][feat] KernelBuilder rehearsal: build kernels while the checkpoint loads For a model a KernelBuilder op family reaches (today fmha_sm100, for MiniMax-M3's MSA attention), the model is first built on dummy weights, a rehearsal runs the warmup's MSA indexer prewarm and attention forwards on them so the ops list and submit their kernels, and the checkpoint then loads while the kernels build. - ModelLoader.load(defer_checkpoint=True) loads what LoadFormat.DUMMY loads into the tensors materialized from meta; load_deferred_checkpoint restores the module tree as created, restores the creation layout with pre_reload_weights, checks every parameter against its creation shape and dtype, and runs load's AUTO fill, post-load hooks and finalization. - ModelEngine.rehearse / load_deferred_checkpoint (the torch.compile wrapper is removed for the load and put back after, as for a weight update); DecoderWarmup.rehearse; pyexecutor/rehearsal.py runs it on throwaway estimation-sized KV cache managers. - KernelBuilder.rehearsal(): a build pass that does not wait; wants_rehearsal() over the _FAMILIES table. fmha_sm100 runs in a rehearsal are replays that compile and launch nothing. - msa_prefill: refresh the NVFP4 K/V scale buffer in inference mode when a forward created it there. - TLLM_KERNEL_BUILDER_REHEARSAL=0 loads the checkpoint first. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 8 +- .../attention/backends/fmha/msa_prefill.py | 6 +- .../minimax_m3/kernels/msa_kernel_builds.py | 65 ++- tensorrt_llm/_torch/kernel_builder.py | 82 +++- .../engine/runners/decoder/runner.py | 11 + .../engine/runners/decoder/warmup.py | 40 ++ .../_torch/pyexecutor/model_engine.py | 34 +- .../_torch/pyexecutor/model_loader.py | 445 +++++++++++++----- .../_torch/pyexecutor/py_executor_creator.py | 16 + tensorrt_llm/_torch/pyexecutor/rehearsal.py | 55 +++ .../sparse/msa/test_msa_kernel_builds.py | 34 ++ .../unittest/_torch/models/test_minimax_m3.py | 9 +- .../pyexecutor/test_deferred_checkpoint.py | 108 +++++ tests/unittest/_torch/test_kernel_builder.py | 42 ++ 14 files changed, 816 insertions(+), 139 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/rehearsal.py create mode 100644 tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 54aa136176bb..e21bd47bc72e 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -205,13 +205,17 @@ curl http://localhost:8000/server_info ### Kernel builds ahead of first use -During the warmup, an op that compiles its kernels on first use lists them the first time -it runs, and worker threads build them in parallel. Supported ops: MiniMax-M3's MSA +An op that compiles its kernels on first use lists them the first time it runs, and worker +threads build them in parallel. For a model with such an op, the engine first loads dummy +weights and runs the attention warmup forwards on them, then loads the checkpoint while the +kernels build. The startup log shows these as the `rehearsal` and `checkpoint` phases, which +`model_engine_creation_seconds` then does not include. Supported ops: MiniMax-M3's MSA attention (`fmha_sm100`). | 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_THREADS` | `8` | Worker threads per process. | ## Unified Cache for Compiled Artifacts diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py index b31151545d19..2c8d275e931e 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 import functools from types import ModuleType from typing import TYPE_CHECKING, Optional @@ -147,7 +148,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 index a3ad9c646ce7..eabab34fda13 100644 --- 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 @@ -12,8 +12,11 @@ segment of every length in ``_QUERY_LENS``, 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 first run in a build pass also submits the plan, split-KV -reduction and sparse top-k modules, which need nothing from a call. +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 @@ -101,6 +104,16 @@ def plan(self, qo_segment_lens, kv_segment_lens, *args, **kwargs): sub_plan[_PLAN_ARGS] = (args, options) 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 ( @@ -109,12 +122,7 @@ def run(self, q, k, v, plan_info, *args, **kwargs): and isinstance(plan_info, tuple) and not torch.cuda.is_current_stream_capturing() ): - if not self.modules_submitted: - self.modules_submitted = True - for name, library, compile_fn in _MODULES: - builder.submit( - _OP, name, functools.partial(self.build_module, name, library, compile_fn) - ) + self.submit_modules() planned = next( (p[_PLAN_ARGS] for p in plan_info[3:] if isinstance(p, dict) and _PLAN_ARGS in p), None, @@ -136,8 +144,29 @@ def run(self, q, k, v, plan_info, *args, **kwargs): signature, functools.partial(self._list, q, k, v, plan_args, plan_options, 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, run_options) -> List[Tuple[str, Callable]]: """Replay the plan and the run for each query length; return the variants asked for. @@ -229,6 +258,26 @@ def get_reduction_module(self): return self.reduction_fn() +def reaches(model_config: Any) -> 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 diff --git a/tensorrt_llm/_torch/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py index 05720fa83cf6..0c034058ad9f 100644 --- a/tensorrt_llm/_torch/kernel_builder.py +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -8,19 +8,27 @@ 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. +A rehearsal is a build pass over forwards on dummy weights, run before the checkpoint +loads: 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_THREADS`` sets the number of worker threads (default 8). +``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 8). """ from __future__ import annotations import contextlib import dataclasses +import importlib import os import queue import threading import time -from typing import Callable, Dict, Hashable, Iterable, Iterator, List, Optional, Tuple +from typing import Any, Callable, Dict, Hashable, Iterable, Iterator, List, Optional, Tuple import torch @@ -29,10 +37,20 @@ __all__ = ["KernelBuilder"] _ENV = "TLLM_KERNEL_BUILDER" +_REHEARSAL_ENV = "TLLM_KERNEL_BUILDER_REHEARSAL" _THREADS_ENV = "TLLM_KERNEL_BUILDER_THREADS" _DEFAULT_THREADS = 8 _PROGRESS_EVERY_S = 60.0 +# The op families a rehearsal can reach: name -> "module:function". The function takes the +# model config, 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" + ), +} + # A build returns True when it compiled the kernel, False when the kernel was already there. Build = Callable[[], bool] @@ -77,8 +95,16 @@ class KernelBuilder: _instance: Optional["KernelBuilder"] = None _instance_lock = threading.Lock() - def __init__(self, threads: Optional[int] = None, enabled: Optional[bool] = None): + 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: threads = int(os.environ.get(_THREADS_ENV, str(_DEFAULT_THREADS))) self.threads = max(1, threads) @@ -88,6 +114,7 @@ def __init__(self, threads: Optional[int] = None, enabled: Optional[bool] = None self._workers: List[threading.Thread] = [] self._device: Optional[int] = None self._passes = 0 + self._rehearsals = 0 self._pending = 0 self._signatures: set = set() self._states: Dict[Tuple[str, Hashable], str] = {} @@ -104,9 +131,9 @@ def get(cls) -> "KernelBuilder": # Build passes. @contextlib.contextmanager - def build_pass(self, name: str) -> Iterator[None]: - """Ops list their kernels on first sight inside; the outermost pass waits for the - builds at its end and logs what was built.""" + def build_pass(self, name: str, wait: bool = True) -> Iterator[None]: + """Ops list their kernels on first sight inside. 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 @@ -122,13 +149,51 @@ def build_pass(self, name: str) -> Iterator[None]: self._passes -= 1 last = self._passes == 0 if last: - waited = self.join() + 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) -> bool: + """True when the rehearsal is on and an op family in ``_FAMILIES`` reaches the + model. 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)) 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( @@ -236,8 +301,11 @@ def report(self, name: str, waited: float = 0.0) -> List[str]: 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. 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 21daf6504308..418fa154904b 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -304,6 +304,46 @@ def run(self, resource_manager: ResourceManager) -> EagerWorkspaceReclaimer | No 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 MSA indexer prewarm and the attention warmup forwards, run eagerly on + whatever weights the model holds, inside the caller's KernelBuilder rehearsal. + + 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() + can_run_general_warmup = self._agreement.agree_flag( + not self.mapping.has_cp_helix() + and self.guided_decoder is None + and not isinstance(kv_cache_manager, MambaHybridCacheManager) + ) + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner + + try: + # force_eager: no torch.compile traces, whole-model or op-level. + with ( + torch.inference_mode(), + torch.compiler.set_stance("force_eager"), + self._torch_compile_bypass(), + ): + self._prewarm_minimax_m3_msa_indexer(kv_cache_manager) + self._run_attention_warmup(resource_manager, can_run_general_warmup) + 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 _get_max_shape_warmup_requests( self, resource_manager: ResourceManager ) -> list[tuple[int, int]]: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 389f3c8f0242..44d3dbffca8a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -124,7 +124,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 @@ -144,7 +148,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() @@ -258,6 +266,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 @@ -274,7 +283,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 @@ -1078,6 +1089,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 f7c054210ab6..17459044a7ad 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 @@ -50,6 +50,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 @@ -461,6 +462,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, @@ -478,23 +520,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 @@ -518,18 +543,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. @@ -879,6 +895,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] @@ -1115,6 +1168,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 = {} @@ -1286,6 +1340,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. @@ -1293,11 +1348,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) @@ -1389,6 +1448,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): @@ -1449,6 +1509,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 @@ -1469,41 +1531,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)): + 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 @@ -1780,60 +1832,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: @@ -1847,7 +2057,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 cfa7fee3617b..734bbbd7eaf5 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 rehearse from .scheduler import gen_preemption _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS = (90, 100, 103, 107, 120, 121) @@ -939,6 +940,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, @@ -949,6 +955,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 @@ -1344,6 +1354,12 @@ 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() + 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..6f10c9103300 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/rehearsal.py @@ -0,0 +1,55 @@ +# 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 attention warmup +forwards on throwaway KV cache managers inside a KernelBuilder rehearsal: the ops list +their kernels, and the kernels build on worker threads while the checkpoint loads. +""" + +from __future__ import annotations + +import copy +import gc +from typing import TYPE_CHECKING + +import torch + +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() 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 index 15f88ee201b1..afac3ba7f9f0 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -53,6 +53,40 @@ def no_compile(name, params): assert jit._variant_manager._run_template is not None +def test_a_rehearsal_run_lists_and_launches_nothing(fmha_sm100, monkeypatch): + from fmha_sm100 import jit + + def no_compile(name, params): + raise AssertionError(f"compiled {name}") + + monkeypatch.setattr(jit._variant_manager, "_compile_only", no_compile) + 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 diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 67af020c53a5..9d15ef9af611 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -177,8 +177,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) @@ -201,7 +203,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..f6e9ebf7cb00 --- /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: 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 index 9b014a84a049..d6cbd3a61754 100644 --- a/tests/unittest/_torch/test_kernel_builder.py +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -2,10 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """KernelBuilder: builds listed by an op on first sight, 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 KernelBuilder pytestmark = pytest.mark.cpu_only @@ -97,11 +100,50 @@ def slow(): assert (stats.built, stats.waited, stats.waited_s, stats.late) == (1, 1, 5.0, 0) +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() + builder.prebuild("op", "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. + builder.prebuild("op", "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: reached.append(model_config) 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") + assert not builder.wants_rehearsal("dense") + assert reached == ["msa", "dense"] + assert not KernelBuilder(threads=1, enabled=True, rehearse=False).wants_rehearsal("msa") + assert not KernelBuilder(threads=1, enabled=False).wants_rehearsal("msa") + + def test_disabled(): builder = KernelBuilder(threads=1, enabled=False) listed = [] with builder.build_pass("test"): assert not builder.active() builder.prebuild("op", "sig", lambda: listed.append(True) or []) + with builder.rehearsal(): + assert not builder.rehearsing() assert listed == [] assert not builder.submit("op", "k", lambda: True) From 812ac7818c450bf3d9507ead15041d0fb861ac76 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 10 Nov 2026 00:19:17 -0800 Subject: [PATCH 03/11] [TRTLLM-17679][fix] KernelBuilder rehearsal: serve NCCL-window requests from plain memory The rehearsal's eager forwards grew the NCCL window pool with buffers the compiled warmup could not reuse, which left 0.8 GiB less for the KV cache on MiniMax-M3 TP2. Run them under nccl_window_allocation_suppressed(), as the autotuner's parallel profiling does, so the pool grows only with the warmup that sizes it. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../_torch/pyexecutor/engine/runners/decoder/warmup.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 418fa154904b..2ffae93562b8 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -22,6 +22,7 @@ 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.config_utils import is_decode_only from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager @@ -326,11 +327,14 @@ def rehearse(self, resource_manager: ResourceManager) -> None: from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner try: - # force_eager: no torch.compile traces, whole-model or op-level. + # 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._prewarm_minimax_m3_msa_indexer(kv_cache_manager) self._run_attention_warmup(resource_manager, can_run_general_warmup) From fa3dfe0bb9b536c700b726fd825707df65a8dd67 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 10 Nov 2026 02:52:50 -0800 Subject: [PATCH 04/11] [TRTLLM-17679][feat] KernelBuilder: sweep each op's shapes down from the forward's, as the autotuner does An op declares the dimensions its kernels depend on in a BuildConfig of DynamicSpecs, as a TuningConfig declares DynamicTensorSpecs. A call supplies the shape; the op lists the kernels of the shapes its dimensions sweep from it that its call signature has not listed yet. fmha_sm100 sweeps query lengths up to the call's longest segment (1-129, and 256); the MSA indexer's maximum-score pass also sweeps KV splits (1, 2). The rehearsal runs the autotuner warmup's two forwards (a context batch at the maximum token count and one generation request) in place of the attention warmup's batches and the MSA indexer prewarm. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 10 +- .../minimax_m3/kernels/msa_kernel_builds.py | 79 ++++++++---- tensorrt_llm/_torch/kernel_builder.py | 104 +++++++++++---- .../engine/runners/decoder/warmup.py | 120 ++++++++++++------ .../sparse/msa/test_msa_kernel_builds.py | 67 +++++++--- tests/unittest/_torch/test_kernel_builder.py | 63 +++++++-- 6 files changed, 326 insertions(+), 117 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index e21bd47bc72e..0e3059debc1c 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -205,10 +205,12 @@ curl http://localhost:8000/server_info ### Kernel builds ahead of first use -An op that compiles its kernels on first use lists them the first time it runs, and worker -threads build them in parallel. For a model with such an op, the engine first loads dummy -weights and runs the attention warmup forwards on them, then loads the checkpoint while the -kernels build. The startup log shows these as the `rehearsal` and `checkpoint` phases, which +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. 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), then loads the checkpoint while the kernels build. +The startup log shows these as the `rehearsal` and `checkpoint` phases, which `model_engine_creation_seconds` then does not include. Supported ops: MiniMax-M3's MSA attention (`fmha_sm100`). 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 index eabab34fda13..076d44ea5ad9 100644 --- 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 @@ -5,15 +5,16 @@ 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. +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 the first run -of each call signature in a build pass, the plan and the run are replayed for one query -segment of every length in ``_QUERY_LENS``, 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. +``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. @@ -29,15 +30,28 @@ import torch -from tensorrt_llm._torch.kernel_builder import KernelBuilder +from tensorrt_llm._torch.kernel_builder import BuildConfig, DynamicSpec, KernelBuilder from tensorrt_llm.logger import logger _OP = "fmha_sm100" -# One query segment of each length: up to 129 crosses every boundary of the pack factor, -# the single warp group and the 128-row query tile; 256 takes the 256-row tile. -_QUERY_LENS = tuple(range(1, 130)) + (256,) -# Plan options a replay does not pass through: it plans its own lengths and offsets. -_PLAN_SKIP = ("qo_offset", "device", "stream") + + +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. @@ -99,9 +113,10 @@ def __init__(self, module: types.ModuleType, api: types.ModuleType, jit: types.M 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) + sub_plan[_PLAN_ARGS] = (args, options, qo_segment_lens, num_kv_splits) return plan_info def submit_modules(self) -> None: @@ -128,8 +143,10 @@ def run(self, q, k, v, plan_info, *args, **kwargs): None, ) if planned is not None: - plan_args, plan_options = planned + 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())), @@ -138,11 +155,20 @@ def run(self, q, k, v, plan_info, *args, **kwargs): 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, - functools.partial(self._list, q, k, v, plan_args, plan_options, run_options), + 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) @@ -167,10 +193,13 @@ def _rehearse(self, q, k, v, plan_info, *args, **kwargs): tensor.fill_(value) return result - def _list(self, q, k, v, plan_args, plan_options, run_options) -> List[Tuple[str, Callable]]: - """Replay the plan and the run for each query length; return the variants asked for. + 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 length the package rejects is skipped: a call of that shape would fail too. + 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 @@ -178,7 +207,8 @@ def _list(self, q, k, v, plan_args, plan_options, run_options) -> List[Tuple[str found: Dict[str, dict] = {} _replay.found = found try: - for qo_len in _QUERY_LENS: + 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) @@ -187,11 +217,14 @@ def _list(self, q, k, v, plan_args, plan_options, run_options) -> List[Tuple[str 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, **plan_options) + 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 {qo_len} query rows: {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 diff --git a/tensorrt_llm/_torch/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py index 0c034058ad9f..87a81d2ab655 100644 --- a/tensorrt_llm/_torch/kernel_builder.py +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -3,15 +3,18 @@ """``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. On the first call of each call signature inside a build pass, the -op lists every kernel that signature reaches over its runtime shapes; 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. - -A rehearsal is a build pass over forwards on dummy weights, run before the checkpoint -loads: 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. +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. + +A rehearsal is a build pass over the autotuner warmup's forwards on dummy weights, run +before the checkpoint loads: 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 @@ -24,17 +27,30 @@ 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, Optional, Tuple +from typing import ( + Any, + Callable, + Dict, + Hashable, + Iterable, + Iterator, + List, + Mapping, + Optional, + Tuple, + Union, +) import torch from tensorrt_llm.logger import logger -__all__ = ["KernelBuilder"] +__all__ = ["BuildConfig", "DynamicSpec", "KernelBuilder"] _ENV = "TLLM_KERNEL_BUILDER" _REHEARSAL_ENV = "TLLM_KERNEL_BUILDER_REHEARSAL" @@ -57,6 +73,39 @@ _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 @@ -116,7 +165,8 @@ def __init__( self._passes = 0 self._rehearsals = 0 self._pending = 0 - self._signatures: set = set() + # 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] = {} @@ -132,8 +182,8 @@ def get(cls) -> "KernelBuilder": @contextlib.contextmanager def build_pass(self, name: str, wait: bool = True) -> Iterator[None]: - """Ops list their kernels on first sight inside. The outermost pass logs what was - built at its end, after waiting for the builds unless ``wait`` is False.""" + """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 @@ -200,22 +250,29 @@ def prebuild( self, op: str, signature: Hashable, - list_builds: Callable[[], Iterable[Tuple[Hashable, Build]]], + config: BuildConfig, + observed: Mapping[str, int], + list_builds: Callable[[List[Dict[str, int]]], Iterable[Tuple[Hashable, Build]]], ) -> None: - """On the first sight of ``(op, signature)`` in a build pass, submit every build - ``list_builds()`` returns as ``(key, build)``. + """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: - if (op, signature) in self._signatures: + 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 - self._signatures.add((op, signature)) + listed.update(tuple(sorted(shape.items())) for shape in shapes) t0 = time.perf_counter() try: - builds = list(list_builds()) + 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 " @@ -226,11 +283,12 @@ def prebuild( queued = sum(self.submit(op, key, build) for key, build in self._rotated(builds)) with self._lock: stats = self._op_stats(op) - stats.signatures += 1 + stats.signatures += new_signature stats.list_s += list_s logger.info( - f"[KernelBuilder] {op}: {len(builds)} kernels for a new call signature, " - f"{queued} queued ({list_s:.1f}s to list)" + 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: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 2ffae93562b8..d138e2cb3ca3 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -301,15 +301,17 @@ 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. Ops list their kernels for - # KernelBuilder on first sight; the pass waits for the builds at its end. + # 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 MSA indexer prewarm and the attention warmup forwards, run eagerly on + """The autotuner warmup's forwards, run eagerly and outside the autotuner on whatever weights the model holds, inside the caller's KernelBuilder rehearsal. - Every rank calls it. Its requests are released and its workspaces freed. + They supply the largest shapes: each op the forwards reach sweeps its own shapes + down from them. Every rank calls it. Its requests are released and its workspaces + freed. """ kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER @@ -319,10 +321,8 @@ def rehearse(self, resource_manager: ResourceManager) -> None: if self.mapping.cp_size > 1 and self.mapping.cp_config.get("cp_type") != CpType.HELIX: return AutoTuner.get() - can_run_general_warmup = self._agreement.agree_flag( - not self.mapping.has_cp_helix() - and self.guided_decoder is None - and not isinstance(kv_cache_manager, MambaHybridCacheManager) + 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 @@ -335,9 +335,10 @@ def rehearse(self, resource_manager: ResourceManager) -> None: torch.compiler.set_stance("force_eager"), self._torch_compile_bypass(), nccl_window_allocation_suppressed(), + self._graphs.no_cuda_graph(), ): - self._prewarm_minimax_m3_msa_indexer(kv_cache_manager) - self._run_attention_warmup(resource_manager, can_run_general_warmup) + 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 @@ -348,6 +349,36 @@ def rehearse(self, resource_manager: ResourceManager) -> None: 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]]: @@ -1260,25 +1291,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 @@ -1308,18 +1323,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 @@ -1735,3 +1738,42 @@ 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 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 index afac3ba7f9f0..ad7941237f10 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -27,39 +27,72 @@ def fmha_sm100(monkeypatch): return module -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) -def test_listing_replays_every_query_length_without_compiling(fmha_sm100, monkeypatch, dtype): +@pytest.fixture +def no_compile(monkeypatch): from fmha_sm100 import jit - def no_compile(name, params): + def compile_variant(name, params): raise AssertionError(f"listing compiled {name}") - monkeypatch.setattr(jit._variant_manager, "_compile_only", no_compile) + 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, num_kv_splits=1, causal=True) - builds = msa_kernel_builds._package._list( - q, k, k, (num_heads,), options, {"kv_indices": kv_indices, "output_maxscore": False} - ) - names = [name for name, _ in builds] - assert len(names) == len(set(names)) - # A query tile of 128 rows and one of 256 cover at least two variants. - assert len(names) >= 2 + 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_a_rehearsal_run_lists_and_launches_nothing(fmha_sm100, monkeypatch): - from fmha_sm100 import jit +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 no_compile(name, params): - raise AssertionError(f"compiled {name}") - monkeypatch.setattr(jit._variant_manager, "_compile_only", no_compile) +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) diff --git a/tests/unittest/_torch/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py index d6cbd3a61754..6c20c8d600d4 100644 --- a/tests/unittest/_torch/test_kernel_builder.py +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""KernelBuilder: builds listed by an op on first sight, run on worker threads.""" +"""KernelBuilder: builds an op lists for the shapes it sweeps, run on worker threads.""" import sys import threading @@ -9,11 +9,16 @@ import pytest from tensorrt_llm._torch import kernel_builder -from tensorrt_llm._torch.kernel_builder import KernelBuilder +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) @@ -30,14 +35,14 @@ def list_builds(): listed.append(True) return [(name, _recorder(built, name)) for name in ("a", "b", "c")] - builder.prebuild("op", "sig", list_builds) + _prebuild(builder, "sig", list_builds) assert listed == [] # no pass, no listing with builder.build_pass("test"): - builder.prebuild("op", "sig", list_builds) - builder.prebuild("op", "sig", list_builds) + _prebuild(builder, "sig", list_builds) + _prebuild(builder, "sig", list_builds) # A kernel another signature also reaches is built once. - builder.prebuild("op", "sig2", lambda: [("a", _recorder(built, "a"))]) + _prebuild(builder, "sig2", lambda: [("a", _recorder(built, "a"))]) # The pass waited for the builds. assert listed == [True] assert sorted(built) == ["a", "b", "c"] @@ -59,7 +64,7 @@ def exit_(): builds = [("bad", fail), ("exit", exit_), ("good", _recorder(built, "good"))] with builder.build_pass("test"): - builder.prebuild("op", "sig", lambda: builds) + _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" @@ -76,7 +81,7 @@ def list_builds(): raise ValueError("unsupported shape") with builder.build_pass("test"): - builder.prebuild("op", "sig", list_builds) + _prebuild(builder, "sig", list_builds) assert builder.state("op", "anything") is None @@ -100,6 +105,42 @@ def slow(): 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() @@ -111,14 +152,14 @@ def slow(): assert not builder.rehearsing() with builder.rehearsal(): assert builder.rehearsing() and builder.active() - builder.prebuild("op", "sig", lambda: [("k", slow)]) + _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. - builder.prebuild("op", "sig", lambda: pytest.fail("listed twice")) + _prebuild(builder, "sig", lambda: pytest.fail("listed twice")) # The warmup pass waited for it. assert builder.state("op", "k") == "done" @@ -142,7 +183,7 @@ def test_disabled(): listed = [] with builder.build_pass("test"): assert not builder.active() - builder.prebuild("op", "sig", lambda: listed.append(True) or []) + _prebuild(builder, "sig", lambda: listed.append(True) or []) with builder.rehearsal(): assert not builder.rehearsing() assert listed == [] From 596a3b2416eac66c2dfd683dadb735765e003340 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 10 Nov 2026 08:15:33 -0800 Subject: [PATCH 05/11] [TRTLLM-17679][feat] Autotuner build state: autotuned ops build their kernels while the weights load An op that calls choose_one now has three states. Inside autotune() it tunes; inside the new build_kernels() (the rehearsal) choose_one queues on KernelBuilder's worker threads the builds of every profile, runner and tactic its tuning will time, without timing them (AutoTuner.prebuild), and the op returns zeros of its output shape; otherwise it runs its tuned tactic. A runner builds a tactic with TunableRunner.build (by default, it runs it once) and says which builds compile the same kernels with build_key. Only ops that opt in with TuningConfig.parallel_jit_warmup build: DeepGEMM's FP8 block-scale GEMMs, and FlashInfer's MXFP8 GEMM, whose CUTLASS build loads its JIT module without running it. The parallel JIT warmup shares the per-profile build and the interleaved order, and skips builds this process ran. The creator waits for these builds (the autotuner_builds phase) before the KV cache estimation measures free memory. A model whose FP8 block-scale linears run DeepGEMM now rehearses too (the deep_gemm KernelBuilder family; families receive the model). KernelBuilder sizes its threads as the parallel JIT warmup does, and join() can wait for the ops of one prefix. TLLM_KERNEL_BUILDER_AUTOTUNER=0 leaves the autotuned ops' kernels to the autotuner warmup. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 21 +- .../minimax_m3/kernels/msa_kernel_builds.py | 2 +- tensorrt_llm/_torch/autotuner.py | 239 ++++++++++++++++-- .../_torch/custom_ops/torch_custom_ops.py | 44 +++- tensorrt_llm/_torch/kernel_builder.py | 63 +++-- tensorrt_llm/_torch/modules/linear.py | 23 +- .../engine/runners/decoder/warmup.py | 141 ++++++----- .../_torch/pyexecutor/model_loader.py | 2 +- .../_torch/pyexecutor/py_executor_creator.py | 4 +- tensorrt_llm/_torch/pyexecutor/rehearsal.py | 24 +- tests/unittest/_torch/misc/test_autotuner.py | 198 +++++++++++++++ .../pyexecutor/test_deferred_checkpoint.py | 2 +- tests/unittest/_torch/test_kernel_builder.py | 23 +- 13 files changed, 656 insertions(+), 130 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 0e3059debc1c..5fd9d3bc4a1a 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -207,18 +207,25 @@ curl http://localhost:8000/server_info 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. 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), then loads the checkpoint while the kernels build. -The startup log shows these as the `rehearsal` and `checkpoint` phases, which -`model_engine_creation_seconds` then does not include. Supported ops: MiniMax-M3's MSA -attention (`fmha_sm100`). +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_THREADS` | `8` | Worker threads per process. | +| `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 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 index 076d44ea5ad9..33a0696afa2c 100644 --- 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 @@ -291,7 +291,7 @@ def get_reduction_module(self): return self.reduction_fn() -def reaches(model_config: Any) -> bool: +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 diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 2d3254157f0c..acb126a4e2cc 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): @@ -152,11 +159,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. @@ -292,6 +301,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, @@ -949,6 +1000,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. @@ -977,6 +1050,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) @@ -1000,7 +1078,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: @@ -1206,6 +1290,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 @@ -1629,11 +1725,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 @@ -1658,12 +1847,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. @@ -1673,15 +1857,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) @@ -1901,13 +2078,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 @@ -1928,13 +2110,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 @@ -2159,6 +2343,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 5c85e3a308c6..d224537067e5 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 index 87a81d2ab655..d3fa89e71c10 100644 --- a/tensorrt_llm/_torch/kernel_builder.py +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -11,15 +11,21 @@ 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: 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. +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 8). +(default: this rank's share of the CPUs, up to 32, as for the autotuner's parallel JIT +warmup). """ from __future__ import annotations @@ -55,16 +61,17 @@ _ENV = "TLLM_KERNEL_BUILDER" _REHEARSAL_ENV = "TLLM_KERNEL_BUILDER_REHEARSAL" _THREADS_ENV = "TLLM_KERNEL_BUILDER_THREADS" -_DEFAULT_THREADS = 8 _PROGRESS_EVERY_S = 60.0 # The op families a rehearsal can reach: name -> "module:function". The function takes the -# model config, returns True when the model reaches the op, and then queues the builds the -# configuration alone determines. A module is imported only when asked. +# 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. @@ -155,7 +162,12 @@ def __init__( rehearse = os.environ.get(_REHEARSAL_ENV, "1") != "0" self.rehearse = self.enabled and rehearse if threads is None: - threads = int(os.environ.get(_THREADS_ENV, str(_DEFAULT_THREADS))) + 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) @@ -165,6 +177,8 @@ def __init__( 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] = {} @@ -226,9 +240,10 @@ 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) -> bool: - """True when the rehearsal is on and an op family in ``_FAMILIES`` reaches the - model. Each family that does queues the builds its configuration determines.""" + 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 @@ -236,7 +251,7 @@ def wants_rehearsal(self, model_config: Any) -> bool: module_name, function = target.split(":") try: reaches = getattr(importlib.import_module(module_name), function) - wanted = bool(reaches(model_config)) or wanted + 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; " @@ -301,6 +316,7 @@ def submit(self, op: str, key: Hashable, build: Build) -> bool: 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 @@ -337,22 +353,26 @@ def op_compiled(self, op: str, key: Hashable, seconds: float) -> None: # Waiting and reporting. - def join(self, timeout: Optional[float] = None) -> float: - """Wait for every queued build; returns the seconds waited.""" + 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 self._pending > 0: + 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 {self._pending} builds" - ) + 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 = [ @@ -407,6 +427,8 @@ def _work(self) -> None: 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 @@ -419,5 +441,6 @@ def _work(self) -> None: else: stats.present += 1 self._pending -= 1 - if self._pending == 0: - self._idle.notify_all() + 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 96bb4cf1e5ba..af0ad304c1a1 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 @@ -1420,6 +1420,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. @@ -3817,6 +3831,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/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index d138e2cb3ca3..c4baab0d4870 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -15,7 +15,7 @@ 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 @@ -306,12 +306,12 @@ def run(self, resource_manager: ResourceManager) -> EagerWorkspaceReclaimer | No return self._warmup_scheduled(resource_manager, kv_cache_manager) def rehearse(self, resource_manager: ResourceManager) -> None: - """The autotuner warmup's forwards, run eagerly and outside the autotuner on - whatever weights the model holds, inside the caller's KernelBuilder rehearsal. + """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. Every rank calls it. Its requests are released and its workspaces - freed. + 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 @@ -321,6 +321,9 @@ def rehearse(self, resource_manager: ResourceManager) -> None: 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() ) @@ -336,6 +339,7 @@ def rehearse(self, resource_manager: ResourceManager) -> None: 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) @@ -1195,70 +1199,14 @@ def _run_autotuner_warmup( requests instead of a context batch, for workers that never run one (Helix CP and decode-only 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. @@ -1777,3 +1725,68 @@ def _autotuner_warmup_configs( # 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_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 17459044a7ad..c97a99868d55 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -1541,7 +1541,7 @@ def init_meta_tensor(t: torch.Tensor): 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)): + and KernelBuilder.get().wants_rehearsal(config, model)): self._deferred = self._load_rehearsal_weights( model, config, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 734bbbd7eaf5..9ccb970981a0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -67,7 +67,7 @@ from .model_engine import ModelEngine from .model_loader import ModelLoader, construct_checkpoint_loader from .py_executor import PyExecutor -from .rehearsal import rehearse +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) @@ -1359,6 +1359,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): 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 index 6f10c9103300..bfe9f4109372 100644 --- a/tensorrt_llm/_torch/pyexecutor/rehearsal.py +++ b/tensorrt_llm/_torch/pyexecutor/rehearsal.py @@ -3,9 +3,11 @@ """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 attention warmup -forwards on throwaway KV cache managers inside a KernelBuilder rehearsal: the ops list -their kernels, and the kernels build on worker threads while the checkpoint loads. +(``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 @@ -16,6 +18,9 @@ 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 @@ -53,3 +58,16 @@ def rehearse(model_engine: ModelEngine, kv_cache_creator: KvCacheCreator) -> Non 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/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/pyexecutor/test_deferred_checkpoint.py b/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py index f6e9ebf7cb00..998d1f844222 100644 --- a/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py +++ b/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py @@ -60,7 +60,7 @@ def test_a_deferred_load_matches_a_direct_one(checkpoint, monkeypatch): expected = _state(direct) del direct - monkeypatch.setattr(KernelBuilder, "wants_rehearsal", lambda self, model_config: True) + 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) diff --git a/tests/unittest/_torch/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py index 6c20c8d600d4..50fb8634af7e 100644 --- a/tests/unittest/_torch/test_kernel_builder.py +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -167,17 +167,32 @@ def slow(): def test_the_rehearsal_runs_for_a_model_an_op_family_reaches(monkeypatch): reached = [] family = types.ModuleType("family") - family.reaches = lambda model_config: reached.append(model_config) or model_config == "msa" + 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") - assert not builder.wants_rehearsal("dense") - assert reached == ["msa", "dense"] + 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 = [] From 0040b98043c741a1d57aa8abf2d2941d4d47ef94 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 10 Nov 2026 22:59:19 -0800 Subject: [PATCH 06/11] [TRTLLM-17679][feat] KernelBuilder: build an op by running it at each swept shape builds_by_running(config, dims, output) decorates an op whose kernels compile inside its own run. In a rehearsal, each call queues on KernelBuilder's workers one run of the op for every shape its BuildConfig sweeps from the call's, on the call's arguments cut or tiled to that shape along the dimensions that dims names, and returns output (zeros) instead of running. Such an op needs no list of its own. Builds that run kernels say so when submitted (uses_gpu), and the creator waits for those, the autotuner's included, before the KV cache estimation (the gpu_builds phase, formerly autotuner_builds). The workers' CUDA stream is kernel_builder.worker_stream(). Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 9 +- tensorrt_llm/_torch/autotuner.py | 43 ++-- tensorrt_llm/_torch/kernel_builder.py | 206 ++++++++++++++++-- .../_torch/pyexecutor/py_executor_creator.py | 6 +- tensorrt_llm/_torch/pyexecutor/rehearsal.py | 18 +- .../sparse/msa/test_msa_kernel_builds.py | 4 +- tests/unittest/_torch/test_kernel_builder.py | 59 ++++- 7 files changed, 271 insertions(+), 74 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 5fd9d3bc4a1a..0574584f182d 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -215,10 +215,11 @@ 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. +The startup log shows these as the `rehearsal`, `checkpoint` and `gpu_builds` phases, +which `model_engine_creation_seconds` then does not include; `gpu_builds` waits for the +builds that run kernels on the GPU. 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 | |----------------------|---------|--------| diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index acb126a4e2cc..215022194eb1 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -37,7 +37,7 @@ import tensorrt_llm from tensorrt_llm._torch.distributed import Distributed -from tensorrt_llm._torch.kernel_builder import KernelBuilder +from tensorrt_llm._torch.kernel_builder import KernelBuilder, worker_stream 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) @@ -1011,17 +1011,6 @@ def _interleaved(profiles: List[OptimizationProfile], 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. @@ -1763,19 +1752,21 @@ def prebuild(self, custom_op: str, runners: List[TunableRunner], # 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)) + queued = 0 + for p in _interleaved(profiles, builder.threads): + build = partial(self._build_profile, + custom_op, + runners, + p, + inputs, + tuning_config, + None, + kwargs, + ready=ready) + queued += builder.submit(BUILD_OP_PREFIX + custom_op, + (runner_ids, p.get_opt_shapes()), + build, + uses_gpu=True) logger.info(f"[Autotuner] {custom_op}: {queued} of {len(profiles)} " "profiles queued to build") @@ -1795,7 +1786,7 @@ def _build_profile(self, the calling thread's own, after ``ready``, and returns whether anything was built. """ - stream = stream if stream is not None else _thread_build_stream() + stream = stream if stream is not None else worker_stream() if ready is not None: stream.wait_event(ready) built = False diff --git a/tensorrt_llm/_torch/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py index d3fa89e71c10..d9e0906043ff 100644 --- a/tensorrt_llm/_torch/kernel_builder.py +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -11,6 +11,11 @@ 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 whose kernels compile inside its own run, and that may run on several threads at +once, needs no list of its own: ``builds_by_running`` has each of its calls in a rehearsal +queue one run of the op for every shape its ``BuildConfig`` sweeps, on the call's own +arguments cut or tiled to that shape. + 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. @@ -32,6 +37,7 @@ import contextlib import dataclasses +import functools import importlib import itertools import os @@ -48,6 +54,7 @@ List, Mapping, Optional, + Sequence, Tuple, Union, ) @@ -56,7 +63,7 @@ from tensorrt_llm.logger import logger -__all__ = ["BuildConfig", "DynamicSpec", "KernelBuilder"] +__all__ = ["BuildConfig", "DynamicSpec", "KernelBuilder", "builds_by_running", "worker_stream"] _ENV = "TLLM_KERNEL_BUILDER" _REHEARSAL_ENV = "TLLM_KERNEL_BUILDER_REHEARSAL" @@ -76,6 +83,8 @@ # A build returns True when it compiled the kernel, False when the kernel was already there. Build = Callable[[], bool] +# Where a dimension sits in an op's arguments: (positional index or keyword, tensor dimension). +Position = Tuple[Union[int, str], int] _QUEUED, _DONE, _FAILED = "queued", "done", "failed" @@ -171,14 +180,14 @@ def __init__( 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._queue: "queue.Queue[Tuple[str, Hashable, Build, bool]]" = 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] = {} + # Queued and running builds that use the GPU (submit's uses_gpu). + self._pending_gpu = 0 # The shapes listed so far for each (op, call signature). self._listed: Dict[Tuple[str, Hashable], set] = {} self._states: Dict[Tuple[str, Hashable], str] = {} @@ -268,12 +277,13 @@ def prebuild( config: BuildConfig, observed: Mapping[str, int], list_builds: Callable[[List[Dict[str, int]]], Iterable[Tuple[Hashable, Build]]], + uses_gpu: bool = False, ) -> 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. + A failure to list leaves the kernels to the op. ``uses_gpu`` as for ``submit``. """ if not self.active(): return @@ -295,7 +305,9 @@ def prebuild( ) return list_s = time.perf_counter() - t0 - queued = sum(self.submit(op, key, build) for key, build in self._rotated(builds)) + queued = sum( + self.submit(op, key, build, uses_gpu=uses_gpu) for key, build in self._rotated(builds) + ) with self._lock: stats = self._op_stats(op) stats.signatures += new_signature @@ -306,8 +318,12 @@ def prebuild( 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.""" + def submit(self, op: str, key: Hashable, build: Build, uses_gpu: bool = False) -> bool: + """Queue one build; False when ``(op, key)`` was already queued or built. + + ``uses_gpu``: the build runs kernels on tensors of its own, so it holds GPU memory + while it runs (``join(gpu_only=True)`` waits for these). + """ if not self.enabled: return False with self._lock: @@ -316,9 +332,9 @@ def submit(self, op: str, key: Hashable, build: Build) -> bool: 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._pending_gpu += uses_gpu self._start_workers() - self._queue.put((op, key, build)) + self._queue.put((op, key, build, uses_gpu)) return True def state(self, op: str, key: Hashable) -> Optional[str]: @@ -353,12 +369,12 @@ def op_compiled(self, op: str, key: Hashable, seconds: float) -> None: # 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.""" + def join(self, timeout: Optional[float] = None, gpu_only: bool = False) -> float: + """Wait for every queued build, or only for those that use the GPU; returns the + seconds waited.""" t0 = last = time.perf_counter() with self._idle: - while (pending := self._pending_for(prefix)) > 0: + while (pending := self._pending_gpu if gpu_only else self._pending) > 0: now = time.perf_counter() if timeout is not None and now - t0 >= timeout: break @@ -368,11 +384,6 @@ def join(self, timeout: Optional[float] = None, prefix: Optional[str] = None) -> 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 = [ @@ -412,7 +423,7 @@ def _start_workers(self) -> None: def _work(self) -> None: device_set = False while True: - op, key, build = self._queue.get() + op, key, build, uses_gpu = self._queue.get() t0 = time.perf_counter() state = _DONE # Every item is settled, whatever it raises: an unsettled one would hang join(). @@ -441,6 +452,157 @@ def _work(self) -> None: 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._pending_gpu -= uses_gpu + # A wait for the GPU builds can end before the others do. self._idle.notify_all() + + +# Builds that run kernels. + +_worker_streams = threading.local() + + +def worker_stream() -> Optional["torch.cuda.Stream"]: + """The calling thread's own CUDA stream, made on first use; None without CUDA.""" + if not torch.cuda.is_available(): + return None + stream = getattr(_worker_streams, "stream", None) + if stream is None: + stream = _worker_streams.stream = torch.cuda.Stream() + return stream + + +def builds_by_running( + config: BuildConfig, + dims: Mapping[str, Sequence[Position]], + output: Optional[Callable[..., Any]] = None, + name: Optional[str] = None, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorates an op whose kernels compile inside its own run, so that it needs no list + of its own. + + In a rehearsal, each call queues on KernelBuilder's workers one run of the op for every + shape ``config`` sweeps from the call's that no earlier call of its signature queued. + ``dims`` gives, for each of ``config``'s dimensions, where it sits in the arguments; a + run takes the call's arguments with those tensors cut or tiled to its shape. The call + itself then returns ``output(*args, **kwargs)``, the op's output filled with zeros, + when given, and runs otherwise. Outside a rehearsal the op runs as it is. + + The op must be safe to run on several threads at once and must not use NCCL windows. + Its runs use the GPU (``KernelBuilder.join(gpu_only=True)``). + + Args: + config: The dimensions the op's kernels depend on. + dims: Each dimension's (positional index or keyword, tensor dimension) positions; + the call's value is read at the first one. + output: The op's output for the call's arguments, filled with zeros. + name: The op's name in KernelBuilder's log; the function's name by default. + """ + + def decorate(fn: Callable[..., Any]) -> Callable[..., Any]: + op = name or fn.__qualname__ + + @functools.wraps(fn) + def call(*args, **kwargs): + builder = KernelBuilder.get() + if not builder.rehearsing(): + return fn(*args, **kwargs) + observed = { + dim: _argument(args, kwargs, positions[0][0]).shape[positions[0][1]] + for dim, positions in dims.items() + } + builder.prebuild( + op, + _signature(args, kwargs, dims), + config, + observed, + functools.partial(_runs, fn, args, kwargs, dims, _ready_event()), + uses_gpu=True, + ) + if output is not None: + return output(*args, **kwargs) + return fn(*args, **kwargs) + + return call + + return decorate + + +def _argument(args: Sequence[Any], kwargs: Mapping[str, Any], where: Union[int, str]) -> Any: + return args[where] if isinstance(where, int) else kwargs[where] + + +def _signature(args: Sequence[Any], kwargs: Mapping[str, Any], dims) -> Hashable: + """What a call has but its swept dimensions: each tensor's dtype, device and shape with + those dimensions left out, and every other argument.""" + swept = {position for positions in dims.values() for position in positions} + + def tag(where: Union[int, str], value: Any) -> Hashable: + if isinstance(value, torch.Tensor): + shape = tuple(-1 if (where, d) in swept else n for d, n in enumerate(value.shape)) + return (value.dtype, value.device.type, shape) + try: + hash(value) + return value + except TypeError: + return type(value).__name__ + + return ( + tuple(tag(i, value) for i, value in enumerate(args)), + tuple(sorted((key, tag(key, value)) for key, value in kwargs.items())), + ) + + +def _ready_event() -> Optional["torch.cuda.Event"]: + """Marks the work queued so far, which produced the call's arguments.""" + if not torch.cuda.is_available(): + return None + event = torch.cuda.Event() + event.record() + return event + + +def _runs(fn, args, kwargs, dims, ready, shapes) -> List[Tuple[Hashable, Build]]: + return [ + ( + tuple(sorted(shape.items())), + functools.partial(_run, fn, args, kwargs, dims, shape, ready), + ) + for shape in shapes + ] + + +def _run(fn, args, kwargs, dims, shape, ready) -> bool: + """Run ``fn`` once, on the worker's own stream after ``ready``, with the tensors at + ``dims`` cut or tiled to ``shape``; its output is discarded.""" + stream = worker_stream() + if stream is not None and ready is not None: + stream.wait_event(ready) + args, kwargs = list(args), dict(kwargs) + on_stream = torch.cuda.stream(stream) if stream is not None else contextlib.nullcontext() + with torch.inference_mode(), on_stream: + for dim, positions in dims.items(): + for where, axis in positions: + resized = _resized(_argument(args, kwargs, where), axis, shape[dim]) + if isinstance(where, int): + args[where] = resized + else: + kwargs[where] = resized + fn(*args, **kwargs) + if stream is not None: + stream.synchronize() + return True + + +def _resized(tensor: torch.Tensor, axis: int, size: int) -> torch.Tensor: + """A copy of ``tensor`` with ``size`` entries along ``axis``: its leading ones, repeated + as many times as it takes. Copying keeps index values valid.""" + if tensor.shape[axis] == 0: + shape = list(tensor.shape) + shape[axis] = size + return tensor.new_zeros(shape) + if size > tensor.shape[axis]: + repeats = [1] * tensor.dim() + repeats[axis] = -(-size // tensor.shape[axis]) + tensor = tensor.repeat(*repeats) + return tensor.narrow(axis, 0, size).clone(memory_format=torch.contiguous_format) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 9ccb970981a0..94fb4a6a0770 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -67,7 +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 .rehearsal import join_gpu_builds, rehearse from .scheduler import gen_preemption _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS = (90, 100, 103, 107, 120, 121) @@ -1359,8 +1359,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): 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() + with _startup_timer.phase("gpu_builds"): + join_gpu_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 index bfe9f4109372..de455a26c3ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/rehearsal.py +++ b/tensorrt_llm/_torch/pyexecutor/rehearsal.py @@ -6,8 +6,8 @@ (``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. +threads while the checkpoint loads. ``join_gpu_builds`` then waits for the builds that run +kernels on the GPU. """ from __future__ import annotations @@ -20,7 +20,6 @@ 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 @@ -60,14 +59,13 @@ def rehearse(model_engine: ModelEngine, kv_cache_creator: KvCacheCreator) -> Non torch.cuda.empty_cache() -def join_autotuner_builds() -> None: - """Wait for the builds the rehearsal's autotuned ops queued, then free their memory. +def join_gpu_builds() -> None: + """Wait for the rehearsal's builds that run kernels on the GPU, 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. + They run 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) + waited = KernelBuilder.get().join(gpu_only=True) torch.cuda.empty_cache() if waited >= 0.1: - logger.info(f"[KernelBuilder] waited {waited:.1f}s for the autotuner's builds") + logger.info(f"[KernelBuilder] waited {waited:.1f}s for the builds that run kernels") 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 index ad7941237f10..21befdf8c7c7 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -95,7 +95,9 @@ def variants(config, observed): 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) + monkeypatch.setattr( + builder, "submit", lambda op, key, build, uses_gpu=False: 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) diff --git a/tests/unittest/_torch/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py index 50fb8634af7e..7311705149f3 100644 --- a/tests/unittest/_torch/test_kernel_builder.py +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -7,9 +7,15 @@ import types import pytest +import torch from tensorrt_llm._torch import kernel_builder -from tensorrt_llm._torch.kernel_builder import BuildConfig, DynamicSpec, KernelBuilder +from tensorrt_llm._torch.kernel_builder import ( + BuildConfig, + DynamicSpec, + KernelBuilder, + builds_by_running, +) pytestmark = pytest.mark.cpu_only @@ -180,17 +186,54 @@ def test_the_rehearsal_runs_for_a_model_an_op_family_reaches(monkeypatch): assert not KernelBuilder(threads=1, enabled=False).wants_rehearsal("msa") -def test_join_waits_for_the_builds_of_the_ops_a_prefix_names(): +def test_join_can_wait_only_for_the_builds_that_use_the_gpu(): 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" + builder.submit("compile", "k", lambda: release.wait(timeout=30)) + builder.submit("run", "k", lambda: True, uses_gpu=True) + assert builder.join(timeout=30, gpu_only=True) < 30 + assert builder.state("run", "k") == "done" + assert builder.state("compile", "k") == "queued" release.set() builder.join(timeout=30) - assert builder.state("other", "k") == "done" + assert builder.state("compile", "k") == "done" + + +def test_an_op_built_by_running_runs_once_per_swept_shape_in_a_rehearsal(monkeypatch): + builder = KernelBuilder(threads=2, enabled=True) + monkeypatch.setattr(KernelBuilder, "_instance", builder) + runs = [] + + @builds_by_running( + BuildConfig((DynamicSpec("m", (1, 8)),)), + dims={"m": [(0, 0), ("out", 0)]}, + output=lambda x, w, out: torch.zeros(x.shape[0], w.shape[1]), + name="op", + ) + def op(x, w, out): + runs.append( + (x.clone(), out.shape[0], threading.get_ident(), torch.is_inference_mode_enabled()) + ) + return x @ w + + x = torch.arange(12.0).reshape(6, 2) + w = torch.ones(2, 3) + # Outside a rehearsal the op runs as it is. + assert torch.equal(op(x, w, out=torch.empty(6, 3)), x @ w) + runs.clear() + with builder.rehearsal(): + assert torch.equal(op(x, w, out=torch.empty(6, 3)), torch.zeros(6, 3)) + # A call of the same signature at a swept size queues nothing new. + op(x[:1], w, out=torch.empty(1, 3)) + builder.join(timeout=30) + # The call's own size and the swept ones, each run once, off the calling thread. + assert sorted(run[0].shape[0] for run in runs) == [1, 6, 8] + assert all(run[1] == run[0].shape[0] for run in runs) + assert all(run[2] != threading.get_ident() and run[3] for run in runs) + # A larger size repeats the call's rows. + tiled = next(run[0] for run in runs if run[0].shape[0] == 8) + assert torch.equal(tiled, torch.cat([x, x[:2]])) + assert builder.stats("op").submitted == 3 def test_disabled(): From 47e87252b69f1438c66a87ca60b3c556118f8668 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 10 Nov 2026 23:08:19 -0800 Subject: [PATCH 07/11] [TRTLLM-17679][fix] KernelBuilder rehearsal: run the warmup's generation-only batch where the warmup does Decode-only workers (sparse KV cache offload) now warm up generation shapes only, as Helix CP workers do. The rehearsal takes the same decision from the same function as the warmup, so it never runs a context batch the warmup does not. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../engine/runners/decoder/warmup.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index c4baab0d4870..7e46f706664e 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -325,7 +325,7 @@ def rehearse(self, resource_manager: ResourceManager) -> None: # 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() + self, resource_manager, generation_only=_generation_only(self) ) from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner @@ -512,12 +512,8 @@ def _warmup_scheduled( and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager) ) - # Helix CP and decode-only workers never run a context batch, so every - # phase below warms up pure-generation shapes only. Both follow from the - # configuration, so every rank takes the same branch. - generation_only = self.mapping.has_cp_helix() or is_decode_only( - self.sparse_attention_config - ) + # Every phase below then warms up pure-generation shapes only. + generation_only = _generation_only(self) log_mem_snapshot("warmup/before_warmup") # Compile the DSv4 indexer-Q CuTe DSL kernels before the first @@ -1790,3 +1786,11 @@ def _select_mxfp8_methods(warmup: DecoderWarmup) -> tuple[list, list, list]: and (not compile_all_batches or method.backend == "flashinfer") ] return mxfp8_methods, native_mxfp8_methods, flashinfer_mxfp8_methods + + +def _generation_only(warmup: DecoderWarmup) -> bool: + """Whether the warmup runs pure-generation shapes only: Helix CP and decode-only workers + never run a context batch. Both follow from the configuration, so every rank agrees. A + function of the warmup, not a method: tests run the warmup's methods on stand-in objects. + """ + return warmup.mapping.has_cp_helix() or is_decode_only(warmup.sparse_attention_config) From 4f4ce62c4569a004d66b2fab24ae4adfb28e4f9c Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 11 Nov 2026 00:52:55 -0800 Subject: [PATCH 08/11] [TRTLLM-17679][feat] KernelBuilder: build the modules a configuration determines from model creation KernelBuilder's op families now run for every model as soon as it is created, not only on the rehearsal path, and queue the builds the configuration alone determines on the worker threads: - FlashInfer's sampling module (every model's warmup loads it); - FlashInfer's trtllm-gen FMHA module, for a model with a layer FlashInferTrtllmGenFmha serves and for MiniMax-M3's dense-layer decode; - FlashInfer's trtllm_utils module, when an MXFP8 linear tunes through FlashInfer; - MiniMax-M3's k2q CSR extension, with the MSA modules. Workers only build the FlashInfer modules, under FlashInfer's own lock; the first call loads them. In a rehearsal, the trtllm-gen FMHA calls and the NVFP4 sparse prefill run nothing and leave zeros, so the rehearsal does not wait for these builds. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 14 +++-- .../backends/fmha/flashinfer_trtllm_gen.py | 54 ++++++++++++++++++- .../attention/backends/fmha/msa_prefill.py | 7 +++ .../minimax_m3/kernels/msa_kernel_builds.py | 26 +++++++-- .../kernels/trtllm_gen_dense_decode.py | 6 +++ tensorrt_llm/_torch/flashinfer_utils.py | 30 +++++++++++ tensorrt_llm/_torch/kernel_builder.py | 49 ++++++++++------- tensorrt_llm/_torch/modules/linear.py | 29 ++++++++++ .../_torch/pyexecutor/model_loader.py | 5 +- .../pyexecutor/sampler/ops/flashinfer.py | 18 ++++++- .../sparse/msa/test_msa_kernel_builds.py | 24 +++++++++ .../test_flashinfer_trtllm_gen_fmha.py | 46 ++++++++++++++++ .../_torch/modules/test_mxfp8_linear.py | 31 +++++++++++ tests/unittest/_torch/test_kernel_builder.py | 50 +++++++++++++++++ 14 files changed, 360 insertions(+), 29 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 0574584f182d..2ec0d27fffaf 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -212,9 +212,15 @@ states: the autotuner warmup tunes it, a build state builds the kernels of every 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. +Kernels the configuration alone determines start building on worker threads when the model +is created, for every model, and build while the weights load: FlashInfer's sampling, +trtllm-gen FMHA and autotuner modules, and MiniMax-M3's MSA modules with its k2q CSR +extension. The calls that need them find them built, or wait for their builds. + +For a model with an op that lists its kernels as it runs, 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 `gpu_builds` phases, which `model_engine_creation_seconds` then does not include; `gpu_builds` waits for the builds that run kernels on the GPU. Supported ops: MiniMax-M3's MSA attention @@ -224,7 +230,7 @@ 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_REHEARSAL` | `1` | `0` loads the checkpoint first; the kernels a forward lists 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. | diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index bab3490c85bb..ba0847b2dc8b 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -40,20 +40,26 @@ import math from functools import lru_cache -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, List, Optional, Tuple import torch -from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from tensorrt_llm._torch.flashinfer_utils import ( + IS_FLASHINFER_AVAILABLE, + get_env_enable_pdl, + queue_flashinfer_module_build, +) if IS_FLASHINFER_AVAILABLE: import flashinfer from tensorrt_llm._torch.attention.backends.interface import ( + AttentionBackend, AttentionForwardArgs, AttentionInputType, ) from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams +from tensorrt_llm._torch.kernel_builder import KernelBuilder from tensorrt_llm._utils import get_sm_version, torch_dtype_to_binding from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal import thop @@ -938,10 +944,21 @@ def _missing_fused_nanobind_ops() -> List[str]: ) return [op for op in required_ops if not hasattr(thop, op)] + @staticmethod + def _rehearsing(params: FmhaParams) -> bool: + """In a KernelBuilder rehearsal no trtllm-gen kernel runs and the output holds zeros: + their module builds on a KernelBuilder worker (``queue_fmha_module_build``).""" + if not KernelBuilder.get().rehearsing(): + return False + params.output.zero_() + return True + def run_context( self, params: FmhaParams, ) -> None: + if self._rehearsing(params): + return attention_input = params.qkv_input if params.qkv_input is not None else params.query_input if attention_input is None: raise RuntimeError("trtllm-gen context requires QKV or query input.") @@ -1103,6 +1120,8 @@ def run_generation( self, params: FmhaParams, ) -> None: + if self._rehearsing(params): + return attention_input = params.qkv_input if params.qkv_input is not None else params.query_input if attention_input is None: raise RuntimeError("trtllm-gen generation requires QKV or query input.") @@ -1275,6 +1294,8 @@ def run_mla_generation( params: FmhaParams, ) -> None: """MLA generation decode using flashinfer MLA kernel.""" + if self._rehearsing(params): + return attn = params.attn meta = params.meta fwd = params.fwd @@ -1371,3 +1392,32 @@ def run_mla_generation( uses_shared_paged_kv_idx=self.USE_SHARED_PAGED_KV_IDX, multi_ctas_kv_counter_buffer=self._get_multi_ctas_kv_counter_buffer(), ) + + +def queue_trtllm_gen_fmha_module_build() -> bool: + """Queue a build of FlashInfer's trtllm-gen FMHA module on a KernelBuilder worker thread. + + The worker only builds it: loading it also installs FlashInfer's cubin loader, which the + thread that runs the kernels does on their first call. + """ + if not IS_FLASHINFER_AVAILABLE: + return False + from flashinfer.jit import gen_trtllm_gen_fmha_module + + return queue_flashinfer_module_build("trtllm_gen_fmha", gen_trtllm_gen_fmha_module) + + +def queue_fmha_module_build(model_config: Any, model: Any) -> bool: + """KernelBuilder op family: the trtllm-gen FMHA module, for a model with an attention + layer this library serves, built from when the model is created. It needs no forward: + False.""" + del model_config + if model is not None and any( + isinstance(fmha, FlashInferTrtllmGenFmha) + for module in model.modules() + for value in vars(module).values() + if isinstance(value, AttentionBackend) + for fmha in getattr(getattr(value, "_fmha_manager", None), "fmha_libs", ()) + ): + queue_trtllm_gen_fmha_module_build() + return False diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py index 2c8d275e931e..555ee9d6f051 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py @@ -344,6 +344,13 @@ def run_msa_prefill_gqa( ) if row_first != 0: raise ValueError("NVFP4 sparse prefill must cover the context prefix") + from tensorrt_llm._torch.kernel_builder import KernelBuilder + + if KernelBuilder.get().rehearsing(): + # A rehearsal runs none of it: the k2q CSR extension it loads builds on a + # KernelBuilder worker (msa_kernel_builds). The output holds zeros. + out_view.zero_() + return k_scale, v_scale = _aligned_nvfp4_dequant_scales(attn, kv_scale_quant_orig) run_msa_nvfp4_sparse_gqa( q_view, 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 index 33a0696afa2c..7703b9645fbd 100644 --- 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 @@ -13,8 +13,10 @@ 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. +reduction and sparse top-k modules, and the k2q CSR extension the sparse prefill loads, +need nothing from a call: ``reaches`` submits them, or else the first run in a build pass. +``reaches`` also queues FlashInfer's trtllm-gen FMHA module, on which M3's dense layers +decode. In a rehearsal every run is such a replay: it compiles and launches nothing, and its outputs hold zeros and an empty maximum score. @@ -64,6 +66,8 @@ def _query_len_buckets(longest: int) -> Tuple[int, ...]: ("reduction", "fmha_sm100_reduction.so", "_do_compile_reduction"), ("sparse_topk", "sparse_topk_select.so", "_do_compile_sparse_topk"), ) +# The package module that compiles the sparse prefill's k2q CSR extension on import. +_K2Q_CSR_MODULE = "src.sm100.build_k2q_csr" _replay = threading.local() # .found: the variants a replay on this thread asked for _package: Optional["_Package"] = None # set by install @@ -128,6 +132,7 @@ def submit_modules(self) -> None: builder.submit( _OP, name, functools.partial(self.build_module, name, library, compile_fn) ) + builder.submit(_OP, _K2Q_CSR_MODULE, _import_k2q_csr) def run(self, q, k, v, plan_info, *args, **kwargs): builder = KernelBuilder.get() @@ -291,9 +296,23 @@ def get_reduction_module(self): return self.reduction_fn() +def _import_k2q_csr() -> bool: + """The k2q CSR extension compiles when its module is first imported; an import the + sparse prefill makes meanwhile waits on the module's import lock and then reuses it.""" + import importlib + + t0 = time.perf_counter() + importlib.import_module(_K2Q_CSR_MODULE) + # Loading a built extension takes milliseconds. + return time.perf_counter() - t0 >= 1.0 + + 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.""" + are then submitted, with the trtllm-gen FMHA module its dense layers decode on.""" + from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + queue_trtllm_gen_fmha_module_build, + ) from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig @@ -308,6 +327,7 @@ def reaches(model_config: Any, model: Any = None) -> bool: if _package is None: return False _package.submit_modules() + queue_trtllm_gen_fmha_module_build() return True diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index fb6e33b3612c..47e8925c35c3 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -28,6 +28,7 @@ from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( fp8_e4m3_saturate, ) +from tensorrt_llm._torch.kernel_builder import KernelBuilder from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers from .msa_utils import check_decode_span_shape @@ -271,6 +272,11 @@ def minimax_m3_trtllm_gen_dense_decode( int(seq_lens.shape[0]), decode_query_len, ) + if KernelBuilder.get().rehearsing(): + # A rehearsal runs no trtllm-gen kernel: their module builds on a KernelBuilder + # worker (msa_kernel_builds.reaches). The output holds zeros. + output.zero_() + return kv_pool, subpages_per_slot = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") num_heads = int(q.shape[1]) diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index 7c478ff936aa..4e1503508418 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -1,6 +1,8 @@ +import functools import os import platform import traceback +from typing import Any, Callable from ..logger import logger @@ -30,3 +32,31 @@ def get_env_enable_pdl() -> bool: from .flashinfer_cubin_pool import install_shared_cubin_pool install_shared_cubin_pool() + + +def queue_flashinfer_module_build(name: str, gen_spec: Callable[[], + Any]) -> bool: + """Queue a build of the FlashInfer JIT module ``gen_spec()`` describes on a + KernelBuilder worker thread; False when nothing was queued. + + The worker only builds, under the lock FlashInfer's own ``build_and_load`` + takes, so the first call that needs the module finds it built, or waits on + that lock, and loads it on its own thread. + """ + if not IS_FLASHINFER_AVAILABLE: + return False + from .kernel_builder import KernelBuilder + + return KernelBuilder.get().submit( + "flashinfer", name, functools.partial(_build_flashinfer_module, + gen_spec)) + + +def _build_flashinfer_module(gen_spec: Callable[[], Any]) -> bool: + if os.environ.get("FLASHINFER_DISABLE_JIT"): + return False + spec = gen_spec() + if getattr(spec, "is_aot", False) or spec.is_compiled: + return False + spec.build(need_lock=True) + return True diff --git a/tensorrt_llm/_torch/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py index d9e0906043ff..b1ceb65b4063 100644 --- a/tensorrt_llm/_torch/kernel_builder.py +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -20,15 +20,18 @@ build state (``autotuner.build_kernels``), ``choose_one`` queues here the builds of every shape, runner and tactic the op's tuning will time. +An op whose kernels the configuration alone determines needs no forward: its op family in +``_FAMILIES`` queues their builds when the model is created, for every model. + 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. +the weights load. It runs for a model one of whose op families in ``_FAMILIES`` lists +kernels from a forward. ``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 +``TLLM_KERNEL_BUILDER_REHEARSAL=0`` keeps the builds but not the rehearsal: those a forward +lists 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). """ @@ -70,15 +73,22 @@ _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. +# The op families: name -> "module:function". The function takes the model config and the +# model as created, queues the builds of the family's kernels that the configuration alone +# determines, and returns True when a rehearsal would list more of them from a forward. 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", + "flashinfer_sampling": ( + "tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer:queue_sampling_module_build" + ), + "flashinfer_trtllm_gen_fmha": ( + "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen:queue_fmha_module_build" + ), + "flashinfer_autotune": "tensorrt_llm._torch.modules.linear:queue_flashinfer_autotune_build", } # A build returns True when it compiled the kernel, False when the kernel was already there. @@ -250,23 +260,26 @@ def rehearsing(self) -> bool: 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: + """Called once ``model`` is created from ``model_config``, before its weights load. + + Each op family in ``_FAMILIES`` queues the builds the configuration alone + determines. True when the rehearsal is on and a family would list more of its + kernels from a forward. + """ + if not self.enabled: 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 + queue_builds = getattr(importlib.import_module(module_name), function) + wanted = bool(queue_builds(model_config, model)) or wanted + except Exception as e: # noqa: BLE001 - the family's kernels build on first use logger.warning( - f"[KernelBuilder] {name}: could not tell whether the model reaches it; " - f"its kernels build at the warmup. {type(e).__name__}: {e}" + f"[KernelBuilder] {name}: could not queue its builds for the model; its " + f"kernels build on first use. {type(e).__name__}: {e}" ) - return wanted + return self.rehearse and wanted # Registration. diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index af0ad304c1a1..30f82b2c77a0 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3626,6 +3626,35 @@ def flashinfer_mxfp8_autotune() -> Iterator[None]: _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.reset(token) +def queue_flashinfer_autotune_build(model_config, model) -> bool: + """KernelBuilder op family: FlashInfer's ``trtllm_utils`` module, whose + delay kernel FlashInfer's autotuner loads when the warmup tunes the MXFP8 + linears that dispatch through FlashInfer (``flashinfer_mxfp8_autotune``), + built on a worker thread from when the model is created. Queued for a + model with MXFP8 linears that use FlashInfer already, or that the warmup + may switch to it (``_use_flashinfer_mxfp8_decode_graph_default``). It + needs no forward: False.""" + del model_config + if model is None: + return False + modules = list(model.modules()) + methods = [ + module.quant_method for module in modules + if isinstance(getattr(module, "quant_method", None), MXFP8LinearMethod) + ] + uses_flashinfer = any(method.needs_flashinfer_autotune + for method in methods) + may_switch = any( + getattr(module, "_use_flashinfer_mxfp8_decode_graph_default", False) + for module in modules) + if methods and (uses_flashinfer or may_switch): + from flashinfer.jit.tllm_utils import gen_trtllm_utils_module + + from ..flashinfer_utils import queue_flashinfer_module_build + queue_flashinfer_module_build("trtllm_utils", gen_trtllm_utils_module) + return False + + @contextmanager def flashinfer_mxfp8_decode_graph_capture() -> Iterator[None]: """Enable auto-dispatched FlashInfer calls only for decode graph capture.""" diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index c97a99868d55..b412dc1f0afb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -1398,6 +1398,9 @@ def load( validate_lm_head_output_dtype(model, self.llm_args.lm_head_output_dtype) + # Queues the kernel builds the configuration alone determines; they + # run on KernelBuilder's worker threads while the weights load. + wants_rehearsal = KernelBuilder.get().wants_rehearsal(config, model) requires_standard_hf_loading = any( getattr(module, "_requires_standard_hf_loading", False) @@ -1541,7 +1544,7 @@ def init_meta_tensor(t: torch.Tensor): 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)): + and wants_rehearsal): self._deferred = self._load_rehearsal_weights( model, config, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index 560426a5b100..fd02cb9da97b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -31,7 +31,11 @@ import torch -from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from tensorrt_llm._torch.flashinfer_utils import ( + IS_FLASHINFER_AVAILABLE, + get_env_enable_pdl, + queue_flashinfer_module_build, +) from tensorrt_llm.logger import logger if TYPE_CHECKING: @@ -402,6 +406,18 @@ def warmup_sampling_module() -> None: ) +def queue_sampling_module_build(model_config: Any, model: Any) -> bool: + """KernelBuilder op family: the sampling module ``warmup_sampling_module`` loads at every + model's warmup, built on a worker thread from when the model is created. It needs no + forward: False.""" + del model_config, model + if IS_FLASHINFER_AVAILABLE: + from flashinfer.jit.sampling import gen_sampling_module + + queue_flashinfer_module_build("sampling", gen_sampling_module) + return False + + _topk_module_ready: Optional[bool] = None 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 index 21befdf8c7c7..674d3e054979 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -118,10 +118,34 @@ def test_a_rehearsal_run_lists_and_launches_nothing(fmha_sm100, no_compile, monk 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} + modules.add(msa_kernel_builds._K2Q_CSR_MODULE) assert len([key for key in submitted if key not in modules]) >= 2 assert torch.equal(out, torch.zeros_like(out)) +def test_an_msa_model_queues_the_modules_its_calls_load(fmha_sm100, monkeypatch): + from types import SimpleNamespace + + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig + + builder = kernel_builder.KernelBuilder.get() + submitted = [] + monkeypatch.setattr( + builder, "submit", lambda op, key, build, uses_gpu=False: submitted.append(key) or True + ) + monkeypatch.setattr(msa_kernel_builds._package, "modules_submitted", False) + config = SimpleNamespace( + sparse_attention_config=MiniMaxM3SparseAttentionConfig(implementation="msa") + ) + assert msa_kernel_builds.reaches(config) + expected = {name for name, _, _ in msa_kernel_builds._MODULES} + expected.add(msa_kernel_builds._K2Q_CSR_MODULE) + if IS_FLASHINFER_AVAILABLE: + expected.add("trtllm_gen_fmha") + assert set(submitted) == expected + + def test_kernels_are_built_once_where_the_package_loads_them(fmha_sm100, monkeypatch, tmp_path): from fmha_sm100 import jit diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py index fac6b0eb5359..969f551349af 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py @@ -308,3 +308,49 @@ def test_sm107_serves_only_extra_page_sizes( assert supported is expected, reason if not expected: assert "SM107" in reason + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("run", ["run_context", "run_generation", "run_mla_generation"]) +def test_a_rehearsal_runs_no_trtllm_gen_kernel(monkeypatch: pytest.MonkeyPatch, run: str) -> None: + from tensorrt_llm._torch import kernel_builder + + builder = kernel_builder.KernelBuilder(threads=1, enabled=True) + monkeypatch.setattr(kernel_builder.KernelBuilder, "_instance", builder) + fmha = FlashInferTrtllmGenFmha(FakeAttention()) + # Nothing but the output: any step past it would fail on the missing inputs. + params = SimpleNamespace(output=torch.ones(4)) + with builder.rehearsal(): + getattr(fmha, run)(params) + assert torch.equal(params.output, torch.zeros(4)) + + +@pytest.mark.cpu_only +def test_a_model_this_library_serves_queues_its_module_build( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.attention.backends.fmha import flashinfer_trtllm_gen + from tensorrt_llm._torch.attention.backends.interface import AttentionBackend + + queued = [] + monkeypatch.setattr( + flashinfer_trtllm_gen, "queue_trtllm_gen_fmha_module_build", lambda: queued.append(True) + ) + + def model(*fmha_libs): + backend = object.__new__(AttentionBackend) + backend._fmha_manager = SimpleNamespace(fmha_libs=list(fmha_libs)) + layer = torch.nn.Module() + layer.attn = backend + root = torch.nn.Module() + root.layer = layer + return root + + assert not flashinfer_trtllm_gen.queue_fmha_module_build( + None, model(FallbackFmha(FakeAttention())) + ) + assert queued == [] + assert not flashinfer_trtllm_gen.queue_fmha_module_build( + None, model(FallbackFmha(FakeAttention()), FlashInferTrtllmGenFmha(FakeAttention())) + ) + assert queued == [True] diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index 4efee07ab4da..9e17c834b5f2 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -922,3 +922,34 @@ def test_mxfp8_fused_scales_follow_weight_rows_per_tp_rank(weight_mode, shard_ro # Dynamic MXFP8 activations cost about 1% relative error; a scale row # on the wrong weight row is off by at least a factor of two. assert rel < 0.05, (rank, float(rel)) + + +@pytest.mark.cpu_only +def test_flashinfer_autotuner_module_is_queued_for_a_model_that_tunes_through_flashinfer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tensorrt_llm._torch.flashinfer_utils as flashinfer_utils + + monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) + queued = [] + monkeypatch.setattr( + flashinfer_utils, + "queue_flashinfer_module_build", + lambda name, gen_spec: queued.append(name) or True, + ) + + def model(flashinfer_decode_graphs: bool) -> torch.nn.Module: + layer = torch.nn.Module() + layer.quant_method = MXFP8LinearMethod() + root = torch.nn.Module() + root.layer = layer + if flashinfer_decode_graphs: + root._use_flashinfer_mxfp8_decode_graph_default = True + return root + + # The native backend never tunes through FlashInfer. + assert not linear_module.queue_flashinfer_autotune_build(None, model(False)) + assert queued == [] + # A model the warmup may switch to FlashInfer tunes through its autotuner. + assert not linear_module.queue_flashinfer_autotune_build(None, model(True)) + assert queued == ["trtllm_utils"] diff --git a/tests/unittest/_torch/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py index 7311705149f3..80091ed29932 100644 --- a/tests/unittest/_torch/test_kernel_builder.py +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -186,6 +186,56 @@ def test_the_rehearsal_runs_for_a_model_an_op_family_reaches(monkeypatch): assert not KernelBuilder(threads=1, enabled=False).wants_rehearsal("msa") +def test_the_builds_a_configuration_determines_start_without_a_rehearsal(monkeypatch): + family = types.ModuleType("family") + builder = KernelBuilder(threads=1, enabled=True, rehearse=False) + built = [] + + def queue_builds(model_config, model): + builder.submit("module", model_config, _recorder(built, model_config)) + return False + + family.queue_builds = queue_builds + monkeypatch.setitem(sys.modules, "family", family) + monkeypatch.setattr(kernel_builder, "_FAMILIES", {"module": "family:queue_builds"}) + assert not builder.wants_rehearsal("config", "model") + builder.join(timeout=30) + assert built == ["config"] + assert builder.state("module", "config") == "done" + + +def test_a_flashinfer_module_builds_once_only_where_nothing_is_built(monkeypatch): + from tensorrt_llm._torch import flashinfer_utils + + builder = KernelBuilder(threads=1, enabled=True) + monkeypatch.setattr(KernelBuilder, "_instance", builder) + monkeypatch.setattr(flashinfer_utils, "IS_FLASHINFER_AVAILABLE", True) + built = [] + + class Spec: + def __init__(self, name, is_aot=False, is_compiled=False): + self.name, self.is_aot, self.is_compiled = name, is_aot, is_compiled + + def build(self, need_lock=False): + built.append((self.name, need_lock)) + + specs = {"cold": {}, "aot": {"is_aot": True}, "warm": {"is_compiled": True}} + for name, flags in specs.items(): + assert flashinfer_utils.queue_flashinfer_module_build( + name, lambda name=name, flags=flags: Spec(name, **flags) + ) + assert not flashinfer_utils.queue_flashinfer_module_build("cold", lambda: Spec("cold")) + builder.join(timeout=30) + # Under the lock FlashInfer's own build_and_load takes. + assert built == [("cold", True)] + stats = builder.stats("flashinfer") + assert (stats.built, stats.present) == (1, 2) + monkeypatch.setenv("FLASHINFER_DISABLE_JIT", "1") + flashinfer_utils.queue_flashinfer_module_build("disabled", lambda: Spec("disabled")) + builder.join(timeout=30) + assert built == [("cold", True)] + + def test_join_can_wait_only_for_the_builds_that_use_the_gpu(): builder = KernelBuilder(threads=2, enabled=True) release = threading.Event() From 0efbb841ba090394bf002c06a76c60102b13d2e7 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:05:55 -0800 Subject: [PATCH 09/11] [TRTLLM-17679][fix] KernelBuilder rehearsal: load the checkpoint first where the KV cache estimation is off Without the estimation the rehearsal has no KV cache to run on, so a deferred checkpoint only makes the model load twice. With TRTLLM_SKIP_KV_CACHE_ESTIMATION=1, context parallelism, the VANILLA attention backend or a decode-only worker, the model engine now loads its checkpoint first, as without KernelBuilder. On Kimi-K3's Helix generation servers the deferred load held 1.02 GiB more memory outside torch per rank, which left the KV cache 0.82 GiB smaller. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../_torch/pyexecutor/py_executor_creator.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 94fb4a6a0770..00fe1969f88b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -943,8 +943,15 @@ def allocation_scope(current_stage: ExecutorMemoryType): # 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). + # multimodal encoders' CUDA graphs), nor where the KV cache estimation + # the rehearsal runs on is off: with TRTLLM_SKIP_KV_CACHE_ESTIMATION=1, + # context parallelism or the VANILLA backend (as in + # KvCacheCreator.try_prepare_estimation), or on a decode-only worker. mm_config = llm_args.multimodal_config + kv_cache_estimation_off = (skip_est or "cp_type" in mapping.cp_config + or llm_args.attn_backend == "VANILLA" + or is_decode_only( + llm_args.sparse_attention_config)) model_engine = ModelEngine( model_path=checkpoint_dir, llm_args=llm_args, @@ -958,7 +965,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): defer_checkpoint=(dwdp_manager is None and dense_dwdp_manager is None and (mm_config is None - or not mm_config.encoder_cuda_graph)), + or not mm_config.encoder_cuda_graph) + and not kv_cache_estimation_off), ) is_generation = model_engine.model.model_config.is_generation From 66f335d5161b38e03bdd766222eec469169ca86f Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:42:30 -0800 Subject: [PATCH 10/11] [TRTLLM-17679][feat] KernelBuilder rehearsal: rehearse wherever the checkpoint can load after it The rehearsal ran only for a model an op family claimed (MSA attention, DeepGEMM linears). It now runs wherever the checkpoint can load after it, unless TLLM_KERNEL_BUILDER_REHEARSAL=0, so every op that lists its kernels from a forward builds ahead without registering, the autotuner's build state included (FlashInfer's MXFP8 GEMM). The checkpoint is deferred only for a model whose runner rehearses (a decoder runner with context warmups): not for encoder-only, pooling, multimodal-encoder or encoder-decoder models, whose rehearsal would run nothing. The op families only queue the builds the configuration alone determines (KernelBuilder.model_created, for every model as it is created); deep_gemm's family only answered whether to rehearse and goes. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- docs/source/llm-api/index.md | 12 ++++--- .../backends/fmha/flashinfer_trtllm_gen.py | 6 ++-- .../minimax_m3/kernels/msa_kernel_builds.py | 20 +++++------ .../kernels/trtllm_gen_dense_decode.py | 2 +- tensorrt_llm/_torch/kernel_builder.py | 32 +++++++---------- tensorrt_llm/_torch/modules/linear.py | 22 ++---------- .../_torch/pyexecutor/model_engine.py | 6 ++-- .../_torch/pyexecutor/model_loader.py | 24 ++++++++++--- .../_torch/pyexecutor/py_executor_creator.py | 2 +- tensorrt_llm/_torch/pyexecutor/rehearsal.py | 12 +++---- .../pyexecutor/sampler/ops/flashinfer.py | 6 ++-- .../sparse/msa/test_msa_kernel_builds.py | 2 +- .../test_flashinfer_trtllm_gen_fmha.py | 6 ++-- .../_torch/modules/test_mxfp8_linear.py | 4 +-- .../pyexecutor/test_deferred_checkpoint.py | 34 ++++++++++++++++-- tests/unittest/_torch/test_kernel_builder.py | 36 ++++++++----------- 16 files changed, 118 insertions(+), 108 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 2ec0d27fffaf..d7a08d3d6294 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -217,10 +217,14 @@ is created, for every model, and build while the weights load: FlashInfer's samp trtllm-gen FMHA and autotuner modules, and MiniMax-M3's MSA modules with its k2q CSR extension. The calls that need them find them built, or wait for their builds. -For a model with an op that lists its kernels as it runs, 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. +Where the checkpoint can load after them, 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 the +forwards listed build. It loads the checkpoint first with load formats other than AUTO, MX +checkpoints, host-resident weights, a failed meta-device init, sleep mode, the MoE load +balancer, DWDP, the multimodal encoders' CUDA graphs, Modeling V2 targets, encoder-only, +pooling, multimodal-encoder and encoder-decoder models, and where the KV cache estimation is +off. The startup log shows these as the `rehearsal`, `checkpoint` and `gpu_builds` phases, which `model_engine_creation_seconds` then does not include; `gpu_builds` waits for the builds that run kernels on the GPU. Supported ops: MiniMax-M3's MSA attention diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index ba0847b2dc8b..b87bf1cbc7fa 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -1407,10 +1407,9 @@ def queue_trtllm_gen_fmha_module_build() -> bool: return queue_flashinfer_module_build("trtllm_gen_fmha", gen_trtllm_gen_fmha_module) -def queue_fmha_module_build(model_config: Any, model: Any) -> bool: +def queue_fmha_module_build(model_config: Any, model: Any) -> None: """KernelBuilder op family: the trtllm-gen FMHA module, for a model with an attention - layer this library serves, built from when the model is created. It needs no forward: - False.""" + layer this library serves, built from when the model is created.""" del model_config if model is not None and any( isinstance(fmha, FlashInferTrtllmGenFmha) @@ -1420,4 +1419,3 @@ def queue_fmha_module_build(model_config: Any, model: Any) -> bool: for fmha in getattr(getattr(value, "_fmha_manager", None), "fmha_libs", ()) ): queue_trtllm_gen_fmha_module_build() - return False 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 index 7703b9645fbd..0abbea1a20d1 100644 --- 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 @@ -14,9 +14,9 @@ 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, and the k2q CSR extension the sparse prefill loads, -need nothing from a call: ``reaches`` submits them, or else the first run in a build pass. -``reaches`` also queues FlashInfer's trtllm-gen FMHA module, on which M3's dense layers -decode. +need nothing from a call: ``queue_builds`` submits them, or else the first run in a build +pass. ``queue_builds`` also queues FlashInfer's trtllm-gen FMHA module, on which M3's dense +layers decode. In a rehearsal every run is such a replay: it compiles and launches nothing, and its outputs hold zeros and an empty maximum score. @@ -307,9 +307,10 @@ def _import_k2q_csr() -> bool: return time.perf_counter() - t0 >= 1.0 -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, with the trtllm-gen FMHA module its dense layers decode on.""" +def queue_builds(model_config: Any, model: Any = None) -> None: + """KernelBuilder op family: for a model that runs MiniMax-M3's MSA attention on + fmha_sm100, submit its modules and the trtllm-gen FMHA module its dense layers decode + on.""" from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( queue_trtllm_gen_fmha_module_build, ) @@ -320,15 +321,14 @@ def reaches(model_config: Any, model: Any = None) -> bool: config = getattr(model_config, "sparse_attention_config", None) if not isinstance(config, MiniMaxM3SparseAttentionConfig) or config.implementation != "msa": - return False + return if not is_sm_100f() or not msa_package_available(): - return False + return require_msa_module() if _package is None: - return False + return _package.submit_modules() queue_trtllm_gen_fmha_module_build() - return True def install(module: types.ModuleType) -> None: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index 47e8925c35c3..e165c78be158 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -274,7 +274,7 @@ def minimax_m3_trtllm_gen_dense_decode( ) if KernelBuilder.get().rehearsing(): # A rehearsal runs no trtllm-gen kernel: their module builds on a KernelBuilder - # worker (msa_kernel_builds.reaches). The output holds zeros. + # worker (msa_kernel_builds.queue_builds). The output holds zeros. output.zero_() return diff --git a/tensorrt_llm/_torch/kernel_builder.py b/tensorrt_llm/_torch/kernel_builder.py index b1ceb65b4063..4766bad7e9b4 100644 --- a/tensorrt_llm/_torch/kernel_builder.py +++ b/tensorrt_llm/_torch/kernel_builder.py @@ -21,13 +21,12 @@ shape, runner and tactic the op's tuning will time. An op whose kernels the configuration alone determines needs no forward: its op family in -``_FAMILIES`` queues their builds when the model is created, for every model. +``_FAMILIES`` queues their builds when the model is created (``model_created``). 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 one of whose op families in ``_FAMILIES`` lists -kernels from a forward. +the weights load. It runs wherever the checkpoint can load after it (``ModelLoader.load``). ``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: those a forward @@ -73,15 +72,14 @@ _THREADS_ENV = "TLLM_KERNEL_BUILDER_THREADS" _PROGRESS_EVERY_S = 60.0 -# The op families: name -> "module:function". The function takes the model config and the -# model as created, queues the builds of the family's kernels that the configuration alone -# determines, and returns True when a rehearsal would list more of them from a forward. A +# The op families with kernels the configuration alone determines: name -> "module:function". +# The function takes the model config and the model as created and queues those builds. A # module is imported only when asked. _FAMILIES = { "fmha_sm100": ( - "tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_kernel_builds:reaches" + "tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_kernel_builds" + ":queue_builds" ), - "deep_gemm": "tensorrt_llm._torch.modules.linear:deep_gemm_reaches", "flashinfer_sampling": ( "tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer:queue_sampling_module_build" ), @@ -259,27 +257,21 @@ 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: - """Called once ``model`` is created from ``model_config``, before its weights load. - - Each op family in ``_FAMILIES`` queues the builds the configuration alone - determines. True when the rehearsal is on and a family would list more of its - kernels from a forward. - """ + def model_created(self, model_config: Any, model: Any = None) -> None: + """Each op family in ``_FAMILIES`` queues the builds the configuration alone + determines; called once ``model`` is created from ``model_config``, before its + weights load.""" if not self.enabled: - return False - wanted = False + return for name, target in _FAMILIES.items(): module_name, function = target.split(":") try: - queue_builds = getattr(importlib.import_module(module_name), function) - wanted = bool(queue_builds(model_config, model)) or wanted + getattr(importlib.import_module(module_name), function)(model_config, model) except Exception as e: # noqa: BLE001 - the family's kernels build on first use logger.warning( f"[KernelBuilder] {name}: could not queue its builds for the model; its " f"kernels build on first use. {type(e).__name__}: {e}" ) - return self.rehearse and wanted # Registration. diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 30f82b2c77a0..8da0697e75e4 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1420,20 +1420,6 @@ 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. @@ -3626,17 +3612,16 @@ def flashinfer_mxfp8_autotune() -> Iterator[None]: _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.reset(token) -def queue_flashinfer_autotune_build(model_config, model) -> bool: +def queue_flashinfer_autotune_build(model_config, model) -> None: """KernelBuilder op family: FlashInfer's ``trtllm_utils`` module, whose delay kernel FlashInfer's autotuner loads when the warmup tunes the MXFP8 linears that dispatch through FlashInfer (``flashinfer_mxfp8_autotune``), built on a worker thread from when the model is created. Queued for a model with MXFP8 linears that use FlashInfer already, or that the warmup - may switch to it (``_use_flashinfer_mxfp8_decode_graph_default``). It - needs no forward: False.""" + may switch to it (``_use_flashinfer_mxfp8_decode_graph_default``).""" del model_config if model is None: - return False + return modules = list(model.modules()) methods = [ module.quant_method for module in modules @@ -3652,7 +3637,6 @@ def queue_flashinfer_autotune_build(model_config, model) -> bool: from ..flashinfer_utils import queue_flashinfer_module_build queue_flashinfer_module_build("trtllm_utils", gen_trtllm_utils_module) - return False @contextmanager diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 44d3dbffca8a..008e320dfe38 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -126,9 +126,9 @@ def __init__( 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``.""" + """``defer_checkpoint``: build the engine on dummy weights for a KernelBuilder + rehearsal, where the loader can (``ModelLoader.load``); ``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 diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index b412dc1f0afb..aba58cffe8ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -483,6 +483,17 @@ def _fill_dummy_quantized(tensor: torch.Tensor, values.uniform_(low, high, generator=generator).to(tensor.dtype)) +def runner_rehearses(model: torch.nn.Module, llm_args) -> bool: + """Whether the runner ``model`` starts with runs the KernelBuilder rehearsal's + forwards: a decoder runner with context warmups (``resolve_runner_type``), not + an encoder-only, pooling, multimodal-encoder or encoder-decoder one.""" + from .engine.runners import DecoderRunner, resolve_runner_type + + runner_cls = resolve_runner_type(model, llm_args) + return (issubclass(runner_cls, DecoderRunner) + and runner_cls._context_warmups_supported) + + def fill_rehearsal_values(tensors: Iterable[torch.Tensor], low: float = -1e-3, high: float = 1e-3, @@ -1348,8 +1359,9 @@ 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` + defer_checkpoint: Where the checkpoint can load later as it would + now (below) and the KernelBuilder rehearsal is on, return the + model with dummy weights instead; `has_deferred_checkpoint` then tells, and `load_deferred_checkpoint` loads the checkpoint. Returns: @@ -1400,7 +1412,7 @@ def load( self.llm_args.lm_head_output_dtype) # Queues the kernel builds the configuration alone determines; they # run on KernelBuilder's worker threads while the weights load. - wants_rehearsal = KernelBuilder.get().wants_rehearsal(config, model) + KernelBuilder.get().model_created(config, model) requires_standard_hf_loading = any( getattr(module, "_requires_standard_hf_loading", False) @@ -1537,14 +1549,16 @@ def init_meta_tensor(t: torch.Tensor): # 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. + # module tree is the Modeling V1 one the refit hooks cover. And + # only where the model's runner rehearses. 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 wants_rehearsal): + and KernelBuilder.get().rehearse + and runner_rehearses(model, self.llm_args)): self._deferred = self._load_rehearsal_weights( model, config, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 00fe1969f88b..eb74ca54cc00 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -940,7 +940,7 @@ 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 + # The model is built on dummy weights for a KernelBuilder rehearsal; 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), nor where the KV cache estimation diff --git a/tensorrt_llm/_torch/pyexecutor/rehearsal.py b/tensorrt_llm/_torch/pyexecutor/rehearsal.py index de455a26c3ab..90ce7dd5f347 100644 --- a/tensorrt_llm/_torch/pyexecutor/rehearsal.py +++ b/tensorrt_llm/_torch/pyexecutor/rehearsal.py @@ -2,12 +2,12 @@ # 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_gpu_builds`` then waits for the builds that run -kernels on the GPU. +Where the checkpoint can load later (``ModelLoader.load``), 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_gpu_builds`` then waits for the +builds that run kernels on the GPU. """ from __future__ import annotations diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index fd02cb9da97b..18ea471b26ce 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -406,16 +406,14 @@ def warmup_sampling_module() -> None: ) -def queue_sampling_module_build(model_config: Any, model: Any) -> bool: +def queue_sampling_module_build(model_config: Any, model: Any) -> None: """KernelBuilder op family: the sampling module ``warmup_sampling_module`` loads at every - model's warmup, built on a worker thread from when the model is created. It needs no - forward: False.""" + model's warmup, built on a worker thread from when the model is created.""" del model_config, model if IS_FLASHINFER_AVAILABLE: from flashinfer.jit.sampling import gen_sampling_module queue_flashinfer_module_build("sampling", gen_sampling_module) - return False _topk_module_ready: Optional[bool] = None 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 index 674d3e054979..3a1db1ed64fc 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_kernel_builds.py @@ -138,7 +138,7 @@ def test_an_msa_model_queues_the_modules_its_calls_load(fmha_sm100, monkeypatch) config = SimpleNamespace( sparse_attention_config=MiniMaxM3SparseAttentionConfig(implementation="msa") ) - assert msa_kernel_builds.reaches(config) + msa_kernel_builds.queue_builds(config) expected = {name for name, _, _ in msa_kernel_builds._MODULES} expected.add(msa_kernel_builds._K2Q_CSR_MODULE) if IS_FLASHINFER_AVAILABLE: diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py index 969f551349af..edb2be177851 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py @@ -346,11 +346,9 @@ def model(*fmha_libs): root.layer = layer return root - assert not flashinfer_trtllm_gen.queue_fmha_module_build( - None, model(FallbackFmha(FakeAttention())) - ) + flashinfer_trtllm_gen.queue_fmha_module_build(None, model(FallbackFmha(FakeAttention()))) assert queued == [] - assert not flashinfer_trtllm_gen.queue_fmha_module_build( + flashinfer_trtllm_gen.queue_fmha_module_build( None, model(FallbackFmha(FakeAttention()), FlashInferTrtllmGenFmha(FakeAttention())) ) assert queued == [True] diff --git a/tests/unittest/_torch/modules/test_mxfp8_linear.py b/tests/unittest/_torch/modules/test_mxfp8_linear.py index 9e17c834b5f2..b213cdd92bb6 100644 --- a/tests/unittest/_torch/modules/test_mxfp8_linear.py +++ b/tests/unittest/_torch/modules/test_mxfp8_linear.py @@ -948,8 +948,8 @@ def model(flashinfer_decode_graphs: bool) -> torch.nn.Module: return root # The native backend never tunes through FlashInfer. - assert not linear_module.queue_flashinfer_autotune_build(None, model(False)) + linear_module.queue_flashinfer_autotune_build(None, model(False)) assert queued == [] # A model the warmup may switch to FlashInfer tunes through its autotuner. - assert not linear_module.queue_flashinfer_autotune_build(None, model(True)) + linear_module.queue_flashinfer_autotune_build(None, model(True)) assert queued == ["trtllm_utils"] diff --git a/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py b/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py index 998d1f844222..5450722cf0ed 100644 --- a/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py +++ b/tests/unittest/_torch/pyexecutor/test_deferred_checkpoint.py @@ -5,12 +5,14 @@ import pytest import torch +from tensorrt_llm._torch import kernel_builder 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, + runner_rehearses, ) from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.mapping import Mapping @@ -54,13 +56,21 @@ def _state(model): return {name: tensor.detach().clone() for name, tensor in model.state_dict().items()} +def _kernel_builder(monkeypatch, rehearse): + # No op family: the families would queue real kernel builds. + monkeypatch.setattr(kernel_builder, "_FAMILIES", {}) + monkeypatch.setattr( + KernelBuilder, "_instance", KernelBuilder(threads=1, enabled=True, rehearse=rehearse) + ) + + @cuda def test_a_deferred_load_matches_a_direct_one(checkpoint, monkeypatch): + _kernel_builder(monkeypatch, rehearse=True) _, 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) @@ -75,11 +85,31 @@ def test_a_deferred_load_matches_a_direct_one(checkpoint, monkeypatch): @cuda -def test_no_deferral_for_a_model_no_op_family_reaches(checkpoint): +def test_no_deferral_with_the_rehearsal_off(checkpoint, monkeypatch): + _kernel_builder(monkeypatch, rehearse=False) loader, _ = _load(checkpoint, defer_checkpoint=True) assert not loader.has_deferred_checkpoint +@pytest.mark.cpu_only +def test_only_a_decoder_runner_with_context_warmups_rehearses(): + from types import SimpleNamespace + + def model(**config): + config = {"is_generation": True, "is_encoder_decoder": False, **config} + return SimpleNamespace(model_config=SimpleNamespace(**config)) + + def args(**flags): + return SimpleNamespace(**{"mm_encoder_only": False, "encode_only": False, **flags}) + + assert runner_rehearses(model(), args()) + # Encoder-decoder: no context warmups. Pooling, encode-only, multimodal encoder: no decoder. + assert not runner_rehearses(model(is_encoder_decoder=True), args()) + assert not runner_rehearses(model(is_generation=False), args()) + assert not runner_rehearses(model(), args(encode_only=True)) + assert not runner_rehearses(model(), args(mm_encoder_only=True)) + + @pytest.mark.cpu_only def test_the_module_tree_goes_back_as_created(): draft = torch.nn.Module() diff --git a/tests/unittest/_torch/test_kernel_builder.py b/tests/unittest/_torch/test_kernel_builder.py index 80091ed29932..fb084fd437c8 100644 --- a/tests/unittest/_torch/test_kernel_builder.py +++ b/tests/unittest/_torch/test_kernel_builder.py @@ -170,38 +170,30 @@ def slow(): 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_the_builds_a_configuration_determines_start_without_a_rehearsal(monkeypatch): - family = types.ModuleType("family") +def test_each_op_family_queues_its_builds_when_the_model_is_created(monkeypatch): builder = KernelBuilder(threads=1, enabled=True, rehearse=False) - built = [] + asked, built = [], [] def queue_builds(model_config, model): + asked.append((model_config, model)) builder.submit("module", model_config, _recorder(built, model_config)) - return False + family = types.ModuleType("family") family.queue_builds = queue_builds monkeypatch.setitem(sys.modules, "family", family) monkeypatch.setattr(kernel_builder, "_FAMILIES", {"module": "family:queue_builds"}) - assert not builder.wants_rehearsal("config", "model") + # The builds need only the configuration, so they start with the rehearsal off too. + builder.model_created("config", "model") builder.join(timeout=30) + assert asked == [("config", "model")] assert built == ["config"] - assert builder.state("module", "config") == "done" + # A family that cannot queue its builds leaves its kernels to their first use. + family.queue_builds = lambda model_config, model: 1 / 0 + builder.model_created("config", "model") + # Turned off, KernelBuilder asks no family. + family.queue_builds = queue_builds + KernelBuilder(threads=1, enabled=False).model_created("other", "model") + assert asked == [("config", "model")] def test_a_flashinfer_module_builds_once_only_where_nothing_is_built(monkeypatch): From 0d7a201220388ea4bfd35511d38e0ff7f3fa1f3e Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Thu, 12 Nov 2026 00:55:52 -0800 Subject: [PATCH 11/11] [TRTLLM-17679][feat] KernelBuilder: build MiniMax-M3's sparse decode kernels for every batch size ahead MiniMax-M3's Triton sparse decode compiles one kernel for each launch geometry the GB300-measured tables pick per local batch (split-K chunks, softmax recurrence, warps and stages, the paired producer), plus the split-K merge. They compiled on the main thread at each batch's first launch: 20 compiles and 12.4 s per rank on M3 ctx, most of them during CUDA graph capture. In a KernelBuilder build pass, a call of the op now lists the compiles of every batch of up to max_num_requests requests with its other arguments, and KernelBuilder's workers compile each one with Triton's compile-only warmup, given dtypes for the tensors. In a rehearsal the call then leaves its output zero. MsaDecodeFmha passes the metadata's max_num_requests. On M3 ctx (GB300, TP2), the main thread's sparse decode compiles drop from 13.6 s to 2.3 s per rank, and startup from 572 s to 562 s. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../attention/backends/fmha/msa_decode.py | 1 + .../kernels/triton_sparse_decode.py | 562 ++++++++++++------ .../msa/test_minimax_m3_sparse_attn_decode.py | 75 +++ 3 files changed, 457 insertions(+), 181 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py index c6c6af4d08cd..b143c64ff09b 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py @@ -216,6 +216,7 @@ def _run_sparse( sm_scale=(head_dim**-0.5) / float(attn.q_scaling), output=params.output.view(num_tokens, attn.num_heads, head_dim), decode_query_len=params.input_seq_length, + max_num_requests=int(params.meta.max_num_requests), **nvfp4_args, ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py index 0150347026a1..fbc357dfd929 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py @@ -36,13 +36,16 @@ from __future__ import annotations +import functools import os -from typing import Optional +import time +from typing import Any, NamedTuple, Optional import torch import triton import triton.language as tl +from tensorrt_llm._torch.kernel_builder import BuildConfig, DynamicSpec, KernelBuilder from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers from .msa_utils import check_decode_span_shape @@ -1001,6 +1004,312 @@ def _check_nvfp4_inputs( ) +# The op's name in KernelBuilder, and its kernels' names in a build's key. +_BUILD_OP = "m3_sparse_decode" +_KERNEL_NAMES = { + _gqa_sparse_decode_kernel: "decode", + _gqa_sparse_decode_pair_kernel: "pair_decode", + _merge_topk_attn_out_kernel: "merge", +} + + +class _Operand(NamedTuple): + """A tensor argument of the kernels and the strides passed for it. In a build, the + value is the tensor's dtype, which Triton compiles for as a 16-byte-aligned tensor.""" + + value: Any + strides: tuple + + +class _Plan(NamedTuple): + """The launch geometry the op chooses for one batch.""" + + query_group_size: int + num_topk_chunks: int + use_linear_softmax: bool + launch_options: dict + merge_launch_options: dict + + +class _DecodeCall(NamedTuple): + """A call's arguments but the top-k table and the batch they depend on.""" + + num_heads: int + num_kv_heads: int + gqa_group_size: int + head_dim: int + max_topk: int + decode_query_len: int + sm_scale: float + use_scale: bool + widen_q: bool + compute_dtype: Any + use_pdl: bool + kv_nvfp4: bool + scale_cols: int + capability: Optional[tuple[int, int]] + q: _Operand + k: _Operand + v: _Operand + kv_scale: _Operand + k_sc: _Operand + v_sc: _Operand + k_gs: _Operand + v_gs: _Operand + block_table: _Operand + seq_lens: _Operand + output: _Operand + + +_TENSOR_FIELDS = ( + "q", + "k", + "v", + "kv_scale", + "k_sc", + "v_sc", + "k_gs", + "v_gs", + "block_table", + "seq_lens", + "output", +) +# Every other field but sm_scale, which Triton does not specialize on. +_SCALAR_FIELDS = tuple( + name for name in _DecodeCall._fields if name not in _TENSOR_FIELDS and name != "sm_scale" +) + + +def _plan(call: _DecodeCall, total_q: int, num_topk_chunks: Optional[int] = None) -> _Plan: + shape = dict( + total_q=total_q, + num_kv_heads=call.num_kv_heads, + gqa_group_size=call.gqa_group_size, + max_topk=call.max_topk, + decode_query_len=call.decode_query_len, + capability=call.capability, + ) + query_group_size = _sm100f_nvfp4_query_group_size(**shape) if call.kv_nvfp4 else 1 + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = _sm100f_nvfp4_num_topk_chunks(**shape) if call.kv_nvfp4 else None + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, call.num_kv_heads, call.max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + use_linear_softmax = call.kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax(**shape) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = _sm100f_nvfp4_launch_options(**shape) if call.kv_nvfp4 else {} + merge_launch_options = _sm100f_nvfp4_merge_launch_options(**shape) if call.kv_nvfp4 else {} + return _Plan( + query_group_size, num_topk_chunks, use_linear_softmax, launch_options, merge_launch_options + ) + + +def _launches( + call: _DecodeCall, plan: _Plan, total_q: int, topk: _Operand, o: _Operand, lse: _Operand +) -> list: + """The kernels the op launches for ``plan``, as ``(kernel, grid, args, kwargs)``: the + producer, then the split-K merge unless one chunk writes the output directly.""" + pair = plan.query_group_size == 2 + producer = _gqa_sparse_decode_pair_kernel if pair else _gqa_sparse_decode_kernel + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": plan.query_group_size, "MAX_TOPK": call.max_topk} if pair else {} + ) + pdl_launch = {"launch_pdl": True} if call.use_pdl else {} + launches = [ + ( + producer, + ((total_q // plan.query_group_size) * plan.num_topk_chunks, call.num_kv_heads), + ( + call.q.value, + call.k.value, + call.v.value, + call.kv_scale.value, + call.k_sc.value, + call.v_sc.value, + call.k_gs.value, + call.v_gs.value, + topk.value, + o.value, + lse.value, + call.block_table.value, + call.seq_lens.value, + total_q, + call.gqa_group_size, + call.head_dim, + call.max_topk, + call.sm_scale, + call.decode_query_len, + *call.q.strides[:3], + *call.k.strides[:4], + *call.v.strides[:4], + *call.k_sc.strides[:2], + *call.v_sc.strides[:2], + *topk.strides[:3], + *o.strides, + *lse.strides, + call.block_table.strides[0], + ), + dict( + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=plan.num_topk_chunks, + USE_SCALE=call.use_scale, + WIDEN_Q=call.widen_q, + COMPUTE_DTYPE=call.compute_dtype, + USE_PDL=call.use_pdl, + KV_NVFP4=call.kv_nvfp4, + USE_LINEAR_SOFTMAX=plan.use_linear_softmax, + SCALE_COLS=call.scale_cols, + DIRECT_OUTPUT=plan.num_topk_chunks == 1, + **pair_kernel_options, + **pdl_launch, + **plan.launch_options, + ), + ) + ] + if plan.num_topk_chunks != 1: + launches.append( + ( + _merge_topk_attn_out_kernel, + (total_q, call.num_heads), + ( + o.value, + lse.value, + call.output.value, + call.head_dim, + *o.strides, + *lse.strides, + *call.output.strides[:3], + ), + dict( + NUM_TOPK_CHUNKS=plan.num_topk_chunks, + USE_PDL=call.use_pdl, + **pdl_launch, + **plan.merge_launch_options, + ), + ) + ) + return launches + + +def _specialization(value: Any) -> Any: + """What Triton compiles a runtime argument for: a tensor's dtype, an integer's + equality to one and divisibility by 16, a float's kind.""" + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, bool): + return value + if isinstance(value, int): + return (value == 1, value % 16 == 0, value >= 2**31) + if isinstance(value, float): + return "float" + return str(value) + + +def _batch_operands(call: _DecodeCall, topk: _Operand, total_q: int, local_total_q: int): + """The top-k table, partial output and LSE of a batch of ``local_total_q`` query + tokens, as dtypes, for a call that has ``total_q``.""" + strides = topk.strides + # The selector's head-major table spans the step's query tokens; another layout + # does not depend on them. + if strides[0] == total_q * strides[1]: + strides = (local_total_q * strides[1], *strides[1:]) + heads, head_dim = call.num_heads, call.head_dim + o = _Operand(torch.float32, (local_total_q * heads * head_dim, heads * head_dim, head_dim, 1)) + lse = _Operand(torch.float32, (local_total_q * heads, heads, 1)) + return _Operand(topk.value.dtype, strides), o, lse + + +def _list_builds(call: _DecodeCall, topk: _Operand, total_q: int, shapes) -> list: + """One build per distinct kernel compile of the batches in ``shapes``, which + ``prebuild`` takes as ``(key, build)``.""" + operands = [getattr(call, name) for name in _TENSOR_FIELDS] + [topk] + if any(op.value.data_ptr() % 16 for op in operands): + # A build compiles for aligned tensors; this call would not use its kernels. + return [] + device = torch.cuda.current_device() + for kernel in _KERNEL_NAMES: + # Triton makes a kernel's per-device cache at its first call on the device. Make + # it on this thread: the workers' first compiles would otherwise each make one, + # and the kernels of all but the last would be lost. + caches = getattr(getattr(kernel, "fn", None), "device_caches", None) + if caches is not None: + _ = caches[device] + as_dtypes = call._replace( + **{ + name: _Operand(getattr(call, name).value.dtype, getattr(call, name).strides) + for name in _TENSOR_FIELDS + } + ) + builds = {} + for shape in shapes: + local_total_q = shape["local_batch"] * call.decode_query_len + plan = _plan(as_dtypes, local_total_q) + if local_total_q % plan.query_group_size or call.decode_query_len % plan.query_group_size: + continue + batch_topk, o, lse = _batch_operands(as_dtypes, topk, total_q, local_total_q) + if plan.num_topk_chunks == 1: + out = as_dtypes.output + o, lse = _Operand(out.value, (0, *out.strides[:3])), _Operand(out.value, (0, 0, 0)) + for launch in _launches(as_dtypes, plan, local_total_q, batch_topk, o, lse): + kernel, _, args, kwargs = launch + # The runtime arguments as Triton specializes them; the constexprs and + # launch options by value. + key = ( + _KERNEL_NAMES[kernel], + tuple(_specialization(arg) for arg in args), + tuple(sorted((name, str(value)) for name, value in kwargs.items())), + ) + builds.setdefault(key, functools.partial(_compile, launch)) + return list(builds.items()) + + +def _compile(launch: tuple) -> bool: + """Compile one launch's kernel without running it; True when that took a compile.""" + kernel, grid, args, kwargs = launch + t0 = time.perf_counter() + kernel.warmup(*args, grid=grid, **kwargs) + # A kernel this process or Triton's cache already holds loads in milliseconds. + return time.perf_counter() - t0 >= 0.05 + + +def _queue_builds( + call: _DecodeCall, topk: _Operand, total_q: int, max_num_requests: Optional[int] +) -> None: + """In a KernelBuilder build pass, queue the compiles that later calls with this + call's other arguments launch, for every batch of up to ``max_num_requests``.""" + local_batch = total_q // call.decode_query_len + largest = max(local_batch, max_num_requests or local_batch) + # The batch-dependent top-k stride is left out, and so is the batch itself. + operands = [getattr(call, name) for name in _TENSOR_FIELDS] + operands.append(_Operand(topk.value, topk.strides[1:])) + signature = ( + tuple( + (op.value.dtype, op.value.data_ptr() % 16 == 0, tuple(map(_specialization, op.strides))) + for op in operands + ), + tuple(str(getattr(call, name)) for name in _SCALAR_FIELDS), + ) + KernelBuilder.get().prebuild( + _BUILD_OP, + signature, + BuildConfig((DynamicSpec("local_batch", tuple(range(1, largest + 1))),)), + {"local_batch": local_batch}, + functools.partial(_list_builds, call, topk, total_q), + ) + + @torch.no_grad() def minimax_m3_sparse_attn_decode( q: torch.Tensor, # [total_q, num_heads, head_dim] @@ -1019,6 +1328,7 @@ def minimax_m3_sparse_attn_decode( k_global_scale: Optional[torch.Tensor] = None, v_global_scale: Optional[torch.Tensor] = None, num_topk_chunks: Optional[int] = None, + max_num_requests: Optional[int] = None, ) -> None: """Block-sparse GQA decode attention, written into output in place. @@ -1037,6 +1347,11 @@ def minimax_m3_sparse_attn_decode( num_topk_chunks overrides the split-K factor and exists for tests: the merged result must not depend on it. + + In a KernelBuilder build pass, a call also queues on KernelBuilder's worker + threads the compiles of the kernels that calls with its other arguments launch, + for every batch of up to max_num_requests requests (only its own without it). + In a rehearsal, it then leaves output zero and launches nothing. """ total_q, num_heads, head_dim = q.shape num_kv_heads = int(k_paged.shape[1]) @@ -1096,200 +1411,85 @@ def minimax_m3_sparse_attn_decode( else: compute_dtype = tl.float32 - query_group_size = ( - _sm100f_nvfp4_query_group_size( - total_q=total_q, - num_kv_heads=num_kv_heads, - gqa_group_size=gqa_group_size, - max_topk=max_topk, - decode_query_len=decode_query_len, - capability=capability, - ) - if kv_nvfp4 - else 1 + def operand(tensor: torch.Tensor) -> _Operand: + return _Operand(tensor, tuple(tensor.stride())) + + call = _DecodeCall( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + head_dim=head_dim, + max_topk=max_topk, + decode_query_len=decode_query_len, + sm_scale=sm_scale, + use_scale=use_scale, + widen_q=widen_q, + compute_dtype=compute_dtype, + use_pdl=_pdl_enabled(), + kv_nvfp4=kv_nvfp4, + scale_cols=scale_cols, + capability=capability if kv_nvfp4 else None, + q=operand(q), + k=operand(k_paged), + v=operand(v_paged), + kv_scale=operand(scale_arg), + k_sc=operand(k_sc_arg), + v_sc=operand(v_sc_arg), + k_gs=operand(k_gs_arg), + v_gs=operand(v_gs_arg), + block_table=operand(block_table), + seq_lens=operand(seq_lens), + output=operand(output), ) - if num_topk_chunks is None: - if query_group_size == 2: - measured_chunks = 8 - else: - measured_chunks = ( - _sm100f_nvfp4_num_topk_chunks( - total_q=total_q, - num_kv_heads=num_kv_heads, - gqa_group_size=gqa_group_size, - max_topk=max_topk, - decode_query_len=decode_query_len, - capability=capability, - ) - if kv_nvfp4 - else None - ) - num_topk_chunks = ( - measured_chunks - if measured_chunks is not None - else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + topk = operand(topk_idx) + builder = KernelBuilder.get() + if ( + num_topk_chunks is None + and builder.active() + and not torch.cuda.is_current_stream_capturing() + ): + _queue_builds(call, topk, total_q, max_num_requests) + if builder.rehearsing(): + # A rehearsal's outputs are discarded, and its kernels build on KernelBuilder's + # worker threads. + output.zero_() + return + + plan = _plan(call, total_q, num_topk_chunks) + if total_q % plan.query_group_size or decode_query_len % plan.query_group_size: + raise ValueError( + f"query group {plan.query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." ) - elif num_topk_chunks & (num_topk_chunks - 1): - raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") - - direct_output = num_topk_chunks == 1 - if direct_output: + if plan.num_topk_chunks == 1: # With one chunk, acc_o is already normalized and final. Point the # producer at the caller's output and compile away the LSE store. # Keeping a real dummy pointer preserves Triton's pointer contract. - o_arg = output - lse_arg = output - o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) - lse_strides = (0, 0, 0) + o = _Operand(output, (0, output.stride(0), output.stride(1), output.stride(2))) + lse = _Operand(output, (0, 0, 0)) else: # Persistent arena rather than torch.empty, so the partials keep one # address across CUDA graph replays. FP32 keeps the split-K factor from # perturbing the merged result. reserve = torch.cuda.is_current_stream_capturing() - o_arg = get_memory_buffers().get_buffer( - [num_topk_chunks, total_q, num_heads, head_dim], - torch.float32, - buffer_name="m3_sparse_decode_o_partial", - reserve_buffer=reserve, - ) - lse_arg = get_memory_buffers().get_buffer( - [num_topk_chunks, total_q, num_heads], - torch.float32, - buffer_name="m3_sparse_decode_lse_partial", - reserve_buffer=reserve, - ) - o_strides = tuple(o_arg.stride()) - lse_strides = tuple(lse_arg.stride()) - - use_pdl = _pdl_enabled() - pdl_launch = {"launch_pdl": True} if use_pdl else {} - use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( - total_q=total_q, - num_kv_heads=num_kv_heads, - gqa_group_size=gqa_group_size, - max_topk=max_topk, - decode_query_len=decode_query_len, - capability=capability, - ) - if query_group_size == 2: - launch_options = {"num_warps": 4, "num_stages": 1} - else: - launch_options = ( - _sm100f_nvfp4_launch_options( - total_q=total_q, - num_kv_heads=num_kv_heads, - gqa_group_size=gqa_group_size, - max_topk=max_topk, - decode_query_len=decode_query_len, - capability=capability, + o = operand( + get_memory_buffers().get_buffer( + [plan.num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, ) - if kv_nvfp4 - else {} ) - - # The pair path merges the two sorted selector rows on device. Equal blocks - # share one K/V dequantization; divergent blocks update only their owning - # query, so the optimization does not rely on runtime selection equality. - if total_q % query_group_size or decode_query_len % query_group_size: - raise ValueError( - f"query group {query_group_size} must divide total_q={total_q} and " - f"decode_query_len={decode_query_len}." - ) - - producer_kernel = ( - _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel - ) - pair_kernel_options = ( - {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} - if query_group_size == 2 - else {} - ) - producer_kernel[ - (total_q // query_group_size) * num_topk_chunks, - num_kv_heads, - ]( - q, - k_paged, - v_paged, - scale_arg, - k_sc_arg, - v_sc_arg, - k_gs_arg, - v_gs_arg, - topk_idx, - o_arg, - lse_arg, - block_table, - seq_lens, - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - q.stride(0), - q.stride(1), - q.stride(2), - k_paged.stride(0), - k_paged.stride(1), - k_paged.stride(2), - k_paged.stride(3), - v_paged.stride(0), - v_paged.stride(1), - v_paged.stride(2), - v_paged.stride(3), - k_sc_arg.stride(0), - k_sc_arg.stride(1), - v_sc_arg.stride(0), - v_sc_arg.stride(1), - topk_idx.stride(0), - topk_idx.stride(1), - topk_idx.stride(2), - *o_strides, - *lse_strides, - block_table.stride(0), - BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, - NUM_TOPK_CHUNKS=num_topk_chunks, - USE_SCALE=use_scale, - WIDEN_Q=widen_q, - COMPUTE_DTYPE=compute_dtype, - USE_PDL=use_pdl, - KV_NVFP4=kv_nvfp4, - USE_LINEAR_SOFTMAX=use_linear_softmax, - SCALE_COLS=scale_cols, - DIRECT_OUTPUT=direct_output, - **pair_kernel_options, - **pdl_launch, - **launch_options, - ) - if not direct_output: - merge_launch_options = ( - _sm100f_nvfp4_merge_launch_options( - total_q=total_q, - num_kv_heads=num_kv_heads, - gqa_group_size=gqa_group_size, - max_topk=max_topk, - decode_query_len=decode_query_len, - capability=capability, + lse = operand( + get_memory_buffers().get_buffer( + [plan.num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, ) - if kv_nvfp4 - else {} - ) - _merge_topk_attn_out_kernel[(total_q, num_heads)]( - o_arg, - lse_arg, - output, - head_dim, - *o_strides, - *lse_strides, - output.stride(0), - output.stride(1), - output.stride(2), - NUM_TOPK_CHUNKS=num_topk_chunks, - USE_PDL=use_pdl, - **pdl_launch, - **merge_launch_options, ) + for kernel, grid, args, kwargs in _launches(call, plan, total_q, topk, o, lse): + kernel[grid](*args, **kwargs) __all__ = [ diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py index e26270df980c..7261db789b50 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py @@ -965,6 +965,81 @@ def test_nvfp4_sparse_decode_fp8_q_matches_reference(): torch.testing.assert_close(out.float(), expected, rtol=3e-2, atol=3e-2) +@skip_not_sm100 +def test_a_rehearsal_has_the_kernels_of_every_batch_size_built_for_the_later_launches(monkeypatch): + """One call in a KernelBuilder rehearsal launches nothing and has KernelBuilder's + workers compile the kernels of every batch of up to max_num_requests requests; each + batch then launches without a compile and matches the reference.""" + import threading + from collections import defaultdict + + import triton + + from tensorrt_llm._torch import kernel_builder + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import ( + triton_sparse_decode, + ) + + # The op's kernels start with nothing compiled in this process. + for kernel in triton_sparse_decode._KERNEL_NAMES: + monkeypatch.setattr(kernel.fn, "device_caches", defaultdict(kernel.fn.create_binder)) + compiled_on = [] + + def listener(*, src, metadata, metadata_group, times, cache_hit): + compiled_on.append(threading.current_thread().name) + + monkeypatch.setattr(triton.knobs.compilation, "listener", listener) + builder = kernel_builder.KernelBuilder(threads=4, enabled=True, rehearse=True) + monkeypatch.setattr(kernel_builder.KernelBuilder, "_instance", builder) + + max_batch = 6 + case = _make_nvfp4_inputs( + [700] * max_batch, num_kv_heads=4, group=16, decode_query_len=4, seed=29 + ) + q = case.q.to(torch.float8_e4m3fn) + + def run(local_batch, **kwargs): + total_q = local_batch * case.decode_query_len + out = torch.full((total_q, *q.shape[1:]), float("nan"), device="cuda", dtype=torch.bfloat16) + minimax_m3_sparse_attn_decode( + q[:total_q].contiguous(), + case.k_paged, + case.v_paged, + case.topk_idx[:, :total_q].contiguous(), + case.block_table[:local_batch], + case.seq_lens[:local_batch], + sm_scale=HEAD_DIM**-0.5, + output=out, + decode_query_len=case.decode_query_len, + **case.scales, + **kwargs, + ) + return out + + with builder.rehearsal(): + assert (run(1, max_num_requests=max_batch) == 0).all() + builder.join() + stats = builder.stats("m3_sparse_decode") + assert stats.submitted > 1 and stats.failed == 0 + assert compiled_on and all(name.startswith("tllm-kernel-builder") for name in compiled_on) + + compiled_on.clear() + for local_batch in range(1, max_batch + 1): + total_q = local_batch * case.decode_query_len + expected = _reference_sparse_decode( + q[:total_q].to(torch.bfloat16), + case.k_ref, + case.v_ref, + case.topk_idx[:, :total_q], + case.block_table[:local_batch], + case.seq_lens[:local_batch], + sm_scale=HEAD_DIM**-0.5, + decode_query_len=case.decode_query_len, + ) + torch.testing.assert_close(run(local_batch).float(), expected, rtol=3e-2, atol=3e-2) + assert compiled_on == [] + + @skip_not_sm100 def test_nvfp4_paged_scale_layouts_are_the_ones_the_kernel_reads(): """Lock the two block-scale layouts and the nibble order.