[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 |
|
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
|
|
PR_Github #64388 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64388 [ skip ] completed with state |
…swap helper Behavior-neutral refactors deferred from NVIDIA#17269: - Extract the duplicated FP8 weight-read module swap (nested _swap in the MoE-MLP and MLA converters, plus the inline KDA o_proj conversion) into a module-level _swap_linear_to_fp8_weight_read helper. - Split the ~450-line KimiLinearForCausalLM.load_weights into focused methods: _validate_checkpoint_keys, _load_trunk_params, _load_expert_slices, and _finalize_weight_load, with load_weights as a short orchestrator. Code moved verbatim; no functional changes. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…dling Item 1.2 of the PR NVIDIA#17269 deferred-cleanup bucket. (a) Record the decision to keep SiTu in the trtllm-gen-local ActType_TrtllmGen enum rather than adding it to the shared ActivationType, which mirrors the cutlass enum in common.h and would then need a cutlass member no cutlass kernel implements. Documented at the enum and left the enums unchanged. (b) Route SiTu and SwiGLU per-expert alpha/beta through a single _gemm1_activation_params accessor at the op call instead of duplicating the is_situ_activation ternary. Behavior-identical: same tensors, same op slot. The deeper storage merge (eliminating the separate situ_alpha/situ_beta buffers by reusing swiglu_alpha/swiglu_beta) is left as a TODO because 'swiglu_alpha is not None' gates quant-method selection and validation on the numeric path; it needs GPU parity revalidation before it can land. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
12 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.