[None][refactor] Adopt per-family native libraries for K3 modeling v2

cpp/tensorrt_llm/CMakeLists.txt

@@ -149,11 +149,6 @@ 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)
@@ -199,7 +194,6 @@ 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,8 +34,6 @@ 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,12 +76,6 @@ 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,9 +133,6 @@ 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

@@ -163,6 +163,7 @@ tests/unittest/_torch/modeling_v2/
   test_modeling_v2_routing.py    what modeling_v2_resolve does (no GPU)
   _validating.py                 the fixtures that interpose is_valid and arm reference
   <category>/test_modeling_v2_<entry>.py                      shared entries
+  models/<family>/test_<entry>.py                            family-owned entries
   models/<family>/<target>/<category>/test_modeling_v2_<entry>.py   owned entries
   comm/_<entry>_op_matrix.py   the collectives' rank bodies: allgather, reducescatter
                                and allreduce at 2/4/8 ranks in one node, mnnvl_allreduce
@@ -181,8 +182,10 @@ 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.
+target's `ops/`. K3 tests and K3-specific helpers live directly in the flat
+`tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/` directory. Test-list entries
+assign them to CPU, single-GPU, or multi-GPU stages; CBTS uses the standard
+family-path selection.
 Claims discovery covers both family-owned and target-owned entries, and import
 isolation rejects sibling-target and unrelated-family dependencies.
 

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

@@ -1,11 +1,32 @@
+#
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
 # All rights reserved. SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+#
+
+add_modeling_v2_family(kimi_k3_vl)
+
+add_modeling_v2_shared(
+  kimi_k3_vl
+  ARCHS 100f
+  SOURCES *.cu *.cpp
+  DEPS internal_cutlass)
 
-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")
+
+add_modeling_v2_target(kimi_k3_vl kimi_k3_mxfp4__sm_100__tp16_moetp4ep4
+                       ARCHS 100f)
+add_modeling_v2_target(kimi_k3_vl kimi_k3_mxfp4__sm_100__tp16_moetp16ep1
+                       ARCHS 100f)

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu

@@ -13,9 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "k3AttnResDecode.h"
+#include "kimi_k3_vl/csrc/k3AttnResDecode.h"
 #include "tensorrt_llm/common/cudaUtils.h"
 #include "tensorrt_llm/common/envUtils.h"
+#include "tensorrt_llm/modeling_v2/common/namespace.h"
 
 #include <cfloat>
 #include <cooperative_groups.h>
@@ -712,7 +713,7 @@ static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_r
 
 }
 
-TRTLLM_NAMESPACE_BEGIN
+TRTLLM_MV2_NAMESPACE_BEGIN
 
 namespace kernels::kimiK3DecodeAttnRes
 {
@@ -778,4 +779,4 @@ void invokeAttnResAddRmsNormFwd(AttnResFwdParams const& params, cudaStream_t str
 
 }
 
-TRTLLM_NAMESPACE_END
+TRTLLM_MV2_NAMESPACE_END

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h

@@ -17,13 +17,14 @@
 #pragma once
 
 #include "tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h"
+#include "tensorrt_llm/modeling_v2/common/namespace.h"
 
-TRTLLM_NAMESPACE_BEGIN
+TRTLLM_MV2_NAMESPACE_BEGIN
 
 namespace kernels::kimiK3DecodeAttnRes
 {
 
-struct AttnResFwdParams : kimiK3AttnRes::AttnResFwdParams
+struct AttnResFwdParams : ::tensorrt_llm::kernels::kimiK3AttnRes::AttnResFwdParams
 {
     bool earlyTrigger = false;
 };
@@ -33,4 +34,4 @@ void invokeAttnResAddRmsNormFwd(AttnResFwdParams const& params, cudaStream_t str
 
 }
 
-TRTLLM_NAMESPACE_END
+TRTLLM_MV2_NAMESPACE_END

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResOp.cpp

@@ -13,7 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "k3AttnResDecode.h"
+#include "kimi_k3_vl/csrc/k3AttnResDecode.h"
+#include "tensorrt_llm/modeling_v2/common/namespace.h"
+#include "tensorrt_llm/modeling_v2/common/torchLibrary.h"
 
 #include <ATen/Functions.h>
 #include <ATen/cuda/CUDAContext.h>
@@ -25,7 +27,7 @@
 #include <optional>
 #include <tuple>
 
-TRTLLM_NAMESPACE_BEGIN
+TRTLLM_MV2_NAMESPACE_BEGIN
 
 namespace torch_ext
 {
@@ -200,9 +202,9 @@ std::tuple<at::Tensor, at::Tensor> k3_attn_res_add_rmsnorm_fwd(at::Tensor layer_
 
 }
 
-TRTLLM_NAMESPACE_END
+TRTLLM_MV2_NAMESPACE_END
 
-TORCH_LIBRARY_FRAGMENT(trtllm, m)
+TRTLLM_MV2_LIBRARY_FRAGMENT(m)
 {
     m.def(
         "k3_attn_res_rmsnorm_fwd(Tensor layer_residual, Tensor block_residual, "
@@ -214,8 +216,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "float rms_eps, float output_rms_eps, bool early_trigger=False) -> (Tensor, Tensor)");
 }
 
-TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+TRTLLM_MV2_LIBRARY_IMPL(CUDA, m)
 {
-    m.impl("k3_attn_res_rmsnorm_fwd", &tensorrt_llm::torch_ext::k3_attn_res_rmsnorm_fwd);
-    m.impl("k3_attn_res_add_rmsnorm_fwd", &tensorrt_llm::torch_ext::k3_attn_res_add_rmsnorm_fwd);
+    m.impl("k3_attn_res_rmsnorm_fwd", &TRTLLM_MV2_NS::torch_ext::k3_attn_res_rmsnorm_fwd);
+    m.impl("k3_attn_res_add_rmsnorm_fwd", &TRTLLM_MV2_NS::torch_ext::k3_attn_res_add_rmsnorm_fwd);
 }

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu

@@ -13,8 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "k3MnnvlAllreduce.h"
+#include "kimi_k3_vl/csrc/k3MnnvlAllreduce.h"
 #include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/modeling_v2/common/namespace.h"
 #include <cooperative_groups.h>
 #include <cstddef>
 #include <cstdint>
@@ -36,10 +37,11 @@
 #include "tensorrt_llm/common/tllmDataType.h"
 #include "tensorrt_llm/kernels/quantization.cuh"
 
-TRTLLM_NAMESPACE_BEGIN
+TRTLLM_MV2_NAMESPACE_BEGIN
 
 namespace kernels::kimiK3MnnvlOneShot
 {
+using namespace ::tensorrt_llm::kernels;
 
 using tensorrt_llm::common::isNegZero;
 using tensorrt_llm::common::isLamportDirty;
@@ -826,4 +828,4 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params)
 
 }
 
-TRTLLM_NAMESPACE_END
+TRTLLM_MV2_NAMESPACE_END

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h

@@ -17,13 +17,14 @@
 #pragma once
 
 #include "tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h"
+#include "tensorrt_llm/modeling_v2/common/namespace.h"
 
-TRTLLM_NAMESPACE_BEGIN
+TRTLLM_MV2_NAMESPACE_BEGIN
 
 namespace kernels::kimiK3MnnvlOneShot
 {
 
-struct AllReduceFusionParams : mnnvl::AllReduceFusionParams
+struct AllReduceFusionParams : ::tensorrt_llm::kernels::mnnvl::AllReduceFusionParams
 {
     bool earlyTrigger = false;
 };
@@ -32,4 +33,4 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params);
 
 }
 
-TRTLLM_NAMESPACE_END
+TRTLLM_MV2_NAMESPACE_END

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlCommOp.cpp

@@ -14,11 +14,13 @@
  * limitations under the License.
  */
 
-#include "mnnvlAllGatherKernels.h"
-#include "mnnvlAllreduceAttnRes.h"
+#include "kimi_k3_vl/csrc/mnnvlAllGatherKernels.h"
+#include "kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h"
 #include "tensorrt_llm/common/mcastDevMemUtils.h"
 #include "tensorrt_llm/common/tllmDataType.h"
 #include "tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h"
+#include "tensorrt_llm/modeling_v2/common/namespace.h"
+#include "tensorrt_llm/modeling_v2/common/torchLibrary.h"
 #include "tensorrt_llm/runtime/mcastDeviceMemory.h"
 
 #include <ATen/cuda/CUDAContext.h>
@@ -27,7 +29,7 @@
 #include <cstdint>
 #include <vector>
 
-TRTLLM_NAMESPACE_BEGIN
+TRTLLM_MV2_NAMESPACE_BEGIN
 
 namespace torch_ext
 {
@@ -91,7 +93,7 @@ std::vector<torch::Tensor> mnnvlAllReduceAttnRes(torch::Tensor const& input,
     params.output = normOut.mutable_data_ptr();
     params.stream = at::cuda::getCurrentCUDAStream(input.get_device());
 
-    tensorrt_llm::kernels::kimiK3Mnnvl::AttnResEpilogueParams epilogue{};
+    TRTLLM_MV2_NS::kernels::kimiK3Mnnvl::AttnResEpilogueParams epilogue{};
     epilogue.blockResidual = block_residual.const_data_ptr();
     epilogue.resWeight = res_weight.const_data_ptr();
     epilogue.rmsWeight = rms_weight.const_data_ptr();
@@ -100,18 +102,18 @@ std::vector<torch::Tensor> mnnvlAllReduceAttnRes(torch::Tensor const& input,
     epilogue.outputRmsEps = static_cast<float>(output_rms_eps);
     epilogue.numCandidates = static_cast<int>(block_residual.size(0)) + 1;
 
-    tensorrt_llm::kernels::kimiK3Mnnvl::oneshotAllreduceAttnResOp(params, epilogue);
+    TRTLLM_MV2_NS::kernels::kimiK3Mnnvl::oneshotAllreduceAttnResOp(params, epilogue);
     return {normOut, prefixOut};
 }
 
 namespace
 {
 
-tensorrt_llm::kernels::kimiK3Mnnvl::AllGatherSplitParams makeAllGatherSplitParams(torch::Tensor const& input,
+TRTLLM_MV2_NS::kernels::kimiK3Mnnvl::AllGatherSplitParams makeAllGatherSplitParams(torch::Tensor const& input,
     int64_t bf16_columns, int64_t world_size, torch::Tensor& comm_buffer, torch::Tensor& buffer_flags,
     torch::Tensor& bf16Out, torch::Tensor& fp32Out)
 {
-    namespace k3Mnnvl = tensorrt_llm::kernels::kimiK3Mnnvl;
+    namespace k3Mnnvl = TRTLLM_MV2_NS::kernels::kimiK3Mnnvl;
     auto* mcast_mem = tensorrt_llm::common::findMcastDevMemBuffer(comm_buffer.data_ptr());
     TORCH_CHECK(
         mcast_mem != nullptr, "[mnnvlAllGatherSplit] comm_buffer must be obtained from a mcastBuffer instance.");
@@ -159,15 +161,15 @@ std::vector<torch::Tensor> mnnvlAllGatherSplit(torch::Tensor const& input, int64
     torch::Tensor fp32Out;
     auto const params
         = makeAllGatherSplitParams(input, bf16_columns, world_size, comm_buffer, buffer_flags, bf16Out, fp32Out);
-    tensorrt_llm::kernels::kimiK3Mnnvl::mnnvlAllGatherSplitOp(params);
+    TRTLLM_MV2_NS::kernels::kimiK3Mnnvl::mnnvlAllGatherSplitOp(params);
     return {bf16Out, fp32Out};
 }
 
 }
 
-TRTLLM_NAMESPACE_END
+TRTLLM_MV2_NAMESPACE_END
 
-TORCH_LIBRARY_FRAGMENT(trtllm, m)
+TRTLLM_MV2_LIBRARY_FRAGMENT(m)
 {
     m.def(
         "mnnvl_allreduce_attn_res(Tensor input, Tensor? prefix_sum, Tensor block_residual, Tensor res_weight, "
@@ -178,8 +180,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "Tensor(b!) buffer_flags) -> Tensor[]");
 }
 
-TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+TRTLLM_MV2_LIBRARY_IMPL(CUDA, m)
 {
-    m.impl("mnnvl_allreduce_attn_res", &tensorrt_llm::torch_ext::mnnvlAllReduceAttnRes);
-    m.impl("mnnvl_allgather_split", &tensorrt_llm::torch_ext::mnnvlAllGatherSplit);
+    m.impl("mnnvl_allreduce_attn_res", &TRTLLM_MV2_NS::torch_ext::mnnvlAllReduceAttnRes);
+    m.impl("mnnvl_allgather_split", &TRTLLM_MV2_NS::torch_ext::mnnvlAllGatherSplit);
 }