Feat/golden prairie multimodal - #17043
Closed
WeiHaocheng wants to merge 145 commits into
Closed
Conversation
…ch backend Signed-off-by: Tao Li <tali@nvidia.com>
[None][feat] Kimi K3 (golden prairie) end-to-end support in the PyTorch backend See merge request ftp/tekit-golden-prairie!2 Signed-off-by: Tao Li <tali@nvidia.com>
- sanity_kimi_k3.sbatch / run_gsm8k.sbatch: decouple the flashinfer JIT/HF cache from the optimization-work checkout via a new KIMI_K3_CACHE_DIR (default: submit dir); the opt-work dir is now mounted read-only so any user can point at another user's checkout. Launch trtllm-llmapi-launch by its in-repo path so no venv console script is required. - run_kimi.py: submit wrapper with sanity (4-GPU/4-layer pipeline check), sanity-full (16-GPU) and gsm8k (16-GPU eval) modes. Resolves all inputs from a workspace directory (default: parent of the repo checkout, override via KIMI_K3_* / IMAGE env vars), preflights that every path and the flashinfer snapshot setup exist before submitting, and supports --dry-run. Signed-off-by: Brian Nguyen <brnguyen@aws-cmh-slurm-1-vscode-01.cm.cluster> Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
- sbatch scripts: drop all user-specific path defaults; the three inputs (KIMI_K3_OPT_WORK_DIR, KIMI_K3_MODEL_DIR, IMAGE) are now required and fail fast with a pointer to run_kimi.py if unset. - run_kimi.py: resolve inputs from ~/.config/kimi-bringup.ini (override location via KIMI_BRINGUP_CONFIG); minimal config is a single workspace path, with per-input keys and env vars taking precedence. Missing or incomplete config prints one-time setup instructions. - kimi-bringup.ini.example: annotated example config. - README: document the config-file flow as the preferred entry point. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…irie' [None][feat] Kimi K3: portable sbatch launch + Python submit wrapper See merge request ftp/tekit-golden-prairie!6 Signed-off-by: Brian Nguyen (TensorRT) <brnguyen@nvidia.com>
lm-eval computes metrics only after every response is collected, so a multi-hour trtllm-eval run gives no quality signal until the very end — there is no way to tell whether a long job is on track or should be killed early. Add _RunningScoreTracker to LmEvalWrapper.generate_until: as each response completes, apply the owning task's filter ensembles and process_results to a throwaway copy of the instance and log a running mean every N responses, e.g. Partial scores after 200/1319 responses (estimate, 0~100): exact_match,strict-match ~ 93.50 | exact_match,flexible-extract ~ 94.00 Opt-in via TLLM_EVAL_PARTIAL_SCORES_EVERY=N (unset/0 = off, no behavior change). Env-var driven so it uniformly covers every lm-eval-backed task without per-task CLI plumbing. Scoring happens on copies — harness state and the final score are unaffected — and any scoring failure permanently disables the tracker for the run instead of failing the eval. Estimates are meaningful for per-sample-decomposable metrics (e.g. exact_match); the multimodal wrapper's generate_until override does not implement partial scoring yet. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…rie' [None][feat] Log running metric estimates during long lm-eval runs See merge request ftp/tekit-golden-prairie!8 Signed-off-by: Brian Nguyen (TensorRT) <brnguyen@nvidia.com>
Signed-off-by: Tao Li <tali@nvidia.com>
Source integrate kda decode kernel as thop op See merge request ftp/tekit-golden-prairie!5 Signed-off-by: Tao Li <tali@nvidia.com>
…ng param bugs Signed-off-by: rosong11 <rosong@nvidia.com>
…lden_prairie' [feat]integrate moe cubins fix trtllm valid dim mismatch and trtllm-gen wrong param bugs See merge request ftp/tekit-golden-prairie!4 Signed-off-by: Song Rong <rosong@nvidia.com>
…to the repo Signed-off-by: Tao Li <tali@nvidia.com>
DO-NOT-MERGE: integrate the KDA prefill and atten residual kernels into the repo See merge request ftp/tekit-golden-prairie!7 Signed-off-by: Tao Li <tali@nvidia.com>
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
[None][feat] B200 IB DEP32 Pipe Cleaning See merge request ftp/tekit-golden-prairie!14 Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
Signed-off-by: Song Rong <rosong@nvidia.com>
…prairie' [None][test] Add KIMI_K3_ADP knob to the K3 sanity for DEP runs See merge request ftp/tekit-golden-prairie!16 Signed-off-by: Song Rong <rosong@nvidia.com>
…worker SAWorker never committed accepted verification-step states for SSM/hybrid models: verification writes per-step states to the cache manager's SpeculativeState scratch buffers (live pools are read-only during verify), and nothing promoted the accepted step — silent state corruption on any hybrid model running standalone SA. Add the update_mamba_states call right after on-device acceptance, mirroring the other one-engine workers (dflash/eagle3); isinstance-gated so pure-attention models are unaffected. Model-agnostic; candidate for a standalone upstream PR. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ne + KDA multi-token verify) - KimiLinearForCausalLM adopts SpecDecOneEngineForCausalLM so the SA worker attaches via the standard one-engine path; guard admits SA (draft-weight-free; MTP/Eagle stay blocked pending a head checkpoint). - SpecDecOneEngineForCausalLM gains optional hidden_size/vocab_size for composite (vision-wrapper) configs that keep them in text_config. - KDA multi-token verification: generation batches with T=1+draft_len tokens/request run per-step conv + fused_recurrent_kda over gathered state copies into SpeculativeState scratch (live pools untouched); T derived per batch (dynamic draft length). MLA needs metadata only. - Exact KDA verify-vs-sequential-decode unit test (random weights). Validated on GB300 (control-calibrated logits certification, 52-prompt sets): no-spec identity gate; 3-layer pure-KDA and 4-layer KDA+MLA parity; full-model certification (30/52 bit-identical, near-tie-only divergences); GSM8K 96.66 +-0.49 vs 96.74 no-spec reference; tokens/step up to 2.97 (draft_len=4 + global pool). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
sanity_kimi_k3.py: KIMI_K3_SPEC_MODE=sa with draft-len/matching/global- pool knobs; cross-process baseline JSON; one-sided logits certification (one-engine samplers emit no per-token logprobs — SpecSamplerBase is tokens-only; baseline-side near-tie classification carries the gate); no-spec regression mode. run_kimi.py: preflight materializes the MoE snapshot's missing flashinfer/_build_meta.py from version.txt (HACK-marked: without it the private-snapshot path reports 0.0.0+unknown and spec-worker construction fails the >=0.6.4 gate). Eval configs for trtllm-eval with SA; 'legacy' MoE-path terminology renamed to 'private-snapshot'. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
4-layer SA logits-parity test (auto-skips without the checkpoint) and the KDA verify unit test, both in l0_gb300_multi_gpus. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Upstream enabled CUDA graphs + overlap by default for K3; a parity baseline must execute in the same regime as the eager spec run it is compared against (and graph capture also faults at truncated scale on the private-snapshot path — unpinned host copy). Parity mode pins KIMI_K3_CUDA_GRAPHS=0 for both instances; standalone no-spec sanity keeps the upstream defaults. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
[None][feat] Kimi K3: Suffix Automaton speculative decoding See merge request ftp/tekit-golden-prairie!10 Signed-off-by: Brian Nguyen (TensorRT) <brnguyen@nvidia.com>
Signed-off-by: Pengbo Wang <sephw@nvidia.com>
…n_prairie' [None][perf] Enable kda optimized decode kernels See merge request ftp/tekit-golden-prairie!17 Signed-off-by: Pengbo Wang <sephw@nvidia.com>
…DP ranks Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…prairie' [None][fix] Keep chunked-MoE size vectors identical across attention-DP ranks See merge request ftp/tekit-golden-prairie!15 Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
… unittest tree Signed-off-by: Tao Li <tali@nvidia.com>
WeiHaocheng
requested review from
FrankD412,
Tabrizian,
allisonlim-nv,
brnguyen2,
chang-l,
kaiyux,
kris1025,
leslie-fang25,
liji-nv,
lowsfer,
mikeiovine,
pengbowang-nv,
reasonsolo,
ruodil,
schetlur-nv,
xxi-nv and
yufeiwu-nv
and removed request for
a team
July 30, 2026 04:07
Collaborator
|
[by Codex] @VALLIS-NERIA Friendly review reminder: this PR is awaiting your review. Thanks! |
Collaborator
Author
|
It is a mistake push. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
9 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Test Coverage
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.