[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 #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>
…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>
…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>
…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>
…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>
…mm plumbing Deferred cleanup from PR NVIDIA#17269: - Move K3-specific runtime modules (kimi_k3_mla, kimi_k3_moe) into modeling_kimi_linear.py per the modeling_xxx.py convention; delete the empty module dirs and repoint unit-test imports. kimi_kda stays standalone. - Remove the unused communication_method parameter chain (create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy); the only caller passed None and TRTLLM_FORCE_COMM_METHOD already forces a strategy. - Use a tuple for KimiLinearConfig.keys_to_ignore_at_inference so the class-level default can't be mutated in place. No functional change. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…mm plumbing Deferred cleanup from PR NVIDIA#17269: - Move K3-specific runtime modules (kimi_k3_mla, kimi_k3_moe) into modeling_kimi_linear.py per the modeling_xxx.py convention; delete the empty module dirs and repoint unit-test imports. kimi_kda stays standalone. - Remove the unused communication_method parameter chain (create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy); the only caller passed None and TRTLLM_FORCE_COMM_METHOD already forces a strategy. - Use a tuple for KimiLinearConfig.keys_to_ignore_at_inference so the class-level default can't be mutated in place. No functional change. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…plumbing Deferred cleanup from PR NVIDIA#17269, tracked in TRTLLM-15177. Redo of PR kimi_k3_moe module (via the shared modules/situ.py), so only the remaining items are carried over here. 1. Inline the K3-specific kimi_k3_mla module into modeling_kimi_linear.py, matching the per-model modeling_xxx.py convention (e.g. DeepSeek-V3). The moved code is unchanged apart from dropping a redundant local torch import and following the file's Linear-as-TrtllmLinear alias. kimi_kda stays a standalone module (general enough to warrant it). 2. Remove the unused communication_method parameter chain create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy. The only caller (modeling_kimi_linear.py) passed None, and TRTLLM_FORCE_COMM_METHOD already provides strategy forcing. The unit test covering the forwarding is deleted with it. 3. Tuple default for KimiLinearConfig.keys_to_ignore_at_inference, so the class-level default cannot be mutated in place. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…plumbing Deferred cleanup from PR NVIDIA#17269, tracked in TRTLLM-15177. Redo of PR kimi_k3_moe module (via the shared modules/situ.py), so only the remaining items are carried over here. 1. Inline the K3-specific kimi_k3_mla module into modeling_kimi_linear.py, matching the per-model modeling_xxx.py convention (e.g. DeepSeek-V3). The moved code is unchanged apart from dropping a redundant local torch import and following the file's Linear-as-TrtllmLinear alias. kimi_kda stays a standalone module (general enough to warrant it). 2. Remove the unused communication_method parameter chain create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy. The only caller (modeling_kimi_linear.py) passed None, and TRTLLM_FORCE_COMM_METHOD already provides strategy forcing. The unit test covering the forwarding is deleted with it. 3. Tuple default for KimiLinearConfig.keys_to_ignore_at_inference, so the class-level default cannot be mutated in place. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…plumbing Deferred cleanup from PR NVIDIA#17269, tracked in TRTLLM-15177. Redo of PR kimi_k3_moe module (via the shared modules/situ.py), so only the remaining items are carried over here. 1. Inline the K3-specific kimi_k3_mla module into modeling_kimi_linear.py, matching the per-model modeling_xxx.py convention (e.g. DeepSeek-V3). The moved code is unchanged apart from dropping a redundant local torch import and following the file's Linear-as-TrtllmLinear alias. kimi_kda stays a standalone module (general enough to warrant it). 2. Remove the unused communication_method parameter chain create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy. The only caller (modeling_kimi_linear.py) passed None, and TRTLLM_FORCE_COMM_METHOD already provides strategy forcing. The unit test covering the forwarding is deleted with it. 3. Tuple default for KimiLinearConfig.keys_to_ignore_at_inference, so the class-level default cannot be mutated in place. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…plumbing Deferred cleanup from PR NVIDIA#17269, tracked in TRTLLM-15177. Redo of PR kimi_k3_moe module (via the shared modules/situ.py), so only the remaining items are carried over here. 1. Inline the K3-specific kimi_k3_mla module into modeling_kimi_linear.py, matching the per-model modeling_xxx.py convention (e.g. DeepSeek-V3). The moved code is unchanged apart from dropping a redundant local torch import and following the file's Linear-as-TrtllmLinear alias. kimi_kda stays a standalone module (general enough to warrant it). 2. Remove the unused communication_method parameter chain create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy. The only caller (modeling_kimi_linear.py) passed None, and TRTLLM_FORCE_COMM_METHOD already provides strategy forcing. The unit test covering the forwarding is deleted with it. 3. Tuple default for KimiLinearConfig.keys_to_ignore_at_inference, so the class-level default cannot be mutated in place. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…rs and decouple from HfWeightLoader Follow-up to the Kimi K3 / KimiLinear support in NVIDIA#17269. That PR shipped a load path validated end-to-end on a real 1.5 TB K3 checkpoint, but it reached the lazy-loading behavior through model-name conditionals in shared code and configured K3 loading through a block of environment variables. This removes both, without changing runtime behavior. 1. Generalize lazy-safetensors into a selectable LoadFormat. Opening safetensors shards and streaming only the rank-local slices, instead of materializing the full checkpoint in host RAM, becomes LoadFormat.LAZY_SAFETENSORS, reachable via load_format="lazy_safetensors" for any HF checkpoint. The sharded-vs-consolidated file selection is handled inside the format. K3 opts in by declaring it as its default in get_model_defaults; an explicit user load_format still wins. 2. Decouple the shared HfWeightLoader from K3. The Kimi/K3 model-name branch (_is_kimi_k3_checkpoint and its dispatch) is removed; the shared loader now dispatches on the resolved LoadFormat, not on any model name. No kimi/k3 string checks remain in weight_loader.py. 3. Move the K3 FP8 weight-read env vars onto the config surface. These 5 knobs decide whether a replicated K3 projection is read from an FP8 (e4m3, 128x128 block-scale) copy of its weights instead of BF16, and they are consumed on the checkpoint-loading path: load_weights keeps the FP8_PB_WO checkpoint pairs only when the read is enabled, and the post-load conversion swaps the modules. They therefore move to QuantConfig. Resolution lives in the new, CPU-unit-testable kimi_k3_knobs.py (config value wins, else the deprecated env var honored with a one-time warning, else the historical default). Deprecated env var -> config equivalent: KIMI_K3_FP8_WEIGHT_READ -> quant_config.kimi_k3_fp8_weight_read KIMI_K3_FP8_WEIGHT_READ_KDA -> quant_config.kimi_k3_fp8_weight_read_kda KIMI_K3_FP8_WEIGHT_READ_MLA -> quant_config.kimi_k3_fp8_weight_read_mla KIMI_K3_FP8_WEIGHT_READ_GATE_UP -> quant_config.kimi_k3_fp8_weight_read_gate_up KIMI_K3_KDA_GLUE_FP8 -> quant_config.kimi_k3_kda_glue_fp8 The env vars are still honored for back-compat, emitting a one-time deprecation warning; an explicit config value takes precedence when both are set. The four remaining K3 env vars (TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ, KIMI_K3_MLA_MAX_POSITIONS, KIMI_K3_FUSED_ATTN_RES, KIMI_K3_ROUTER_BF16) are left as they are. They configure the forward path -- the latent MoE GEMM selection, the identity-RoPE table size, the fused attn_res op, and the router GEMM dtype -- not loading or quantization, and none of them is required for K3 to load. Migrating them, and deciding which of them belong on the public API surface at all, is tracked separately. Test coverage: - tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py exercises the LAZY_SAFETENSORS format on a synthetic multi-shard fixture (sharded and consolidated selection, lazy slice materialization, dtype/shape/values, FP8 weights and scale tensors), no large checkpoint required. - tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py covers the 5 gates: config defaults, sub-gate narrowing, the is_sm_100f() arch gate, deprecated-env back-compat, config-over-env precedence, warn-once, and a zero-env/zero-config completeness check. - API-stability reference updated for the new QuantConfig fields. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…rs and decouple from HfWeightLoader Follow-up to the Kimi K3 / KimiLinear support in NVIDIA#17269. That PR shipped a load path validated end-to-end on a real 1.5 TB K3 checkpoint, but it reached the lazy-loading behavior through model-name conditionals in shared code and configured K3 loading through a block of environment variables. This removes both, without changing runtime behavior. 1. Generalize lazy-safetensors into a selectable LoadFormat. Opening safetensors shards and streaming only the rank-local slices, instead of materializing the full checkpoint in host RAM, becomes LoadFormat.LAZY_SAFETENSORS, reachable via load_format="lazy_safetensors" for any HF checkpoint. The sharded-vs-consolidated file selection is handled inside the format. K3 opts in by declaring it as its default in get_model_defaults; an explicit user load_format still wins. 2. Decouple the shared HfWeightLoader from K3. The Kimi/K3 model-name branch (_is_kimi_k3_checkpoint and its dispatch) is removed; the shared loader now dispatches on the resolved LoadFormat, not on any model name. No kimi/k3 string checks remain in weight_loader.py. 3. Move the K3 FP8 weight-read env vars onto the config surface. These 5 knobs decide whether a replicated K3 projection is read from an FP8 (e4m3, 128x128 block-scale) copy of its weights instead of BF16, and they are consumed on the checkpoint-loading path: load_weights keeps the FP8_PB_WO checkpoint pairs only when the read is enabled, and the post-load conversion swaps the modules. They therefore move to QuantConfig. Resolution lives in the new, CPU-unit-testable kimi_k3_knobs.py (config value wins, else the deprecated env var honored with a one-time warning, else the historical default). Deprecated env var -> config equivalent: KIMI_K3_FP8_WEIGHT_READ -> quant_config.kimi_k3_fp8_weight_read KIMI_K3_FP8_WEIGHT_READ_KDA -> quant_config.kimi_k3_fp8_weight_read_kda KIMI_K3_FP8_WEIGHT_READ_MLA -> quant_config.kimi_k3_fp8_weight_read_mla KIMI_K3_FP8_WEIGHT_READ_GATE_UP -> quant_config.kimi_k3_fp8_weight_read_gate_up KIMI_K3_KDA_GLUE_FP8 -> quant_config.kimi_k3_kda_glue_fp8 The env vars are still honored for back-compat, emitting a one-time deprecation warning; an explicit config value takes precedence when both are set. The four remaining K3 env vars (TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ, KIMI_K3_MLA_MAX_POSITIONS, KIMI_K3_FUSED_ATTN_RES, KIMI_K3_ROUTER_BF16) are left as they are. They configure the forward path -- the latent MoE GEMM selection, the identity-RoPE table size, the fused attn_res op, and the router GEMM dtype -- not loading or quantization, and none of them is required for K3 to load. Migrating them, and deciding which of them belong on the public API surface at all, is tracked separately. Test coverage: - tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py exercises the LAZY_SAFETENSORS format on a synthetic multi-shard fixture (sharded and consolidated selection, lazy slice materialization, dtype/shape/values, FP8 weights and scale tensors), no large checkpoint required. - tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py covers the 5 gates: config defaults, sub-gate narrowing, the is_sm_100f() arch gate, deprecated-env back-compat, config-over-env precedence, warn-once, and a zero-env/zero-config completeness check. - API-stability reference updated for the new QuantConfig fields. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
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.