[None][fix] Warm up flashinfer sampling module during engine warmup

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)
 

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,