Stabilize perf warmup

tests/performance_tests/client/static_benchmark.py

@@ -2,8 +2,9 @@
 """Static throughput/latency benchmark against an OpenAI-compatible completions server.
 
 Fires --batch-size requests simultaneously via asyncio.gather, waits for all to
-finish, and reports throughput, latency (avg/p50/p99), and TPOT. Iterates over
-warmup + timed batches and emits a JSON results file consumable by
+finish, and reports throughput, latency (avg/p50/p99), and TPOT. Warmup batches
+are widened to cover every data-parallel worker when needed. Timed batches keep
+the requested batch size and emit a JSON results file consumable by
 compare_to_baseline.py.
 
 Hits the server's POST /v1/completions endpoint directly via aiohttp — no
@@ -92,10 +93,15 @@ async def _run_batch(
     url: str,
     prompts: list[str],
     iter_start_index: int,
+    request_count: int | None = None,
 ) -> tuple[list[int], list[int], list[float], float]:
-    """Fire batch_size requests in parallel. Cycles through `prompts` deterministically
-    starting at `iter_start_index` so each timed iteration sees the same prompt
-    distribution (reduces run-to-run variance for gsm8k mode)."""
+    """Fire requests in parallel, defaulting to the measured batch size.
+
+    Cycles through `prompts` deterministically starting at `iter_start_index`
+    so each timed iteration sees the same prompt distribution (reduces
+    run-to-run variance for gsm8k mode).
+    """
+    request_count = args.batch_size if request_count is None else request_count
     t0 = time.perf_counter()
     results = await asyncio.gather(
         *[
@@ -107,7 +113,7 @@ async def _run_batch(
                 args.num_output_tokens,
                 args.temperature,
             )
-            for i in range(args.batch_size)
+            for i in range(request_count)
         ]
     )
     wall = time.perf_counter() - t0
@@ -122,8 +128,14 @@ def _percentile(sorted_values: list[float], pct: float) -> float:
     return sorted_values[idx]
 
 
+def _get_warmup_batch_size(batch_size: int, data_parallel_size: int) -> int:
+    """Keep batch-shape warmup while issuing enough requests to cover DP workers."""
+    return max(batch_size, data_parallel_size)
+
+
 async def main(args: argparse.Namespace) -> dict:
     url = f"{args.server_url.rstrip('/')}/completions"
+    warmup_batch_size = _get_warmup_batch_size(args.batch_size, args.data_parallel_size)
 
     if args.dataset == "gsm8k":
         prompts = _load_gsm8k_prompts()
@@ -140,26 +152,34 @@ async def main(args: argparse.Namespace) -> dict:
     print(f"Dataset         : {prompt_source}")
     print(f"Output tokens   : {args.num_output_tokens}")
     print(f"Warmup iters    : {args.num_warmup_iters}")
+    print(f"Warmup batch    : {warmup_batch_size}")
     print(f"Timed iters     : {args.num_iters}", flush=True)
 
     connector = aiohttp.TCPConnector(limit=0)
     async with aiohttp.ClientSession(connector=connector) as session:
-        cursor = 0
+        warmup_cursor = 0
         for i in range(args.num_warmup_iters):
-            print(f"\nWarmup {i + 1}/{args.num_warmup_iters}...", flush=True)
-            await _run_batch(session, args, url, prompts, cursor)
-            cursor += args.batch_size
+            print(
+                f"\nWarmup {i + 1}/{args.num_warmup_iters} (batch={warmup_batch_size})...",
+                flush=True,
+            )
+            await _run_batch(
+                session, args, url, prompts, warmup_cursor, request_count=warmup_batch_size
+            )
+            warmup_cursor += warmup_batch_size
 
         all_wall: list[float] = []
         all_output_tokens: list[int] = []
         all_input_tokens: list[int] = []
         all_latencies: list[float] = []
+        # Keep the timed prompt sequence stable when widening warmup batches.
+        timed_cursor = args.num_warmup_iters * args.batch_size
 
         for i in range(args.num_iters):
             input_counts, output_counts, latencies, wall = await _run_batch(
-                session, args, url, prompts, cursor
+                session, args, url, prompts, timed_cursor
             )
-            cursor += args.batch_size
+            timed_cursor += args.batch_size
             total_out = sum(output_counts)
             all_wall.append(wall)
             all_output_tokens.append(total_out)
@@ -226,6 +246,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--num-output-tokens", type=int, default=128)
     parser.add_argument("--temperature", type=float, default=0.0)
     parser.add_argument("--num-warmup-iters", type=int, default=2)
+    parser.add_argument(
+        "--data-parallel-size",
+        type=int,
+        default=1,
+        help="Number of coordinator-addressable data-parallel workers. Warmup batches "
+        "use at least this many concurrent requests; timed batch size is unchanged.",
+    )
     parser.add_argument("--num-iters", type=int, default=5)
     parser.add_argument(
         "--output-json",

tests/performance_tests/shell_test_utils/run_perf_test.sh

@@ -95,13 +95,17 @@ mapfile -t BATCH_SIZES < <("$YQ" '.BATCH_SIZES[]' "$CONFIG_PATH")
 # and MoE-with-DP=1 picks up EP correctly.
 GROUP_SIZE=$((DP > EP ? DP : EP))
 WORLD_SIZE=$((TP * PP * GROUP_SIZE))
+# The inference coordinator uses the dense-model DP group. EP-only configs
+# therefore expose GROUP_SIZE workers even when the YAML DP value is one.
+COORDINATOR_WORKERS=$GROUP_SIZE
 ARGS_FILE="$PERF_DIR/server/model_args/${MODEL}.args"
 if [[ ! -f "$ARGS_FILE" ]]; then
     echo "[run_perf_test] error: model args file $ARGS_FILE not found" >&2
     exit 2
 fi
 
 echo "[run_perf_test] MODEL=$MODEL  TP=$TP PP=$PP DP=$DP EP=$EP  world_size=$WORLD_SIZE  dataset=$DATASET"
+echo "[run_perf_test] coordinator workers: $COORDINATOR_WORKERS"
 echo "[run_perf_test] ISL=$NUM_INPUT_TOKENS  OSL=$NUM_OUTPUT_TOKENS"
 echo "[run_perf_test] batch sizes: ${BATCH_SIZES[*]}"
 
@@ -257,6 +261,7 @@ for BS in "${BATCH_SIZES[@]}"; do
         --num-input-tokens "$NUM_INPUT_TOKENS" \
         --num-output-tokens "$NUM_OUTPUT_TOKENS" \
         --num-warmup-iters "$NUM_WARMUP_ITERS" \
+        --data-parallel-size "$COORDINATOR_WORKERS" \
         --num-iters "$NUM_TIMED_ITERS" \
         --output-json "$RESULTS_JSON" \
         2>&1 | tee -a "$RESULTS_ROOT/benchmark.log"

tests/performance_tests/test_cases/gpt/gpt_583m_perf/baseline_values.json

@@ -2,47 +2,51 @@
   "h100": {
     "batch_1": {
       "batch_size": 1,
-      "num_input_tokens": 512,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
       "num_output_tokens": 128,
       "num_iters": 5,
-      "throughput_tok_per_sec": 22.08203581766323,
-      "avg_latency_ms": 5796.511190757155,
-      "p50_latency_ms": 5842.958671972156,
-      "p99_latency_ms": 5919.582245871425,
-      "tpot_ms_per_tok": 45.28567964734975
+      "throughput_tok_per_sec": 44.399582392645556,
+      "avg_latency_ms": 2882.8655768185854,
+      "p50_latency_ms": 2883.104130625725,
+      "p99_latency_ms": 2886.1329462379217,
+      "tpot_ms_per_tok": 22.522734361700714
     },
     "batch_8": {
       "batch_size": 8,
-      "num_input_tokens": 512,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
       "num_output_tokens": 128,
       "num_iters": 5,
-      "throughput_tok_per_sec": 357.49352020243185,
-      "avg_latency_ms": 2808.0778209026903,
-      "p50_latency_ms": 2813.233459368348,
-      "p99_latency_ms": 2896.668652072549,
-      "tpot_ms_per_tok": 22.378027986269444
+      "throughput_tok_per_sec": 343.99702071382734,
+      "avg_latency_ms": 2918.7685920856893,
+      "p50_latency_ms": 2911.766432225704,
+      "p99_latency_ms": 3012.841146439314,
+      "tpot_ms_per_tok": 23.25601536722388
     },
     "batch_32": {
       "batch_size": 32,
-      "num_input_tokens": 512,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
       "num_output_tokens": 128,
       "num_iters": 5,
-      "throughput_tok_per_sec": 1432.3391490750541,
-      "avg_latency_ms": 2812.452972715255,
-      "p50_latency_ms": 2819.5638693869114,
-      "p99_latency_ms": 2865.5092362314463,
-      "tpot_ms_per_tok": 22.341077544842847
+      "throughput_tok_per_sec": 1378.968944067799,
+      "avg_latency_ms": 2920.1371596893296,
+      "p50_latency_ms": 2913.53677585721,
+      "p99_latency_ms": 2975.6032899022102,
+      "tpot_ms_per_tok": 23.2057437824551
     },
     "batch_128": {
       "batch_size": 128,
-      "num_input_tokens": 512,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
       "num_output_tokens": 128,
       "num_iters": 5,
-      "throughput_tok_per_sec": 5643.249306634135,
-      "avg_latency_ms": 2839.432980850688,
-      "p50_latency_ms": 2846.7628210783005,
-      "p99_latency_ms": 2900.5435090512037,
-      "tpot_ms_per_tok": 22.681967966491356
+      "throughput_tok_per_sec": 5414.296145922958,
+      "avg_latency_ms": 2964.334140153369,
+      "p50_latency_ms": 2969.226948916912,
+      "p99_latency_ms": 3022.6969085633755,
+      "tpot_ms_per_tok": 23.641115400823765
     }
   }
 }

tests/unit_tests/test_static_benchmark.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import argparse
+import sys
+from unittest import mock
+
+import pytest
+
+from tests.performance_tests.client import static_benchmark
+
+
+@pytest.mark.parametrize(
+    ("batch_size", "data_parallel_size", "expected"),
+    [(1, 8, 8), (8, 8, 8), (32, 8, 32), (1, 4, 4), (128, 4, 128), (1, 1, 1)],
+)
+def test_get_warmup_batch_size(batch_size, data_parallel_size, expected):
+    assert static_benchmark._get_warmup_batch_size(batch_size, data_parallel_size) == expected
+
+
+def test_parse_args_preserves_single_worker_default(monkeypatch):
+    monkeypatch.setattr(sys, "argv", ["static_benchmark.py"])
+
+    assert static_benchmark.parse_args().data_parallel_size == 1
+
+
+@pytest.mark.asyncio
+async def test_run_batch_request_count_override(monkeypatch):
+    single_request = mock.AsyncMock(return_value=(512, 128, 0.1))
+    monkeypatch.setattr(static_benchmark, "_single_request", single_request)
+    args = argparse.Namespace(
+        batch_size=1, model="gpt_583m", num_output_tokens=128, temperature=0.0
+    )
+
+    inputs, outputs, latencies, _ = await static_benchmark._run_batch(
+        mock.sentinel.session,
+        args,
+        "http://localhost:5000/v1/completions",
+        ["prompt 0", "prompt 1"],
+        iter_start_index=1,
+        request_count=8,
+    )
+
+    assert single_request.await_count == 8
+    assert [call.args[3] for call in single_request.await_args_list] == [
+        "prompt 1",
+        "prompt 0",
+        "prompt 1",
+        "prompt 0",
+        "prompt 1",
+        "prompt 0",
+        "prompt 1",
+        "prompt 0",
+    ]
+    assert inputs == [512] * 8
+    assert outputs == [128] * 8
+    assert latencies == [0.1] * 8
+
+    single_request.reset_mock()
+    await static_benchmark._run_batch(
+        mock.sentinel.session,
+        args,
+        "http://localhost:5000/v1/completions",
+        ["prompt"],
+        iter_start_index=0,
+    )
+    single_request.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_main_widens_only_warmup_batches_and_preserves_timed_prompts(monkeypatch):
+    calls = []
+
+    async def fake_run_batch(session, args, url, prompts, iter_start_index, request_count=None):
+        count = args.batch_size if request_count is None else request_count
+        calls.append((iter_start_index, count))
+        return [512] * count, [128] * count, [0.1] * count, 1.0
+
+    class FakeClientSession:
+        async def __aenter__(self):
+            return mock.sentinel.session
+
+        async def __aexit__(self, exc_type, exc_value, traceback):
+            return False
+
+    monkeypatch.setattr(static_benchmark, "_run_batch", fake_run_batch)
+    monkeypatch.setattr(static_benchmark.aiohttp, "TCPConnector", mock.Mock())
+    monkeypatch.setattr(
+        static_benchmark.aiohttp, "ClientSession", mock.Mock(return_value=FakeClientSession())
+    )
+    args = argparse.Namespace(
+        server_url="http://localhost:5000/v1",
+        model="gpt_583m",
+        batch_size=1,
+        data_parallel_size=8,
+        dataset="synthetic",
+        num_input_tokens=512,
+        num_output_tokens=128,
+        temperature=0.0,
+        num_warmup_iters=2,
+        num_iters=2,
+    )
+
+    summary = await static_benchmark.main(args)
+
+    assert calls == [(0, 8), (8, 8), (2, 1), (3, 1)]
+    assert summary["batch_size"] == 1