[TRTLLM-14813][feat] Add Kimi K3 (KimiLinear) model - #17269
Conversation
|
/bot run |
|
PR_Github #63876 [ run ] triggered by Bot. Commit: |
|
PR_Github #63876 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63898 [ run ] triggered by Bot. Commit: |
|
PR_Github #63898 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63930 [ run ] triggered by Bot. Commit: |
|
PR_Github #63930 [ run ] completed with state
|
e40834b to
1365842
Compare
|
/bot run |
|
PR_Github #64074 [ run ] triggered by Bot. Commit: |
|
PR_Github #64074 [ run ] completed with state
|
1365842 to
d7cd8ee
Compare
|
/bot run --skip-test |
|
PR_Github #64103 [ run ] triggered by Bot. Commit: |
|
/bot run --skip-test |
|
PR_Github #64111 [ run ] triggered by Bot. Commit: |
|
PR_Github #64103 [ run ] completed with state |
|
PR_Github #64111 [ run ] completed with state |
|
/bot skip --comment "Full pipeline 52001 on the pre-rebase head failed only on AutoDeploy MoE unit tests already failing on main and waived there (nvbugs 6564714, waives merged via #17162), plus cascade aborts from those stages. The branch was then rebased onto current main (which carries those waives) with byte-identical content, adding only a fake-impl registration for the new custom op; that fix passed the H100 unit-test stage on PR #17225 (pipeline 52037), and a build-only run on this exact head passed (pipeline 52033)." |
|
PR_Github #64129 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64129 [ skip ] completed with state |
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
f092cd2 to
82fab8c
Compare
|
/bot run |
VALLIS-NERIA
left a comment
VALLIS-NERIA
left a comment
There was a problem hiding this comment.
Looks good. The comment is non-blocking
|
PR_Github #64257 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64318 [ run ] triggered by Bot. Commit: |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #64318 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64354 [ run ] triggered by Bot. Commit: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (1)
234-241: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winDo not treat filesystem errors as a missing Kimi configuration.
Line 234 can return
Falsewhenstat()cannot access the path, even ifconfig.jsonexists. (docs.python.org) If one rank takes this result, it enters the eager path and executes the collectives at Lines 213-217, while another rank enters the lazy path. The job can deadlock.Open the file directly and catch only
FileNotFoundError. Propagate other filesystem errors.Proposed fix
config_path = os.path.join(checkpoint_dir, "config.json") - if not os.path.isfile(config_path): - return False - with open(config_path) as f: - model_type = json.load(f).get("model_type") + try: + with open(config_path) as f: + model_type = json.load(f).get("model_type") + except FileNotFoundError: + return False return model_type in ("kimi_k3", "kimi_linear")As per coding guidelines, catch the narrowest exception possible.
🤖 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/checkpoints/hf/weight_loader.py` around lines 234 - 241, Update the Kimi configuration check around config_path and model_type to open config.json directly, catching only FileNotFoundError to return False for a genuinely missing file. Remove the os.path.isfile precheck so permission, I/O, and other filesystem errors propagate consistently across ranks.Source: Coding guidelines
🧹 Nitpick comments (1)
tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py (1)
315-315: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse PEP 604 optional annotations.
Change the two annotations to
torch.Tensor | None. The project supports Python 3.10+, and the coding guidelines prefer|unions.🤖 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/custom_ops/cpp_custom_ops.py` at line 315, Update the two Optional[torch.Tensor] annotations in the affected function signature to use torch.Tensor | None, and remove the now-unused Optional import if applicable.Sources: Coding guidelines, Learnings
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 234-241: Update the Kimi configuration check around config_path
and model_type to open config.json directly, catching only FileNotFoundError to
return False for a genuinely missing file. Remove the os.path.isfile precheck so
permission, I/O, and other filesystem errors propagate consistently across
ranks.
---
Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py`:
- Line 315: Update the two Optional[torch.Tensor] annotations in the affected
function signature to use torch.Tensor | None, and remove the now-unused
Optional import if applicable.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 307b17e2-eccc-4be5-8c10-ca19ac2bea3a
📒 Files selected for processing (5)
tensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytensorrt_llm/_torch/models/modeling_kimi_linear.pytensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.pytensorrt_llm/_torch/modules/mla.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py
- tensorrt_llm/_torch/modules/mla.py
- tensorrt_llm/_torch/models/modeling_kimi_linear.py
|
PR_Github #64354 [ run ] completed with state
|
|
py_executor changes and disagg changes LGTM. |
CI status on head
|
|
|
||
| class KimiLinearConfig(PretrainedConfig): | ||
| model_type = "kimi_linear" | ||
| keys_to_ignore_at_inference = ["past_key_values"] |
There was a problem hiding this comment.
Would it still work if this is a tuple instead of a list? Otherwise this is globally mutable through e.g. append.
| model_type = json.load(f).get("model_type") | ||
| return model_type in ("kimi_k3", "kimi_linear") | ||
|
|
||
| def _load_lazy_safetensors( |
There was a problem hiding this comment.
Could be left as a followup (although please create a JIRA ticket + reference here in TODO comment?) if the existing code doesn't allow for this: should we have a Kimi K3 specific weight loader that encapsulates all this? Otherwise my fear is that this HfWeightLoader class eventually becomes a hodge-podge of if / else for various model-specific business. If the logic in here is generic enough, we could even have it be a HfLazyWeightLoader.
| ``forward``), with no FP8 dequant path. Returns the number of projections | ||
| converted. | ||
| """ | ||
| import gc |
There was a problem hiding this comment.
Nit: consider making the import module-level (it's a standard library package).
|
|
||
| count = 0 | ||
|
|
||
| def _swap(parent: nn.Module, attr: str) -> None: |
There was a problem hiding this comment.
This local function seems to be duplicated in several functions. Any chance for it to be made a common helper function?
| self.routed_expert_up_proj = nn.Linear( | ||
| self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype | ||
| ) | ||
| assert getattr(cfg, "latent_moe_use_norm", False), ( |
There was a problem hiding this comment.
Nit: move this earlier in the function (raise earlier)?
| ) | ||
| rt = KimiKDARuntime(cfg, layer_idx=0).to("cuda") | ||
| gen = torch.Generator(device="cuda").manual_seed(seed) | ||
| with torch.no_grad(): |
There was a problem hiding this comment.
Nit: this whole function could be wrapped in @torch.no_grad()? Saves an indentation layer.
|
|
||
|
|
||
| def _make_runtime(seed): | ||
| from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime |
| cache[:, :, : W - 1] = conv_pool[:, section * d : (section + 1) * d, 1:].float() | ||
| return cache | ||
|
|
||
| return SimpleNamespace( |
There was a problem hiding this comment.
Nit: is there anyway we could use the actual object instead of this SimpleNamespace? It's usually better to use something as close to the actual object instead of mocks.
Same comment throughout the rest of the file.
| ).to(torch.bfloat16) | ||
|
|
||
| ok = True | ||
| with torch.no_grad(): |
There was a problem hiding this comment.
Nit: @torch.no_grad() the entire function?
| assert ok | ||
|
|
||
|
|
||
| if __name__ == "__main__": |
There was a problem hiding this comment.
Nit: please remove. We never execute these tests as executables, and only through pytest anyway.
11 participants
Description
Adds the Kimi K3 (KimiLinear) model: hybrid KDA/MLA layer stack with the
K3 MoE blocks, HF checkpoint config and weight loading, executor
integration (config utils, resource manager, mamba cache-manager sizing
for KDA state, mapping), chunked prefill enabled by default with KV-cache
block reuse as an opt-in, and quant-config handling for the K3 checkpoint
format. Model-level defaults enable the fused attention-residual path,
CUDA graphs, and the CuTe DSL kernels from the dependency PRs.
Ships with the K3 module packages the model builds on — the KDA mixer +
kernel glue (
modules/kimi_kda), the K3 SITU MoE block/gate/MLP(
modules/kimi_k3_moe), the fused attention-residual op wrapper(
modules/kimi_k3_attn_res), and the K3 MLA module as a thin wrapperover the general MLA path (
modules/kimi_k3_mla, TRTLLM-14811). Alsoincludes fused_moe SiTu activation + communication_method support and
the SiTu entry in the python
ActType_TrtllmGenenum mirror.Dependencies (all merged to main)
attention-residual kernel
All dependency PRs have landed; this branch is now a single model commit
on top of current main.
Notes
default.
excluded and will follow separately.
Test Coverage
PR Checklist
[TRTLLM-14813][feat]conventionimport tensorrt_llm+ model module import clean; unit suites green — kimi_kda 35 passed/2 skipped, attn_res 10 passed, situ_moe 23 passed, KDA modeling parity 14 passedDev Engineer Review
kda_decode.get_valid_configs()argument mismatch in FP8 block-scale MoE fallback handling. PR#17362addresses this unrelated failure.QA Engineer Review
test_update_quant_config_from_compressed_tensors_mxfp4_with_fp8_kv_cache.test-db/orqa/files.