Perf tests

tests/performance_tests/client/data/README.md

@@ -0,0 +1,14 @@
+# GSM8K prompt subset
+
+This directory contains a 256-example subset of the GSM8K dataset from
+OpenAI's grade-school-math repository, used as realistic prompt input for
+inference performance benchmarks.
+
+- Original dataset: https://github.com/openai/grade-school-math
+- License: MIT
+
+`gsm8k_prompts.jsonl` holds one prompt per line as `{"prompt": "..."}`,
+drawn from the dataset's test split. Only the question text is retained;
+answers and chain-of-thought solutions are stripped because the benchmark
+generates a fixed `NUM_OUTPUT_TOKENS` worth of tokens and discards the
+content.

tests/performance_tests/client/static_benchmark.py

@@ -9,6 +9,19 @@
 Hits the server's POST /v1/completions endpoint directly via aiohttp — no
 `openai` python client dep required. Vendored and adapted from
 /Users/shanmugamr/inference-bench/static_benchmark.py.
+
+Prompt source:
+  --dataset synthetic   (default for dense models) — uses `"hello " * N` of
+                        fixed length `--num-input-tokens`. Cheap, deterministic,
+                        but every request is identical → not representative for
+                        MoE/hybrid models because routers/dispatchers see the
+                        same token-to-expert assignment each time.
+  --dataset gsm8k       Loads prompts from data/gsm8k_prompts.jsonl (vendored
+                        from openai/gsm8k test split, 256 prompts). Each
+                        request in a batch gets the next prompt, cycling.
+                        Required for MoE models per reviewer feedback on
+                        PR #4917 — synthetic input gives misleading perf for
+                        anything with token-dependent routing.
 """
 
 import argparse
@@ -19,17 +32,36 @@
 
 import aiohttp
 
+_GSM8K_PATH = Path(__file__).parent / "data" / "gsm8k_prompts.jsonl"
+
+
+def _load_gsm8k_prompts() -> list[str]:
+    """Load the vendored gsm8k prompts. Returns the list of question strings."""
+    if not _GSM8K_PATH.exists():
+        raise FileNotFoundError(f"gsm8k prompt file not found at {_GSM8K_PATH}")
+    prompts: list[str] = []
+    for line in _GSM8K_PATH.read_text().splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        obj = json.loads(line)
+        text = obj.get("prompt") or obj.get("problem") or obj.get("question") or obj.get("text")
+        if text:
+            prompts.append(text)
+    if not prompts:
+        raise ValueError(f"no prompts loaded from {_GSM8K_PATH}")
+    return prompts
+
 
 async def _single_request(
     session: aiohttp.ClientSession,
     url: str,
     model: str,
     prompt: str,
-    expected_input_tokens: int,
     num_output_tokens: int,
     temperature: float,
-) -> tuple[int, float]:
-    """POST /v1/completions and return (tokens_generated, latency_s)."""
+) -> tuple[int, int, float]:
+    """POST /v1/completions and return (input_tokens, output_tokens, latency_s)."""
     payload = {
         "model": model,
         "prompt": prompt,
@@ -45,41 +77,44 @@ async def _single_request(
         body = await resp.json()
     latency = time.perf_counter() - t0
     usage = body.get("usage", {})
-    actual_input = usage.get("prompt_tokens")
+    actual_input = usage.get("prompt_tokens") or 0
     actual_output = usage.get("completion_tokens")
-    if actual_input is not None:
-        assert (
-            actual_input == expected_input_tokens
-        ), f"Expected {expected_input_tokens} prompt tokens, server saw {actual_input}."
     if actual_output is not None:
         assert (
             actual_output == num_output_tokens
         ), f"Expected {num_output_tokens} output tokens, got {actual_output}."
-    return actual_output or num_output_tokens, latency
+    return actual_input, actual_output or num_output_tokens, latency
 
 
 async def _run_batch(
-    session: aiohttp.ClientSession, args: argparse.Namespace, url: str, prompt: str
-) -> tuple[list[int], list[float], float]:
+    session: aiohttp.ClientSession,
+    args: argparse.Namespace,
+    url: str,
+    prompts: list[str],
+    iter_start_index: int,
+) -> 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)."""
     t0 = time.perf_counter()
     results = await asyncio.gather(
         *[
             _single_request(
                 session,
                 url,
                 args.model,
-                prompt,
-                args.num_input_tokens,
+                prompts[(iter_start_index + i) % len(prompts)],
                 args.num_output_tokens,
                 args.temperature,
             )
-            for _ in range(args.batch_size)
+            for i in range(args.batch_size)
         ]
     )
     wall = time.perf_counter() - t0
-    token_counts = [r[0] for r in results]
-    latencies = [r[1] for r in results]
-    return token_counts, latencies, wall
+    input_counts = [r[0] for r in results]
+    output_counts = [r[1] for r in results]
+    latencies = [r[2] for r in results]
+    return input_counts, output_counts, latencies, wall
 
 
 def _percentile(sorted_values: list[float], pct: float) -> float:
@@ -89,52 +124,69 @@ def _percentile(sorted_values: list[float], pct: float) -> float:
 
 async def main(args: argparse.Namespace) -> dict:
     url = f"{args.server_url.rstrip('/')}/completions"
-    prompt = ("hello " * args.num_input_tokens).strip()
+
+    if args.dataset == "gsm8k":
+        prompts = _load_gsm8k_prompts()
+        prompt_source = f"gsm8k ({len(prompts)} prompts)"
+    elif args.dataset == "synthetic":
+        prompts = [("hello " * args.num_input_tokens).strip()]
+        prompt_source = f"synthetic (ISL={args.num_input_tokens})"
+    else:
+        raise ValueError(f"unknown --dataset {args.dataset!r}; expected 'synthetic' or 'gsm8k'")
 
     print(f"Server          : {args.server_url}")
     print(f"Model           : {args.model}")
     print(f"Batch size      : {args.batch_size}")
-    print(f"Input tokens    : {args.num_input_tokens}")
+    print(f"Dataset         : {prompt_source}")
     print(f"Output tokens   : {args.num_output_tokens}")
     print(f"Warmup iters    : {args.num_warmup_iters}")
     print(f"Timed iters     : {args.num_iters}", flush=True)
 
     connector = aiohttp.TCPConnector(limit=0)
     async with aiohttp.ClientSession(connector=connector) as session:
+        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, prompt)
+            await _run_batch(session, args, url, prompts, cursor)
+            cursor += args.batch_size
 
         all_wall: list[float] = []
-        all_tokens: list[int] = []
+        all_output_tokens: list[int] = []
+        all_input_tokens: list[int] = []
         all_latencies: list[float] = []
 
         for i in range(args.num_iters):
-            token_counts, latencies, wall = await _run_batch(session, args, url, prompt)
-            total_tokens = sum(token_counts)
+            input_counts, output_counts, latencies, wall = await _run_batch(
+                session, args, url, prompts, cursor
+            )
+            cursor += args.batch_size
+            total_out = sum(output_counts)
             all_wall.append(wall)
-            all_tokens.append(total_tokens)
+            all_output_tokens.append(total_out)
+            all_input_tokens.extend(input_counts)
             all_latencies.extend(latencies)
             print(
                 f"Iter {i + 1}/{args.num_iters}: "
                 f"wall={wall * 1000:.0f} ms, "
-                f"throughput={total_tokens / wall:.1f} tok/s, "
+                f"throughput={total_out / wall:.1f} tok/s, "
                 f"avg_latency={sum(latencies) / len(latencies) * 1000:.0f} ms",
                 flush=True,
             )
 
     avg_wall = sum(all_wall) / len(all_wall)
-    avg_tokens = sum(all_tokens) / len(all_tokens)
-    throughput = avg_tokens / avg_wall
+    avg_out = sum(all_output_tokens) / len(all_output_tokens)
+    throughput = avg_out / avg_wall
     sorted_lat = sorted(all_latencies)
     avg_latency_ms = sum(sorted_lat) / len(sorted_lat) * 1000
     p50_latency_ms = _percentile(sorted_lat, 0.50) * 1000
     p99_latency_ms = _percentile(sorted_lat, 0.99) * 1000
     tpot_ms = avg_wall * 1000 / args.num_output_tokens
+    avg_input_tokens = sum(all_input_tokens) / len(all_input_tokens) if all_input_tokens else 0
 
     summary = {
         "batch_size": args.batch_size,
-        "num_input_tokens": args.num_input_tokens,
+        "dataset": args.dataset,
+        "num_input_tokens_avg": avg_input_tokens,
         "num_output_tokens": args.num_output_tokens,
         "num_iters": args.num_iters,
         "throughput_tok_per_sec": throughput,
@@ -158,7 +210,19 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--server-url", default="http://localhost:5000/v1")
     parser.add_argument("--model", default="")
     parser.add_argument("--batch-size", type=int, default=32)
-    parser.add_argument("--num-input-tokens", type=int, default=512)
+    parser.add_argument(
+        "--dataset",
+        choices=["synthetic", "gsm8k"],
+        default="synthetic",
+        help="Prompt source. 'synthetic' uses 'hello '*N (deterministic, MoE-misleading); "
+        "'gsm8k' uses the vendored 256-prompt gsm8k subset (real tokens, recommended for MoE).",
+    )
+    parser.add_argument(
+        "--num-input-tokens",
+        type=int,
+        default=512,
+        help="Synthetic-prompt length in tokens. Ignored when --dataset=gsm8k.",
+    )
     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)

tests/performance_tests/server/model_args/gpt_16b.args

@@ -44,5 +44,20 @@
 --no-use-tokenizer-model-from-checkpoint-args
 --no-load-optim
 --inference-ckpt-non-strict
+# Inference perf flags. NOTE: reviewer asked for the full kitchen-sink list
+# (--transformer-impl inference_optimized, --cuda-graph-impl local, --cuda-
+# graph-scope full_iteration_inference, --inference-dynamic-batching-num-cuda-
+# graphs -1, --inference-grouped-gemm-backend vllm, --moe-shared-expert-
+# overlap, --inference-moe-token-dispatcher-type nvls). Three blockers:
+#   1. `inference_optimized` rejects --swiglu, which this model uses.
+#   2. The vllm-fused-MoE / nvls dispatcher / shared-expert overlap args all
+#      require inference_optimized, so they're out by transitive closure.
+#   3. CUDA graphs at full_iteration_inference scope crash during capture
+#      because the alltoall MoE dispatcher does a d2h event.synchronize()
+#      inside its forward (token_dispatcher.py:_maybe_dtoh_and_synchronize)
+#      which is illegal in a capturing stream.
+# Until a graph-friendly MoE dispatcher is available for the transformer_engine
+# path, we keep the safe subset: chunked prefill.
 --transformer-impl transformer_engine
 --distributed-backend nccl
+--enable-chunked-prefill

tests/performance_tests/server/model_args/hybrid_nanov3_3b.args

@@ -0,0 +1,33 @@
+# Model-specific args for the nanov3 3B hybrid MoE checkpoint (TP=1 PP=1 EP=8 — 8 GPUs).
+# Sourced from tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_ep8_nanov3_chunked_prefill/model_config.yaml
+--load ${CHECKPOINT_LOAD_PATH}/model/nemotron6/3b_hybrid_moe/checkpoints/phase2_lc_reinit_emb/
+--tokenizer-model ${CHECKPOINT_LOAD_PATH}/model/nemotron6/tokenizers/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json
+--tokenizer-type TikTokenizer
+--tiktoken-pattern v2
+--use-mcore-models
+--model-provider mamba
+--expert-model-parallel-size 8
+--expert-tensor-parallel-size 1
+--use-checkpoint-args
+--no-use-tokenizer-model-from-checkpoint-args
+--dist-ckpt-strictness log_unexpected
+--ckpt-format torch_dist
+--ckpt-fully-parallel-load
+--ckpt-assume-constant-structure
+--no-load-optim
+--moe-router-score-function sigmoid
+--moe-router-enable-expert-bias
+--moe-router-topk-scaling-factor 2.5
+--moe-router-dtype fp32
+--bf16
+--attention-backend flash
+--no-create-attention-mask-in-dataloader
+--mamba-inference-conv-states-dtype fp32
+--mamba-inference-ssm-states-dtype fp32
+--transformer-impl inference_optimized
+--inference-moe-token-dispatcher-type nvls
+--enable-chunked-prefill
+--cuda-graph-impl local
+--inference-cuda-graph-scope block
+--inference-dynamic-batching-num-cuda-graphs -1
+--distributed-backend nccl