From cae41e317ace0306a3854fea72988f1c7693780f Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 18:31:32 -0700 Subject: [PATCH] [TRTLLM-14814][chore] Add SA speculative-decoding eval config for Kimi K3 Suffix-automaton (SA) speculative decoding is supported for the Kimi K3 example, but there was no example configuration to evaluate it with: examples/kimi_k3 has eval configs for the default and block-reuse modes only, so an accuracy run of the SA path had to be assembled by hand. Add examples/kimi_k3/eval_extra_llm_options_sa.yaml (the base eval config plus the SA-specific keys: max_batch_size 8, a matching CUDA-graph batch size, overlap scheduler off, chunked prefill off, and speculative_config with decoding_type SA / max_draft_len 2), and a --sa flag on run_gsm8k_kimi_k3.sbatch that selects it and turns on the speculative-decoding acceptance summary (TLLM_EVAL_SPEC_STATS). --sa and --reuse are mutually exclusive. Document the flag in the example README. Configuration only; no library code changes. Signed-off-by: Brian Nguyen --- examples/kimi_k3/README.md | 17 +++++++- .../kimi_k3/eval_extra_llm_options_sa.yaml | 41 +++++++++++++++++++ examples/kimi_k3/run_gsm8k_kimi_k3.sbatch | 25 +++++++++-- 3 files changed, 79 insertions(+), 4 deletions(-) create mode 100644 examples/kimi_k3/eval_extra_llm_options_sa.yaml diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index e58c618b5628..2ac4c475088c 100644 --- a/examples/kimi_k3/README.md +++ b/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. diff --git a/examples/kimi_k3/eval_extra_llm_options_sa.yaml b/examples/kimi_k3/eval_extra_llm_options_sa.yaml new file mode 100644 index 000000000000..22a25324f81f --- /dev/null +++ b/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 diff --git a/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch b/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch index 0526f1c7c512..441247201612 100644 --- a/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch +++ b/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;