diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 8f246f744287..60c5f724dbf8 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -1135,6 +1135,10 @@ def getPytestBaseCommandLine( extraInternalEnv += " CBTS_COVERAGE_CONFIG=${coverageConfigFile}" extraInternalEnv += " CBTS_MARKER_FILE=${outputPath}/cbts_current_test.txt" } + // Skip TRTLLM-Gen FMHA JIT warmup for non-performance tests. + if (!stageName.contains("-Perf-") && !stageName.contains("-PerfSanity-")) { + extraInternalEnv += " TLLM_ENABLE_TRTLLM_GEN_FMHA_JIT_WARMUP=0" + } // Container port allocation environment variables for avoiding port conflicts def portEnvVars = "" diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d83f51a1bbe3..fd97b8620d28 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1280,6 +1280,10 @@ def _run_attention_warmup(self, if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata): return + if os.environ.get("TLLM_ENABLE_TRTLLM_GEN_FMHA_JIT_WARMUP", "1") != "1": + logger.info("Skipping TRTLLM-Gen FMHA JIT warmup") + return + @contextlib.contextmanager def trtllm_gen_fmha_jit_warmup(): previous = self._trtllm_gen_jit_warmup