diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 17b2dc5a6c2a..9cbb8fc496f8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -90,6 +90,7 @@ PeftCacheManager, ResourceManager, ResourceManagerType) from .sampler import SampleStateTensors +from .sampler.ops.flashinfer import warmup_sampling_module from .scheduler import ScheduledRequests from .trace_log_utils import log_mem_snapshot @@ -1154,6 +1155,12 @@ def warmup(self, resource_manager: ResourceManager) -> None: Orchestrates the warmup process by calling specialized warmup methods for torch.compile, the autotuner, and CUDA graphs. """ + # Ahead of the early returns below, since it holds regardless of why + # warmup is skipped: only the advanced-sampling CUDA graph capture pass + # exercises the non-greedy sampler, so with cuda_graph_config=None + # flashinfer's sampling kernels would be JIT-built mid-serving. + warmup_sampling_module() + kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index a0e1c7c44322..6ee643336eb7 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -31,6 +31,7 @@ import torch from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from tensorrt_llm.logger import logger if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode @@ -353,6 +354,30 @@ def sample_from_logits_op( return sampling_from_probs_op(probs, seed=seed, offset=offset) +@_compiler_disable +def warmup_sampling_module() -> None: + """Build flashinfer's sampling kernels now instead of on first use. + + flashinfer ships these kernels as source, so the first sampling call runs + nvcc inline (~90s with a cold cache). Doing that inside the executor loop + stalls the rank past the hang detector's threshold. Cheap no-op once built. + + Best-effort, like the neighbouring pre-JIT warmups: this runs ahead of every + guard in ``warmup()``, so a broken JIT toolchain must not abort startup for + deployments that never sample non-greedily. Ones that do will raise at the + real call site instead. + """ + if not IS_FLASHINFER_AVAILABLE: + return + try: + flashinfer.sampling.get_sampling_module() + except Exception as e: # noqa: BLE001 + logger.warning( + "flashinfer sampling module prewarm failed; it will be built " + f"lazily on first use. {type(e).__name__}: {e}" + ) + + @torch.compile(options={"max-autotune": True}) def sampling_batch_spec_dec_one_model_for_rejection( logits: torch.Tensor,