[TRTLLM-14814][chore] Add SA speculative-decoding eval config for Kimi K3

examples/kimi_k3/README.md

@@ -97,6 +97,21 @@ approximately:
 The expected average accuracy is approximately 96.47. Small differences are
 possible with different checkpoint or dependency revisions.
 
+To evaluate with suffix-automaton (SA) speculative decoding, add `--sa`:
+
+```bash
+sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \
+    --model /path/to/kimi-k3-checkpoint \
+    --image /path/to/tensorrt-llm-container.sqsh \
+    --sa
+```
+
+This selects `eval_extra_llm_options_sa.yaml` (see Current limitations
+below for what SA changes) and logs a speculative-decoding acceptance
+summary at the end of the run. SA is lossless, so the scores should match
+the non-SA run within noise. Speedup is workload-dependent, proportional
+to the n-gram repetition in the generated output.
+
 For serving performance, use the standard sweep under
 `examples/kimi_k3/perf_sweep/` (this supersedes the older
 `run_serving_benchmark_kimi_k3.sbatch` single-recipe benchmark). It submits
@@ -192,4 +207,4 @@ default cache manager.
   and the TEP16/TEP8 latency recipes are unaffected. Tracked as
   TRTLLM-14904.
 - FP8 KV cache (`kv_cache_config.dtype: fp8`) is not yet supported.
-- Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). Combining speculation with disaggregated serving is not yet supported.
+- Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). For evaluation, use `eval_extra_llm_options_sa.yaml` (the `--sa` flag of the GSM8K job): that configuration runs with the overlap scheduler off, `max_batch_size` 8, and a matching CUDA-graph `max_batch_size`. Combining speculation with disaggregated serving is not yet supported.

examples/kimi_k3/eval_extra_llm_options_sa.yaml

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Suffix-automaton (SA) speculative-decoding evaluation config. Identical
+# to eval_extra_llm_options.yaml except for the SA-specific keys, each
+# marked below: max_batch_size, cuda_graph_config.max_batch_size,
+# disable_overlap_scheduler, enable_chunked_prefill, and
+# speculative_config. Keep every other key in sync with the base file when
+# editing either.
+tensor_parallel_size: 16
+enable_attention_dp: true
+moe_expert_parallel_size: 16
+# SA-specific: SA allocates per-request speculative state; the evaluated
+# configuration uses a batch size of 8.
+max_batch_size: 8
+max_num_tokens: 8192
+max_seq_len: 8192
+trust_remote_code: true
+# SA-specific: the evaluated SA configuration runs with the overlap
+# scheduler off.
+disable_overlap_scheduler: true
+# SA-specific: the evaluated SA configuration runs without chunked
+# prefill (the base config enables it).
+enable_chunked_prefill: false
+cuda_graph_config:
+  # SA-specific: keep the CUDA-graph batch size equal to max_batch_size
+  # above. SA is compatible with CUDA graphs. enable_padding is left at
+  # its default (off) here; the base config turns it on.
+  max_batch_size: 8
+moe_config:
+  max_num_tokens: 33024
+  use_low_precision_moe_combine: true
+kv_cache_config:
+  enable_block_reuse: false
+  free_gpu_memory_fraction: 0.25
+  tokens_per_block: 64
+# SA-specific: standalone suffix-automaton speculative decoding. SA is
+# lossless, so GSM8K scores should match the non-SA run within noise.
+speculative_config:
+  decoding_type: SA
+  max_draft_len: 2

examples/kimi_k3/run_gsm8k_kimi_k3.sbatch

@@ -10,10 +10,16 @@
 #       --model /path/to/kimi-k3-checkpoint \
 #       --image /path/to/tensorrt-llm-container.sqsh
 #
+# Pass --sa to evaluate with suffix-automaton (SA) speculative decoding
+# (selects eval_extra_llm_options_sa.yaml, which carries the SA-required
+# max_batch_size of 8).
+#
 # Pass --reuse to additionally enable KV-cache block reuse
 # (eval_extra_llm_options_reuse.yaml); chunked prefill is enabled by
 # default.
 #
+# --sa and --reuse are mutually exclusive.
+#
 # REPO defaults to the submit directory (override by exporting REPO), and
 # must contain the built TensorRT-LLM checkout. The #SBATCH partition and
 # account below are placeholders — edit them or override on the sbatch
@@ -32,16 +38,19 @@
 set -euo pipefail
 
 usage() {
-    echo "Usage: sbatch $0 --model PATH --image PATH [--reuse]"
+    echo "Usage: sbatch $0 --model PATH --image PATH [--sa | --reuse]"
 }
 
 MODEL=""
 CONTAINER_IMAGE=""
 MODE=default
+# Speculative-decoding acceptance stats: on for --sa, off otherwise.
+SPEC_STATS_DEFAULT=0
 while [[ $# -gt 0 ]]; do
     case "$1" in
-        --reuse)
-            MODE=reuse
+        --sa|--reuse)
+            [[ "$MODE" == default ]] || { echo "error: --sa and --reuse are mutually exclusive" >&2; usage >&2; exit 2; }
+            MODE=${1#--}
             shift
             ;;
         --model)
@@ -84,8 +93,14 @@ REPO=${REPO:-$SLURM_SUBMIT_DIR}
 # Per-mode LLM options. All engine options live in the selected yaml
 # (single source of truth); trtllm-eval lets explicit CLI flags override
 # --config, so none are duplicated here.
+# - sa: suffix-automaton speculative decoding, which needs its own LLM
+#   options (overlap scheduler off, max_batch_size 8; see the yaml).
 # - reuse: KV-cache block reuse (enable_block_reuse: true in the yaml).
 case "$MODE" in
+    sa)
+        EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_sa.yaml
+        SPEC_STATS_DEFAULT=1
+        ;;
     reuse)
         EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_reuse.yaml
         ;;
@@ -143,6 +158,10 @@ srun --mpi=pmix \
         export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\"
         # Cap in-flight requests: may cost some throughput vs submit-all, but yields steady partial scores (early failure signal) instead of one burst at the end (0 = off).
         export TLLM_EVAL_MAX_IN_FLIGHT=\"\${TLLM_EVAL_MAX_IN_FLIGHT:-0}\"
+        # Speculative-decoding acceptance summary (acceptance length /
+        # rate) at eval end; on by default for --sa, off otherwise (an
+        # explicit env value still wins either way).
+        export TLLM_EVAL_SPEC_STATS=\"\${TLLM_EVAL_SPEC_STATS:-$SPEC_STATS_DEFAULT}\"
 
         # Import tensorrt_llm from $REPO, not from wherever the venv's
         # in-place install points (for git-worktree submits those differ;