Skip to content

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

Open
brnguyen2 wants to merge 1 commit into
NVIDIA:mainfrom
brnguyen2:k3/sa-eval-config
Open

[TRTLLM-14814][chore] Add SA speculative-decoding eval config for Kimi K3#17455
brnguyen2 wants to merge 1 commit into
NVIDIA:mainfrom
brnguyen2:k3/sa-eval-config

Conversation

@brnguyen2

@brnguyen2 brnguyen2 commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Description

Suffix-automaton (SA) speculative decoding is supported for the Kimi K3
example, but examples/kimi_k3/ has eval configurations for the default and
block-reuse modes only. There is no example config for the SA path, so an
accuracy run of it has to be assembled by hand. This adds one.

  • New examples/kimi_k3/eval_extra_llm_options_sa.yaml: the base
    eval_extra_llm_options.yaml plus the SA-specific keys — max_batch_size: 8,
    a matching CUDA-graph max_batch_size, overlap scheduler off, chunked
    prefill off, and
    speculative_config: {decoding_type: SA, max_draft_len: 2}.
  • run_gsm8k_kimi_k3.sbatch: new --sa flag that selects that config and
    enables the speculative-decoding acceptance summary
    (TLLM_EVAL_SPEC_STATS) by default for that mode. --sa and --reuse
    are mutually exclusive; an explicit TLLM_EVAL_SPEC_STATS still wins.
  • examples/kimi_k3/README.md: documents --sa and what the SA
    configuration changes.

Configuration/docs only; no library code changes.

The config keys were checked against the current args schema rather than
copied blindly: speculative_config resolves through the decoding_type
discriminated union to SADecodingConfig in
tensorrt_llm/llmapi/llm_args.py, max_draft_len: 2 satisfies its
validator, and global_pool_size is left unset so the
global_pool_size >= max_batch_size check does not apply. Every other key
is already present in the checked-in sibling configs. One deliberate
deviation from the base config: cuda_graph_config.enable_padding is left
at its default (off) rather than true, matching the configuration SA was
evaluated with; this is called out in a comment in the file.

Still missing for examples/kimi_k3/, out of scope here: there is no
example eval configuration for the DFlash speculative-decoding path, which
is scaffolding only at this point.

Test Coverage

No new automated tests: this adds an example YAML config, an sbatch flag,
and README text, none of which are exercised by CI. Verified locally that
the new YAML parses, that its keys and value constraints match the current
LlmArgs/SADecodingConfig schema, and that the modified sbatch passes
bash -n and selects the expected config file per mode. The accuracy claim
it enables (SA is lossless, so GSM8K should match the non-SA run within
noise) is what the config is intended to let users check.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

  • Added SA speculative-decoding configuration for the Kimi K3 GSM8K evaluation.
  • Added --sa mode with --sa/--reuse mutual exclusion.
  • Preserved explicit TLLM_EVAL_SPEC_STATS environment settings.
  • Updated README documentation.
  • YAML and shell validation passed.
  • No library code or public API changes.
  • No test-list changes.

QA Engineer Review

No test changes.

…i 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 <brnguyen@nvidia.com>
@brnguyen2
brnguyen2 requested a review from a team as a code owner August 10, 2026 01:32
@brnguyen2
brnguyen2 requested review from arysef and chang-l August 10, 2026 01:32
@coderabbitai

coderabbitai Bot commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b1df7b3b-7d51-45c7-9923-ed664f237eb3

📥 Commits

Reviewing files that changed from the base of the PR and between d16d01f and cae41e3.

📒 Files selected for processing (3)
  • examples/kimi_k3/README.md
  • examples/kimi_k3/eval_extra_llm_options_sa.yaml
  • examples/kimi_k3/run_gsm8k_kimi_k3.sbatch

Walkthrough

Adds suffix-automaton speculative decoding support to the Kimi K3 GSM8K evaluation launcher. The change adds SA-specific configuration, mutually exclusive launcher modes, acceptance-statistics control, and usage documentation.

Changes

Suffix-automaton GSM8K evaluation

Layer / File(s) Summary
SA evaluation configuration
examples/kimi_k3/eval_extra_llm_options_sa.yaml, examples/kimi_k3/README.md
Adds SA evaluation settings for parallelism, batching, scheduling, CUDA graphs, KV cache, and draft length. Documents evaluation limitations.
SA launcher integration
examples/kimi_k3/run_gsm8k_kimi_k3.sbatch, examples/kimi_k3/README.md
Adds the --sa mode, rejects simultaneous --sa and --reuse, selects the SA configuration, and controls speculative-decoding statistics through TLLM_EVAL_SPEC_STATS. Documents the command and expected acceptance output.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and concisely identifies the SA speculative-decoding evaluation configuration added for Kimi K3.
Description check ✅ Passed The description explains the problem, solution, affected files, validation performed, test coverage, and checklist status.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@brnguyen2

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64915 [ run ] triggered by Bot. Commit: cae41e3 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64915 [ run ] completed with state SUCCESS. Commit: cae41e3
/LLM/main/L0_MergeRequest_PR pipeline #52754 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@brnguyen2

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64954 [ run ] triggered by Bot. Commit: cae41e3 Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants