[None][chore] Fix Kimi_k25 with spec dec

[ziyixiong-nv]

Summary by CodeRabbit

Release Notes

• New Features

  • Draft model support properties now accessible for advanced configurations

• Improvements

  • Explicit multimodal parameter handling in the forward interface for clearer configuration
  • Enhanced vision embedding computation for multimodal requests
  • Improved token processing for inputs containing media elements

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-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.

[coderabbitai]

📝 Walkthrough

Walkthrough

KimiK25ForConditionalGeneration gains draft-related pass-through properties (draft_config, draft_model, load_draft_weights) and its forward method now explicitly accepts multimodal_params as an optional argument. The multimodal path computes vision embeddings, validates placeholder token presence, constructs token mappings, filters kwargs keys, and propagates remaining kwargs into the underlying LLM call.

Changes

Model Interface and Multimodal Path Updates

Layer / File(s)	Summary
Draft property delegation tensorrt_llm/_torch/models/modeling_kimi_k25.py	KimiK25ForConditionalGeneration exposes draft_config, draft_model, and load_draft_weights as properties that delegate to the wrapped self.llm.
Multimodal forward refactoring and kwargs passthrough tensorrt_llm/_torch/models/modeling_kimi_k25.py	The forward method signature is updated to accept multimodal_params explicitly. When multimodal inputs are present, vision embeddings are computed for the first attn_metadata.num_contexts requests, embeddings are discarded if no media placeholder tokens exist in input_ids, mm_token_ids are constructed, and mm_token_indices/text_token_indices are filtered from fuse_kwargs before multimodal embedding fusion. The final self.llm.forward call includes **kwargs.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Suggested reviewers

• ZhanruiSunCh
• reasonsolo

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description contains only the repository template with no actual content filling the required sections: Description, Test Coverage, and PR Checklist remain empty.	Complete the Description section explaining the issue and solution, list relevant tests in Test Coverage, and check off applicable PR Checklist items before merging.
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the specific model (Kimi_k25) and the purpose (fix with spec dec), directly matching the code changes to the KimiK25ForConditionalGeneration model.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tensorrt_llm/_torch/models/modeling_kimi_k25.py (1)

1584-1595: ⚡ Quick win

Add explicit return type annotations to the new delegated properties.

These new property methods currently have no return annotations, which weakens static typing on this public interface.

♻️ Suggested update

 `@property`
-    def draft_config(self):
+    def draft_config(self) -> Any:
         return self.llm.draft_config
 
 `@property`
-    def draft_model(self):
+    def draft_model(self) -> Any:
         return self.llm.draft_model
 
 `@property`
-    def load_draft_weights(self):
+    def load_draft_weights(self) -> Any:
         return self.llm.load_draft_weights

As per coding guidelines "Static type checking via mypy is opt-in by submodule and highly recommended; always annotate Python functions with return types (use None if no return); make return type explicit."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_kimi_k25.py` around lines 1584 - 1595,
The three new property getters draft_config, draft_model, and load_draft_weights
lack return type annotations; update each `@property` method to include an
explicit return type that matches the delegated attribute on self.llm (e.g., the
actual DraftConfig/DraftModel types or the callable signature for
load_draft_weights). Import the concrete types from the module that defines the
llm interface or use typing.Any as a temporary fallback, and annotate the
methods as: def draft_config(self) -> <Type>, def draft_model(self) -> <Type>,
def load_draft_weights(self) -> <Type>. Ensure the annotations reflect the
public API types used elsewhere.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_k25.py`:
- Around line 1645-1679: The code currently assumes input_ids exists and only
strips mm_token_indices/text_token_indices from kwargs in the branch where
mm_embeds are kept, which can cause a crash when input_ids is None or allow
unwanted kwargs to reach fuse_input_embeds; update the block around
multimodal_params/mm_embeds so you first check input_ids is not None before
checking for placeholder tokens (use self._media_placeholder_token_id and guard
the int((input_ids == placeholder_id).sum().item()) call), and always sanitize
fuse_kwargs by removing "mm_token_indices" and "text_token_indices" from kwargs
(regardless of whether embeddings were found) before calling fuse_input_embeds;
refer to multimodal_params, mm_embeds, mm_token_ids, fuse_kwargs,
fuse_input_embeds and _media_placeholder_token_id when making the changes.

---

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_k25.py`:
- Around line 1584-1595: The three new property getters draft_config,
draft_model, and load_draft_weights lack return type annotations; update each
`@property` method to include an explicit return type that matches the delegated
attribute on self.llm (e.g., the actual DraftConfig/DraftModel types or the
callable signature for load_draft_weights). Import the concrete types from the
module that defines the llm interface or use typing.Any as a temporary fallback,
and annotate the methods as: def draft_config(self) -> <Type>, def
draft_model(self) -> <Type>, def load_draft_weights(self) -> <Type>. Ensure the
annotations reflect the public API types used elsewhere.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2efd1c67-8b9e-4226-9d7f-fc1194379e83

📥 Commits

Reviewing files that changed from the base of the PR and between 3b8387c and 7b29a1e.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/models/modeling_kimi_k25.py

[ziyixiong-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49557 [ run ] triggered by Bot. Commit: 216c906 Link to invocation

[tensorrt-cicd]

PR_Github #49557 [ run ] completed with state SUCCESS. Commit: 216c906
/LLM/main/L0_MergeRequest_PR pipeline #39185 completed with status: 'SUCCESS'

CI Report

Link to invocation

[tianyuxbear]

/bot run --extra-stage GB200-4_GPUs-PyTorch-PerfSanity-Post-Merge-1

[tensorrt-cicd]

PR_Github #49659 [ run ] triggered by Bot. Commit: 216c906 Link to invocation

[tensorrt-cicd]

PR_Github #49659 [ run ] completed with state SUCCESS. Commit: 216c906
/LLM/main/L0_MergeRequest_PR pipeline #39273 completed with status: 'SUCCESS'

CI Report

Link to invocation