[TRTLLM-14813][feat] Add Kimi K3 (KimiLinear) model by brnguyen2 · Pull Request #17269 · NVIDIA/TensorRT-LLM

brnguyen2 · 2026-08-05T00:13:08Z

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 wrapper
over the general MLA path (modules/kimi_k3_mla, TRTLLM-14811). Also
includes fused_moe SiTu activation + communication_method support and
the SiTu entry in the python ActType_TrtllmGen enum mirror.

Dependencies (all merged to main)

[TRTLLM-14701][feat] Update trtllm-gen batchedGemm kernel drop for Kimi K3 MoE #17190 — K3 MoE trtllm-gen batchedGemm kernels (TRTLLM-14701)
[TRTLLM-14702][feat] Integrate Kimi K3 KDA decode kernel #17054 — KDA decode kernel + torch op (TRTLLM-14702)
[None][feat] add CuteDSL FP8/FP16 MLA decode attention fmha lib #15138 — CuteDSL FP8/FP16 MLA decode fmha lib
[None][feat] Add Kimi K3 KDA prefill/MTP decode CuTe DSL kernels and fused attention-residual kernel #17225 — KDA prefill + MTP-verify CuTe DSL kernels, fused
attention-residual kernel

All dependency PRs have landed; this branch is now a single model commit
on top of current main.

Notes

FP8 KV cache support is not part of this PR; bf16 KV is the validated
default.
Later K3 perf work (fused-A-GEMM mxfp8 dispatch etc.) is intentionally
excluded and will follow separately.

Test Coverage

tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py
tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py
tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py
tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py
tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py

PR Checklist

PR title follows the [TRTLLM-14813][feat] convention
Sanity rerun on the updated base: build green on a Blackwell (sm_100 family) node; import 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 passed
Dependencies merged; full CI run kicked on the main-based single-commit head

Dev Engineer Review

Added Kimi K3 support across configuration, model execution, KDA, MLA, MoE, checkpoint loading, cache management, quantization, mapping, and runtime integration.
Added Blackwell-specific fused kernels with validated PyTorch fallbacks.
Added SiTu activation and MXFP4 expert loading support.
Added optional output reuse to kda_decode.
Added validation for unsupported K3 configurations, cache combinations, routing methods, quantization settings, and hardware paths.
CI results were mixed. The latest pipeline passed x86_64 and SBSA builds and the SBSA single-GPU stage. Four x86_64 DGX B200 shards failed because of the pre-existing get_valid_configs() argument mismatch in FP8 block-scale MoE fallback handling. PR #17362 addresses this unrelated failure.

QA Engineer Review

Added KDA decode, verification, prefill, and cache soundness tests.
Added attention-residual parity tests.
Added Kimi K3 SiTu MoE tests for routing, activation contracts, kernel execution, numerical parity, weight loading, TP/EP selection, and TP8 execution.
Added test_update_quant_config_from_compressed_tensors_mxfp4_with_fp8_kv_cache.
No test-list files were changed.
The new KDA, attention-residual, and Kimi K3 SiTu MoE tests are not listed in the available test-db/ or qa/ files.
Existing list entries cover related Kimi K2/K2.5 and generic MXFP4 MoE tests.
Verdict: needs follow-up.

brnguyen2 · 2026-08-05T00:32:41Z

/bot run

tensorrt-cicd · 2026-08-05T00:38:48Z

PR_Github #63876 [ run ] triggered by Bot. Commit: 108cfbe Link to invocation

tensorrt-cicd · 2026-08-05T02:06:26Z

PR_Github #63876 [ run ] completed with state SUCCESS. Commit: 108cfbe
/LLM/main/L0_MergeRequest_PR pipeline #51815 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-05T02:15:48Z

/bot run

tensorrt-cicd · 2026-08-05T02:21:59Z

PR_Github #63898 [ run ] triggered by Bot. Commit: e40834b Link to invocation

tensorrt-cicd · 2026-08-05T03:57:10Z

PR_Github #63898 [ run ] completed with state SUCCESS. Commit: e40834b
/LLM/main/L0_MergeRequest_PR pipeline #51840 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-05T04:18:13Z

/bot run

tensorrt-cicd · 2026-08-05T04:23:29Z

PR_Github #63930 [ run ] triggered by Bot. Commit: e40834b Link to invocation

tensorrt-cicd · 2026-08-05T06:00:56Z

PR_Github #63930 [ run ] completed with state SUCCESS. Commit: e40834b
/LLM/main/L0_MergeRequest_PR pipeline #51868 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-05T15:38:59Z

/bot run

tensorrt-cicd · 2026-08-05T15:44:58Z

PR_Github #64074 [ run ] triggered by Bot. Commit: 1365842 Link to invocation

tensorrt-cicd · 2026-08-05T17:58:43Z

PR_Github #64074 [ run ] completed with state SUCCESS. Commit: 1365842
/LLM/main/L0_MergeRequest_PR pipeline #52001 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-05T19:08:59Z

/bot run --skip-test

tensorrt-cicd · 2026-08-05T19:15:04Z

PR_Github #64103 [ run ] triggered by Bot. Commit: d7cd8ee Link to invocation

brnguyen2 · 2026-08-05T19:58:47Z

/bot run --skip-test

tensorrt-cicd · 2026-08-05T20:04:54Z

PR_Github #64111 [ run ] triggered by Bot. Commit: f092cd2 Link to invocation

tensorrt-cicd · 2026-08-05T20:07:45Z

PR_Github #64103 [ run ] completed with state ABORTED. Commit: d7cd8ee

Link to invocation

tensorrt-cicd · 2026-08-05T22:03:06Z

PR_Github #64111 [ run ] completed with state SUCCESS. Commit: f092cd2
/LLM/main/L0_MergeRequest_PR pipeline #52033 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

brnguyen2 · 2026-08-05T23:00:10Z

/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)."

tensorrt-cicd · 2026-08-05T23:06:39Z

PR_Github #64129 [ skip ] triggered by Bot. Commit: f092cd2 Link to invocation

tensorrt-cicd · 2026-08-05T23:32:30Z

PR_Github #64129 [ skip ] completed with state SUCCESS. Commit: f092cd2
Skipping testing for commit f092cd2

Link to invocation

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>

brnguyen2 · 2026-08-06T00:00:00Z

/bot run

tensorrt-cicd · 2026-08-06T17:36:29Z

PR_Github #64354 [ run ] completed with state SUCCESS. Commit: c0f2bb9
/LLM/main/L0_MergeRequest_PR pipeline #52248 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Tabrizian · 2026-08-06T17:37:39Z

py_executor changes and disagg changes LGTM.

brnguyen2 · 2026-08-06T17:54:37Z

CI status on head `c0f2bb9` (pipeline 52248) and why a skip is justified

Full L0 ran on the exact merge head:

Both builds green (x86_64 and SBSA).
SBSA single-GPU: fully green (28/28 stages).
x86_64 single-GPU: 208/216 stages green. The 4 failing stages are all DGX_B200 PyTorch shards.
Multi-GPU stages did not run: they were blocked-collateral of the x86_64 single-GPU failure, not independent failures.

The 4 failures share one root cause, and it is not from this PR

Per the CI failure analysis for pipeline 52248, all four failing shards (DeepSeek V3 Lite FP8 perf sanity, Gemma4 dummy e2e, layer-wise benchmarks, LTX2 example) fail on the same error:

RuntimeError: get_valid_configs() is missing value for argument '_6'.

Root cause: #15297 added a numFusedSharedExpert argument to FP8BlockScaleMoERunner::getValidConfigs and updated the get_valid_tactics call site, but missed the memoized fallback path — get_fallback_tactic() in tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py still builds a 5-element key for the now 6-argument binding. Any FP8 block-scale MoE forward that misses the autotuner profiling cache fails hard. The bug is readable straight from main's source and is deterministic once that path is hit; nothing in the path is Kimi-related, and this PR does not modify the failing file. The one-line fix is up as #17362.

All four affected tests are non-Kimi tests exercising FP8 block-scale MoE on B200, which is exactly the blast radius that bug predicts.

Prior run on the previous head

Pipeline 52214 (head b379d30) was pure infrastructure: the test artifact download hit repeated 503s and then a truncated tarball, so no tests executed. The delta between that head and c0f2bb9 is a review-feedback commit (comment cleanups, an eps fallback, defensive guards) that has never been implicated in any test failure.

Additional coverage beyond this pipeline

The stacked serving/spec-dec branch (#17327), which contains this PR's full diff, passed 652 Kimi unit tests plus a speculative-decoding-vs-baseline logits-parity integration run and GSM8K accuracy checks on multi-GPU hardware.

Known residual gap, stated for the record: the multi-GPU L0 stages were blocked by the failure above and so did not run in this pipeline. The multi-GPU validation cited above covers the Kimi paths this PR adds; shared-code risk to other models is bounded by the fully-green single-GPU coverage of the same files.

brnguyen2 · 2026-08-06T17:55:02Z

/bot skip --comment "All 4 failures in pipeline 52248 trace to a single pre-existing main bug (FP8 MoE fallback-tactic signature mismatch from #15297, fixed by #17362) in code this PR does not modify; see previous comment for the full analysis and coverage evidence."

brnguyen2 · 2026-08-06T17:57:49Z

/bot skip --comment "All 4 failures in pipeline 52248 trace to a single pre-existing main bug (FP8 MoE fallback-tactic signature mismatch from #15297, fixed by #17362) in code this PR does not modify; see previous comment for the full analysis and coverage evidence."

tensorrt-cicd · 2026-08-06T22:38:23Z

PR_Github #64388 [ skip ] triggered by Bot. Commit: c883e1d Link to invocation

tensorrt-cicd · 2026-08-06T22:56:43Z

PR_Github #64388 [ skip ] completed with state SUCCESS. Commit: c883e1d
Skipping testing for commit c883e1d

Link to invocation

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

…o_ignore_at_inference Class-level mutable list default could be mutated process-wide via append; a tuple is iteration-only and cannot be shared-mutated. Deferred review nit from PR NVIDIA#17269. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

…iene Deferred review nits from PR NVIDIA#17269: - Hoist stdlib imports (gc, json, ExitStack) and safetensors.safe_open to module level; drop the three redundant local 'import gc' and the function-local json/contextlib/safetensors imports. Remove the json -> _json alias (use json.load directly). - Raise the latent_moe_use_norm precondition assert to the top of KimiK3MoERuntime.__init__, beside the routed_expert_hidden_size assert, so config validation fails before any layer is allocated. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

Deferred review nits from PR NVIDIA#17269: - Build the runtime config from the real KimiLinearConfig instead of a SimpleNamespace mock. - Hoist the KimiKDARuntime import (and the new config import) to module level. - Replace the two 'with torch.no_grad():' blocks with @torch.no_grad() decorators on _make_runtime and the test, dropping an indent level. - Drop the 'if __name__ == "__main__"' runner (pytest-only). The per-layer replay-cache mocks stay SimpleNamespace: they mirror the cache-manager's KDA slot allocation, which has no standalone class to instantiate in a unit test. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

Register the K3 (KimiLinear) unit suites in the single-GPU Blackwell pre-merge list so they run in CI: KDA modeling parity tests, the kimi_kda module suites, the attn-res op test, and the SiTU MoE parity test. Deferred from PR NVIDIA#17269 (tests shipped but absent from any list). qa/ flat-list enablement is tracked separately under TRTLLM-15036. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

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

github-actions Bot assigned brnguyen2 Aug 5, 2026

brnguyen2 force-pushed the k3/14813-kimilinear-model branch from e40834b to 1365842 Compare August 5, 2026 15:38

brnguyen2 force-pushed the k3/14813-kimilinear-model branch from 1365842 to d7cd8ee Compare August 5, 2026 19:08

brnguyen2 force-pushed the k3/14813-kimilinear-model branch from f092cd2 to 82fab8c Compare August 5, 2026 23:59

brnguyen2 marked this pull request as ready for review August 6, 2026 00:05

brnguyen2 requested review from a team as code owners August 6, 2026 00:05

brnguyen2 requested review from jiaganc, liji-nv and rosong11 August 6, 2026 14:50

rosong11 approved these changes Aug 6, 2026

View reviewed changes

Tabrizian approved these changes Aug 6, 2026

View reviewed changes

brnguyen2 enabled auto-merge (squash) August 6, 2026 17:56

Merge branch 'main' into k3/14813-kimilinear-model

c883e1d

coderabbitai Bot mentioned this pull request Aug 6, 2026

[#15853][fix] Gate B12x MoE on CUDA13 CuTe DSL #15851

Open

2ez4bz approved these changes Aug 6, 2026

View reviewed changes

schetlur-nv approved these changes Aug 6, 2026

View reviewed changes

brnguyen2 merged commit 36922c7 into NVIDIA:main Aug 6, 2026
9 of 10 checks passed

coderabbitai Bot mentioned this pull request Aug 7, 2026

[TRTLLM-14955][refactor] Declare MoE backend behaviour instead of comparing classes #17411

Open

1 task

This was referenced Aug 7, 2026

[TRTLLM-15177][chore] Kimi K3 post-merge cleanup: config/import/test hygiene + L0 wiring #17413

Open

[TRTLLM-15177][chore] Consolidate trtllm-gen SiTu activation slot handling #17414

Open

brnguyen2 mentioned this pull request Aug 7, 2026

[None][fix] Kimi K3: eager CUDA-graph buffer allocation and prebuilt fused-verify constants #17421

Merged

Uh oh!

Conversation

brnguyen2 commented Aug 5, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Dependencies (all merged to main)

Notes

Test Coverage

PR Checklist

Dev Engineer Review

QA Engineer Review

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

brnguyen2 commented Aug 6, 2026

Uh oh!

tensorrt-cicd commented Aug 6, 2026

Uh oh!

Tabrizian commented Aug 6, 2026

Uh oh!

brnguyen2 commented Aug 6, 2026

CI status on head c0f2bb9 (pipeline 52248) and why a skip is justified

The 4 failures share one root cause, and it is not from this PR

Prior run on the previous head

Additional coverage beyond this pipeline

Uh oh!

brnguyen2 commented Aug 6, 2026

Uh oh!

brnguyen2 commented Aug 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

brnguyen2 commented Aug 5, 2026 •

edited by coderabbitai Bot

Loading

CI status on head `c0f2bb9` (pipeline 52248) and why a skip is justified