[#11755][feat] AutoDeploy onboarding agent + Kimi K2.5 AD modeling code - #11780
Merged
bmarimuthu-nv merged 8 commits intoMar 5, 2026
Merged
Conversation
bmarimuthu-nv
force-pushed
the
bala/model-onboarding-agent
branch
from
February 27, 2026 19:16
7106ace to
a473b47
Compare
bmarimuthu-nv
changed the title
bmarimuthu-nv
force-pushed
the
bala/model-onboarding-agent
branch
from
February 27, 2026 19:41
a473b47 to
9bc5a89
Compare
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
Author
|
@CodeRabbit summary |
Contributor
✅ Actions performedSummary regeneration triggered. |
Contributor
📝 WalkthroughWalkthroughAdds export-ready Kimi-K2/K2.5 model support, MLA cache dtype consistency changes, comprehensive unit tests, deployment config, and onboarding/debugging documentation for AutoDeploy model workflows and reviews. Changes
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
✅ Actions performedSummary regeneration triggered. |
bmarimuthu-nv
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (5)
.claude/agents/ad-debug-agent.md (1)
30-40: Clean up repeated typo’d instruction line.“Remember to use you your own tools ...” appears twice and contains a typo. Keep a single corrected instance to reduce noise.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.claude/agents/ad-debug-agent.md around lines 30 - 40, Remove the duplicated, typo-containing instruction line "Remember to use you your own tools - Read, Grep, Glob, Bash, Edit, Write" and replace it with a single corrected instance "Remember to use your own tools - Read, Grep, Glob, Bash, Edit, Write"; update the .claude/agents/ad-debug-agent.md content where that phrase appears so only the corrected line remains (search for the duplicated string to locate both occurrences) and ensure surrounding spacing/newlines remain consistent with the file's formatting.examples/auto_deploy/model_registry/configs/kimi_k2.yaml (1)
1-2: Clarify config purpose: validation vs. production.The comment mentions "minimum layers for validation: 1 dense + 2 MoE = 3 total", suggesting this is a validation/testing config rather than a full production deployment config. Consider either:
- Renaming the file to
kimi_k2_validation.yamlto clarify intent, or- Updating the comment to explain this is intentionally a minimal config for CI/validation purposes
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/auto_deploy/model_registry/configs/kimi_k2.yaml` around lines 1 - 2, The config file kimi_k2.yaml appears to be a minimal/validation-only setup (comment "minimum layers for validation: 1 dense + 2 MoE = 3 total"); clarify intent by either renaming the file to kimi_k2_validation.yaml or updating the top-of-file comment to explicitly state this is a minimal CI/validation/testing config (not production) so readers know the reduced layer counts are intentional.tensorrt_llm/_torch/auto_deploy/models/custom/modeling_kimi_k2.py (2)
440-442: Prefix unused shape variables with underscore.
bszandseq_lenare extracted but unused. Prefix with_to indicate intentional non-use.🔧 Suggested fix
def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - bsz, seq_len, hidden_dim = hidden_states.shape + _bsz, _seq_len, hidden_dim = hidden_states.shape hidden_states_flat = hidden_states.view(-1, hidden_dim)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_kimi_k2.py` around lines 440 - 442, In the forward method of class/model in modeling_kimi_k2.py, rename the unused extracted shape variables bsz and seq_len to _bsz and _seq_len (i.e., change "bsz, seq_len, hidden_dim = hidden_states.shape" to "_bsz, _seq_len, hidden_dim = hidden_states.shape") so the linter/reader knows they are intentionally unused; keep the rest of the function (including hidden_states_flat = hidden_states.view(-1, hidden_dim)) unchanged.
253-259: Consider documenting unusedseq_lenparameter.The
seq_lenparameter is unused because full cached tables are returned for export compatibility. Consider adding a note in the docstring or using_seq_lento make this explicit.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_kimi_k2.py` around lines 253 - 259, The forward method's seq_len parameter is unused (forward in modeling_kimi_k2.py returns full cached tables _ad_cos_cached and _ad_sin_cached) so either document that behavior or mark the parameter as intentionally unused to satisfy linters and readers; update the forward docstring to mention that seq_len is ignored because full cached tables are returned for export compatibility, or change the parameter name to _seq_len (or assign seq_len to a dummy _ = seq_len) inside forward to indicate intentional nonuse, referencing the forward method and the returned _ad_cos_cached and _ad_sin_cached symbols.tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py (1)
530-532: Consider conditional casting for consistency and performance.Unlike
torch_backend_mla.pywhich uses conditional checks before casting, this code unconditionally calls.to(). While.to()is a no-op when dtypes match, it still incurs function call overhead and may create unnecessary tensor views. Consider aligning with the pattern intorch_backend_mla.py:🔧 Suggested change
- # Cast to cache dtype for writes (no-op when dtypes already match). - compressed_kv_for_cache = compressed_kv_flat.to(ckv_cache.dtype) - kpe_for_cache = kpe_flat.to(kpe_cache.dtype) + # Cast to cache dtype for writes if needed. + cache_dtype = ckv_cache.dtype + compressed_kv_for_cache = compressed_kv_flat + kpe_for_cache = kpe_flat + if compressed_kv_flat.dtype != cache_dtype: + compressed_kv_for_cache = compressed_kv_flat.to(cache_dtype) + if kpe_flat.dtype != cache_dtype: + kpe_for_cache = kpe_flat.to(cache_dtype)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py` around lines 530 - 532, The unconditional casts on compressed_kv_flat.to(ckv_cache.dtype) and kpe_flat.to(kpe_cache.dtype) should be made conditional to avoid unnecessary overhead; in the block where compressed_kv_for_cache and kpe_for_cache are set, check if compressed_kv_flat.dtype != ckv_cache.dtype and only then assign compressed_kv_flat.to(ckv_cache.dtype) otherwise reuse compressed_kv_flat, and similarly for kpe_flat vs kpe_cache.dtype (mirroring the pattern used in torch_backend_mla.py) so you avoid extra .to() calls and potential tensor view creation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.claude/agents/ad-debug-agent.md:
- Around line 17-21: The example run command is malformed (broken flag
`--args.yaml-`, stray `extra` token and `5.log`) so update the example to a
single executable command: set the AD_DUMP_GRAPHS_DIR environment variable, run
python examples/auto_deploy/build_and_run_ad.py with the --model <MODEL_HF_ID>
flag and the correct args flag (replace the broken `--args.yaml-`/`extra` with
the proper flag and argument, e.g. `--args.yaml-extra` followed by
examples/auto_deploy/model_registry/configs/<CONFIG_YAML_FILE>), and pipe
stdout/stderr to tee <LOG_FILE> (remove the stray `5.log`).
In @.claude/agents/onboard-reviewer.md:
- Line 89: The fenced code block shown as triple backticks (``` ) in the onboard
reviewer guidance needs a language tag to satisfy markdown linting; change the
opening fence from ``` to a tagged fence such as ```text (or another appropriate
language) so the fenced output example is explicitly labeled and renders
correctly in markdown; update the single fenced block in the document (the
example surrounded by ``` ) to use ```text.
- Around line 65-68: The table in section F has duplicate row IDs and
inconsistent column counts; rename the duplicated F3 entries to unique IDs
(e.g., F2/F3/F4 sequence) and ensure every row has the same number of
pipe-separated cells (fill the missing cell for the row that currently ends
after "Non isNan or isInf checks - these are smoke tests..." or move that text
into the appropriate column). Update the ID references so they are unique (F1,
F2, F3, F4) and verify the header/column count matches all rows to restore
deterministic PASS/FAIL reporting.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_kimi_k2_modeling.py`:
- Around line 730-754: The test asserts that RoPE weights are de-interleaved by
a pre-hook but no such hook exists; fix by either registering a load_state_dict
pre-hook on KimiK2ForCausalLM (or its base KimiK2PreTrainedModel) that calls
_deinterleave_attention_weights on the incoming state_dict before
load_state_dict, or modify the test to call
_deinterleave_attention_weights(hf_state_dict, config) immediately prior to
custom_model.load_state_dict(hf_state_dict) so the weights are transformed
consistently with the block/layer tests; reference KimiK2ForCausalLM,
KimiK2PreTrainedModel, _deinterleave_attention_weights,
register_load_state_dict_pre_hook, and load_state_dict when making the change.
---
Nitpick comments:
In @.claude/agents/ad-debug-agent.md:
- Around line 30-40: Remove the duplicated, typo-containing instruction line
"Remember to use you your own tools - Read, Grep, Glob, Bash, Edit, Write" and
replace it with a single corrected instance "Remember to use your own tools -
Read, Grep, Glob, Bash, Edit, Write"; update the
.claude/agents/ad-debug-agent.md content where that phrase appears so only the
corrected line remains (search for the duplicated string to locate both
occurrences) and ensure surrounding spacing/newlines remain consistent with the
file's formatting.
In `@examples/auto_deploy/model_registry/configs/kimi_k2.yaml`:
- Around line 1-2: The config file kimi_k2.yaml appears to be a
minimal/validation-only setup (comment "minimum layers for validation: 1 dense +
2 MoE = 3 total"); clarify intent by either renaming the file to
kimi_k2_validation.yaml or updating the top-of-file comment to explicitly state
this is a minimal CI/validation/testing config (not production) so readers know
the reduced layer counts are intentional.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py`:
- Around line 530-532: The unconditional casts on
compressed_kv_flat.to(ckv_cache.dtype) and kpe_flat.to(kpe_cache.dtype) should
be made conditional to avoid unnecessary overhead; in the block where
compressed_kv_for_cache and kpe_for_cache are set, check if
compressed_kv_flat.dtype != ckv_cache.dtype and only then assign
compressed_kv_flat.to(ckv_cache.dtype) otherwise reuse compressed_kv_flat, and
similarly for kpe_flat vs kpe_cache.dtype (mirroring the pattern used in
torch_backend_mla.py) so you avoid extra .to() calls and potential tensor view
creation.
In `@tensorrt_llm/_torch/auto_deploy/models/custom/modeling_kimi_k2.py`:
- Around line 440-442: In the forward method of class/model in
modeling_kimi_k2.py, rename the unused extracted shape variables bsz and seq_len
to _bsz and _seq_len (i.e., change "bsz, seq_len, hidden_dim =
hidden_states.shape" to "_bsz, _seq_len, hidden_dim = hidden_states.shape") so
the linter/reader knows they are intentionally unused; keep the rest of the
function (including hidden_states_flat = hidden_states.view(-1, hidden_dim))
unchanged.
- Around line 253-259: The forward method's seq_len parameter is unused (forward
in modeling_kimi_k2.py returns full cached tables _ad_cos_cached and
_ad_sin_cached) so either document that behavior or mark the parameter as
intentionally unused to satisfy linters and readers; update the forward
docstring to mention that seq_len is ignored because full cached tables are
returned for export compatibility, or change the parameter name to _seq_len (or
assign seq_len to a dummy _ = seq_len) inside forward to indicate intentional
nonuse, referencing the forward method and the returned _ad_cos_cached and
_ad_sin_cached symbols.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
.claude/agents/ad-debug-agent.md.claude/agents/onboard-reviewer.md.claude/skills/ad-model-onboard/SKILL.mdexamples/auto_deploy/model_registry/configs/kimi_k2.yamltensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.pytensorrt_llm/_torch/auto_deploy/models/custom/__init__.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_kimi_k2.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_kimi_k2_modeling.py
Collaborator
|
PR_Github #37115 [ run ] triggered by Bot. Commit: |
Contributor
ℹ️ Recent review infoConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro 📒 Files selected for processing (9)
📝 WalkthroughWalkthroughIntroduces the Kimi-K2.5 VLM model to AutoDeploy with custom PyTorch implementations, comprehensive onboarding documentation and validation procedures. Adds model definitions with MLA/MoE architecture, configuration file, updated dtype handling in ops, and extensive test coverage. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37127 [ run ] triggered by Bot. Commit: |
suyoggupta
reviewed
Feb 27, 2026
lucaslie
reviewed
Feb 27, 2026
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37133 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37133 [ run ] completed with state
|
lucaslie
reviewed
Mar 3, 2026
suyoggupta
approved these changes
Mar 3, 2026
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
bmarimuthu-nv
force-pushed
the
bala/model-onboarding-agent
branch
from
March 3, 2026 18:03
20fe8f2 to
603aff0
Compare
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37546 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37546 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37554 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37554 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37574 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37574 [ run ] completed with state
|
StanleySun639
approved these changes
Mar 4, 2026
Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37629 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37629 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --reuse-test --disable-fail-fast |
Collaborator
|
PR_Github #37727 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37727 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --reuse-test |
Collaborator
|
PR_Github #37755 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37755 [ run ] completed with state |
dominicshanshan
pushed a commit
to dominicshanshan/TensorRT-LLM
that referenced
this pull request
Mar 9, 2026
…ing code (NVIDIA#11780) Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
This was referenced Mar 11, 2026
tianyuz-nv
pushed a commit
to wanqian-nv/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
…ing code (NVIDIA#11780) Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
limin2021
pushed a commit
to limin2021/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
…ing code (NVIDIA#11780) Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com>
1 task
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
6 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.
Summary by CodeRabbit
New Features
Bug Fixes
Tests
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)
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.