[None][test] Add a warmup pass before perf-sanity measured clients

tests/integration/defs/perf/README.md

@@ -32,6 +32,9 @@ If a CI run fails `sanity_perf_check.py`, it will upload a patch file as an arti
 
 This patch will only update the metrics that had a difference which was over the accepted threshold. The patch will also remove/add metrics according to the removed or added tests.
 
+## Warmup pass
+Before measuring, `test_perf.py` warms up with one un-measured wave of the benchmark's own traffic (its concurrency, or `min(reqs, maxbs)` without one), so JIT kernels for the real prefill shapes compile outside the measured window: `bench` cases raise `trtllm-bench --warmup` to that wave, and `serve` cases send it with other prompts before each client. `TRTLLM_PERF_WARMUP_PASS=0` turns it off.
+
 ## Running locally
 Given a `target_perf_csv_path` you can compare it to another perf csv file.
 First make sure you install the dependencies:

tests/integration/defs/perf/README_test_perf_sanity.md

@@ -338,6 +338,17 @@ queue (nvbugs/6564113). Notes:
   connection-warmup iterations; start the range later to profile only measured
   iterations.
 
+### Warmup pass
+
+Before each measured client of a normal aggregated, `ctx_only` or disaggregated `e2e`
+test, the harness sends one un-measured wave (`concurrency` requests) of the client's
+own `benchmark_serving` command with `--seed` + 1, so JIT kernels for the real prefill
+shapes compile outside the measured window. `gen_only` modes, AgentX and accuracy-only
+clients get none. Its log (`trtllm-benchmark-warmup-pass.*.log`) is not parsed, and its
+records stay out of `d_tb_*` and the device step time; with the time breakdown on, only
+the first measured client gets one. `TRTLLM_PERF_WARMUP_PASS=0` turns it off (locally,
+via `EXTRA_CONTAINER_EXPORTS`).
+
 ### The optional instrumentation modifier
 
 The three shapes that read a disaggregated config (2, 3 and 4 above) take one
@@ -399,8 +410,9 @@ to run after the writers are done. Only the **gen** workers announce that
 workers and the disaggregated server do not. `wait_for_perf_metrics_files` therefore
 polls the discovered set until its total size holds still for
 `PERF_METRICS_SETTLE_SECONDS`, bounded by `PERF_METRICS_SETTLE_TIMEOUT`, before
-anything is read, and then compares the largest file's complete-record count against
-the client's `--num-prompts`. Both a still-growing set at the timeout and a census
+anything is read, and then compares the census (the complete records of one server
+role, summed over its files) against the client's `--num-prompts`. Both a
+still-growing set at the timeout and a census
 shortfall are logged as warnings, not failures — a nearly-complete file still yields
 usable statistics, and the numbers are not regression-gated. This matters because a
 truncated read is otherwise **invisible**: all 108 fields are still populated and the
@@ -409,13 +421,14 @@ mid-record) costs only that line, never the whole file: dropping the disagg serv
 file would silently reroute the per-request groups to same-role worker fallbacks,
 which produce plausible values for the wrong phase.
 
-The servers also append records of the e2e connection warmup's requests to these
-files. The harness therefore waits until the largest file holds a record for every
-successful warmup request and the files hold still, then notes each file's line count; the aggregation skips those leading lines and leaves
+The servers also append records of the e2e connection warmup's and the warmup pass's
+requests to these files. The harness therefore waits until the census holds a record
+for every successful warmup request and the files hold still, then notes each
+file's line count; the aggregation skips those leading lines and leaves
 them out of the census. Unlike the measured run's census, this one is a failure: if
 the count is not reached within `PERF_METRICS_SETTLE_TIMEOUT`, the test fails rather
 than mixing late warmup records into the breakdown. The log line
-`Time breakdown: excluded connection-warmup records from {...}` prints the counts as
+`Time breakdown: excluded warmup records from {...}` prints the counts as
 JSON. To reproduce the uploaded numbers offline, run
 `time_breakdown_metrics.py --mode e2e --output-dir <test output dir>/perf_metrics
 --drop-warmup-request --leading-lines-to-skip '<that JSON>'`.

tests/integration/defs/perf/test_perf.py

@@ -24,7 +24,7 @@
 import os
 import re
 import shutil
-from typing import Dict, List, NamedTuple
+from typing import Dict, List, NamedTuple, Optional
 
 import pytest
 import yaml
@@ -40,6 +40,8 @@
 from .utils import (AbstractPerfScriptTestClass, PerfBenchScriptTestCmds,
                     PerfMetricType, PerfServeScriptTestCmds,
                     generate_test_nodes)
+from .warmup_pass import (WarmupPass, build_warmup_pass_cmd, one_wave_requests,
+                          warmup_pass_enabled)
 
 if not hasattr(re, "Pattern"):
     re.Pattern = type(re.compile(""))
@@ -49,6 +51,14 @@
 
 TIMING_CACHE_DIR = os.environ.get("TIMING_CACHE_DIR", "")
 
+# trtllm-bench throughput's --warmup default.
+TRTLLM_BENCH_DEFAULT_WARMUP = 2
+
+
+class NotEnoughSamplesError(ValueError):
+    """A real-text source holds fewer usable articles than a dataset file asks for."""
+
+
 # Models served with the Nemotron reasoning/tool parsers, a long max model len
 # and a longer server-start timeout. "nano-v3" and "nemotron-v3" are two
 # registered names for the same reasoning parser, so Lightning belongs here too.
@@ -1115,6 +1125,29 @@ def get_prepare_data_command(self, engine_dir, input_len,
 
         return data_cmd
 
+    def _get_bench_warmup_requests(self) -> int:
+        """Return trtllm-bench's --warmup for this case, or 0 for its default.
+
+        trtllm-bench runs these requests in its own process before the measured
+        ones, at the same concurrency, and reports none of them. With the warmup
+        pass on, they are one wave of the measured traffic (see warmup_pass.py),
+        and never fewer than without it. They are drawn from the measured
+        dataset, though, so a case whose config reuses KV cache blocks keeps the
+        fewer requests it has without the pass: otherwise its measured requests
+        would hit the warmup's cached prefixes.
+        """
+        without_pass = 2 * self._config.num_gpus if self._config.num_gpus > 1 else 0
+        if not warmup_pass_enabled():
+            return without_pass
+        kv_cache_config = self._get_model_yaml_config().get(
+            'kv_cache_config') or {}
+        if kv_cache_config.get('enable_block_reuse'):
+            return without_pass
+        wave = one_wave_requests(self._config.concurrency,
+                                 self._config.num_reqs,
+                                 self._config.max_batch_size)
+        return max(without_pass or TRTLLM_BENCH_DEFAULT_WARMUP, wave)
+
     def get_trtllm_bench_command(self, engine_dir):
         model_dir = self.get_trtllm_bench_model()
         model_name = self._config.model_name
@@ -1149,8 +1182,9 @@ def get_trtllm_bench_command(self, engine_dir):
             benchmark_cmd += [f"--pp={self._config.pp_size}"]
         if self._config.streaming == "streaming":
             benchmark_cmd += [f"--streaming"]
-        if self._config.num_gpus > 1:
-            benchmark_cmd += [f"--warmup={2 * self._config.num_gpus}"]
+        warmup_requests = self._get_bench_warmup_requests()
+        if warmup_requests:
+            benchmark_cmd += [f"--warmup={warmup_requests}"]
 
         # Add extra-llm-api-config.yml for the PyTorch backend.
         if self._config.backend == "pytorch":
@@ -1227,9 +1261,20 @@ def get_trtllm_serve_server_command(self, engine_dir):
 
         return server_cmd
 
-    def generate_trtllm_custom_dataset(self, dst_dataset_path: str,
-                                       input_len: int, output_len: int,
-                                       dataset_source: str):
+    def generate_trtllm_custom_dataset(self,
+                                       dst_dataset_path: str,
+                                       input_len: int,
+                                       output_len: int,
+                                       dataset_source: str,
+                                       num_requests: Optional[int] = None,
+                                       skip_requests: int = 0):
+        """Write a real-text dataset of fixed-length prompts.
+
+        Args:
+            num_requests: Samples to write; the case's request count if None.
+            skip_requests: Leading usable articles to skip, so that another
+                file can hold prompts this one does not.
+        """
         # Currently only support cnn_dailymail dataset source.
         if dataset_source != "cnn_dailymail":
             raise ValueError(
@@ -1260,8 +1305,9 @@ def generate_trtllm_custom_dataset(self, dst_dataset_path: str,
         if not os.path.exists(os.path.dirname(dst_dataset_path)):
             os.makedirs(os.path.dirname(dst_dataset_path), exist_ok=True)
 
-        num_reqs = self._config.num_reqs
+        num_reqs = self._config.num_reqs if num_requests is None else num_requests
         req_count = 0
+        num_skipped = 0
         with open(dst_dataset_path, "w", encoding="utf-8") as f:
             for req in dataset:
                 article = req.get("article")
@@ -1281,6 +1327,9 @@ def generate_trtllm_custom_dataset(self, dst_dataset_path: str,
                         prompt_ids = (prompt_ids * repeats)[:input_len]
                     elif len(prompt_ids) > input_len:
                         prompt_ids = prompt_ids[:input_len]
+                if num_skipped < skip_requests:
+                    num_skipped += 1
+                    continue
                 prompt_text = tokenizer.decode(prompt_ids,
                                                skip_special_tokens=False)
                 if self._config.model_name in DEEPSEEK_V41_MODELS:
@@ -1322,7 +1371,7 @@ def generate_trtllm_custom_dataset(self, dst_dataset_path: str,
                     break
 
         if req_count < num_reqs:
-            raise ValueError(
+            raise NotEnoughSamplesError(
                 f"Cannot sample enough requests from cnn_dailymail: requested={num_reqs}, sampled={req_count}"
             )
         print_info(f"Generated {req_count} samples from {dataset_source} to "
@@ -1416,6 +1465,52 @@ def get_trtllm_serve_client_command(self,
             client_cmd += ["--trust-remote-code"]
         return client_cmd
 
+    def _get_serve_warmup_pass(self, engine_dir: str, input_len: int,
+                               output_len: int,
+                               client_cmd: List[str]) -> Optional[WarmupPass]:
+        """Return the warmup pass before a serve client, or None; see warmup_pass.py.
+
+        A real-text client gets a prompt file of its own, with the articles
+        after the measured ones. The measured file holds exactly the measured
+        prompts, so another shuffle of it would resend them, and KV cache block
+        reuse would then serve their prefixes to the measured requests. If the
+        source has no articles to spare, the client gets no pass.
+        """
+        if not warmup_pass_enabled():
+            return None
+        num_requests = one_wave_requests(self._config.concurrency,
+                                         self._config.num_reqs,
+                                         self._config.max_batch_size)
+        if num_requests < 1:
+            return None
+        dataset_source = SPEC_DEC_REAL_DATASET_MODELS.get(
+            self._config.model_name)
+        if dataset_source is not None:
+            warmup_dataset_path = os.path.join(
+                engine_dir,
+                f"dataset_custom_{input_len}_{output_len}_warmup_pass.jsonl")
+            try:
+                self.generate_trtllm_custom_dataset(
+                    warmup_dataset_path,
+                    input_len,
+                    output_len,
+                    dataset_source=dataset_source,
+                    num_requests=num_requests,
+                    skip_requests=self._config.num_reqs)
+            except NotEnoughSamplesError as error:
+                print_warning(
+                    f"Warmup pass: none for this client, since {dataset_source} "
+                    f"has no {num_requests} articles to spare after the "
+                    f"{self._config.num_reqs} measured ones: {error}")
+                return None
+            client_cmd = self.get_trtllm_serve_client_command(
+                engine_dir,
+                input_len,
+                output_len,
+                real_dataset_path=warmup_dataset_path)
+        return WarmupPass(build_warmup_pass_cmd(client_cmd, num_requests),
+                          num_requests)
+
     def get_commands(self):
         num_gpus = self._config.num_gpus
 
@@ -1428,6 +1523,7 @@ def get_commands(self):
         if self._config.runtime == "serve":
             server_cmd = self.get_trtllm_serve_server_command(engine_dir)
             client_cmds = []
+            warmup_passes = []
             data_cmds = []
             for bs in self._config.batch_sizes:
                 for len_idx, input_len in enumerate(self._config.input_lens):
@@ -1453,6 +1549,9 @@ def get_commands(self):
                         output_len,
                         real_dataset_path=real_dataset_path)
                     client_cmds.append(client_cmd)
+                    warmup_passes.append(
+                        self._get_serve_warmup_pass(engine_dir, input_len,
+                                                    output_len, client_cmd))
             server_env = os.environ.copy()
             if self._config.model_name.startswith("qwen3.8_flash_next_"):
                 server_env["TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD"] = "1"
@@ -1490,7 +1589,8 @@ def get_commands(self):
                                            client_cmds=client_cmds,
                                            data_cmds=data_cmds,
                                            server_env=server_env,
-                                           server_timeout=server_timeout)
+                                           server_timeout=server_timeout,
+                                           warmup_passes=warmup_passes)
 
         # PyTorch loads model checkpoints directly.
         build_cmd = []