[TRTLLM-14814][chore] Add SA speculative-decoding eval config for Kimi K3 - #17455
[TRTLLM-14814][chore] Add SA speculative-decoding eval config for Kimi K3#17455brnguyen2 wants to merge 1 commit into
Conversation
…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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
WalkthroughAdds 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. ChangesSuffix-automaton GSM8K evaluation
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run |
|
PR_Github #64915 [ run ] triggered by Bot. Commit: |
|
PR_Github #64915 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64954 [ run ] triggered by Bot. Commit: |
Description
Suffix-automaton (SA) speculative decoding is supported for the Kimi K3
example, but
examples/kimi_k3/has eval configurations for the default andblock-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.
examples/kimi_k3/eval_extra_llm_options_sa.yaml: the baseeval_extra_llm_options.yamlplus the SA-specific keys —max_batch_size: 8,a matching CUDA-graph
max_batch_size, overlap scheduler off, chunkedprefill off, and
speculative_config: {decoding_type: SA, max_draft_len: 2}.run_gsm8k_kimi_k3.sbatch: new--saflag that selects that config andenables the speculative-decoding acceptance summary
(
TLLM_EVAL_SPEC_STATS) by default for that mode.--saand--reuseare mutually exclusive; an explicit
TLLM_EVAL_SPEC_STATSstill wins.examples/kimi_k3/README.md: documents--saand what the SAconfiguration changes.
Configuration/docs only; no library code changes.
The config keys were checked against the current args schema rather than
copied blindly:
speculative_configresolves through thedecoding_typediscriminated union to
SADecodingConfigintensorrt_llm/llmapi/llm_args.py,max_draft_len: 2satisfies itsvalidator, and
global_pool_sizeis left unset so theglobal_pool_size >= max_batch_sizecheck does not apply. Every other keyis already present in the checked-in sibling configs. One deliberate
deviation from the base config:
cuda_graph_config.enable_paddingis leftat its default (off) rather than
true, matching the configuration SA wasevaluated with; this is called out in a comment in the file.
Still missing for
examples/kimi_k3/, out of scope here: there is noexample 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/SADecodingConfigschema, and that the modified sbatch passesbash -nand selects the expected config file per mode. The accuracy claimit 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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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
--samode with--sa/--reusemutual exclusion.TLLM_EVAL_SPEC_STATSenvironment settings.QA Engineer Review
No test changes.