[TRTLLM-17679][feat] KernelBuilder and the autotuner's build state: build kernels on worker threads while the weights load

docs/source/llm-api/index.md

@@ -203,6 +203,41 @@ curl http://localhost:8000/server_info
 }
 ```
 
+### Kernel builds ahead of first use
+
+An op that compiles its kernels on first use lists them when it runs, for every shape it
+sweeps down from that call's and has not listed yet, as the autotuner sweeps its tuning
+buckets; worker threads build them in parallel. An op that the autotuner tunes has three
+states: the autotuner warmup tunes it, a build state builds the kernels of every shape,
+runner and tactic its tuning will time on the same worker threads without timing them, and
+serving runs the tuned tactic.
+
+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.
+
+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
+(`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 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. |
+
 ## Unified Cache for Compiled Artifacts
 
 Set `TRTLLM_CACHE_DIR` before importing TensorRT-LLM to place its compiled

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,30 @@ 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) -> 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."""
+    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()

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,
         )
 

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]
@@ -340,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,