[TRTLLM-14814][feat] Kimi K3 serving parsers, chat template, and speculative decoding (suffix automaton + DFlash scaffold) - #17327
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #64149 [ run ] triggered by Bot. Commit: |
|
PR_Github #64149 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64252 [ run ] triggered by Bot. Commit: |
|
PR_Github #64252 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64342 [ run ] triggered by Bot. Commit: |
|
Internal validation summary for this PR's SA and parser claims (full-size Kimi K3 checkpoint, Blackwell nodes, branch head 174f9f6 — before the latest test-scaffolding fix commit, which does not affect these paths):
One scoped issue found, not blocking SA/parser correctness above: with a named/forced |
|
PR_Github #64342 [ run ] completed with state
|
…k-decode path
The TRTLLM-14814 DSpark scaffold added a guard in
DFlashForCausalLM.__init__ rejecting any dflash_config with causal=true.
Legacy DFlash drafter configs (e.g. Laguna) also declare causal=true;
their causality is implemented by the legacy decode path
(_sliding_layers_causal in the Laguna subclass), so the unconditional
guard broke construction of pre-existing drafters
(TestLagunaXS_2_1::test_{bf16,fp8,nvfp4}_dflash).
Only reject causal=true when the config actually selects the dspark
convention: projector_type == 'dspark' or any dspark feature field
(shift_label, use_swa, markov_rank, use_confidence_head) enabled.
Add hw-agnostic unit tests: legacy causal config constructs; dspark +
causal (via feature flags or projector_type alone) still raises.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…follow-up PRs test_kimi_k3_dflash_scaffold.py loaded examples/kimi_k3/make_synthetic_dflash_drafter.py at module level, so on branches without that generator (it ships with the examples PR) the whole module failed at collection — taking down every unittest stage that collects unittest/_torch/speculative/hw_agnostic (l0_cpu and the l0_h100 sampler/speculative stage). Load it conditionally and skip the generator-schema tests when it is absent; the DFlashDecodingConfig, KimiLinearModel, and DFlashSpecMetadata tests keep running. test_kimi_k3_disagg_parity_selftest subprocess-runs kimi_k3_disagg_parity.py, which ships with the disagg parity PR; skip the self-test while the harness is absent. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…arsers The new kimi_k3 reasoning parser and tool parser register into ReasoningParserFactory / ToolParserFactory, which feed the trtllm-serve --reasoning_parser / --tool_parser Click Choice lists pinned by tests/unittest/api_stability/references/trtllm_serve_cli.yaml. Add kimi_k3 to both recorded choice lists. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
- kimi_k3_sa_harness.py: validate KIMI_K3_SPEC_PARITY values, check baseline/prompt-set length match, and guard baseline-side logprob alignment before indexing in _compare_logits_parity - modeling_speculative.py: explicit slice | None annotation for vocab_slice (RUF013) - test_dflash_accept_stats.py: catch ImportError instead of Exception in the torch-less fallback loader; use tmp_path instead of /tmp Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…d KV managers Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…TODO Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…e graph-safe append The clamp keeps padded/warmup rows in bounds but targets a real block 0; TRTLLM-15199 tracks excluding invalid rows or reserving a scratch block. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…aceholders and K3 MoE backend pinning - modeling_speculative.py: explain why the DSpark markov/confidence head attributes are plain None placeholders instead of pre-allocated parameters (checkpoint-dependent shapes, absent on most checkpoints). - kimi_k3_sa_harness.py / test_kimi_k3_specdec.py: KimiK3MoERuntime pins the routed MoE backend to TRTLLM, so the KIMI_K3_MOE_BACKEND=VANILLA default was a no-op; drop it and correct the comments (parity is unaffected since baseline and spec runs share the same backend). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…rimental The recorder is opt-in and DFlash/DSpark-specific; log a warning at activation so enabling it with other speculative decoding methods is visibly unsupported rather than silently recording nothing. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…rop chat-completions dependency note, mark SA speculation supported for aggregated serving Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
45f2312 to
f9376ca
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #64750 [ run ] triggered by Bot. Commit: |
|
PR_Github #64735 [ run ] completed with state |
|
/bot skip --comment "Union-green across runs on this head and its immediate predecessor. Current run (pipeline 52600, commit f9376ca): SBSA single-GPU 28/28 green, A100X green, 214/218 x86 stages green; sole red is DGX_H100-PyTorch-5, which the test harness itself classified as transient infra (monitor lost contact with a live job: "this is not a test failure") - that stage passed on the previous head 45f2312 (pipeline 52568), whose only delta to this head is a rebase onto main plus a doc-only commit. 52568 additionally had all six DGX_H100 shards green; its single red (A100X NIXL test) is an intermittent failure that passed on a same-commit rerun elsewhere and is green on this run. DGX_H100-PyTorch-4 multimodal failures in 52600 were node-level CUDA errors (cudaErrorMapBufferObjectFailed) in a subsystem this PR does not touch and auto-retried green." |
|
/bot skip --comment "Re-issuing skip for the rebased head fac670d (content unchanged vs f9376ca, which this skip rationale was established for): union-green across pipeline 52600 (SBSA 28/28, A100X green, 214/218 x86 stages; sole red harness-classified as transient infra and green on the prior head) and pipeline 52568 (all six DGX_H100 shards green; single red an intermittent test that passed on a same-commit rerun and in 52600)." |
|
PR_Github #64762 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64750 [ run ] completed with state |
|
PR_Github #64763 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64762 [ skip ] completed with state |
|
PR_Github #64763 [ skip ] completed with state |
17 participants
Description
Stages the Kimi K3 serving and speculative-decoding layers on top of the model PR:
TRTLLM-14812 — reasoning parser, tool parser, chat template (
ec71a9cc67): K3 reasoning parser, K3 tool parser + factory registration, OpenAI-server special-token decoding config for the parsers, and tokenizer-native chat-template dispatch, with unit tests for each.TRTLLM-14814 — speculative decoding (
174f9f6a37): suffix-automaton spec-dec for K3 with acceptance statistics, the DFlash draft-model scaffold (Markov-head helpers inmodeling_speculative.py; the trained K3 DFlash draft checkpoint is still in progress and not required by this PR), KDA replay-cache seeding for generation-side restarts, and the TRTLLM-14706 graph-safe latent-cache append fix it depends on. Unit tests plus an SA logits-parity harness.Base
Notes
Test Coverage
PR Checklist
Dev Engineer Review
reasoning_contentcorrectly.tool_choice: raw model text can appear intool_calls[0].argumentsinstead of extracted JSON arguments. Auto tool choice was not validated.QA Engineer Review
test_kimi_k3_sa_specdec_logits_paritytest_kimi_k3_disagg_parity_selftesttests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.test-db/,qa/, orwaives.txtchanges were included.