test: add async scheduling inference coverage

tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_async_sched/model_config.yaml

@@ -0,0 +1,59 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+TEST_TYPE: frozen-start
+MODE: inference
+MODEL_ARGS:
+  --tiktoken-pattern: v2
+  --use-mcore-models: true
+  --tokenizer-type: TikTokenizer
+  --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/mcore_mistral/nemo_minitron-0.5b/v1/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json
+  --auto-detect-ckpt-format: true
+  --max-tokens-to-oom: 3600000
+  --inference-max-seq-length: 4096
+  --attention-backend: flash
+  --use-checkpoint-args: true
+  --micro-batch-size: 1
+  --no-load-optim: true
+  --no-use-tokenizer-model-from-checkpoint-args: true
+  --timing-log-level: 0
+  --load: ${CHECKPOINT_LOAD_PATH}/model/mcore_mistral/nemo_minitron-0.5b/v1
+  --distributed-backend: nccl
+  --log-interval: 1
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --deterministic-mode: true
+  --ckpt-format: torch_dist
+  --bf16: true
+  --log-memory-to-tensorboard: true
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --num-layers: 24
+  --hidden-size: 1152
+  --num-attention-heads: 16
+  --max-position-embeddings: 1024
+  --seq-length: 1024
+  --temperature: 1.0
+  --top_k: 1
+  # Async scheduling only supports greedy sampling (top_k=1, top_p=0.0) and does
+  # not support log probabilities, stop words, chunked prefill, or prefix
+  # caching (see dynamic_engine._validate_async_sched_support_for_request).
+  --inference-dynamic-batching-buffer-size-gb: 20
+  --inference-dynamic-batching-async-sched-mode: async
+  --dist-ckpt-strictness: log_unexpected
+  --inference-ckpt-non-strict: true # To handle the extra_state errors
+  --output-path: ${INFERENCE_OUTPUT_PATH}
+  --output-every-n-results: 32
+  --prompt-file: ${DATA_PATH}/text/sharegpt-vicuna/filtered/processed.jsonl
+  --prompt-file-num-truncate: 128 # originally 1024
+  --num-tokens-to-generate: 128 # originally 512
+  --incoming-requests-per-step: 32
+  --termination-id: -1
+  --inference-repeat-n: 3
+  --inference-logging-step-interval: 1
+METRICS:
+  - "generated_tokens"

tests/performance_tests/shell_test_utils/run_perf_test.sh

@@ -88,6 +88,12 @@ NUM_TIMED_ITERS=$("$YQ" '.NUM_TIMED_ITERS // 5' "$CONFIG_PATH")
 # hybrid models should use 'gsm8k' — synthetic input gives misleading
 # perf because every token is identical (uniform expert routing, hot KV).
 DATASET=$("$YQ" '.DATASET // "synthetic"' "$CONFIG_PATH")
+# Async prefill scheduling for dynamic batching. When true, the server is
+# launched with --inference-dynamic-batching-async-sched-mode async (overlaps
+# the prefill scheduler with GPU compute). Requires greedy sampling / no
+# logprobs / no stop words — the static benchmark client already satisfies
+# these (temperature 0.0, ignore_eos, no stop tokens, no logprobs requested).
+ASYNC_SCHED=$("$YQ" '.ASYNC_SCHED // false' "$CONFIG_PATH")
 mapfile -t BATCH_SIZES < <("$YQ" '.BATCH_SIZES[]' "$CONFIG_PATH")
 
 # For MoE configs, expert-parallelism is orthogonal to DP and reshapes the
@@ -194,6 +200,20 @@ SERVER_COMMON_ARGS=(
     --host 0.0.0.0
 )
 
+# Enable async prefill scheduling when the test case opts in. Async scheduling
+# requires materialize_only_last_token_logits=True. run_dynamic_text_generation_server
+# force-sets return_log_probs=True (for echo/loglikelihood support), which would
+# flip materialize_only_last_token_logits to False; passing --skip-prompt-log-probs
+# keeps it True (materialize = not(return_log_probs and not skip_prompt_log_probs)).
+# The perf client never requests prompt logprobs, so this is a no-op for the metrics.
+if [[ "$ASYNC_SCHED" == "true" ]]; then
+    echo "[run_perf_test] async scheduling enabled (--inference-dynamic-batching-async-sched-mode async --skip-prompt-log-probs)"
+    SERVER_COMMON_ARGS+=(
+        --inference-dynamic-batching-async-sched-mode async
+        --skip-prompt-log-probs
+    )
+fi
+
 (
     cd "$ROOT_DIR"
     uv run --no-sync python -m torch.distributed.run \

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

@@ -0,0 +1,52 @@
+{
+  "h100": {
+    "batch_1": {
+      "batch_size": 1,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 42.75601962213067,
+      "avg_latency_ms": 2993.6806965619326,
+      "p50_latency_ms": 2992.9271759465337,
+      "p99_latency_ms": 3003.761636093259,
+      "tpot_ms_per_tok": 23.388519530999474
+    },
+    "batch_8": {
+      "batch_size": 8,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 344.0237043263471,
+      "avg_latency_ms": 2920.4657254274935,
+      "p50_latency_ms": 2921.021580696106,
+      "p99_latency_ms": 2982.457813806832,
+      "tpot_ms_per_tok": 23.254211554012727
+    },
+    "batch_32": {
+      "batch_size": 32,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 1372.033198644746,
+      "avg_latency_ms": 2927.9737344244495,
+      "p50_latency_ms": 2925.4276445135474,
+      "p99_latency_ms": 3030.2867460995913,
+      "tpot_ms_per_tok": 23.32305080635706
+    },
+    "batch_128": {
+      "batch_size": 128,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 5395.836458577838,
+      "avg_latency_ms": 2958.10645565507,
+      "p50_latency_ms": 2953.0187863856554,
+      "p99_latency_ms": 3037.2949857264757,
+      "tpot_ms_per_tok": 23.72199398232624
+    }
+  }
+}

tests/performance_tests/test_cases/gpt/gpt_583m_perf_async_sched/model_config.yaml

@@ -0,0 +1,38 @@
+# Inference perf test: 583M mcore-mistral checkpoint, TP=1 PP=1 DP=8 (8 GPUs),
+# async prefill scheduling enabled.
+#
+# Same 583M checkpoint / DP=8 dynamic-batching path as gpt_583m_perf, but the
+# server runs with --inference-dynamic-batching-async-sched-mode async (the
+# scheduler overlaps prefill setup with GPU compute). Measures the throughput /
+# latency of the async scheduling path exercised by the functional test
+# gpt_dynamic_inference_tp1_pp1_583m_async_sched.
+#
+# Async scheduling requires greedy sampling / no logprobs / no stop words; the
+# static benchmark client already satisfies these (temperature 0.0, ignore_eos).
+#
+# Baseline values are recorded by `RECORD_BASELINE=1 run_perf_test.sh ...`
+# and compared on subsequent runs with TOLERANCE_PCT tolerance.
+
+MODEL: gpt_583m
+TP: 1
+PP: 1
+DP: 8
+ASYNC_SCHED: true
+NUM_INPUT_TOKENS: 512
+NUM_OUTPUT_TOKENS: 128
+NUM_WARMUP_ITERS: 2
+NUM_TIMED_ITERS: 5
+BATCH_SIZES:
+  - 1
+  - 8
+  - 32
+  - 128
+TOLERANCE_PCT: 10
+# p99 omitted on purpose: with NUM_TIMED_ITERS=5 it is the max of 5 samples,
+# not a real percentile, so it produces flaky regressions even when throughput
+# / avg / p50 are stable. p99 is still recorded in results.json for visibility.
+METRICS:
+  - throughput_tok_per_sec
+  - avg_latency_ms
+  - p50_latency_ms
+  - tpot_ms_per_tok

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

@@ -0,0 +1,40 @@
+{
+  "gb200": {
+    "batch_8": {
+      "batch_size": 8,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 306.77571598122546,
+      "avg_latency_ms": 3292.265848722309,
+      "p50_latency_ms": 3284.7761889570393,
+      "p99_latency_ms": 3369.791687990073,
+      "tpot_ms_per_tok": 26.07768341249539
+    },
+    "batch_32": {
+      "batch_size": 32,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 1217.121447486444,
+      "avg_latency_ms": 3322.6687765843963,
+      "p50_latency_ms": 3319.4968919851817,
+      "p99_latency_ms": 3425.9593110182323,
+      "tpot_ms_per_tok": 26.29154228288826
+    },
+    "batch_128": {
+      "batch_size": 128,
+      "dataset": "synthetic",
+      "num_input_tokens_avg": 512.0,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 4573.446310856698,
+      "avg_latency_ms": 3390.523879592547,
+      "p50_latency_ms": 3371.0599309997633,
+      "p99_latency_ms": 3771.179543051403,
+      "tpot_ms_per_tok": 27.987646798464993
+    }
+  }
+}

tests/performance_tests/test_cases/gpt/gpt_583m_perf_async_sched_gb200_4gpu/model_config.yaml

@@ -0,0 +1,36 @@
+# Inference perf test: 583M mcore-mistral checkpoint, TP=1 PP=1 DP=4 (4 GPUs),
+# async prefill scheduling enabled.
+#
+# GB200 single-node variant of gpt_583m_perf_async_sched. GB200 nodes have a
+# 4-GPU/node limit, so the DP=8 configuration used on H100 doesn't fit on a
+# single GB200 node. This is a separate test (different world size, different
+# baseline) — not a multi-node port of the DP=8 case.
+#
+# The server runs with --inference-dynamic-batching-async-sched-mode async.
+# Async scheduling requires greedy sampling / no logprobs / no stop words; the
+# static benchmark client already satisfies these (temperature 0.0, ignore_eos).
+
+MODEL: gpt_583m
+TP: 1
+PP: 1
+DP: 4
+ASYNC_SCHED: true
+NUM_INPUT_TOKENS: 512
+NUM_OUTPUT_TOKENS: 128
+# 2 warmup iters can leave the first timed iteration cold on GB200, poisoning
+# the mean/tail metrics; warm up more so timing starts at steady state.
+NUM_WARMUP_ITERS: 5
+NUM_TIMED_ITERS: 5
+# batch_size=1 is omitted: at single-stream the 4-GPU (DP=4) deployment is
+# barely utilized, so per-iteration jitter dominates and the small-sample p50
+# latency is too noisy to gate on. Larger batches are stable perf signals.
+BATCH_SIZES:
+  - 8
+  - 32
+  - 128
+TOLERANCE_PCT: 10
+METRICS:
+  - throughput_tok_per_sec
+  - avg_latency_ms
+  - p50_latency_ms
+  - tpot_ms_per_tok

tests/test_utils/recipes/gb200/gpt-dynamic-inference.yaml

@@ -0,0 +1,65 @@
+type: basic
+format_version: 1
+maintainers: [mcore]
+loggers: [stdout]
+spec:
+  name: '{test_case}_{environment}_{platforms}'
+  model: gpt
+  build: mcore-pyt-{environment}
+  nodes: 1
+  gpus: 4
+  n_repeat: 1
+  platforms: dgx_gb200
+  script_setup: |
+    set -euo pipefail
+    unset https_proxy
+    echo "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" | tee -a /root/.netrc
+
+    # Checkout latest
+    cd /opt
+    rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm
+    git init
+    git remote add origin $MCORE_REPO
+    git fetch origin '+refs/merge-requests/*:refs/remotes/merge-requests/*'
+    git fetch origin $MCORE_MR_COMMIT
+    git checkout $MCORE_MR_COMMIT
+    git rev-parse HEAD
+    # Checkout backwards-ref
+    cd /opt
+    rm -rf /opt/megatron-lm-legacy; mkdir megatron-lm-legacy; cd megatron-lm-legacy
+    git init
+    git remote add origin $MCORE_REPO
+    git fetch origin $MCORE_BACKWARDS_COMMIT
+    git checkout $MCORE_BACKWARDS_COMMIT
+    git rev-parse HEAD
+    rm -rf megatron; cp -a /opt/megatron-lm/megatron ./
+  script: |-
+    set -euo pipefail
+    ls
+    cd /opt/megatron-lm
+    export GPUS_PER_NODE={gpus}
+
+    ARGUMENTS=(
+        "CHECKPOINT_LOAD_PATH=/mnt/artifacts"
+        "CHECKPOINT_SAVE_PATH=/tmp/checkpoints"
+        "DATA_PATH=/mnt/artifacts"
+        "DATA_CACHE_PATH=/workspace/data/cache"
+        "TRAINING_SCRIPT_PATH=examples/inference/advanced/gpt_dynamic_inference.py"
+        "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
+        "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
+        "OUTPUT_PATH={assets_dir}"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
+        "N_REPEAT={n_repeat}"
+        "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE:-}}"
+        "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS:-}}"
+    )
+
+    bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh ${{ARGUMENTS[@]}}
+
+products:
+  - test_case: [gpt_dynamic_inference_tp1_pp1_583m_async_sched]
+    products:
+      - environment: [dev]
+        scope: [mr]
+        platforms: [dgx_gb200]

tests/test_utils/recipes/gb200/gpt-perf-dp4.yaml

@@ -40,7 +40,7 @@ spec:
     GPUS_PER_NODE=4 bash ./tests/performance_tests/shell_test_utils/run_perf_test.sh ${{ARGUMENTS[@]}}
 
 products:
-  - test_case: [gpt_583m_perf_gb200_4gpu]
+  - test_case: [gpt_583m_perf_gb200_4gpu, gpt_583m_perf_async_sched_gb200_4gpu]
     products:
       - environment: [dev]
         scope: [mr]

tests/test_utils/recipes/h100/gpt-dynamic-inference.yaml

@@ -178,3 +178,8 @@ products:
       - environment: [dev]
         scope: [mr]
         platforms: [dgx_h100]
+  - test_case: [gpt_dynamic_inference_tp1_pp1_583m_async_sched]
+    products:
+      - environment: [dev]
+        scope: [mr]
+        platforms: [dgx_h100]

tests/test_utils/recipes/h100/gpt-perf-dp8.yaml

@@ -38,7 +38,7 @@ spec:
     GPUS_PER_NODE=8 bash ./tests/performance_tests/shell_test_utils/run_perf_test.sh ${{ARGUMENTS[@]}}
 
 products:
-  - test_case: [gpt_583m_perf]
+  - test_case: [gpt_583m_perf, gpt_583m_perf_async_sched]
     products:
       - environment: [dev]
         scope: [mr]