[None][feat] Consolidate Kimi K3 native modeling work into main (#19990, #19991)

cpp/tensorrt_llm/CMakeLists.txt

@@ -149,6 +149,11 @@ set_target_properties(
                                  CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 
 add_subdirectory(kernels)
+set(KIMI_K3_MODEL_CSRC
+    "${CMAKE_CURRENT_SOURCE_DIR}/../../tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc"
+)
+add_subdirectory("${KIMI_K3_MODEL_CSRC}"
+                 "${CMAKE_CURRENT_BINARY_DIR}/kimi_k3_model")
 add_subdirectory(runtime)
 
 set(BATCH_MANAGER_TARGET tensorrt_llm_batch_manager_static)
@@ -194,6 +199,7 @@ set(TRTLLM_LINK_LIBS
     mamba2_mtp_ssm_cache_src
     kda_decode_src
     kimi_k3_attn_res_src
+    kimi_k3_model_src
     ws_layernorm_src
     fusedGatedRMSNormQuant_src
     fpA_intB_gemm_src

cpp/tensorrt_llm/modeling_v2/CMakeLists.txt

@@ -34,6 +34,8 @@ file(
   GLOB family_lists
   LIST_DIRECTORIES false
   CONFIGURE_DEPENDS "${modeling_v2_models_dir}/*/csrc/CMakeLists.txt")
+# K3 objects and torch bindings are linked through the stock libraries.
+list(FILTER family_lists EXCLUDE REGEX "/kimi_k3_vl/csrc/CMakeLists[.]txt$")
 list(SORT family_lists)
 
 # A family removed or renamed since an earlier configure leaves its build

cpp/tensorrt_llm/thop/CMakeLists.txt

@@ -76,6 +76,12 @@ add_th_common_family(
   nvfp4MlaKvCacheGatherOp.cpp
   trtllmGenQKVProcessOp.cpp)
 
+target_sources(
+  th_common_attention
+  PRIVATE "${KIMI_K3_MODEL_CSRC}/k3AttnResOp.cpp"
+          "${KIMI_K3_MODEL_CSRC}/k3MnnvlCommOp.cpp"
+          "${KIMI_K3_MODEL_CSRC}/k3MnnvlFusionOp.cpp")
+
 add_th_common_family(
   th_common_gemm_quant
   cutlassScaledMM.cpp

jenkins/scripts/cbts/cbts/rules/modeling_v2.py

@@ -133,6 +133,9 @@ def _is_mv2_claim(path: str) -> bool:
     fixture, so the rule keeps claiming it and re-runs the stages (safe
     over-run).
     """
+    # K3 native code and fakes reach stock imports, so baseline selection applies.
+    if path.startswith(_MV2_SRC_PREFIXES[0] + "models/kimi_k3_vl/"):
+        return False
     if not (path.startswith(_MV2_SRC_PREFIXES) or path in _CPP_BUILD_FILES):
         return False
     if is_docs_path(path):

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -6,10 +6,11 @@ built-in model zoo rather than inside it.
 Where `_torch/models/modeling_deepseekv3.py` is one class serving V3, V3-Lite,
 R1 and V3.2 across every GPU generation and parallel topology,
 `_torch/_experimental/modeling_v2/models/deepseek_v3/` is one flat forward per (checkpoint,
-GPU architecture, parallel topology) triple — assembled only from `catalog/`
-entries, sharing nothing with its siblings, and trusted through accuracy gates
-instead of shared abstractions. The one-to-one correspondence between
-`models/<x>/` here and `modeling_<x>.py` there is the point of the exercise.
+GPU architecture, parallel topology) triple — assembled from certified `OpWrapper`
+entries and trusted through accuracy gates. Model-family ops and kernels may be
+shared by that family's targets; target-only code remains inside its target.
+The one-to-one correspondence between `models/<x>/` here and `modeling_<x>.py`
+there is the point of the exercise.
 
 ## Using it
 
@@ -114,6 +115,8 @@ _router_index.py    architectures[0] -> routing module. Small, stable, test-guar
 _native.py          loads a target's native code (see "Native code" below)
 models/<family>/
   routing.py        one forward-reading decision tree per architecture family
+  ops/<category>/   entries shared only by this family
+  kernels/          family-owned CuTe DSL packages
   fake_ops.py       fakes of the family's shared native ops, if it has any
   csrc/             the family's CMakeLists.txt and the native code its targets share
   <checkpoint>__<gpu arch>__<parallel>/
@@ -131,6 +134,19 @@ torch, `is_valid` names the inputs it accepts and answers wrongly, and `CELLS`
 wrapper is stateless and `__call__` reads nothing off `self`, so a call site
 shows exactly what the kernel receives; `catalog/_op.py` states the rest.
 
+Kimi K3's `csrc/` is compiled into the existing native libraries through CMake
+OBJECT targets and `target_sources`; it has no runtime extension loader. Its
+one-shot MNNVL specialization and fused attention-residual decode kernels are
+family-owned. The generic attention-residual, persistent fallback and two-shot
+MNNVL kernels remain stock, as do the shared native schemas and workspace
+ownership. Rebuild the native libraries when these model-owned sources change.
+
+K3's all-reduce adapter keeps its one-shot ceiling and entry-release setting on
+the selected model's modules. Stock PDL releases after rank reduction; K3 can
+release immediately after its entry dependency wait. Neither release implies
+output readiness: a dependent must wait before reading producer outputs or
+reusing its workspace. Shared kernels and defaults are unchanged.
+
 Accuracy anchors are not kept here. They live where every other model's do,
 `tests/integration/defs/accuracy/references/`, read by the gates CI runs.
 
@@ -146,7 +162,8 @@ tests/unittest/_torch/modeling_v2/
   comm/_<entry>_op_matrix.py   the collectives' rank bodies: allgather, reducescatter
                                and allreduce at 2/4/8 ranks in one node, mnnvl_allreduce
                                at 8 ranks over two trays (srun only, see its docstring)
-  comm/_rank_job.py            starts one of those and asserts on its exit code
+  comm/_lockstep.py            the stateful matrices' launcher and shared reference
+  comm/_rank_job.py            starts one at a world size and asserts on its exit code
 ```
 
 That split is not a preference; it is where this repo's CI collects from, and
@@ -155,13 +172,65 @@ an in-package test is on no list. Each test is parametrized off its entry's
 and derives, so the certification is re-proved on every commit rather than
 recorded once.
 
+Kimi K3's shared TP4×EP4 / TP16×EP1 wrappers and CuTe packages live once under
+`models/kimi_k3_vl/{ops,kernels}/`. The TP16×EP1 M1/M2 MoE engines live under
+that target's `kernels/k3_moe_engines/`; target-only wrappers live under their
+target's `ops/`. Their existing tests retain their paths under the category
+and `tests/unittest/_torch/kernels/kimi_k3/` trees so CI selectors do not change.
+Claims discovery covers both family-owned and target-owned entries, and import
+isolation rejects sibling-target and unrelated-family dependencies.
+
+The Markov chain declares its Lamport buffers, rotation flags and optional
+KV-length rewind as mutable custom-op inputs. Its no-rewind call uses a
+separate cached length buffer, allocated before CUDA-graph capture.
+
+Model-owned native fake implementations register with the standard native-op
+initialization, without importing a target. KDA verify consumes the cache
+manager's bf16 convolution pools and independent per-slot strides directly;
+standalone fp32 convolution pools remain supported. Pool views are never
+copied or cast on the call path.
+
+KDA prefill takes raw packed BF16 Q/K/V, gate and beta logits with mandatory
+int64 sequence offsets, FP32 gate parameters and indexed BF16/FP32 state rows.
+Normalization and gate/beta activation run inside the shared prefill op; the
+model-owned wrapper uses its current 16-argument contract, including optional
+sequence-order and chunk-prefix metadata.
+
+KDA verify's sequential FP32 FMA recurrence and main's chunked 3xTF32
+tensor-core recurrence agree numerically, not bit-for-bit. With the same
+external gate, 48 SM100 rounds across dense and strided pools measured maximum
+state divergence of 6.31e-6 relative to the largest reference-state entry.
+The cross-implementation gate is 1e-5; the existing 2e-2 output and 1e-3
+FP64-reference state bands remain unchanged. Repeated same-implementation
+states and outputs, convolution caches and graph replay remain bit-exact.
+
+MLA launch admission queries the launch stream's owning context resources,
+including converted green contexts, without querying stream execution during
+graph capture. Invalid resources still reject the launch. Wide MoE states use
+the canonical decode build for up to eight tokens, sharing their preallocated
+scratch through prefix views; larger steps retain the wide build.
+
+
 The collectives' rank bodies sit in the tests tree with everything else
 that only tests run. Both halves are started by file path -- the launcher must
 not import `tensorrt_llm`, because that calls `MPI_Init` and an
 MPI-initialized process cannot start `mpirun`, and the ranks reach the catalog
 by absolute import -- so neither needs a package to live in. Neither those
 file names nor their `check_*` bodies match pytest's collection patterns:
-each is one fixed 4-rank sequence that cannot run as independent cases.
+each is one fixed sequence over one job's ranks that cannot run as
+independent cases. CI runs them at 4 ranks on one node. A stateful entry's
+launcher also takes `--world-size N --launcher srun`, which runs the same
+body as one of N ranks started across nodes: that is how its multi-node
+receipt is recorded.
+
+Stateful matrices compare ordered per-tensor SHA-256 fingerprints of raw
+bytes together with shape and dtype across ranks. The large KDA-pool tests
+check free CUDA memory before allocating their full pools, including 1 GiB
+of headroom for weights, small reference pools and compilation.
+
+Routing and claims tests run on the CPU list; GPU lists exclude their
+`cpu_only` marker. KDA post-convolution cells exercise both channel-major and
+transposed token-major inputs accepted by the shared kernel.
 
 Identity is the directory name, and it carries all three segments:
 `gpt_oss_120b__sm_10x__default`. They were three nested directories once, which
@@ -171,6 +240,11 @@ carries the same triple; `test_modeling_v2_claims.py` asserts they agree.
 
 ### Native code
 
+Kimi K3 retains its stock-linked `kimi_k3_model_src` objects, `libth_common` torch bindings,
+and existing operator names and fake registrations. Its `csrc/` directory is excluded from
+per-family library discovery. The per-family mechanism below applies to the other native
+families.
+
 A target's C++/CUDA lives with its Python code: `<target>/csrc/`, and the family's
 `csrc/` for code several targets share. Each family builds into one torch-op library,
 `libth_modeling_v2_<family>.so`, installed next to `libth_common.so`; nothing of it is
@@ -347,7 +421,13 @@ Perf is measured, never gated.
 
 **Every entry the gpt-oss and deepseek default targets call is certified on
 sm_100 and sm_103, and both targets are gated on both.** The qwen4_exp target
-and the entries it owns are certified on sm_103.
+and the entries it owns are certified on sm_103. Kimi K3's sm_100 operation
+receipts remain in its adjacent records; its generic-operation regressions
+exercise the relocated entries.
+
+The K3 targets use the engine-selected `KimiK3HybridCacheManager`; its
+model-owned replay state supplies the pending accepted-draft counts consumed
+by fused acceptance and KDA verification.
 
 The default targets' 21 entries -- 20 shared, 1 (`fused_moe`) owned by deepseek's default target -- each drive
 their cells through `__call__` and `reference`. Where that runs depends on the
@@ -359,6 +439,11 @@ entry:
 * `mnnvl_allreduce` needs 8 ranks over two trays. Its matrix runs only by
   hand under srun (see its docstring), and no CI job re-runs it.
 
+Matrix subprocesses use the launcher's interpreter and inherited Python search
+path to import the installed TensorRT-LLM package. Install a matching wheel or
+complete editable build before running them; an unbuilt checkout must not be
+prepended to `PYTHONPATH` in place of the installed package.
+
 Certification is per architecture: `ARCHS` says
 where the cells were driven, and a missing architecture reads as unknown,
 which is the honest state.

tensorrt_llm/_torch/_experimental/modeling_v2/models/__init__.py

@@ -14,4 +14,5 @@
 from .deepseek_v3 import routing as deepseek_v3_routing  # noqa: F401
 from .glm_moe_dsa import routing as glm_moe_dsa_routing  # noqa: F401
 from .gpt_oss import routing as gpt_oss_routing  # noqa: F401
+from .kimi_k3_vl import routing as kimi_k3_vl_routing  # noqa: F401
 from .qwen4_exp import routing as qwen4_exp_routing  # noqa: F401

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""KimiK3ForConditionalGeneration targets."""

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/.clang-format

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+BasedOnStyle: InheritParentConfig
+FixNamespaceComments: false

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/CMakeLists.txt

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved. SPDX-License-Identifier: Apache-2.0
+
+add_library(
+  kimi_k3_model_src OBJECT k3AttnResDecode.cu k3MnnvlAllreduce.cu
+                           mnnvlAllGatherKernels.cu mnnvlAllreduceAttnRes.cu)
+set_cuda_architectures(kimi_k3_model_src 100f)
+set_property(TARGET kimi_k3_model_src PROPERTY POSITION_INDEPENDENT_CODE ON)
+set_property(TARGET kimi_k3_model_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
+set_source_files_properties(k3AttnResDecode.cu PROPERTIES COMPILE_OPTIONS
+                                                          "--use_fast_math")