From 212610da28ac3606b4cc48ad01353bcc054a1d27 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Wed, 28 Oct 2026 10:14:39 -0700 Subject: [PATCH 01/11] [None][refactor] isolate K3 modeling v2 native operators by family Signed-off-by: Tao Li --- cpp/tensorrt_llm/CMakeLists.txt | 6 -- cpp/tensorrt_llm/modeling_v2/CMakeLists.txt | 2 - cpp/tensorrt_llm/thop/CMakeLists.txt | 6 -- .../scripts/cbts/cbts/rules/modeling_v2.py | 3 - .../models/kimi_k3_vl/csrc/CMakeLists.txt | 33 +++++-- .../models/kimi_k3_vl/csrc/k3AttnResDecode.cu | 8 +- .../models/kimi_k3_vl/csrc/k3AttnResDecode.h | 6 +- .../models/kimi_k3_vl/csrc/k3AttnResOp.cpp | 16 ++-- .../kimi_k3_vl/csrc/k3MnnvlAllreduce.cu | 8 +- .../models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h | 6 +- .../models/kimi_k3_vl/csrc/k3MnnvlCommOp.cpp | 28 +++--- .../kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp | 18 ++-- .../kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu | 8 +- .../kimi_k3_vl/csrc/mnnvlAllGatherKernels.h | 6 +- .../kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.cu | 8 +- .../kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h | 6 +- .../{ops/_native.py => fake_ops.py} | 14 ++- .../__init__.py | 4 + .../modeling.py | 5 +- .../__init__.py | 4 + .../modeling.py | 5 +- .../models/kimi_k3_vl/ops/__init__.py | 4 +- .../ops/comm/mnnvl_allgather_split.py | 3 +- .../kimi_k3_vl/ops/comm/mnnvl_allreduce.py | 3 +- .../ops/comm/mnnvl_allreduce_attn_res.py | 3 +- .../ops/comm/mnnvl_fusion_allreduce.py | 3 +- .../ops/norm/attn_res_rmsnorm_fwd.py | 3 +- .../_torch/custom_ops/cpp_custom_ops.py | 2 - .../comm/_mnnvl_allgather_split_op_matrix.py | 4 +- .../kimi_k3_vl/test_modeling_v2_k3_native.py | 43 +++++++++ .../test_modeling_v2_attn_res_rmsnorm_fwd.py | 7 +- .../test_modeling_v2_native_claims.py | 91 ++----------------- .../scripts/cbts/test_modeling_v2_rule.py | 6 +- 33 files changed, 191 insertions(+), 181 deletions(-) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/{ops/_native.py => fake_ops.py} (89%) create mode 100644 tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 1e13b43ef27a..92fb45dc57cf 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/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 diff --git a/cpp/tensorrt_llm/modeling_v2/CMakeLists.txt b/cpp/tensorrt_llm/modeling_v2/CMakeLists.txt index aeb3965a035a..3785811be8cf 100644 --- a/cpp/tensorrt_llm/modeling_v2/CMakeLists.txt +++ b/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 diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 43485b14f4ce..fa7c0a458fa2 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/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 diff --git a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py index 1da3b6e1ed88..fee9b67e260b 100644 --- a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py +++ b/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): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/CMakeLists.txt b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/CMakeLists.txt index deb70b340b09..06d46e32eb89 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/CMakeLists.txt +++ b/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) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu index 22eaee9115e2..18b065fe5321 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu +++ b/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 #include @@ -712,10 +713,11 @@ 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 { +using namespace ::tensorrt_llm::kernels; template static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStream_t stream) @@ -778,4 +780,4 @@ void invokeAttnResAddRmsNormFwd(AttnResFwdParams const& params, cudaStream_t str } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h index 3f4285866974..0b57f05219fd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h @@ -17,11 +17,13 @@ #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 { +using namespace ::tensorrt_llm::kernels; struct AttnResFwdParams : kimiK3AttnRes::AttnResFwdParams { @@ -33,4 +35,4 @@ void invokeAttnResAddRmsNormFwd(AttnResFwdParams const& params, cudaStream_t str } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResOp.cpp b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResOp.cpp index 4425feef0590..ed7167ee127c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResOp.cpp +++ b/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 #include @@ -25,7 +27,7 @@ #include #include -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace torch_ext { @@ -200,9 +202,9 @@ std::tuple 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); } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu index 697b2ad3bfe8..7f015ae19923 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu +++ b/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 #include #include @@ -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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h index f80fd7b4a7c8..f6132b02dcd9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h @@ -17,11 +17,13 @@ #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 { +using namespace ::tensorrt_llm::kernels; struct AllReduceFusionParams : mnnvl::AllReduceFusionParams { @@ -32,4 +34,4 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params); } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlCommOp.cpp b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlCommOp.cpp index 98405a00df61..234f42454099 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlCommOp.cpp +++ b/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 @@ -27,7 +29,7 @@ #include #include -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace torch_ext { @@ -91,7 +93,7 @@ std::vector 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 mnnvlAllReduceAttnRes(torch::Tensor const& input, epilogue.outputRmsEps = static_cast(output_rms_eps); epilogue.numCandidates = static_cast(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 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); } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp index 32a177d32566..c07561130299 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp @@ -14,10 +14,12 @@ * 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/mcastDevMemUtils.h" #include "tensorrt_llm/kernels/customAllReduceKernels.h" #include "tensorrt_llm/kernels/quantization.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" +#include "tensorrt_llm/modeling_v2/common/torchLibrary.h" #include "tensorrt_llm/runtime/mcastDeviceMemory.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/thUtils.h" @@ -28,7 +30,7 @@ using tensorrt_llm::kernels::AllReduceFusionOp; -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace torch_ext { @@ -174,7 +176,7 @@ std::vector k3MnnvlFusionAllReduce(torch::Tensor& input, torch::o } } - auto allreduce_params = tensorrt_llm::kernels::kimiK3MnnvlOneShot::AllReduceFusionParams(); + auto allreduce_params = TRTLLM_MV2_NS::kernels::kimiK3MnnvlOneShot::AllReduceFusionParams(); allreduce_params.nRanks = mcast_mem->getWorldSize(); allreduce_params.rank = mcast_mem->getRank(); allreduce_params.dType = dtype; @@ -219,7 +221,7 @@ std::vector k3MnnvlFusionAllReduce(torch::Tensor& input, torch::o if (numTokens * hiddenDim * allreduce_params.nRanks * input.itemsize() <= static_cast(one_shot_max_bytes)) { - tensorrt_llm::kernels::kimiK3MnnvlOneShot::oneshotAllreduceFusionOp(allreduce_params); + TRTLLM_MV2_NS::kernels::kimiK3MnnvlOneShot::oneshotAllreduceFusionOp(allreduce_params); } else { @@ -244,9 +246,9 @@ std::vector k3MnnvlFusionAllReduce(torch::Tensor& input, torch::o } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END -TORCH_LIBRARY_FRAGMENT(trtllm, m) +TRTLLM_MV2_LIBRARY_FRAGMENT(m) { m.def( "k3_mnnvl_fusion_allreduce(Tensor input, Tensor? gamma, Tensor? residual, " @@ -254,7 +256,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor? scale=None, int fusion_op=0, int one_shot_max_bytes=1048576, bool early_trigger=False) -> Tensor[]"); } -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +TRTLLM_MV2_LIBRARY_IMPL(CUDA, m) { - m.impl("k3_mnnvl_fusion_allreduce", &tensorrt_llm::torch_ext::k3MnnvlFusionAllReduce); + m.impl("k3_mnnvl_fusion_allreduce", &TRTLLM_MV2_NS::torch_ext::k3MnnvlFusionAllReduce); } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu index bb9e2d82ecfc..e253cc3f90e5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu @@ -14,16 +14,18 @@ * limitations under the License. */ -#include "mnnvlAllGatherKernels.h" +#include "kimi_k3_vl/csrc/mnnvlAllGatherKernels.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/lamportUtils.cuh" -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { +using namespace ::tensorrt_llm::kernels; using tensorrt_llm::common::isLamportDirty; using tensorrt_llm::common::LamportFlags; @@ -157,4 +159,4 @@ void mnnvlAllGatherSplitOp(AllGatherSplitParams const& params) } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h index bafc45bf7d6a..c6122471603f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h @@ -17,15 +17,17 @@ #pragma once #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" #include #include #include -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { +using namespace ::tensorrt_llm::kernels; struct AllGatherSplitParams { @@ -50,4 +52,4 @@ void mnnvlAllGatherSplitOp(AllGatherSplitParams const& params); } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.cu index 01b9d5278d81..63a609517992 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.cu +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.cu @@ -13,8 +13,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "mnnvlAllreduceAttnRes.h" +#include "kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h" #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" #include #include #include @@ -37,10 +38,11 @@ #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/quantization.cuh" -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { +using namespace ::tensorrt_llm::kernels; using mnnvl::AllReduceFusionParams; using tensorrt_llm::common::isNegZero; @@ -683,4 +685,4 @@ void oneshotAllreduceAttnResOp(AllReduceFusionParams const& params, AttnResEpilo } } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h index e30052506361..43d4ca90d54b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h @@ -17,11 +17,13 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" -TRTLLM_NAMESPACE_BEGIN +TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { +using namespace ::tensorrt_llm::kernels; struct AttnResEpilogueParams { @@ -38,4 +40,4 @@ void oneshotAllreduceAttnResOp(mnnvl::AllReduceFusionParams const& params, AttnR } -TRTLLM_NAMESPACE_END +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/_native.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/fake_ops.py similarity index 89% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/_native.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/fake_ops.py index 8cf7ec483b86..fd0671e3b90f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/_native.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/fake_ops.py @@ -7,8 +7,12 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils +from ..._native import namespace_of -@torch.library.register_fake("trtllm::k3_mnnvl_fusion_allreduce") +_NS = namespace_of(__name__) + + +@torch.library.register_fake(f"{_NS}::k3_mnnvl_fusion_allreduce") def _( input, gamma, @@ -53,7 +57,7 @@ def _( return [torch.empty_like(input)] -@torch.library.register_fake("trtllm::k3_attn_res_rmsnorm_fwd") +@torch.library.register_fake(f"{_NS}::k3_attn_res_rmsnorm_fwd") def _( layer_residual: torch.Tensor, block_residual: torch.Tensor, @@ -67,7 +71,7 @@ def _( return torch.empty_like(layer_residual) -@torch.library.register_fake("trtllm::k3_attn_res_add_rmsnorm_fwd") +@torch.library.register_fake(f"{_NS}::k3_attn_res_add_rmsnorm_fwd") def _( layer_residual: torch.Tensor, layer_residual_add: torch.Tensor, @@ -82,7 +86,7 @@ def _( return torch.empty_like(layer_residual), torch.empty_like(layer_residual) -@torch.library.register_fake("trtllm::mnnvl_allgather_split") +@torch.library.register_fake(f"{_NS}::mnnvl_allgather_split") def _(input, bf16_columns, world_size, comm_buffer, buffer_flags): num_tokens, columns = input.shape return [ @@ -91,7 +95,7 @@ def _(input, bf16_columns, world_size, comm_buffer, buffer_flags): ] -@torch.library.register_fake("trtllm::mnnvl_allreduce_attn_res") +@torch.library.register_fake(f"{_NS}::mnnvl_allreduce_attn_res") def _(input, prefix_sum, block_residual, res_weight, rms_weight, output_rms_weight, rms_eps, output_rms_eps, comm_buffer, buffer_flags): # fmt: skip return [torch.empty_like(input), torch.empty_like(input)] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py index 2b24a35badfd..eba167c488f1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py @@ -1,3 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 16 x moe_ep 1.""" + +from ...._native import load_target_ops + +native_ops = load_target_ops(__name__) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index 0ba1801fb9f0..2649dfc1ea5f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -98,6 +98,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_mla.decode_view import ( k3_mla_decode_view, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.attention.k3_drafter_attn_qknorm import ( k3_drafter_attn_qknorm, ) @@ -647,7 +648,7 @@ def _apply_attn_res_rmsnorm_fused( _note_attn_res_fusion("attn_res+norm", False, M, H, N) return None try: - attn_res_rmsnorm_op = torch.ops.trtllm.k3_attn_res_rmsnorm_fwd + attn_res_rmsnorm_op = native_ops.k3_attn_res_rmsnorm_fwd except (AttributeError, RuntimeError): return None layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() @@ -721,7 +722,7 @@ def _apply_attn_res_add_rmsnorm_fused( _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) return None try: - attn_res_add_rmsnorm_op = torch.ops.trtllm.k3_attn_res_add_rmsnorm_fwd + attn_res_add_rmsnorm_op = native_ops.k3_attn_res_add_rmsnorm_fwd except (AttributeError, RuntimeError): return None layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py index 474e865afce4..727d94347449 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py @@ -1,3 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4.""" + +from ...._native import load_target_ops + +native_ops = load_target_ops(__name__) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index 6a5f6dfd42c1..7b6d5597b484 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -96,6 +96,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_mla.decode_view import ( k3_mla_decode_view, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.attention.k3_drafter_attn_qknorm import ( k3_drafter_attn_qknorm, ) @@ -641,7 +642,7 @@ def _apply_attn_res_rmsnorm_fused( _note_attn_res_fusion("attn_res+norm", False, M, H, N) return None try: - attn_res_rmsnorm_op = torch.ops.trtllm.k3_attn_res_rmsnorm_fwd + attn_res_rmsnorm_op = native_ops.k3_attn_res_rmsnorm_fwd except (AttributeError, RuntimeError): return None layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() @@ -715,7 +716,7 @@ def _apply_attn_res_add_rmsnorm_fused( _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) return None try: - attn_res_add_rmsnorm_op = torch.ops.trtllm.k3_attn_res_add_rmsnorm_fwd + attn_res_add_rmsnorm_op = native_ops.k3_attn_res_add_rmsnorm_fwd except (AttributeError, RuntimeError): return None layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py index 81a61039f370..7e0d1509ec00 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py @@ -1,4 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from . import _native as _native +from ...._native import load_family_ops + +native_ops = load_family_ops(__name__) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py index f6c544132102..7a67c64c6882 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py @@ -9,6 +9,7 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from .mnnvl_workspace import MnnvlWorkspace @@ -136,7 +137,7 @@ def __call__( f"mnnvl_allgather_split: the call needs {need} bytes per Lamport buffer, the workspace has " f"{workspace.buffer_bytes}" ) - bf16_out, fp32_out = torch.ops.trtllm.mnnvl_allgather_split( + bf16_out, fp32_out = native_ops.mnnvl_allgather_split( input, bf16_columns, workspace.world_size, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py index 9270b4fa00a8..f7bbbb185b9e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py @@ -3,6 +3,7 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm._torch.distributed import AllReduceParams from tensorrt_llm._torch.distributed.ops import ( MNNVLAllReduce, @@ -56,7 +57,7 @@ def forward( is_fusion = fusion_op != AllReduceFusionOp.NONE if is_fusion and fusion_op not in MNNVLAllReduce.SUPPORTED_FUSION_OPS: return None - outputs = torch.ops.trtllm.k3_mnnvl_fusion_allreduce( + outputs = native_ops.k3_mnnvl_fusion_allreduce( input, all_reduce_params.norm_weight, all_reduce_params.residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py index 52e3c786b96c..6d2f6acf3dd6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py @@ -9,6 +9,7 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from .mnnvl_workspace import MnnvlWorkspace @@ -140,7 +141,7 @@ def __call__( output_rms_eps: float, workspace: MnnvlWorkspace, ) -> Tuple[torch.Tensor, torch.Tensor]: - normed, updated = torch.ops.trtllm.mnnvl_allreduce_attn_res( + normed, updated = native_ops.mnnvl_allreduce_attn_res( input, prefix_sum, block_residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py index 8ee91b6ea636..15ceccb84068 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py @@ -9,6 +9,7 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm.functional import AllReduceFusionOp from .mnnvl_workspace import MnnvlWorkspace @@ -199,7 +200,7 @@ def __call__( f"not {workspace.buffer_bytes}" ) fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM if fused else AllReduceFusionOp.NONE - outputs = torch.ops.trtllm.k3_mnnvl_fusion_allreduce( + outputs = native_ops.k3_mnnvl_fusion_allreduce( input, norm_weight, residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py index 1eaaf48da12b..5ebd3fcd9b66 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py @@ -6,6 +6,7 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops class AttnResRmsnormFwd(OpWrapper): @@ -93,7 +94,7 @@ def __call__( rms_eps: float, output_rms_eps: float, ) -> torch.Tensor: - return torch.ops.trtllm.k3_attn_res_rmsnorm_fwd( + return native_ops.k3_attn_res_rmsnorm_fwd( layer_residual, block_residual, res_weight, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 427e863acada..1a1afaf02929 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -15,8 +15,6 @@ def _register_fake(): - from .._experimental.modeling_v2.models.kimi_k3_vl.ops import \ - _native as _native @torch.library.register_fake("trtllm::allreduce") def allreduce( diff --git a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py index 9aff8e6820f4..77fb30abe4f0 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py @@ -35,6 +35,8 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops + sys.path.insert(0, str(Path(__file__).resolve().parent)) import _lockstep as ls # noqa: E402 @@ -436,7 +438,7 @@ def check_unsupported_calls_raise_on_every_rank() -> None: expect_refusal(lambda: allgather_split(rows.bfloat16(), 8, WS_A), RuntimeError, "bf16 rows") expect_refusal(lambda: allgather_split(rows[:0], 8, WS_A), RuntimeError, "no rows") expect_refusal( - lambda: torch.ops.trtllm.mnnvl_allgather_split( + lambda: native_ops.mnnvl_allgather_split( rows, 8, R.world + 1, WS_A.comm_buffer(torch.bfloat16), WS_A.buffer_flags ), RuntimeError, diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py new file mode 100644 index 000000000000..5c9a4af7f1c0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise K3 family packaging, both target declarations, dispatch and fake registration.""" + +import importlib + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="K3 native code requires SM100-family GPUs", +) + + +def test_k3_family_native_ops(): + package = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl" + for target in ( + "kimi_k3_mxfp4__sm_100__tp16_moetp4ep4", + "kimi_k3_mxfp4__sm_100__tp16_moetp16ep1", + ): + importlib.import_module(f"{package}.{target}") + ops = load_family_ops(package) + # A single candidate has softmax weight one; unit RMS weights leave ones unchanged. + residual = torch.ones((2, 1, 7168), device="cuda", dtype=torch.bfloat16) + blocks = residual.new_empty((0, 2, 1, 7168)) + weight = residual.new_ones((7168,)) + actual = ops.k3_attn_res_rmsnorm_fwd(residual, blocks, weight, weight, weight, 1e-6, 1e-6) + torch.testing.assert_close(actual, residual, rtol=0, atol=0) + with FakeTensorMode() as mode: + fake = ops.k3_attn_res_rmsnorm_fwd( + mode.from_tensor(residual), + mode.from_tensor(blocks), + mode.from_tensor(weight), + mode.from_tensor(weight), + mode.from_tensor(weight), + 1e-6, + 1e-6, + ) + assert fake.shape == actual.shape and fake.dtype == actual.dtype diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py index 529d918d874f..8b3777490fec 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py @@ -17,6 +17,7 @@ import pytest import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.norm.attn_res_rmsnorm_fwd import ( attn_res_rmsnorm_fwd, ) @@ -154,15 +155,13 @@ def test_k3_fused_add_preserves_rounding_and_graph_inputs( layer, block, res_weight, rms_weight, output_weight = inputs addend = layer.roll(1, dims=-1) * 0.25 args = (layer, addend, block, res_weight, rms_weight, output_weight, RMS_EPS, RMS_EPS) - updated, normalized = torch.ops.trtllm.k3_attn_res_add_rmsnorm_fwd( - *args, early_trigger=early_trigger - ) + updated, normalized = native_ops.k3_attn_res_add_rmsnorm_fwd(*args, early_trigger=early_trigger) for replay in (False, True): if replay: torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - updated, normalized = torch.ops.trtllm.k3_attn_res_add_rmsnorm_fwd( + updated, normalized = native_ops.k3_attn_res_add_rmsnorm_fwd( *args, early_trigger=early_trigger ) layer.copy_(layer.roll(3, dims=-1)) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py index 99e3c2a7a4a0..b6087e13faff 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py @@ -9,8 +9,7 @@ nothing, a target never depending on a sibling, and CBTS testing a family-only change with that family's tests alone (``jenkins/scripts/cbts/cbts/rules/modeling_v2.py``). The tests here guard the -rules that keep it true for per-family libraries; stock-linked K3 retains its -explicit compatibility path. None imports a target or needs a GPU. +rules that keep it true for per-family libraries. None imports a target or needs a GPU. The paths are the source checkout's, found from this file: these are claims about the tree a PR changes, C++ included, not about an installed wheel. @@ -22,7 +21,6 @@ import re import shutil import subprocess -import sys from pathlib import Path import pytest @@ -42,8 +40,6 @@ pytest.mark.skipif(not _INFRA.is_dir(), reason="needs the source checkout"), ] -_STOCK_LINKED_FAMILIES = frozenset({"kimi_k3_vl"}) - _NATIVE_SUFFIXES = (".cu", ".cuh", ".h", ".hpp", ".inl", ".cpp") # A target name's sm segment that names a family of architectures. _SM_GROUPS = {"10x": ((10, 0), (10, 3))} @@ -86,8 +82,6 @@ def _families() -> dict[str, Path]: """family -> its csrc/ directory, from every models//csrc/CMakeLists.txt.""" families = {} for cmake in sorted(_MODELS.glob("*/csrc/CMakeLists.txt")): - if cmake.parent.parent.name in _STOCK_LINKED_FAMILIES: - continue calls = _calls(cmake.read_text(), "add_modeling_v2_family") assert len(calls) == 1 and len(calls[0]) == 1, ( f"{cmake.relative_to(_REPO)} must call add_modeling_v2_family() exactly once" @@ -245,31 +239,16 @@ def test_native_code_stays_in_its_own_unit(): _STOCK_FAMILY_PATH = re.compile(r"_experimental/modeling_v2/models/[a-z]") -_STOCK_COMPATIBILITY_REFERENCES = { - "tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py": ( - "from .._experimental.modeling_v2.models.kimi_k3_vl.ops import", - ), - "cpp/tensorrt_llm/CMakeLists.txt": ("_experimental/modeling_v2/models/kimi_k3_vl/csrc",), -} - - -def _stock_source(path: Path) -> str: - """Read stock code with its explicit stock-linked K3 references removed.""" - text = path.read_text(errors="ignore") - for reference in _STOCK_COMPATIBILITY_REFERENCES.get(path.relative_to(_REPO).as_posix(), ()): - text = text.replace(reference, "") - return text - - def test_stock_does_not_reach_modeling_v2(): - """Except for stock-linked K3 bindings, stock Python reaches modeling_v2 only through the resolver (no target + """Stock Python reaches modeling_v2 only through the resolver (no target import, op call or fake), and no stock CMake, C++ or packaging file names a family: a family reaches the build only through cpp/tensorrt_llm/modeling_v2/'s discovery, so adding one edits none of them.""" offenders = [ str(path.relative_to(_REPO)) for path in sorted((_REPO / "tensorrt_llm").rglob("*.py")) - if not path.is_relative_to(_PY) and _STOCK_PY_REFERENCE.search(_stock_source(path)) + if not path.is_relative_to(_PY) + and _STOCK_PY_REFERENCE.search(path.read_text(errors="ignore")) ] for root in ("cpp/tensorrt_llm", "cpp/include", "cpp/tests", "cpp/cmake"): for path in sorted((_REPO / root).rglob("*")): @@ -282,10 +261,10 @@ def test_stock_does_not_reach_modeling_v2(): pattern = _STOCK_FAMILY_PATH else: continue - if pattern.search(_stock_source(path)): + if pattern.search(path.read_text(errors="ignore")): offenders.append(str(path.relative_to(_REPO))) for name in ("cpp/CMakeLists.txt", "setup.py", "scripts/build_wheel.py"): - if _STOCK_FAMILY_PATH.search(_stock_source(_REPO / name)): + if _STOCK_FAMILY_PATH.search((_REPO / name).read_text()): offenders.append(name) assert not offenders, f"stock code reaches modeling_v2 targets or families: {offenders}" @@ -409,61 +388,3 @@ def test_cmake_and_python_agree_on_names_and_architectures(tmp_path): assert _native._namespace("gpt_oss", "gpt_oss_120b__sm_10x__default") == ( "mv2_gpt_oss_gpt_oss_120b_sm_10x_default" ) - - -@pytest.mark.parametrize( - "relative, reference, pattern", - [ - ( - "tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py", - "from .._experimental.modeling_v2.models.kimi_k3_vl.ops import _native", - _STOCK_PY_REFERENCE, - ), - ( - "cpp/tensorrt_llm/CMakeLists.txt", - '"_experimental/modeling_v2/models/kimi_k3_vl/csrc"', - _STOCK_FAMILY_PATH, - ), - ], -) -def test_stock_reference_exception_is_k3_only(tmp_path, monkeypatch, relative, reference, pattern): - monkeypatch.setattr(sys.modules[__name__], "_REPO", tmp_path) - path = tmp_path / relative - path.parent.mkdir(parents=True) - path.write_text(reference) - assert not pattern.search(_stock_source(path)) - path.write_text(reference + "\n" + reference.replace("kimi_k3_vl", "other_family")) - assert pattern.search(_stock_source(path)) - other = tmp_path / "unlisted.py" - other.write_text(reference) - assert pattern.search(_stock_source(other)) - - -@pytest.mark.skipif(shutil.which("cmake") is None, reason="needs cmake") -def test_stock_linked_k3_is_not_discovered_twice(tmp_path): - project = tmp_path / "cpp" - project.mkdir() - models = tmp_path / "tensorrt_llm/_torch/_experimental/modeling_v2/models" - for family, content in ( - ("kimi_k3_vl", "add_library(kimi_k3_model_src INTERFACE)"), - ("other_family", "add_custom_target(discovered_native_family)"), - ): - csrc = models / family / "csrc" - csrc.mkdir(parents=True) - (csrc / "CMakeLists.txt").write_text(content + "\n") - (project / "CMakeLists.txt").write_text( - "cmake_minimum_required(VERSION 3.18)\n" - "project(discovery LANGUAGES NONE)\n" - f'list(APPEND CMAKE_MODULE_PATH "{_CMAKE_MODULE.parent}")\n' - "add_library(kimi_k3_model_src INTERFACE)\n" - f'add_subdirectory("{_INFRA}" "${{CMAKE_BINARY_DIR}}/native")\n' - "if(NOT TARGET discovered_native_family)\n" - ' message(FATAL_ERROR "native family discovery was disabled")\n' - "endif()\n" - ) - subprocess.run( - ["cmake", "-S", str(project), "-B", str(tmp_path / "build")], - check=True, - capture_output=True, - text=True, - ) diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index a499a583d097..23f85c2d4755 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -159,7 +159,7 @@ def test_scope(rule): @pytest.mark.parametrize("relative", ["csrc/kernel.cu", "ops/_native.py", "routing.py"]) -def test_stock_linked_k3_uses_baseline_selection(rule, relative): +def test_k3_uses_modeling_v2_selection(rule, relative): path = f"{_PY}/models/kimi_k3_vl/{relative}" - assert not _is_mv2_claim(path) - assert _apply(rule, path) is None + assert _is_mv2_claim(path) + assert _apply(rule, path).scope == "modelingv2only" From 7fa132b9ad3b42e4b2359f4a47b5bf2ee6c7c657 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Wed, 28 Oct 2026 10:25:36 -0700 Subject: [PATCH 02/11] [None][fix] qualify stock helpers used by K3 family operators Signed-off-by: Tao Li --- .../modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu | 1 - .../modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h | 3 +-- .../modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h | 3 +-- .../modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp | 5 +++-- .../models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu | 1 - .../models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h | 1 - .../models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h | 4 ++-- .../kimi_k3_vl/kernels/k3_sandwich/k3_sandwich_kernel.py | 2 +- .../modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/op.py | 2 +- .../models/kimi_k3_vl/ops/comm/mnnvl_workspace.py | 2 +- tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py | 5 +++-- tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py | 2 +- tests/unittest/scripts/cbts/test_modeling_v2_rule.py | 2 +- 13 files changed, 15 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu index 18b065fe5321..263f5083961b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.cu @@ -717,7 +717,6 @@ TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3DecodeAttnRes { -using namespace ::tensorrt_llm::kernels; template static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStream_t stream) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h index 0b57f05219fd..ba7abecb2901 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3AttnResDecode.h @@ -23,9 +23,8 @@ TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3DecodeAttnRes { -using namespace ::tensorrt_llm::kernels; -struct AttnResFwdParams : kimiK3AttnRes::AttnResFwdParams +struct AttnResFwdParams : ::tensorrt_llm::kernels::kimiK3AttnRes::AttnResFwdParams { bool earlyTrigger = false; }; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h index f6132b02dcd9..101b17bc65d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.h @@ -23,9 +23,8 @@ TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3MnnvlOneShot { -using namespace ::tensorrt_llm::kernels; -struct AllReduceFusionParams : mnnvl::AllReduceFusionParams +struct AllReduceFusionParams : ::tensorrt_llm::kernels::mnnvl::AllReduceFusionParams { bool earlyTrigger = false; }; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp index c07561130299..20fb4a77cf52 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/k3MnnvlFusionOp.cpp @@ -163,10 +163,11 @@ std::vector k3MnnvlFusionAllReduce(torch::Tensor& input, torch::o std::vector quantShape(input.sizes().begin(), input.sizes().end()); quantShape.back() /= 2; - quantOut = at::detail::empty_cuda(quantShape, FLOAT4_E2M1X2, input.device(), std::nullopt); + quantOut = at::detail::empty_cuda( + quantShape, ::tensorrt_llm::torch_ext::FLOAT4_E2M1X2, input.device(), std::nullopt); scaleOut = at::detail::empty_cuda({tensorrt_llm::computeSwizzledLayoutSFSize(numTokens, hiddenDim / kSfVecSize)}, - SF_DTYPE, input.device(), std::nullopt); + ::tensorrt_llm::torch_ext::SF_DTYPE, input.device(), std::nullopt); } else { diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu index e253cc3f90e5..de4a372af599 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.cu @@ -25,7 +25,6 @@ TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { -using namespace ::tensorrt_llm::kernels; using tensorrt_llm::common::isLamportDirty; using tensorrt_llm::common::LamportFlags; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h index c6122471603f..e89150d78e60 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllGatherKernels.h @@ -27,7 +27,6 @@ TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { -using namespace ::tensorrt_llm::kernels; struct AllGatherSplitParams { diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h index 43d4ca90d54b..05b8d7267ed2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/csrc/mnnvlAllreduceAttnRes.h @@ -23,7 +23,6 @@ TRTLLM_MV2_NAMESPACE_BEGIN namespace kernels::kimiK3Mnnvl { -using namespace ::tensorrt_llm::kernels; struct AttnResEpilogueParams { @@ -36,7 +35,8 @@ struct AttnResEpilogueParams int numCandidates; }; -void oneshotAllreduceAttnResOp(mnnvl::AllReduceFusionParams const& params, AttnResEpilogueParams const& epilogue); +void oneshotAllreduceAttnResOp( + ::tensorrt_llm::kernels::mnnvl::AllReduceFusionParams const& params, AttnResEpilogueParams const& epilogue); } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/k3_sandwich_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/k3_sandwich_kernel.py index 2af8a6c5d67a..3f1f75bf95bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/k3_sandwich_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/k3_sandwich_kernel.py @@ -19,7 +19,7 @@ updated = bf16(prefix + bf16(sum_r partial_r)) (or the sum alone without a prefix) normed = KimiK3RMSNorm(attn_res(snapshots..., updated)) -as ``o_proj`` followed by ``trtllm::mnnvl_allreduce_attn_res``; the reduction and the epilogue follow +as ``o_proj`` followed by ``mv2_kimi_k3_vl::mnnvl_allreduce_attn_res``; the reduction and the epilogue follow ``oneshotAllreduceAttnResKernel``'s arithmetic (rank chunks of 8, the same statistics, summation orders and roundings), so the outputs are bit-identical to that pair. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/op.py index 87c65c76f3b0..26e6ce2ffd72 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_sandwich/op.py @@ -341,7 +341,7 @@ def k3_sandwich_oproj( src_slab: Optional[torch.Tensor] = None, src_buf: int = 0, ) -> List[torch.Tensor]: - """``(normed, updated)`` of ``o_proj`` followed by ``trtllm::mnnvl_allreduce_attn_res``. + """``(normed, updated)`` of ``o_proj`` followed by ``mv2_kimi_k3_vl::mnnvl_allreduce_attn_res``. ``core`` bf16 [M, 768] is this rank's attention output, ``o_weight`` [7168, 768] its o_proj slice; ``prefix`` [M, 7168] (or None), ``block_residual`` [S, M, 7168] the S valid snapshots, the weights [7168]; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_workspace.py index a030e1e02602..2a7d9fa4fbb7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_workspace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_workspace.py @@ -4,7 +4,7 @@ A state type, not an entry: it launches nothing per call. Its constructor is collective and eager, the target builds one in ``post_load_weights`` (before any CUDA-graph capture) and passes it to every MNNVL -one-shot op of that group (``comm/mnnvl_allreduce_attn_res``, and ``trtllm::k3_mnnvl_fusion_allreduce`` given +one-shot op of that group (``comm/mnnvl_allreduce_attn_res``, and ``mv2_kimi_k3_vl::k3_mnnvl_fusion_allreduce`` given the same ``comm_buffer`` / ``buffer_flags``). """ diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py b/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py index a8a1c6d3f8db..5a4e4dc4a36e 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py +++ b/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py @@ -14,12 +14,13 @@ # limitations under the License. """The Kimi K3 decode collectives on the MNNVL all-reduce workspace, one process per GPU over the TP group of this run (4 on one GB200 tray), at every M in 1..8 and at 16, 32, 64 tokens: - trtllm::mnnvl_allreduce_attn_res (over the stock MNNVL all-reduce's workspace): against the unfused path it replaces, + mv2_kimi_k3_vl::mnnvl_allreduce_attn_res (over the stock MNNVL all-reduce's workspace): + against the unfused path it replaces, the MNNVL all-reduce then trtllm::attn_res_add_rmsnorm_fwd (attn_res_rmsnorm_fwd without a prefix sum): the updated prefix sum bit for bit, the normed rows within 1e-2 (max |d| / max |ref|), both against an fp32 port of the attention-residual selection within 2e-2; 0, 1, 3, 8 and 11 snapshots, with and without the prefix; the reference all-reduce sent one-shot (the fused op's order); - trtllm::mnnvl_allgather_split (over the same workspace): bit for bit against the same gather on the host (bf16 + mv2_kimi_k3_vl::mnnvl_allgather_split (over the same workspace): bit for bit against the same gather on the host (bf16 columns rounded to nearest, -0.0 arriving as +0.0), at the K3 sharded MoE head's per-rank widths for TP16 (224 bf16 + 56 fp32 columns) and TP4 (896 + 224), interleaved with all-reduces on the same Lamport rotation; each with run-to-run identical bits, the same bits on every rank, each M's rows bit-identical to the same rows of the diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py b/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py index 68120ff910c7..a19212d4cbae 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py +++ b/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py @@ -16,7 +16,7 @@ one kernel, M <= 8) at the Kimi K3 TP16 per-rank shapes, one process per GPU over the TP group of this run (4 on one GB200 tray; the per-rank shapes do not depend on the group size, the reduction is 4-way instead of 16-way), at every M in 1..8: - oproj against o_proj (cuBLAS, rows as in an 8-row call) -> trtllm::mnnvl_allreduce_attn_res, bit for bit, + oproj against o_proj (cuBLAS, rows as in an 8-row call) -> mv2_kimi_k3_vl::mnnvl_allreduce_attn_res, bit for bit, 0 / 1 / 3 / 8 snapshots, with and without the prefix sum; tail against the tail in torch (fp32 accumulators, the latent RMS on the latent one) -> mnnvl_allreduce_attn_res (fp32 tolerance: another summation order), and with the DSpark capture tap (the pre-norm mixture against diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 23f85c2d4755..7e1e30dcfcad 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -158,7 +158,7 @@ def test_scope(rule): assert _COMM in _selected(result) # unreadable family falls back to all modeling_v2 -@pytest.mark.parametrize("relative", ["csrc/kernel.cu", "ops/_native.py", "routing.py"]) +@pytest.mark.parametrize("relative", ["csrc/kernel.cu", "fake_ops.py", "routing.py"]) def test_k3_uses_modeling_v2_selection(rule, relative): path = f"{_PY}/models/kimi_k3_vl/{relative}" assert _is_mv2_claim(path) From 4b81089c077f14fcd83edbcf8f942ec55947e250 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Wed, 28 Oct 2026 12:52:46 -0700 Subject: [PATCH 03/11] [None][fix] preserve K3 lazy imports and native test coverage Signed-off-by: Tao Li --- .../scripts/cbts/cbts/rules/modeling_v2.py | 20 +++++++++++-- .../modeling.py | 12 ++------ .../modeling.py | 12 ++------ .../models/kimi_k3_vl/ops/__init__.py | 4 --- .../ops/comm/mnnvl_allgather_split.py | 5 ++-- .../kimi_k3_vl/ops/comm/mnnvl_allreduce.py | 4 +-- .../ops/comm/mnnvl_allreduce_attn_res.py | 5 ++-- .../ops/comm/mnnvl_fusion_allreduce.py | 5 ++-- .../ops/norm/attn_res_rmsnorm_fwd.py | 5 ++-- .../kernels/kimi_k3/test_k3_mnnvl_comm.py | 11 +++++-- .../kernels/kimi_k3/test_k3_sandwich.py | 9 +++++- .../comm/_mnnvl_allgather_split_op_matrix.py | 9 ++++-- .../kimi_k3_vl/test_modeling_v2_k3_native.py | 8 +++++ .../test_modeling_v2_attn_res_rmsnorm_fwd.py | 15 ++++++---- .../test_modeling_v2_kimi_k3_construction.py | 17 +++++++---- .../test_modeling_v2_kimi_k3_decode_step.py | 29 +++++++++++-------- .../test_modeling_v2_native_claims.py | 4 +++ .../scripts/cbts/test_modeling_v2_rule.py | 18 ++++++++++++ 18 files changed, 125 insertions(+), 67 deletions(-) diff --git a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py index fee9b67e260b..8898f86cf19e 100644 --- a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py +++ b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py @@ -107,6 +107,11 @@ _UNIT_ROOT = "unittest/_torch/modeling_v2" _ACCURACY_ROOT = "accuracy/modeling_v2" +# Some family tests also exercise stock/CuTe wrappers and live outside models/. +_FAMILY_TEST_PREFIXES = { + "kimi_k3_vl": ("unittest/_torch/kernels/kimi_k3", f"{_UNIT_ROOT}/comm"), +} + # Substrings that mark a test entry as modeling_v2. All are unambiguous: # - "unittest/_torch/modeling_v2/" → the op and target unit tests # - "accuracy/modeling_v2/" → the per-model accuracy gates @@ -141,8 +146,13 @@ def _is_mv2_claim(path: str) -> bool: def _entry_is_mv2(entry: str) -> bool: - return entry.split()[0].rstrip("/") == _UNIT_ROOT or any( - p in entry for p in _MV2_ENTRY_PATTERNS + if entry.split()[0].rstrip("/") == _UNIT_ROOT or any(p in entry for p in _MV2_ENTRY_PATTERNS): + return True + target = _entry_target(entry) + return any( + _target_in_filter_subtree(target, prefix) + for prefixes in _FAMILY_TEST_PREFIXES.values() + for prefix in prefixes ) @@ -254,7 +264,11 @@ def _prefixes_for(self, path: str) -> Optional[set[str]]: gates = self._gates_under(node) if gates is None: return None - return {f"{_UNIT_ROOT}/models/{node.as_posix()}"} | gates + return ( + {f"{_UNIT_ROOT}/models/{node.as_posix()}"} + | gates + | set(_FAMILY_TEST_PREFIXES.get(node.parts[0], ())) + ) def apply(self, pr: PRInputs) -> Optional[RuleResult]: claimed = {f for f in pr.changed_files if _is_mv2_claim(f)} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index 2649dfc1ea5f..4c31137a567f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -95,10 +95,10 @@ import torch from torch import nn +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_mla.decode_view import ( k3_mla_decode_view, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.attention.k3_drafter_attn_qknorm import ( k3_drafter_attn_qknorm, ) @@ -647,10 +647,7 @@ def _apply_attn_res_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("attn_res+norm", False, M, H, N) return None - try: - attn_res_rmsnorm_op = native_ops.k3_attn_res_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None + attn_res_rmsnorm_op = load_family_ops(__name__).k3_attn_res_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() output = attn_res_rmsnorm_op( @@ -721,10 +718,7 @@ def _apply_attn_res_add_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) return None - try: - attn_res_add_rmsnorm_op = native_ops.k3_attn_res_add_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None + attn_res_add_rmsnorm_op = load_family_ops(__name__).k3_attn_res_add_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() addend_kernel = addend.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index 7b6d5597b484..3f72eb890168 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -93,10 +93,10 @@ import torch from torch import nn +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_mla.decode_view import ( k3_mla_decode_view, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.attention.k3_drafter_attn_qknorm import ( k3_drafter_attn_qknorm, ) @@ -641,10 +641,7 @@ def _apply_attn_res_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("attn_res+norm", False, M, H, N) return None - try: - attn_res_rmsnorm_op = native_ops.k3_attn_res_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None + attn_res_rmsnorm_op = load_family_ops(__name__).k3_attn_res_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() output = attn_res_rmsnorm_op( @@ -715,10 +712,7 @@ def _apply_attn_res_add_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) return None - try: - attn_res_add_rmsnorm_op = native_ops.k3_attn_res_add_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None + attn_res_add_rmsnorm_op = load_family_ops(__name__).k3_attn_res_add_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() addend_kernel = addend.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py index 7e0d1509ec00..52a7a9daf028 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/__init__.py @@ -1,6 +1,2 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 - -from ...._native import load_family_ops - -native_ops = load_family_ops(__name__) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py index 7a67c64c6882..b85a1e5fa509 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py @@ -7,9 +7,8 @@ import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from .mnnvl_workspace import MnnvlWorkspace @@ -137,7 +136,7 @@ def __call__( f"mnnvl_allgather_split: the call needs {need} bytes per Lamport buffer, the workspace has " f"{workspace.buffer_bytes}" ) - bf16_out, fp32_out = native_ops.mnnvl_allgather_split( + bf16_out, fp32_out = load_family_ops(__name__).mnnvl_allgather_split( input, bf16_columns, workspace.world_size, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py index f7bbbb185b9e..c240e2f20f67 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py @@ -3,7 +3,7 @@ import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch.distributed import AllReduceParams from tensorrt_llm._torch.distributed.ops import ( MNNVLAllReduce, @@ -57,7 +57,7 @@ def forward( is_fusion = fusion_op != AllReduceFusionOp.NONE if is_fusion and fusion_op not in MNNVLAllReduce.SUPPORTED_FUSION_OPS: return None - outputs = native_ops.k3_mnnvl_fusion_allreduce( + outputs = load_family_ops(__name__).k3_mnnvl_fusion_allreduce( input, all_reduce_params.norm_weight, all_reduce_params.residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py index 6d2f6acf3dd6..7c22f5cc3d3a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py @@ -7,9 +7,8 @@ import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from .mnnvl_workspace import MnnvlWorkspace @@ -141,7 +140,7 @@ def __call__( output_rms_eps: float, workspace: MnnvlWorkspace, ) -> Tuple[torch.Tensor, torch.Tensor]: - normed, updated = native_ops.mnnvl_allreduce_attn_res( + normed, updated = load_family_ops(__name__).mnnvl_allreduce_attn_res( input, prefix_sum, block_residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py index 15ceccb84068..4ca9abcd2e71 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py @@ -7,9 +7,8 @@ import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops from tensorrt_llm.functional import AllReduceFusionOp from .mnnvl_workspace import MnnvlWorkspace @@ -200,7 +199,7 @@ def __call__( f"not {workspace.buffer_bytes}" ) fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM if fused else AllReduceFusionOp.NONE - outputs = native_ops.k3_mnnvl_fusion_allreduce( + outputs = load_family_ops(__name__).k3_mnnvl_fusion_allreduce( input, norm_weight, residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py index 5ebd3fcd9b66..db19905ac37d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py @@ -4,9 +4,8 @@ import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops class AttnResRmsnormFwd(OpWrapper): @@ -94,7 +93,7 @@ def __call__( rms_eps: float, output_rms_eps: float, ) -> torch.Tensor: - return native_ops.k3_attn_res_rmsnorm_fwd( + return load_family_ops(__name__).k3_attn_res_rmsnorm_fwd( layer_residual, block_residual, res_weight, diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py b/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py index 5a4e4dc4a36e..e70e6e1f46ed 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py +++ b/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py @@ -62,6 +62,13 @@ def namespace(): return namespace() +def _k3_native(): + """Resolve the family namespace inside each MPI worker.""" + from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops + + return load_family_ops("tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl") + + WORLD = 4 H, LATENT, EXPERTS = 7168, 3584, 896 EPS, OUT_EPS = 1e-6, 1e-5 @@ -136,7 +143,7 @@ def _stock_attn_res(mnnvl, input, prefix_sum, block_residual, res_weight, rms_we workspace = get_or_scale_allreduce_mnnvl_workspace( mnnvl.mapping, mnnvl.dtype, buffer_size_bytes=one_shot_bytes ) - normed, updated = _trtllm().mnnvl_allreduce_attn_res( + normed, updated = _k3_native().mnnvl_allreduce_attn_res( input, prefix_sum, block_residual, res_weight, rms_weight, output_rms_weight, rms_eps, output_rms_eps, workspace["uc_buffer"].view(mnnvl.dtype).view(3, -1), workspace["buffer_flags"], ) # fmt: skip @@ -153,7 +160,7 @@ def _stock_allgather_split(mnnvl, input, bf16_columns): workspace = get_or_scale_allreduce_mnnvl_workspace( mnnvl.mapping, mnnvl.dtype, buffer_size_bytes=footprint ) - bf16_out, fp32_out = _trtllm().mnnvl_allgather_split( + bf16_out, fp32_out = _k3_native().mnnvl_allgather_split( input, bf16_columns, mnnvl.mapping.tp_size, workspace["uc_buffer"].view(mnnvl.dtype).view(3, -1), workspace["buffer_flags"], ) # fmt: skip diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py b/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py index a19212d4cbae..4b1b985c61c5 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py +++ b/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py @@ -64,6 +64,13 @@ def namespace(): return namespace() +def _k3_native(): + """Resolve the family namespace inside each MPI worker.""" + from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops + + return load_family_ops("tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl") + + WORLD = 4 H, K_O, LATENT, WIDTH, PAD, ACT = 7168, 768, 3584, 224, 256, 384 PLAIN_K, DOWN_K = 384, 896 @@ -171,7 +178,7 @@ def _attn_res_ar(ctx, partial, prefix, block, res_w, rms_w, out_w): mnnvl.dtype, buffer_size_bytes=num_tokens * hidden_dim * mnnvl.mapping.tp_size * partial.element_size(), ) - return _trtllm().mnnvl_allreduce_attn_res( + return _k3_native().mnnvl_allreduce_attn_res( partial, prefix, block, res_w, rms_w, out_w, EPS, EPS, workspace["uc_buffer"].view(mnnvl.dtype).view(3, -1), workspace["buffer_flags"], ) # fmt: skip diff --git a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py index 77fb30abe4f0..5353d9703919 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py @@ -35,7 +35,7 @@ import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops sys.path.insert(0, str(Path(__file__).resolve().parent)) import _lockstep as ls # noqa: E402 @@ -88,6 +88,9 @@ STATS = {"normed_err": 0.0, "attn_res_err": 0.0} +_FAMILY = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl" + + def buffer_bytes(world: int) -> int: """One Lamport buffer holds the largest call this matrix makes: the all-gather at the widest certified split and 64 tokens, a wide step's [64, 7168] all-reduce sent two-shot, or the attention-residual all-reduce at 16 @@ -438,7 +441,9 @@ def check_unsupported_calls_raise_on_every_rank() -> None: expect_refusal(lambda: allgather_split(rows.bfloat16(), 8, WS_A), RuntimeError, "bf16 rows") expect_refusal(lambda: allgather_split(rows[:0], 8, WS_A), RuntimeError, "no rows") expect_refusal( - lambda: native_ops.mnnvl_allgather_split( + lambda: load_family_ops( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl" + ).mnnvl_allgather_split( rows, 8, R.world + 1, WS_A.comm_buffer(torch.bfloat16), WS_A.buffer_flags ), RuntimeError, diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py index 5c9a4af7f1c0..544c56d97de9 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py @@ -24,6 +24,14 @@ def test_k3_family_native_ops(): ): importlib.import_module(f"{package}.{target}") ops = load_family_ops(package) + for name in ( + "k3_mnnvl_fusion_allreduce", + "k3_attn_res_rmsnorm_fwd", + "k3_attn_res_add_rmsnorm_fwd", + "mnnvl_allgather_split", + "mnnvl_allreduce_attn_res", + ): + assert getattr(ops, name).default._schema is not None, name # A single candidate has softmax weight one; unit RMS weights leave ones unchanged. residual = torch.ones((2, 1, 7168), device="cuda", dtype=torch.bfloat16) blocks = residual.new_empty((0, 2, 1, 7168)) diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py index 8b3777490fec..580ab1f8dc76 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py @@ -17,7 +17,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops import native_ops +from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.norm.attn_res_rmsnorm_fwd import ( attn_res_rmsnorm_fwd, ) @@ -39,6 +39,9 @@ INPUT_NAMES = ("layer_residual", "block_residual", "res_weight", "rms_weight", "output_rms_weight") +_FAMILY = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl" + + def _make_inputs(num_tokens: int, num_candidates: int, seed: int = 0) -> tuple[torch.Tensor, ...]: """Inputs at the scales of main's op test, seeded by the cell.""" torch.manual_seed(seed) @@ -155,15 +158,17 @@ def test_k3_fused_add_preserves_rounding_and_graph_inputs( layer, block, res_weight, rms_weight, output_weight = inputs addend = layer.roll(1, dims=-1) * 0.25 args = (layer, addend, block, res_weight, rms_weight, output_weight, RMS_EPS, RMS_EPS) - updated, normalized = native_ops.k3_attn_res_add_rmsnorm_fwd(*args, early_trigger=early_trigger) + updated, normalized = load_family_ops( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl" + ).k3_attn_res_add_rmsnorm_fwd(*args, early_trigger=early_trigger) for replay in (False, True): if replay: torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - updated, normalized = native_ops.k3_attn_res_add_rmsnorm_fwd( - *args, early_trigger=early_trigger - ) + updated, normalized = load_family_ops( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl" + ).k3_attn_res_add_rmsnorm_fwd(*args, early_trigger=early_trigger) layer.copy_(layer.roll(3, dims=-1)) addend.neg_() graph.replay() diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py index d50944d4a23b..2784d143ce57 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py @@ -8,15 +8,20 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 - modeling as route_a, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 - modeling as route_b, -) +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget from tensorrt_llm.functional import AllReduceStrategy from tensorrt_llm.quantization.mode import QuantAlgo +try: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 + modeling as route_a, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 + modeling as route_b, + ) +except NoTarget as no_target: + pytest.skip(str(no_target), allow_module_level=True) + def _config(moe_tp, moe_ep, split_set=True, spec_config=None, attention_dp=False): """The ModelConfig fields the construction checks read.""" diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py index 15084d158155..68f036a922ce 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py @@ -11,24 +11,29 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 - modeling as route_a, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( # noqa: E501 - DecodeStep, - _attn_res_max_tokens, - decode_step, - latent_push, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 - modeling as route_b, -) +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( KDAIntermediateLayerCache, KDAReplayLayerCache, ) from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import _KdaPrequantizedCore +try: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 + modeling as route_a, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( # noqa: E501 + DecodeStep, + _attn_res_max_tokens, + decode_step, + latent_push, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 + modeling as route_b, + ) +except NoTarget as no_target: + pytest.skip(str(no_target), allow_module_level=True) + spec_worker_a = import_module(f"{route_a.__package__}.spec_worker") spec_worker_b = import_module(f"{route_b.__package__}.spec_worker") spec_accept_op = import_module( diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py index b6087e13faff..0de82b1dffee 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py @@ -303,6 +303,10 @@ def test_native_code_loads_only_when_routing_picks_its_target(): assert _is_docstring_or_future(node), ( f"{init.relative_to(_REPO)}:{node.lineno} -- a family package holds a docstring only" ) + for init in sorted(_MODELS.glob("*/ops/__init__.py")): + assert all(_is_docstring_or_future(node) for node in ast.parse(init.read_text()).body), ( + f"{init.relative_to(_REPO)}: importing unrelated ops must not load a native library" + ) for routing in sorted(_MODELS.glob("*/routing.py")): for node in ast.parse(routing.read_text()).body: assert not (isinstance(node, ast.ImportFrom) and node.level == 1), ( diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 7e1e30dcfcad..2b38815383cc 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -163,3 +163,21 @@ def test_k3_uses_modeling_v2_selection(rule, relative): path = f"{_PY}/models/kimi_k3_vl/{relative}" assert _is_mv2_claim(path) assert _apply(rule, path).scope == "modelingv2only" + + +def test_k3_native_change_keeps_existing_kernel_and_collective_tests(rule): + family = rule._src_root / "models/kimi_k3_vl" + (family / "target").mkdir(parents=True) + (family / "target/modeling.py").touch() + (family / "routing.py").write_text("from .target.modeling import KimiK3\n") + kernel_tests = { + "unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py", + "unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py", + } + matrix = "unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py" + rule.yaml_index.blocks[1].tests.extend([*sorted(kernel_tests), matrix]) + result = _apply(rule, f"{_PY}/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu") + assert kernel_tests | {matrix, _COMM} <= _selected(result) + assert result.scope == "modelingv2only" + other = _selected(_apply(rule, f"{_PY}/models/gpt_oss/csrc/kernel.cu")) + assert not (kernel_tests | {matrix}) & other From 471b115d04a5efda903924186fa0b516d89fe564 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Wed, 28 Oct 2026 13:18:33 -0700 Subject: [PATCH 04/11] [None][test] retain K3 RMSNorm coverage in family CI selection Signed-off-by: Tao Li --- jenkins/scripts/cbts/cbts/rules/modeling_v2.py | 6 +++++- tests/unittest/scripts/cbts/test_modeling_v2_rule.py | 7 +++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py index 8898f86cf19e..c83879d718ab 100644 --- a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py +++ b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py @@ -109,7 +109,11 @@ # Some family tests also exercise stock/CuTe wrappers and live outside models/. _FAMILY_TEST_PREFIXES = { - "kimi_k3_vl": ("unittest/_torch/kernels/kimi_k3", f"{_UNIT_ROOT}/comm"), + "kimi_k3_vl": ( + "unittest/_torch/kernels/kimi_k3", + f"{_UNIT_ROOT}/comm", + f"{_UNIT_ROOT}/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py", + ), } # Substrings that mark a test entry as modeling_v2. All are unambiguous: diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 2b38815383cc..5891e29f317d 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -176,8 +176,11 @@ def test_k3_native_change_keeps_existing_kernel_and_collective_tests(rule): } matrix = "unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py" rule.yaml_index.blocks[1].tests.extend([*sorted(kernel_tests), matrix]) + norm = "unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py" + rule.yaml_index.blocks[2].tests.append(norm) result = _apply(rule, f"{_PY}/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu") - assert kernel_tests | {matrix, _COMM} <= _selected(result) + assert kernel_tests | {matrix, norm, _COMM} <= _selected(result) + assert "B200-PyTorch-1" in result.affected_stages assert result.scope == "modelingv2only" other = _selected(_apply(rule, f"{_PY}/models/gpt_oss/csrc/kernel.cu")) - assert not (kernel_tests | {matrix}) & other + assert not (kernel_tests | {matrix, norm}) & other From 90722f465f7f99af5b15ed6acd65cd265c8551dd Mon Sep 17 00:00:00 2001 From: Tao Li Date: Wed, 28 Oct 2026 13:50:10 -0700 Subject: [PATCH 05/11] [None][test] assert the exact K3 CI selection Signed-off-by: Tao Li --- tests/unittest/scripts/cbts/test_modeling_v2_rule.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 5891e29f317d..56c1780813e0 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -179,7 +179,7 @@ def test_k3_native_change_keeps_existing_kernel_and_collective_tests(rule): norm = "unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py" rule.yaml_index.blocks[2].tests.append(norm) result = _apply(rule, f"{_PY}/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu") - assert kernel_tests | {matrix, norm, _COMM} <= _selected(result) + assert _selected(result) == kernel_tests | {matrix, norm, _COMM, _UNIT, _CLAIMS} assert "B200-PyTorch-1" in result.affected_stages assert result.scope == "modelingv2only" other = _selected(_apply(rule, f"{_PY}/models/gpt_oss/csrc/kernel.cu")) From d0ee18259605a7ed3841f64190a98b199bc8a84a Mon Sep 17 00:00:00 2001 From: Tao Li Date: Wed, 28 Oct 2026 16:36:08 -0700 Subject: [PATCH 06/11] [None][fix] align K3 expert adapters with the shared loader contract Signed-off-by: Tao Li --- .../modeling.py | 11 +++++- .../modeling.py | 11 +++++- .../test_modeling_v2_kimi_k3_construction.py | 37 +++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index 4c31137a567f..0a9f0c607ba8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -834,7 +834,15 @@ def _attn_res_max_tokens(step: DecodeStep) -> int: _K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 -def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: +def _load_packed_mxfp4_expert( + backend: nn.Module, + base: str, + expert_idx: int, + local_slot_id: int, + get_tensor: Callable[[str], torch.Tensor], + get_optional_tensor: Callable[[str], Optional[torch.Tensor]], +) -> None: + del get_optional_tensor # MXFP4 stores no optional per-expert tensors. backend.quant_method.load_packed_mxfp4_expert( backend, global_expert_id=expert_idx, @@ -858,6 +866,7 @@ class _K3ExpertCkptSpec(NamedTuple): loaded_slots_attr: str # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). needs_layer_finalize: bool + optional_suffixes: Tuple[str, ...] = () _K3_EXPERT_CKPT_SPECS = { diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index 3f72eb890168..5c5c87039a14 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -828,7 +828,15 @@ def _attn_res_max_tokens(step: DecodeStep) -> int: _K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 -def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: +def _load_packed_mxfp4_expert( + backend: nn.Module, + base: str, + expert_idx: int, + local_slot_id: int, + get_tensor: Callable[[str], torch.Tensor], + get_optional_tensor: Callable[[str], Optional[torch.Tensor]], +) -> None: + del get_optional_tensor # MXFP4 stores no optional per-expert tensors. backend.quant_method.load_packed_mxfp4_expert( backend, global_expert_id=expert_idx, @@ -852,6 +860,7 @@ class _K3ExpertCkptSpec(NamedTuple): loaded_slots_attr: str # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). needs_layer_finalize: bool + optional_suffixes: Tuple[str, ...] = () _K3_EXPERT_CKPT_SPECS = { diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py index 2784d143ce57..73817109d006 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py @@ -4,6 +4,7 @@ ``tp16_moetp16ep1`` also an explicit expert split and no speculative decoding.""" import types +from unittest.mock import Mock import pytest import torch @@ -100,3 +101,39 @@ def test_tp16_moetp16ep1_keeps_no_per_token_kda_states(monkeypatch): object.__setattr__(model, "layers", [kda, mla]) assert model._latent_push(None, route_b.DecodeStep(1, 1, 1)) assert not model._latent_push(None, route_b.DecodeStep(8, 1, 8)) + + +@pytest.mark.parametrize("target", [route_a, route_b]) +def test_mxfp4_expert_adapter_matches_shared_loader_contract(target): + """Both targets accept the shared loader's optional getter without changing MXFP4 inputs.""" + from tensorrt_llm._torch.models import modeling_kimi_linear as stock + + quant = QuantAlgo.W4A8_MXFP4_MXFP8 + spec = target._k3_expert_ckpt_spec(quant) + reference = stock._k3_expert_ckpt_spec(quant) + assert ( + spec.kinds, + spec.loaded_slots_attr, + spec.needs_layer_finalize, + spec.optional_suffixes, + ) == ( + reference.kinds, + reference.loaded_slots_attr, + reference.needs_layer_finalize, + reference.optional_suffixes, + ) + backend = torch.nn.Module() + load = Mock() + backend.quant_method = types.SimpleNamespace(load_packed_mxfp4_expert=load) + keys = [ + f"experts.7.{weight}.{kind}" for weight in ("w1", "w2", "w3") for kind in reference.kinds + ] + tensors = {key: torch.tensor(index) for index, key in enumerate(keys)} + optional = Mock(side_effect=AssertionError("MXFP4 must not read optional calibration tensors")) + reference.loader(backend, "experts", 7, 3, tensors.__getitem__, optional) + expected = load.call_args + load.reset_mock() + spec.loader(backend, "experts", 7, 3, tensors.__getitem__, optional) + assert load.call_count == 1 + assert load.call_args == expected + optional.assert_not_called() From 4e785cd6e6b2c52253d5be7fa5494f227f226b54 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Thu, 29 Oct 2026 02:16:15 -0700 Subject: [PATCH 07/11] Avoid duplicate modeling v2 test module names Signed-off-by: Tao Li --- ...y => test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py} | 0 ...v1d_fwd.py => test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/{test_modeling_v2_rms_norm_gated_token_major.py => test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py} (100%) rename tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/{test_modeling_v2_causal_conv1d_fwd.py => test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py} (100%) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py From 937bfabad35b19248a2db4ed80a9c405545300fd Mon Sep 17 00:00:00 2001 From: Tao Li Date: Thu, 29 Oct 2026 05:19:09 -0700 Subject: [PATCH 08/11] Load K3 native consumers at module scope Signed-off-by: Tao Li --- .../modeling.py | 8 ++--- .../modeling.py | 8 ++--- .../ops/comm/mnnvl_allgather_split.py | 4 ++- .../kimi_k3_vl/ops/comm/mnnvl_allreduce.py | 4 ++- .../ops/comm/mnnvl_allreduce_attn_res.py | 4 ++- .../ops/comm/mnnvl_fusion_allreduce.py | 4 ++- .../ops/norm/attn_res_rmsnorm_fwd.py | 4 ++- .../integration/test_lists/test-db/l0_cpu.yml | 2 +- .../test_modeling_v2_k3_mnnvl_workspace.py | 33 +++++++++++++++++++ .../test_modeling_v2_mnnvl_workspace.py | 21 ------------ 10 files changed, 57 insertions(+), 35 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mnnvl_workspace.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index 0a9f0c607ba8..e6f5e94f29ca 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -161,6 +161,8 @@ # >>> route B: no KDA verify kernels, and no trtllm::kda_mtp_decode (the built-in KDA verify's) # <<< route B +native_ops = load_family_ops(__name__) + if TYPE_CHECKING: from transformers import PretrainedConfig @@ -647,10 +649,9 @@ def _apply_attn_res_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("attn_res+norm", False, M, H, N) return None - attn_res_rmsnorm_op = load_family_ops(__name__).k3_attn_res_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output = attn_res_rmsnorm_op( + output = native_ops.k3_attn_res_rmsnorm_fwd( layer_kernel, block_kernel, proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), @@ -718,11 +719,10 @@ def _apply_attn_res_add_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) return None - attn_res_add_rmsnorm_op = load_family_ops(__name__).k3_attn_res_add_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() addend_kernel = addend.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - updated_prefix_sum, output = attn_res_add_rmsnorm_op( + updated_prefix_sum, output = native_ops.k3_attn_res_add_rmsnorm_fwd( layer_kernel, addend_kernel, block_kernel, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index 5c5c87039a14..f540e23fecf7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -158,6 +158,8 @@ from .ops.ssm.k3_kda_attn import k3_kda_attn from .ops.ssm.k3_kda_verify import k3_kda_verify +native_ops = load_family_ops(__name__) + if TYPE_CHECKING: from transformers import PretrainedConfig @@ -641,10 +643,9 @@ def _apply_attn_res_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("attn_res+norm", False, M, H, N) return None - attn_res_rmsnorm_op = load_family_ops(__name__).k3_attn_res_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output = attn_res_rmsnorm_op( + output = native_ops.k3_attn_res_rmsnorm_fwd( layer_kernel, block_kernel, proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), @@ -712,11 +713,10 @@ def _apply_attn_res_add_rmsnorm_fused( if M > max_fused_tokens or H != 7168 or N > 12: _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) return None - attn_res_add_rmsnorm_op = load_family_ops(__name__).k3_attn_res_add_rmsnorm_fwd layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() addend_kernel = addend.reshape(M, 1, H).contiguous() block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - updated_prefix_sum, output = attn_res_add_rmsnorm_op( + updated_prefix_sum, output = native_ops.k3_attn_res_add_rmsnorm_fwd( layer_kernel, addend_kernel, block_kernel, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py index b85a1e5fa509..056c0fec7647 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py @@ -12,6 +12,8 @@ from .mnnvl_workspace import MnnvlWorkspace +native_ops = load_family_ops(__name__) + __all__ = ["MnnvlWorkspace", "mnnvl_allgather_split", "required_buffer_bytes"] @@ -136,7 +138,7 @@ def __call__( f"mnnvl_allgather_split: the call needs {need} bytes per Lamport buffer, the workspace has " f"{workspace.buffer_bytes}" ) - bf16_out, fp32_out = load_family_ops(__name__).mnnvl_allgather_split( + bf16_out, fp32_out = native_ops.mnnvl_allgather_split( input, bf16_columns, workspace.world_size, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py index c240e2f20f67..bce500c0bcf0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce.py @@ -14,6 +14,8 @@ from .mnnvl_fusion_allreduce import required_buffer_bytes +native_ops = load_family_ops(__name__) + class K3MnnvlAllReduce(MNNVLAllReduce): def __init__(self, mapping: Mapping, dtype: torch.dtype): @@ -57,7 +59,7 @@ def forward( is_fusion = fusion_op != AllReduceFusionOp.NONE if is_fusion and fusion_op not in MNNVLAllReduce.SUPPORTED_FUSION_OPS: return None - outputs = load_family_ops(__name__).k3_mnnvl_fusion_allreduce( + outputs = native_ops.k3_mnnvl_fusion_allreduce( input, all_reduce_params.norm_weight, all_reduce_params.residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py index 7c22f5cc3d3a..feff2441b0aa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py @@ -12,6 +12,8 @@ from .mnnvl_workspace import MnnvlWorkspace +native_ops = load_family_ops(__name__) + class MnnvlAllreduceAttnRes(OpWrapper): """Return ``(normed, updated)``: ``updated = prefix_sum + allreduce(input)`` over ``workspace``'s TP group (the @@ -140,7 +142,7 @@ def __call__( output_rms_eps: float, workspace: MnnvlWorkspace, ) -> Tuple[torch.Tensor, torch.Tensor]: - normed, updated = load_family_ops(__name__).mnnvl_allreduce_attn_res( + normed, updated = native_ops.mnnvl_allreduce_attn_res( input, prefix_sum, block_residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py index 4ca9abcd2e71..f65bdaad67cd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py @@ -13,6 +13,8 @@ from .mnnvl_workspace import MnnvlWorkspace +native_ops = load_family_ops(__name__) + __all__ = ["MnnvlWorkspace", "mnnvl_fusion_allreduce", "required_buffer_bytes"] @@ -199,7 +201,7 @@ def __call__( f"not {workspace.buffer_bytes}" ) fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM if fused else AllReduceFusionOp.NONE - outputs = load_family_ops(__name__).k3_mnnvl_fusion_allreduce( + outputs = native_ops.k3_mnnvl_fusion_allreduce( input, norm_weight, residual, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py index db19905ac37d..4b990f21abfb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py @@ -7,6 +7,8 @@ from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +native_ops = load_family_ops(__name__) + class AttnResRmsnormFwd(OpWrapper): """Return rmsnorm(attn_res_fwd's bf16 mixture) * output_rms_weight as a new bf16 [T, 1, 7168] tensor. @@ -93,7 +95,7 @@ def __call__( rms_eps: float, output_rms_eps: float, ) -> torch.Tensor: - return load_family_ops(__name__).k3_attn_res_rmsnorm_fwd( + return native_ops.k3_attn_res_rmsnorm_fwd( layer_residual, block_residual, res_weight, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index edf634c48fcf..ebe3d26c04f1 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -80,7 +80,7 @@ l0_cpu: - unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py # l0_b300 selects unittest/_torch/modeling_v2 with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py - - unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py::test_k3_workspace_threshold_is_inclusive_and_two_shot_pads + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mnnvl_workspace.py - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mnnvl_workspace.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mnnvl_workspace.py new file mode 100644 index 000000000000..db1a4a2d94c8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mnnvl_workspace.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""K3 all-reduce workspace sizing at the one-shot boundary.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + +try: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.comm.mnnvl_allreduce import ( + K3MnnvlAllReduce, + ) +except NoTarget as no_target: + pytest.skip(str(no_target), allow_module_level=True) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("world", (4, 16)) +@pytest.mark.parametrize("rows", (1, 3, 8, 16, 33, 64)) +def test_k3_workspace_threshold_is_inclusive_and_two_shot_pads(rows, world): + hidden = 7168 + one_shot = rows * hidden * world * 2 + assert ( + K3MnnvlAllReduce.get_required_workspace_size(rows, hidden, world, torch.bfloat16, one_shot) + == one_shot + ) + assert ( + K3MnnvlAllReduce.get_required_workspace_size( + rows, hidden, world, torch.bfloat16, one_shot - 1 + ) + == 2 * ((rows + world - 1) // world) * world * hidden * 2 + ) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py index f3f130495288..f57edccb88b9 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py @@ -16,9 +16,6 @@ workspace_bytes, workspace_bytes_up_to, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.comm.mnnvl_allreduce import ( - K3MnnvlAllReduce, -) from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce # (hidden, world): the two default targets at their MNNVL sizes. @@ -52,21 +49,3 @@ def test_the_largest_one_shot_call_sets_the_size_below_the_band_edge(): assert workspace_bytes(64, 2880, 8, torch.bfloat16) == 737_280 assert workspace_bytes(22, 2880, 8, torch.bfloat16) == 1_013_760 assert workspace_bytes_up_to(64, 2880, 8, torch.bfloat16) == 1_013_760 - - -@pytest.mark.cpu_only -@pytest.mark.parametrize("world", (4, 16)) -@pytest.mark.parametrize("rows", (1, 3, 8, 16, 33, 64)) -def test_k3_workspace_threshold_is_inclusive_and_two_shot_pads(rows, world): - hidden = 7168 - one_shot = rows * hidden * world * 2 - assert ( - K3MnnvlAllReduce.get_required_workspace_size(rows, hidden, world, torch.bfloat16, one_shot) - == one_shot - ) - assert ( - K3MnnvlAllReduce.get_required_workspace_size( - rows, hidden, world, torch.bfloat16, one_shot - 1 - ) - == 2 * ((rows + world - 1) // world) * world * hidden * 2 - ) From 927480ee7a5044ad0002366feb34df70f13c8606 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Thu, 29 Oct 2026 05:22:16 -0700 Subject: [PATCH 09/11] Organize K3 modeling v2 tests by family Signed-off-by: Tao Li --- .../scripts/cbts/cbts/rules/modeling_v2.py | 24 +--- .../_experimental/modeling_v2/README.md | 7 +- .../ops/attention/k3_mla_attn_vb_out.py | 2 +- .../kimi_k3_vl/ops/attention/k3_mla_qkv.py | 2 +- .../test_lists/test-db/l0_b200.yml | 107 ++++++++++-------- .../test_lists/test-db/l0_b300.yml | 41 ++++++- .../integration/test_lists/test-db/l0_cpu.yml | 8 +- .../test-db/l0_gb200_multi_gpus.yml | 34 +++--- .../test-db/l0_gb300_multi_gpus.yml | 9 ++ .../comm/test_modeling_v2_rank_job.py | 9 +- .../_k3_latent_reduce_op_matrix.py | 3 +- .../kimi_k3_vl}/_k3_moe_engines.py | 0 .../kimi_k3_vl}/_k3_moe_front_op_matrix.py | 3 +- .../kimi_k3_vl}/_k3_sandwich_common.py | 0 .../_k3_sandwich_oproj_op_matrix.py | 3 +- .../_k3_sandwich_plain_op_matrix.py | 3 +- .../_k3_sandwich_tail_op_matrix.py | 3 +- .../{ssm => models/kimi_k3_vl}/_kda_cells.py | 0 .../kimi_k3_vl}/_kda_reference.py | 0 .../_kimi_k3_decode_comm_op_matrix.py | 1 + .../_kimi_k3_route_b_decode_moe_op_matrix.py | 1 + .../_mnnvl_allgather_split_op_matrix.py | 3 +- .../_mnnvl_allreduce_attn_res_op_matrix.py | 3 +- .../_mnnvl_fusion_allreduce_op_matrix.py | 3 +- .../kimi_k3_vl}/test_k3_cluster_waits.py | 0 .../models/kimi_k3_vl}/test_k3_ctm_gemv.py | 0 .../kimi_k3_vl}/test_k3_ctm_gemv_wide.py | 0 .../models/kimi_k3_vl}/test_k3_ctx_kv.py | 0 .../models/kimi_k3_vl}/test_k3_decode_gemv.py | 0 .../kimi_k3_vl}/test_k3_drafter_attn.py | 0 .../models/kimi_k3_vl}/test_k3_embed.py | 0 .../models/kimi_k3_vl}/test_k3_fused_moe.py | 0 .../models/kimi_k3_vl}/test_k3_head_gemv.py | 0 .../models/kimi_k3_vl}/test_k3_kda_attn.py | 0 .../kimi_k3_vl}/test_k3_kda_decode_attn.py | 0 .../kimi_k3_vl}/test_k3_kda_pool_layout.py | 2 +- .../kimi_k3_vl}/test_k3_kda_pools_past_2g.py | 0 .../models/kimi_k3_vl}/test_k3_kda_verify.py | 0 .../kimi_k3_vl}/test_k3_latent_reduce.py | 2 +- .../models/kimi_k3_vl}/test_k3_markov.py | 0 .../models/kimi_k3_vl}/test_k3_mla_attn.py | 0 .../kimi_k3_vl}/test_k3_mla_decode_view.py | 0 .../models/kimi_k3_vl}/test_k3_mla_launch.py | 2 +- .../models/kimi_k3_vl}/test_k3_mla_q.py | 0 .../models/kimi_k3_vl}/test_k3_mnnvl_comm.py | 0 .../models/kimi_k3_vl}/test_k3_moe_front.py | 0 .../kimi_k3_vl}/test_k3_moe_front_geometry.py | 0 .../models/kimi_k3_vl}/test_k3_moe_m1.py | 0 .../models/kimi_k3_vl}/test_k3_moe_m2.py | 0 .../models/kimi_k3_vl}/test_k3_moe_push.py | 0 .../models/kimi_k3_vl}/test_k3_moe_wide.py | 0 .../models/kimi_k3_vl}/test_k3_route_quant.py | 0 .../models/kimi_k3_vl}/test_k3_sandwich.py | 0 .../models/kimi_k3_vl}/test_k3_spec_accept.py | 0 .../test_k3_spec_accept_sharded.py | 0 .../kimi_k3_vl}/test_k3_tcgen05_fences.py | 0 .../test_modeling_v2_attn_res_fwd.py | 2 +- .../test_modeling_v2_attn_res_rmsnorm_fwd.py | 2 +- .../test_modeling_v2_causal_conv1d_fwd.py | 2 +- .../kimi_k3_vl}/test_modeling_v2_chunk_kda.py | 2 +- .../test_modeling_v2_k3_ctm_gemv.py | 2 +- .../test_modeling_v2_k3_ctm_gemv_long.py | 2 +- .../test_modeling_v2_k3_ctm_gemv_swiglu.py | 2 +- .../test_modeling_v2_k3_ctm_gemv_wide.py | 2 +- .../test_modeling_v2_k3_decode_gemv.py | 2 +- .../test_modeling_v2_k3_drafter_attn.py | 2 +- ...test_modeling_v2_k3_drafter_attn_qknorm.py | 2 +- .../test_modeling_v2_k3_embed_norm.py | 2 +- .../test_modeling_v2_k3_head_gemv.py | 2 +- .../test_modeling_v2_k3_kda_attn.py | 2 +- .../test_modeling_v2_k3_kda_decode_attn.py | 2 +- .../test_modeling_v2_k3_kda_verify.py | 2 +- ..._modeling_v2_k3_latent_reduce_op_matrix.py | 7 +- .../test_modeling_v2_k3_mla_attn_vb_out.py | 2 +- .../test_modeling_v2_k3_mla_qkv.py | 2 +- .../kimi_k3_vl}/test_modeling_v2_k3_moe.py | 9 +- ...test_modeling_v2_k3_moe_front_op_matrix.py | 12 +- .../kimi_k3_vl}/test_modeling_v2_k3_moe_m1.py | 2 +- .../kimi_k3_vl}/test_modeling_v2_k3_moe_m2.py | 2 +- .../test_modeling_v2_k3_route_quant.py | 2 +- ...modeling_v2_k3_sandwich_oproj_op_matrix.py | 7 +- ...modeling_v2_k3_sandwich_plain_op_matrix.py | 7 +- ..._modeling_v2_k3_sandwich_tail_op_matrix.py | 7 +- .../test_modeling_v2_k3_situ_mul.py | 2 +- .../test_modeling_v2_kda_decode.py | 2 +- .../test_modeling_v2_kda_post_conv.py | 2 +- .../test_modeling_v2_kda_prefill.py | 2 +- .../test_modeling_v2_kimi_k3_construction.py | 0 ...deling_v2_kimi_k3_decode_comm_op_matrix.py | 7 +- .../test_modeling_v2_kimi_k3_decode_gemv.py | 0 .../test_modeling_v2_kimi_k3_decode_step.py | 0 .../test_modeling_v2_kimi_k3_drafter.py | 0 .../test_modeling_v2_kimi_k3_drift.py | 0 .../test_modeling_v2_kimi_k3_moe_routing.py | 0 ...odeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py | 2 +- ...v2_kimi_k3_route_b_decode_moe_op_matrix.py | 7 +- .../test_modeling_v2_kimi_k3_route_b_moe.py | 0 ...ling_v2_mnnvl_allgather_split_op_matrix.py | 7 +- ...g_v2_mnnvl_allreduce_attn_res_op_matrix.py | 7 +- ...ing_v2_mnnvl_fusion_allreduce_op_matrix.py | 7 +- ..._modeling_v2_rms_norm_gated_token_major.py | 2 +- .../test_modeling_v2_situ_and_mul.py | 2 +- .../scripts/cbts/test_modeling_v2_rule.py | 34 ++++-- 103 files changed, 272 insertions(+), 177 deletions(-) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_k3_latent_reduce_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/_k3_moe_engines.py (100%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_k3_moe_front_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_k3_sandwich_common.py (100%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_k3_sandwich_oproj_op_matrix.py (98%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_k3_sandwich_plain_op_matrix.py (98%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_k3_sandwich_tail_op_matrix.py (98%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/_kda_cells.py (100%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/_kda_reference.py (100%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_kimi_k3_decode_comm_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_kimi_k3_route_b_decode_moe_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_mnnvl_allgather_split_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_mnnvl_allreduce_attn_res_op_matrix.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/_mnnvl_fusion_allreduce_op_matrix.py (99%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_cluster_waits.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_ctm_gemv.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_ctm_gemv_wide.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_ctx_kv.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_decode_gemv.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_drafter_attn.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_embed.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_fused_moe.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_head_gemv.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_kda_attn.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_kda_decode_attn.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_kda_pool_layout.py (98%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_kda_pools_past_2g.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_kda_verify.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_latent_reduce.py (99%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_markov.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_mla_attn.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_mla_decode_view.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_mla_launch.py (99%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_mla_q.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_mnnvl_comm.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_moe_front.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_moe_front_geometry.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_moe_m1.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_moe_m2.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_moe_push.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_moe_wide.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_route_quant.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_sandwich.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_spec_accept.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_spec_accept_sharded.py (100%) rename tests/unittest/_torch/{kernels/kimi_k3 => modeling_v2/models/kimi_k3_vl}/test_k3_tcgen05_fences.py (100%) rename tests/unittest/_torch/modeling_v2/{norm => models/kimi_k3_vl}/test_modeling_v2_attn_res_fwd.py (99%) rename tests/unittest/_torch/modeling_v2/{norm => models/kimi_k3_vl}/test_modeling_v2_attn_res_rmsnorm_fwd.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_causal_conv1d_fwd.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_chunk_kda.py (99%) rename tests/unittest/_torch/modeling_v2/{gemm => models/kimi_k3_vl}/test_modeling_v2_k3_ctm_gemv.py (99%) rename tests/unittest/_torch/modeling_v2/{gemm => models/kimi_k3_vl}/test_modeling_v2_k3_ctm_gemv_long.py (99%) rename tests/unittest/_torch/modeling_v2/{gemm => models/kimi_k3_vl}/test_modeling_v2_k3_ctm_gemv_swiglu.py (99%) rename tests/unittest/_torch/modeling_v2/{gemm => models/kimi_k3_vl}/test_modeling_v2_k3_ctm_gemv_wide.py (99%) rename tests/unittest/_torch/modeling_v2/{gemm => models/kimi_k3_vl}/test_modeling_v2_k3_decode_gemv.py (99%) rename tests/unittest/_torch/modeling_v2/{attention => models/kimi_k3_vl}/test_modeling_v2_k3_drafter_attn.py (99%) rename tests/unittest/_torch/modeling_v2/{attention => models/kimi_k3_vl}/test_modeling_v2_k3_drafter_attn_qknorm.py (99%) rename tests/unittest/_torch/modeling_v2/{norm => models/kimi_k3_vl}/test_modeling_v2_k3_embed_norm.py (99%) rename tests/unittest/_torch/modeling_v2/{gemm => models/kimi_k3_vl}/test_modeling_v2_k3_head_gemv.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_k3_kda_attn.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_k3_kda_decode_attn.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_k3_kda_verify.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_k3_latent_reduce_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{attention => models/kimi_k3_vl}/test_modeling_v2_k3_mla_attn_vb_out.py (99%) rename tests/unittest/_torch/modeling_v2/{attention => models/kimi_k3_vl}/test_modeling_v2_k3_mla_qkv.py (99%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/test_modeling_v2_k3_moe.py (99%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/test_modeling_v2_k3_moe_front_op_matrix.py (53%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/test_modeling_v2_k3_moe_m1.py (98%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/test_modeling_v2_k3_moe_m2.py (98%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/test_modeling_v2_k3_route_quant.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_k3_sandwich_oproj_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_k3_sandwich_plain_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_k3_sandwich_tail_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{activation => models/kimi_k3_vl}/test_modeling_v2_k3_situ_mul.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_kda_decode.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_kda_post_conv.py (99%) rename tests/unittest/_torch/modeling_v2/{ssm => models/kimi_k3_vl}/test_modeling_v2_kda_prefill.py (99%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_construction.py (100%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py (88%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_decode_gemv.py (100%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_decode_step.py (100%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_drafter.py (100%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_drift.py (100%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_moe_routing.py (100%) rename tests/unittest/_torch/modeling_v2/{moe => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py (99%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py (88%) rename tests/unittest/_torch/modeling_v2/{ => models/kimi_k3_vl}/test_modeling_v2_kimi_k3_route_b_moe.py (100%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_mnnvl_allgather_split_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{comm => models/kimi_k3_vl}/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py (83%) rename tests/unittest/_torch/modeling_v2/{norm => models/kimi_k3_vl}/test_modeling_v2_rms_norm_gated_token_major.py (99%) rename tests/unittest/_torch/modeling_v2/{activation => models/kimi_k3_vl}/test_modeling_v2_situ_and_mul.py (99%) diff --git a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py index c83879d718ab..fee9b67e260b 100644 --- a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py +++ b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py @@ -107,15 +107,6 @@ _UNIT_ROOT = "unittest/_torch/modeling_v2" _ACCURACY_ROOT = "accuracy/modeling_v2" -# Some family tests also exercise stock/CuTe wrappers and live outside models/. -_FAMILY_TEST_PREFIXES = { - "kimi_k3_vl": ( - "unittest/_torch/kernels/kimi_k3", - f"{_UNIT_ROOT}/comm", - f"{_UNIT_ROOT}/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py", - ), -} - # Substrings that mark a test entry as modeling_v2. All are unambiguous: # - "unittest/_torch/modeling_v2/" → the op and target unit tests # - "accuracy/modeling_v2/" → the per-model accuracy gates @@ -150,13 +141,8 @@ def _is_mv2_claim(path: str) -> bool: def _entry_is_mv2(entry: str) -> bool: - if entry.split()[0].rstrip("/") == _UNIT_ROOT or any(p in entry for p in _MV2_ENTRY_PATTERNS): - return True - target = _entry_target(entry) - return any( - _target_in_filter_subtree(target, prefix) - for prefixes in _FAMILY_TEST_PREFIXES.values() - for prefix in prefixes + return entry.split()[0].rstrip("/") == _UNIT_ROOT or any( + p in entry for p in _MV2_ENTRY_PATTERNS ) @@ -268,11 +254,7 @@ def _prefixes_for(self, path: str) -> Optional[set[str]]: gates = self._gates_under(node) if gates is None: return None - return ( - {f"{_UNIT_ROOT}/models/{node.as_posix()}"} - | gates - | set(_FAMILY_TEST_PREFIXES.get(node.parts[0], ())) - ) + return {f"{_UNIT_ROOT}/models/{node.as_posix()}"} | gates def apply(self, pr: PRInputs) -> Optional[RuleResult]: claimed = {f for f in pr.changed_files if _is_mv2_claim(f)} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 3956bcdf6725..3493dde7c57c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/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 /test_modeling_v2_.py shared entries + models//test_.py family-owned entries models////test_modeling_v2_.py owned entries comm/__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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py index e6c5e0d96822..3a4291d3343c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py @@ -136,7 +136,7 @@ class K3MlaAttnVbOut(OpWrapper): already raise ValueError before launch. First calls compile outside graph capture. is_valid checks lengths, not whether an earlier call actually wrote the rows or whether another stream is using the workspace. - tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_attn.py additionally + tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py additionally covers every split of up to eight tokens, R x 8 verification, stock cute_dsl_mla_decode_fp16_blackwell parity, 640-element interleaved rows, poison/counter-wrap checks, refusals and optional K3_BASE_TRTLLM identity. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py index 2460df5bc5a6..fe65a88a2b8e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py @@ -58,7 +58,7 @@ class K3MlaQkv(OpWrapper): k3_mla_q returns the same fused_q bits without storing rows; k3_mla_qkv_out stores the same rows densely in bf16 [M, 576]. Query rows are computed in chunks of 8, 16 or 32; every eight-token chunk is bit-identical to its own - call. The op-level test tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_q.py + call. The op-level test tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py covers M 1..64 against the unfused RMSNorm/GEMM/bmm chain, 640-element interleaved rows, short lengths, refusals, and optional K3_BASE_TRTLLM single-request identity. The catalog tests retain the 2e-2 query bound diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4727ba8b8e40..cb250f4d4948 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -164,9 +164,9 @@ l0_b200: - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py # Kimi K3's generic path (prefill, mixed and wide steps) in the modeling_v2 catalog: its own entries # and the sm_100 cells of the shared entries it calls. - - unittest/_torch/modeling_v2/norm/test_modeling_v2_rms_norm_gated_token_major.py - - unittest/_torch/modeling_v2/activation/test_modeling_v2_situ_and_mul.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py - unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py::test_kimi_k3_config - unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py::test_kimi_k3_routed_hidden - unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py @@ -181,39 +181,39 @@ l0_b200: - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py # Kimi K3 decode kernels (CuTe DSL): the GEMVs, the head GEMV and the # embedding, and the cluster-wait and tcgen05-fence lints over them. - - unittest/_torch/kernels/kimi_k3/test_k3_embed.py - - unittest/_torch/kernels/kimi_k3/test_k3_ctm_gemv.py - - unittest/_torch/kernels/kimi_k3/test_k3_ctm_gemv_wide.py - - unittest/_torch/kernels/kimi_k3/test_k3_decode_gemv.py - - unittest/_torch/kernels/kimi_k3/test_k3_head_gemv.py - - unittest/_torch/kernels/kimi_k3/test_k3_cluster_waits.py - - unittest/_torch/kernels/kimi_k3/test_k3_tcgen05_fences.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py # modeling_v2 catalog entries with sm_100 receipts (the K3 ones skip on # other architectures). - - unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py - - unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py - - unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py - - unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py - unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py # The Kimi K3 target's decode GEMVs, LM head and embedding on those entries. - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py # The Kimi K3 target's DSpark drafter on the drafter entries, against the stock block forward. - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py # The Kimi K3 targets' construction checks and the drift check of tp16_moetp16ep1's copied modules. # l0_b300's modeling_v2 entry, the only other one that collects them, is waived (nvbugs/6853741). - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py # tp16_moetp16ep1's MoE decode path: the layers and steps it takes (its kernels: l0_gb200_multi_gpus.yml). - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py # The Kimi K3 targets' generic-path MoE routing, outside the TRTLLM-Gen kernel, against the kernel's own. - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py # KDA runtime state-pool and mixed-batch core coverage. - unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py - unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py @@ -345,14 +345,14 @@ l0_b200: - unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k "test_spec_decode_graph_step_gather_kernel" # Kimi K3 DSpark drafter attention (CuTe DSL, SM 100) and its modeling_v2 catalog entries, plus the sm_100 cells of # fused_qk_norm_rope, which the drafter calls. - - unittest/_torch/kernels/kimi_k3/test_k3_drafter_attn.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py - unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py # Kimi K3 DSpark decode kernels (CuTe DSL, SM 100), one GPU: the acceptance against vocabulary-sharded target # logits and the drafter's context K / V. Their multi-rank tests are in l0_gb200_multi_gpus.yml. - - unittest/_torch/kernels/kimi_k3/test_k3_spec_accept.py - - unittest/_torch/kernels/kimi_k3/test_k3_ctx_kv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py # Vocab-parallel token selection kernels on a single shard. - unittest/_torch/speculative/test_vocab_parallel.py - unittest/_torch/moe/kernels/parallel @@ -423,16 +423,16 @@ l0_b200: # ------------- Prefix-aware scheduling E2E tests --------------- - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke # ------------- Kimi K3 decode MoE kernels and their catalog entries (sm_100) --------------- - - unittest/_torch/kernels/kimi_k3/test_k3_fused_moe.py - - unittest/_torch/kernels/kimi_k3/test_k3_moe_wide.py - - unittest/_torch/kernels/kimi_k3/test_k3_route_quant.py - - unittest/_torch/kernels/kimi_k3/test_k3_moe_front_geometry.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_route_quant.py - - unittest/_torch/kernels/kimi_k3/test_k3_moe_m1.py - - unittest/_torch/kernels/kimi_k3/test_k3_moe_m2.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py # ------------- Visual Gen tests --------------- - unittest/_torch/visual_gen/kernels/test_nvfp4_conv3d.py - unittest/_torch/visual_gen/test_media_decode.py @@ -528,15 +528,22 @@ l0_b200: - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[fl2va] TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[ref2va] TIMEOUT (30) # ------------- Kimi K3 KDA / MLA decode kernels (sm_100) and their modeling_v2 entries --------------- - - unittest/_torch/kernels/kimi_k3/test_k3_kda_attn.py - - unittest/_torch/kernels/kimi_k3/test_k3_kda_decode_attn.py - - unittest/_torch/kernels/kimi_k3/test_k3_kda_pools_past_2g.py - - unittest/_torch/kernels/kimi_k3/test_k3_kda_verify.py - - unittest/_torch/kernels/kimi_k3/test_k3_mla_attn.py - - unittest/_torch/kernels/kimi_k3/test_k3_mla_q.py - - unittest/_torch/modeling_v2/ssm - - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index ffe4b3e8255b..4c421dd99214 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -141,7 +141,46 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. - - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) + - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm --ignore=unittest/_torch/modeling_v2/models/kimi_k3_vl TIMEOUT (30) + # K3 single-GPU and host tests; collective launchers have dedicated multi-GPU entries. + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_step.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ebe3d26c04f1..9864ffab0234 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -45,9 +45,9 @@ l0_cpu: # The GPU lists' attention/ entries run -m "not cpu_only", so hw_agnostic/ runs only here. - unittest/_torch/attention # Kimi K3 MLA decode view from the attention metadata: every test is cpu_only (CPU tensors). - - unittest/_torch/kernels/kimi_k3/test_k3_mla_decode_view.py - - unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py - - unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_decode_view.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_launch.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pool_layout.py - unittest/_torch/compilation/test_auto_multi_stream.py - unittest/_torch/compilation/test_phase_selective_forward.py - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs @@ -79,7 +79,7 @@ l0_cpu: - unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_equality.py - unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py # l0_b300 selects unittest/_torch/modeling_v2 with -m "not cpu_only", so this entry is what runs it. - - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_step.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mnnvl_workspace.py - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 79e79d8848ec..ca9023a1faac 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -77,28 +77,28 @@ l0_gb200_multi_gpus: - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_preserves_cuda_graph_addresses - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_rejects_wrong_group_membership - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_growth_keeps_captured_graphs - - unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py - - unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py # Kimi K3 collective kernels (sm_100) and the catalog entries over caller-owned state, 4 ranks - - unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py - - unittest/_torch/kernels/kimi_k3/test_k3_moe_front.py - - unittest/_torch/kernels/kimi_k3/test_k3_latent_reduce.py - - unittest/_torch/kernels/kimi_k3/test_k3_moe_push.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py # Kimi K3 DSpark decode kernels' cross-rank exchanges (sm_100): the sharded acceptance and the Markov chain, 4 ranks - - unittest/_torch/kernels/kimi_k3/test_k3_spec_accept_sharded.py - - unittest/_torch/kernels/kimi_k3/test_k3_markov.py - - unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_oproj_op_matrix.py - - unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_tail_op_matrix.py - - unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_plain_op_matrix.py - - unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py - - unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py - - unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_latent_reduce_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py - unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py # The Kimi K3 target's decode-path collectives over those entries, 4 ranks - - unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py # tp16_moetp16ep1's MoE decode path: the engines pushing into the latent exchange, 4 ranks - - unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_failure_is_collective_and_bounded diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 58d14ea21ba6..97e375284bc5 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -17,6 +17,15 @@ l0_gb300_multi_gpus: tests: # One ordered sequence in one job: the cases share module-global rank state. - unittest/_torch/modeling_v2/comm + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_identity_vs_mtp3[dep4] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_mtp3_acceptance[modeling_v2-dep4] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_mtp3_acceptance[stock-dep4] diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py index 02a840090e43..51662f609dd8 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py @@ -11,7 +11,8 @@ pytestmark = pytest.mark.cpu_only -def test_rank_job_does_not_shadow_installed_package(tmp_path, monkeypatch): +@pytest.mark.parametrize("family_launcher", [False, True]) +def test_rank_job_does_not_shadow_installed_package(tmp_path, monkeypatch, family_launcher): source_root = tmp_path / "source" comm = source_root / "tests/unittest/_torch/modeling_v2/comm" comm.mkdir(parents=True) @@ -27,7 +28,9 @@ def test_rank_job_does_not_shadow_installed_package(tmp_path, monkeypatch): monkeypatch.setenv("MPI4PY_RC_INITIALIZE", "0") monkeypatch.setenv("TLLM_DISABLE_MPI", "1") result = tmp_path / "package-path.txt" - (comm / "_package_selection_op_matrix.py").write_text( + launcher_dir = comm.parent / "models/example_family" if family_launcher else comm + launcher_dir.mkdir(parents=True, exist_ok=True) + (launcher_dir / "_package_selection_op_matrix.py").write_text( "from pathlib import Path\n" "DEADLINE_S = 30\n" "if __name__ == '__main__':\n" @@ -38,5 +41,5 @@ def test_rank_job_does_not_shadow_installed_package(tmp_path, monkeypatch): spec = importlib.util.spec_from_file_location("_rank_job_package_selection", rank_job_path) rank_job = importlib.util.module_from_spec(spec) spec.loader.exec_module(rank_job) - rank_job.run("package_selection") + rank_job.run("package_selection", directory=launcher_dir if family_launcher else None) assert Path(result.read_text()).resolve() == installed_package diff --git a/tests/unittest/_torch/modeling_v2/comm/_k3_latent_reduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_latent_reduce_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_k3_latent_reduce_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_latent_reduce_op_matrix.py index 7b9cf44f387e..33d0eebddb41 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_k3_latent_reduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_latent_reduce_op_matrix.py @@ -41,6 +41,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "k3_latent_reduce requires CUDA devices" @@ -648,7 +649,7 @@ def _run_one_rank(args) -> int: mnnvl_fusion_allreduce as mnnvl_module, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_context validating = validating_context diff --git a/tests/unittest/_torch/modeling_v2/moe/_k3_moe_engines.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_engines.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/moe/_k3_moe_engines.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_engines.py diff --git a/tests/unittest/_torch/modeling_v2/comm/_k3_moe_front_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_k3_moe_front_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py index cd5e94908e9f..803aaed333e6 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_k3_moe_front_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py @@ -50,6 +50,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "k3_moe_front requires CUDA devices" @@ -946,7 +947,7 @@ def _run_one_rank(args) -> int: k3_moe_front as front, ) - sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating OPS = SimpleNamespace( diff --git a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_common.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_common.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_common.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_common.py diff --git a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_oproj_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_oproj_op_matrix.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_oproj_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_oproj_op_matrix.py index 488ccdad6a32..f8013e2e91ed 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_oproj_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_oproj_op_matrix.py @@ -33,6 +33,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _k3_sandwich_common as cm # noqa: E402 import _lockstep as ls # noqa: E402 @@ -269,7 +270,7 @@ def _run_one_rank(args) -> int: K3SandwichWorkspace, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_context validating = validating_context diff --git a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_plain_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_plain_op_matrix.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_plain_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_plain_op_matrix.py index 98be928ab95b..8967b31e113e 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_plain_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_plain_op_matrix.py @@ -33,6 +33,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _k3_sandwich_common as cm # noqa: E402 import _lockstep as ls # noqa: E402 @@ -183,7 +184,7 @@ def _run_one_rank(args) -> int: K3SandwichWorkspace, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_context validating = validating_context diff --git a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_tail_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_tail_op_matrix.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_tail_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_tail_op_matrix.py index b9fe6a95a5e3..25469ea30ede 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_k3_sandwich_tail_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_sandwich_tail_op_matrix.py @@ -26,6 +26,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _k3_sandwich_common as cm # noqa: E402 import _lockstep as ls # noqa: E402 @@ -249,7 +250,7 @@ def _run_one_rank(args) -> int: K3SandwichWorkspace, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_context validating = validating_context diff --git a/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_cells.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_cells.py diff --git a/tests/unittest/_torch/modeling_v2/ssm/_kda_reference.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_reference.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/ssm/_kda_reference.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_reference.py diff --git a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kimi_k3_decode_comm_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kimi_k3_decode_comm_op_matrix.py index 38da2d4b79ea..a4716c63427a 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kimi_k3_decode_comm_op_matrix.py @@ -65,6 +65,7 @@ from torch import nn sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "the Kimi K3 target's decode collectives require CUDA devices" diff --git a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kimi_k3_route_b_decode_moe_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kimi_k3_route_b_decode_moe_op_matrix.py index afda4d0112f7..88827e952186 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kimi_k3_route_b_decode_moe_op_matrix.py @@ -43,6 +43,7 @@ from torch import nn sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "the Kimi K3 target's MoE decode path requires CUDA devices" diff --git a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_allgather_split_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_allgather_split_op_matrix.py index 5353d9703919..2e3141c6b563 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allgather_split_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_allgather_split_op_matrix.py @@ -38,6 +38,7 @@ from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "mnnvl_allgather_split requires CUDA devices" @@ -713,7 +714,7 @@ def _run_one_rank(args) -> int: mnnvl_fusion_allreduce as reduce_module, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_cm validating = validating_cm diff --git a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allreduce_attn_res_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_allreduce_attn_res_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_mnnvl_allreduce_attn_res_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_allreduce_attn_res_op_matrix.py index ad1dda2deccf..a73f23f69a49 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_allreduce_attn_res_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_allreduce_attn_res_op_matrix.py @@ -28,6 +28,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "mnnvl_allreduce_attn_res requires CUDA devices" @@ -347,7 +348,7 @@ def _run_one_rank(args) -> int: MnnvlWorkspace, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_cm validating = validating_cm diff --git a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_fusion_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_fusion_allreduce_op_matrix.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/comm/_mnnvl_fusion_allreduce_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_fusion_allreduce_op_matrix.py index a1de1dca48e0..b81f261f32fe 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_mnnvl_fusion_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_mnnvl_fusion_allreduce_op_matrix.py @@ -37,6 +37,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _lockstep as ls # noqa: E402 assert torch.cuda.is_available(), "mnnvl_fusion_allreduce requires CUDA devices" @@ -785,7 +786,7 @@ def _run_one_rank(args) -> int: mnnvl_fusion_allreduce as module, ) - sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _validating import validating as validating_cm validating = validating_cm diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_cluster_waits.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_cluster_waits.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_ctm_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_ctm_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_ctm_gemv_wide.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_ctm_gemv_wide.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_ctx_kv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_ctx_kv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_decode_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_drafter_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_drafter_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_embed.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_embed.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_fused_moe.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_fused_moe.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_head_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_head_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_decode_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_decode_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pool_layout.py similarity index 98% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pool_layout.py index 8d05f79c1751..24abc1ae17fa 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_pool_layout.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pool_layout.py @@ -8,7 +8,7 @@ import torch source = ( - Path(__file__).resolve().parents[5] + Path(__file__).resolve().parents[6] / "tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_kda_verify/_pool_layout.py" ) spec = importlib.util.spec_from_file_location("_k3_kda_pool_layout_test", source) diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_pools_past_2g.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_pools_past_2g.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_verify.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_kda_verify.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_latent_reduce.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py similarity index 99% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_latent_reduce.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py index 49faee297b4b..b1eba2a98708 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_latent_reduce.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py @@ -252,7 +252,7 @@ def test_latent_exchange_refuses_unsupported_tp_size(): """The exchange serves TP groups of 4, 8 or 16 ranks: another size raises ValueError before any collective step (the size is the same on every rank of a group), eagerly and under CUDA-graph capture alike. One process, a group of one. Its refusal of capture at a supported size is a collective agreement: the 4-rank matrix - (tests/unittest/_torch/modeling_v2/comm/_k3_latent_reduce_op_matrix.py) checks it.""" + (tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_latent_reduce_op_matrix.py) checks it.""" import tensorrt_llm # noqa: F401 from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_fused_moe import ( latent_op, diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_markov.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_markov.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_decode_view.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_decode_view.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_decode_view.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_decode_view.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_launch.py similarity index 99% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_launch.py index 06017f950a12..56fdda84b681 100644 --- a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_launch.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_launch.py @@ -8,7 +8,7 @@ import pytest source = ( - Path(__file__).resolve().parents[5] + Path(__file__).resolve().parents[6] / "tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_mla/_launch.py" ) spec = importlib.util.spec_from_file_location("_k3_mla_launch_test", source) diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_q.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_mla_q.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_front.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_front.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_front_geometry.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_front_geometry.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_m1.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_m1.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_m2.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_m2.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_push.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_push.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_wide.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_moe_wide.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_route_quant.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_route_quant.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_spec_accept.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_spec_accept.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_spec_accept_sharded.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_spec_accept_sharded.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py diff --git a/tests/unittest/_torch/kernels/kimi_k3/test_k3_tcgen05_fences.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py similarity index 100% rename from tests/unittest/_torch/kernels/kimi_k3/test_k3_tcgen05_fences.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py index 1cec917af382..2f9082d93928 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py @@ -21,7 +21,7 @@ attn_res_fwd, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "attn_res_fwd requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py index 580ab1f8dc76..9e371d4c7a60 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py @@ -22,7 +22,7 @@ attn_res_rmsnorm_fwd, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "attn_res_rmsnorm_fwd requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_causal_conv1d_fwd.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py index 618fed7edd30..c36118083ab4 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_causal_conv1d_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py @@ -12,7 +12,7 @@ ) from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "causal_conv1d_fwd requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py index 79ba6e832536..4404276a2a22 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_chunk_kda.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py @@ -17,7 +17,7 @@ ) from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "chunk_kda requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py index 2d40a3db8e00..82ea0637863d 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py @@ -11,7 +11,7 @@ k3_ctm_gemv, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_ctm_gemv requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py index 4f6e221d9143..20d74d4c1309 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py @@ -11,7 +11,7 @@ k3_ctm_gemv_long, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_ctm_gemv_long requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py index 3be5fc31e3d1..32569113e62d 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py @@ -11,7 +11,7 @@ k3_ctm_gemv_swiglu, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_ctm_gemv_swiglu requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py index c1feac1dcbd4..78f5e5404dda 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py @@ -14,7 +14,7 @@ k3_ctm_gemv_wide, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_ctm_gemv_wide requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py index dc3852629592..537c3278dfb9 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py @@ -12,7 +12,7 @@ k3_decode_gemv, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_decode_gemv requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py index 0186e13d2dce..d98232ab5731 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py @@ -21,7 +21,7 @@ k3_drafter_attn_qknorm, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating # noqa: E402 D = 64 diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py index ff080aa83dfe..02ff7d504612 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py @@ -17,7 +17,7 @@ k3_drafter_attn_qknorm, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating # noqa: E402 D = 64 diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py index 41b643706f3c..b0034836d60d 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py @@ -20,7 +20,7 @@ k3_embed_norm, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_embed_norm requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py index 41eb17a18e84..e02c60f6ccbb 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py @@ -13,7 +13,7 @@ k3_head_gemv, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_head_gemv requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py index 8b7b3a728302..d24373adbbf1 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py @@ -32,7 +32,7 @@ K3KdaBuffers, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating _attn = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py index 19683e0b84ea..8eb55cb10608 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py @@ -37,7 +37,7 @@ kda_decode, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating _attn = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py index 754db075f388..eaa72d0669d7 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py @@ -24,7 +24,7 @@ import pytest import torch -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating k3_kda_verify = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_latent_reduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_latent_reduce_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py index a37c5edfbcf3..924f4e96982a 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_latent_reduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py @@ -6,7 +6,10 @@ caller-owned exchanges, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_k3_latent_reduce_op_matrix() -> None: - _rank_job.run("k3_latent_reduce") + _rank_job.run("k3_latent_reduce", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py index b7955dfc5400..c97a4188f849 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py @@ -54,7 +54,7 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py index d7006ca5a2b4..99018d2c45c7 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py @@ -45,7 +45,7 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py index 4d95bbf4c605..870d3e27ec62 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py @@ -15,11 +15,12 @@ The checks run in file order and share the state objects: a K3MoeState with layers A, B, C (A and C over this rank's experts, B over other experts), a second K3MoeState, a K3MoeWideState with layers A and B. Each check makes its own -first (compiling) call eagerly where it needs one, so it also runs alone (-k). The kernel tests under -tests/unittest/_torch/kernels/kimi_k3/ remain the exhaustive numerics; this file copies what it needs. +first (compiling) call eagerly where it needs one, so it also runs alone (-k). +The test_k3_* modules beside this file provide the exhaustive kernel numerics; +this file copies what it needs. The head_flags build needs the TP group's head workspace and the MoE front: its cells run in moe/k3_moe_front's -4-rank matrix, tests/unittest/_torch/modeling_v2/comm/_k3_moe_front_op_matrix.py. +4-rank matrix, tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py. """ import contextlib @@ -42,7 +43,7 @@ k3_route_quant, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating H, TOP_K, NUM_EXPERTS, SV = 3584, 16, 896, 32 diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py similarity index 53% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py index d19d722ff2f2..f32afe2e80c4 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_front_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py @@ -2,11 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Collected entry point for the k3_moe_front op's certification matrix. -The matrix also certifies moe/k3_moe's cells behind the front (trtllm::k3_moe_front, then the k3_moe entry on a plain -and on a head_flags K3MoeState). It is ``comm/_k3_moe_front_op_matrix.py``: rank bodies live beside ``_lockstep`` and -``_rank_job`` in ``comm/``, and it is its own W-rank launcher (call sequences over caller-owned workspaces and states, -so one job, not independent cases); see ``_rank_job`` for why that is left intact. This file puts ``comm/`` on the -import path to reach ``_rank_job``. +The matrix beside this file also certifies moe/k3_moe's cells behind the front +(trtllm::k3_moe_front, then the k3_moe entry on a plain or head_flags K3MoeState). +The shared comm/_rank_job helper starts its W-rank launcher as one ordered job. """ import sys @@ -15,7 +13,7 @@ import pytest import torch -sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "comm")) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "comm")) import _rank_job # noqa: E402 @@ -29,4 +27,4 @@ def _is_sm100() -> bool: # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_k3_moe_front_op_matrix() -> None: - _rank_job.run("k3_moe_front") + _rank_job.run("k3_moe_front", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py index 9487eb8c60d2..7c88da4eb934 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m1.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py @@ -10,7 +10,7 @@ import _k3_moe_engines as engines import pytest -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating k3_moe_m1 = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py index 493eb1a5272f..4b6afdc98838 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_moe_m2.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py @@ -10,7 +10,7 @@ import _k3_moe_engines as engines import pytest -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating k3_moe_m2 = importlib.import_module( diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_route_quant.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_route_quant.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py index 3ec7f6015732..e0a750bdd4fc 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_k3_route_quant.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py @@ -26,7 +26,7 @@ k3_route_quant, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating E, K, H = 896, 16, 3584 diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_oproj_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_oproj_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py index 60d8ff2afa7a..3224e6b470f6 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_oproj_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py @@ -6,7 +6,10 @@ caller-owned workspace, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_k3_sandwich_oproj_op_matrix() -> None: - _rank_job.run("k3_sandwich_oproj") + _rank_job.run("k3_sandwich_oproj", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_plain_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_plain_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py index 5acc2536f623..1c7d379532fc 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_plain_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py @@ -6,7 +6,10 @@ caller-owned workspace, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_k3_sandwich_plain_op_matrix() -> None: - _rank_job.run("k3_sandwich_plain") + _rank_job.run("k3_sandwich_plain", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_tail_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_tail_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py index 8a08b0b5adb5..c1411bd39d6b 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_k3_sandwich_tail_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py @@ -6,7 +6,10 @@ caller-owned workspace, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_k3_sandwich_tail_op_matrix() -> None: - _rank_job.run("k3_sandwich_tail") + _rank_job.run("k3_sandwich_tail", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py index f7c12ea5db6d..fb85cd558e0b 100644 --- a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py @@ -17,7 +17,7 @@ k3_situ_mul, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "k3_situ_mul requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py index 1f105019ae8a..ec9e1d6012ce 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py @@ -20,7 +20,7 @@ kda_decode, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating pytestmark = pytest.mark.skipif(not kc.sm100(), reason="certified on SM100 only") diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_post_conv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_post_conv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py index bd67ce4ca8cc..ba4c82a5b600 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_post_conv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py @@ -11,7 +11,7 @@ kda_post_conv, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "kda_post_conv requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py index 78f91c4a4e48..6001c86d0fa1 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_prefill.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py @@ -14,7 +14,7 @@ kda_prefill, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "kda_prefill requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py similarity index 88% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py index 74f1c9733040..b52427943efc 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py @@ -7,7 +7,10 @@ target's collective state, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -29,4 +32,4 @@ # The case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_kimi_k3_decode_comm_op_matrix() -> None: - _rank_job.run("kimi_k3_decode_comm") + _rank_job.run("kimi_k3_decode_comm", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_step.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_step.py diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py index b41a007e3403..d91f34f0b578 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py @@ -15,7 +15,7 @@ kimi_k3_noaux_tc_mxfp8_quant, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "kimi_k3_noaux_tc_mxfp8_quant requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py similarity index 88% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py index a7d1b8df1a09..d53299b07c8d 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py @@ -7,7 +7,10 @@ the path's collective state, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -29,4 +32,4 @@ # The case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_kimi_k3_route_b_decode_moe_op_matrix() -> None: - _rank_job.run("kimi_k3_route_b_decode_moe") + _rank_job.run("kimi_k3_route_b_decode_moe", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py index 1c59d2c29088..4c9aa757517c 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py @@ -6,7 +6,10 @@ over caller-owned workspaces, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_mnnvl_allgather_split_op_matrix() -> None: - _rank_job.run("mnnvl_allgather_split") + _rank_job.run("mnnvl_allgather_split", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py index 08836fbb349c..1b2d1676611f 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py @@ -6,7 +6,10 @@ over one caller-owned workspace, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_mnnvl_allreduce_attn_res_op_matrix() -> None: - _rank_job.run("mnnvl_allreduce_attn_res") + _rank_job.run("mnnvl_allreduce_attn_res", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py index 8ed3cab58d6b..f7444af490ba 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py @@ -6,7 +6,10 @@ over caller-owned workspaces, so one job, not independent cases); see ``_rank_job`` for why that is left intact. """ -import _rank_job +from pathlib import Path + +__extra_import_path__ = ["../../comm"] +import _rank_job # noqa: E402 import pytest import torch @@ -19,4 +22,4 @@ # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. @pytest.mark.no_xdist def test_mnnvl_fusion_allreduce_op_matrix() -> None: - _rank_job.run("mnnvl_fusion_allreduce") + _rank_job.run("mnnvl_fusion_allreduce", directory=Path(__file__).resolve().parent) diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_rms_norm_gated_token_major.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py index 1a3831b15ca4..c91d8a1a6559 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py @@ -11,7 +11,7 @@ rms_norm_gated_token_major, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_situ_and_mul.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py similarity index 99% rename from tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_situ_and_mul.py rename to tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py index eaca0e7c038b..62123daf1e24 100644 --- a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_situ_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py @@ -12,7 +12,7 @@ situ_and_mul, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../.."] from _validating import validating assert torch.cuda.is_available(), "situ_and_mul requires a CUDA device" diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 56c1780813e0..b6eaa3f1ee81 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -165,22 +165,38 @@ def test_k3_uses_modeling_v2_selection(rule, relative): assert _apply(rule, path).scope == "modelingv2only" -def test_k3_native_change_keeps_existing_kernel_and_collective_tests(rule): +@pytest.mark.parametrize("relative", ["csrc/k3MnnvlAllreduce.cu", "target/modeling.py"]) +def test_k3_family_paths_select_kernel_collective_and_host_tests(rule, relative): family = rule._src_root / "models/kimi_k3_vl" (family / "target").mkdir(parents=True) (family / "target/modeling.py").touch() (family / "routing.py").write_text("from .target.modeling import KimiK3\n") + prefix = "unittest/_torch/modeling_v2/models/kimi_k3_vl" kernel_tests = { - "unittest/_torch/kernels/kimi_k3/test_k3_mnnvl_comm.py", - "unittest/_torch/kernels/kimi_k3/test_k3_sandwich.py", + f"{prefix}/test_k3_mnnvl_comm.py", + f"{prefix}/test_k3_sandwich.py", } - matrix = "unittest/_torch/modeling_v2/comm/test_modeling_v2_mnnvl_allgather_split_op_matrix.py" + matrix = f"{prefix}/test_modeling_v2_mnnvl_allgather_split_op_matrix.py" + norm = f"{prefix}/test_modeling_v2_attn_res_rmsnorm_fwd.py" + native = f"{prefix}/test_modeling_v2_k3_native.py" + host = f"{prefix}/test_modeling_v2_kimi_k3_decode_step.py" + accuracy = "accuracy/modeling_v2/test_kimi_k3_vl.py::TestKimiK3::test_gsm8k" + # The broad single-GPU entry excludes the flat family directory; its + # single-GPU tests and collective tests have separate stage entries. + rule.yaml_index.blocks[0].tests.remove(_UNIT) + rule.yaml_index.blocks[0].tests.extend([f"{_UNIT} --ignore={prefix}", native, accuracy]) rule.yaml_index.blocks[1].tests.extend([*sorted(kernel_tests), matrix]) - norm = "unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py" rule.yaml_index.blocks[2].tests.append(norm) - result = _apply(rule, f"{_PY}/models/kimi_k3_vl/csrc/k3MnnvlAllreduce.cu") - assert _selected(result) == kernel_tests | {matrix, norm, _COMM, _UNIT, _CLAIMS} - assert "B200-PyTorch-1" in result.affected_stages + rule.yaml_index.blocks[3].tests.append(host) + result = _apply(rule, f"{_PY}/models/kimi_k3_vl/{relative}") + selected = kernel_tests | {matrix, norm, native, host, accuracy, _CLAIMS} + assert _selected(result) == selected + assert result.affected_stages == { + "B300-PyTorch-1", + "B200-PyTorch-1", + "GB300-4_GPUs-PyTorch-1", + "CPU-Generic-x86-1", + } assert result.scope == "modelingv2only" other = _selected(_apply(rule, f"{_PY}/models/gpt_oss/csrc/kernel.cu")) - assert not (kernel_tests | {matrix, norm}) & other + assert not (selected - {_CLAIMS}) & other From da6676877bee0982b8ac1ada2a022116f4469209 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Thu, 29 Oct 2026 07:19:19 -0700 Subject: [PATCH 10/11] Enable K3 SM103 kernel tests and consolidate GB300 CI coverage Signed-off-by: Tao Li --- jenkins/L0_Test.groovy | 2 + .../kimi_k3_vl/kernels/k3_fused_moe/op.py | 7 +- .../kernels/k3_moe_engines/op.py | 7 +- .../ops/moe/k3_moe_m1.py | 2 +- .../ops/moe/k3_moe_m2.py | 2 +- .../ops/ssm/k3_kda_attn.py | 2 +- .../ops/ssm/k3_kda_verify.py | 2 +- .../kimi_k3_vl/ops/activation/k3_situ_mul.py | 2 +- .../kimi_k3_vl/ops/activation/situ_and_mul.py | 2 +- .../ops/attention/k3_drafter_attn.py | 2 +- .../ops/attention/k3_drafter_attn_qknorm.py | 2 +- .../ops/attention/k3_mla_attn_vb_out.py | 2 +- .../kimi_k3_vl/ops/attention/k3_mla_qkv.py | 2 +- .../kimi_k3_vl/ops/comm/k3_latent_reduce.py | 2 +- .../kimi_k3_vl/ops/comm/k3_sandwich_oproj.py | 2 +- .../kimi_k3_vl/ops/comm/k3_sandwich_plain.py | 2 +- .../kimi_k3_vl/ops/comm/k3_sandwich_tail.py | 2 +- .../ops/comm/mnnvl_allgather_split.py | 2 +- .../ops/comm/mnnvl_allreduce_attn_res.py | 2 +- .../ops/comm/mnnvl_fusion_allreduce.py | 2 +- .../models/kimi_k3_vl/ops/gemm/k3_ctm_gemv.py | 2 +- .../kimi_k3_vl/ops/gemm/k3_ctm_gemv_long.py | 2 +- .../kimi_k3_vl/ops/gemm/k3_ctm_gemv_swiglu.py | 2 +- .../kimi_k3_vl/ops/gemm/k3_ctm_gemv_wide.py | 2 +- .../kimi_k3_vl/ops/gemm/k3_decode_gemv.py | 2 +- .../kimi_k3_vl/ops/gemm/k3_head_gemv.py | 2 +- .../models/kimi_k3_vl/ops/moe/k3_moe.py | 2 +- .../models/kimi_k3_vl/ops/moe/k3_moe_front.py | 2 +- .../kimi_k3_vl/ops/moe/k3_route_quant.py | 4 +- .../ops/moe/kimi_k3_noaux_tc_mxfp8_quant.py | 2 +- .../kimi_k3_vl/ops/norm/attn_res_fwd.py | 2 +- .../ops/norm/attn_res_rmsnorm_fwd.py | 2 +- .../kimi_k3_vl/ops/norm/k3_embed_norm.py | 2 +- .../ops/norm/rms_norm_gated_token_major.py | 2 +- .../kimi_k3_vl/ops/ssm/causal_conv1d_fwd.py | 2 +- .../models/kimi_k3_vl/ops/ssm/chunk_kda.py | 2 +- .../kimi_k3_vl/ops/ssm/k3_kda_decode_attn.py | 2 +- .../models/kimi_k3_vl/ops/ssm/kda_decode.py | 2 +- .../kimi_k3_vl/ops/ssm/kda_post_conv.py | 2 +- .../models/kimi_k3_vl/ops/ssm/kda_prefill.py | 2 +- .../test_lists/test-db/l0_b200.yml | 77 +------------------ .../test_lists/test-db/l0_b300.yml | 39 ---------- .../test-db/l0_gb200_multi_gpus.yml | 21 ----- .../test_lists/test-db/l0_gb300.yml | 77 ++++++++++++++++++- .../test-db/l0_gb300_multi_gpus.yml | 31 +++++++- .../models/kimi_k3_vl/_k3_moe_engines.py | 2 +- .../kimi_k3_vl/_k3_moe_front_op_matrix.py | 4 +- .../models/kimi_k3_vl/_kda_cells.py | 4 +- .../models/kimi_k3_vl/test_k3_fused_moe.py | 4 +- .../models/kimi_k3_vl/test_k3_moe_front.py | 2 +- .../models/kimi_k3_vl/test_k3_moe_m1.py | 4 +- .../models/kimi_k3_vl/test_k3_moe_m2.py | 4 +- .../models/kimi_k3_vl/test_k3_moe_push.py | 2 +- .../models/kimi_k3_vl/test_k3_moe_wide.py | 4 +- .../test_modeling_v2_attn_res_fwd.py | 3 +- .../test_modeling_v2_attn_res_rmsnorm_fwd.py | 3 +- .../test_modeling_v2_causal_conv1d_fwd.py | 3 +- .../kimi_k3_vl/test_modeling_v2_chunk_kda.py | 3 +- .../test_modeling_v2_k3_ctm_gemv.py | 3 +- .../test_modeling_v2_k3_ctm_gemv_long.py | 3 +- .../test_modeling_v2_k3_ctm_gemv_swiglu.py | 3 +- .../test_modeling_v2_k3_ctm_gemv_wide.py | 3 +- .../test_modeling_v2_k3_decode_gemv.py | 3 +- .../test_modeling_v2_k3_drafter_attn.py | 6 +- ...test_modeling_v2_k3_drafter_attn_qknorm.py | 6 +- .../test_modeling_v2_k3_embed_norm.py | 3 +- .../test_modeling_v2_k3_head_gemv.py | 3 +- ..._modeling_v2_k3_latent_reduce_op_matrix.py | 4 +- .../test_modeling_v2_k3_mla_attn_vb_out.py | 6 +- .../kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py | 6 +- .../kimi_k3_vl/test_modeling_v2_k3_moe.py | 4 +- ...test_modeling_v2_k3_moe_front_op_matrix.py | 4 +- .../kimi_k3_vl/test_modeling_v2_k3_moe_m1.py | 2 +- .../kimi_k3_vl/test_modeling_v2_k3_moe_m2.py | 2 +- .../test_modeling_v2_k3_route_quant.py | 4 +- ...modeling_v2_k3_sandwich_oproj_op_matrix.py | 4 +- ...modeling_v2_k3_sandwich_plain_op_matrix.py | 4 +- ..._modeling_v2_k3_sandwich_tail_op_matrix.py | 4 +- .../test_modeling_v2_k3_situ_mul.py | 3 +- .../kimi_k3_vl/test_modeling_v2_kda_decode.py | 2 +- .../test_modeling_v2_kda_post_conv.py | 3 +- .../test_modeling_v2_kda_prefill.py | 3 +- ...deling_v2_kimi_k3_decode_comm_op_matrix.py | 4 +- .../test_modeling_v2_kimi_k3_decode_gemv.py | 4 +- .../test_modeling_v2_kimi_k3_drafter.py | 4 +- .../test_modeling_v2_kimi_k3_moe_routing.py | 4 +- ...odeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py | 3 +- ...v2_kimi_k3_route_b_decode_moe_op_matrix.py | 4 +- ...ling_v2_mnnvl_allgather_split_op_matrix.py | 4 +- ...g_v2_mnnvl_allreduce_attn_res_op_matrix.py | 6 +- ...ing_v2_mnnvl_fusion_allreduce_op_matrix.py | 4 +- ..._modeling_v2_rms_norm_gated_token_major.py | 3 +- .../test_modeling_v2_situ_and_mul.py | 3 +- .../scripts/cbts/test_modeling_v2_rule.py | 63 ++++++++++++++- 94 files changed, 315 insertions(+), 258 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 0fded65aa33e..d9075e450ab4 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -6675,7 +6675,9 @@ def launchTestJobs(pipeline, testFilter, globalVars) "GB200-4_GPUs-PyTorch-Post-Merge-1": ["auto:gb200-x4", "l0_gb200_multi_gpus", 1, 2, 4, 1, false, true], "GB200-4_GPUs-PyTorch-Post-Merge-2": ["auto:gb200-x4", "l0_gb200_multi_gpus", 2, 2, 4, 1, false, true], "GB10-PyTorch-Post-Merge-1": ["gb10x-single", "l0_gb10", 1, 1], + "GB300-PyTorch-1": ["auto:gb300-flex", "l0_gb300", 1, 1, 1, 1, true, false], "GB300-4_GPUs-PyTorch-1": ["auto:gb300-x4", "l0_gb300", 1, 1, 4, 1, true, false], + "GB300-4_GPUs-PyTorch-2": ["auto:gb300-x4", "l0_gb300_multi_gpus", 1, 1, 4, 1, true, false], "GB300-4_GPUs-PyTorch-Post-Merge-1": ["auto:gb300-x4", "l0_gb300_multi_gpus", 1, 3, 4, 1, true, false], "GB300-4_GPUs-PyTorch-Post-Merge-2": ["auto:gb300-x4", "l0_gb300_multi_gpus", 2, 3, 4, 1, true, false], "GB300-4_GPUs-PyTorch-Post-Merge-3": ["auto:gb300-x4", "l0_gb300_multi_gpus", 3, 3, 4, 1, true, false], diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_fused_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_fused_moe/op.py index 1895004daef1..c389c73e584e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_fused_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kernels/k3_fused_moe/op.py @@ -87,8 +87,11 @@ def is_supported(w3_w1_weight: torch.Tensor, w3_w1_weight_scale: torch.Tensor, w """Whether these TRTLLM-Gen W4A8_MXFP4_MXFP8 buffers fit the fused kernels. Only metadata is read, so the buffers may still be on the meta device.""" - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - return False, "needs sm_100" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ( + (10, 0), + (10, 3), + ): + return False, "needs sm_100 or sm_103" try: import cutlass # noqa: F401 except ImportError: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/kernels/k3_moe_engines/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/kernels/k3_moe_engines/op.py index 9261aa6064b8..c7b0ce08048c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/kernels/k3_moe_engines/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/kernels/k3_moe_engines/op.py @@ -66,8 +66,11 @@ def m1_supported( logical intermediate; the buffers may hold it zero-padded to a multiple of 128, as TRT-LLM's loader lays out an unaligned shard (e.g. 192 -> 256 at TP16). Only metadata is read, so the buffers may still be on the meta device.""" - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - return False, "needs sm_100" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ( + (10, 0), + (10, 3), + ): + return False, "needs sm_100 or sm_103" try: import cutlass # noqa: F401 except ImportError: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m1.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m1.py index cef20d897af6..bc426bae6122 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m1.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m1.py @@ -20,7 +20,7 @@ class K3MoeM1(OpWrapper): """This rank's routed partial ``[M, 3584]`` bf16 of ``layer``'s experts for the M tokens of its state. Advances the state's workspace by one call: the calls on one state run in one stream order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="TP16 x EP1 single-token weight stream with 192 real columns padded to 256", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m2.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m2.py index 39be66661484..124c9b788c2d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m2.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/ops/moe/k3_moe_m2.py @@ -20,7 +20,7 @@ class K3MoeM2(OpWrapper): """This rank's routed partial ``[2, 3584]`` bf16 of ``layer``'s experts for two tokens. Advances the state's workspace by one call: the calls on one state run in one stream order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="TP16 x EP1 two-token grouped schedule with 192 real columns padded to 256", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_attn.py index b9c23aa13c94..d58854fd1874 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_attn.py @@ -29,7 +29,7 @@ class K3KdaAttn(OpWrapper): [8, 7168]: the projection ``x w^T`` and the verify of ``ssm/k3_kda_verify`` on it, in one launch. Updates the slot's conv caches, state and per-draft states in place; advances ``buffers`` by one launch.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS = ( Cell( why="TP16 golden token plus seven drafts: two successive requests, two layers, pending through 0..7", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_verify.py index 663e53fd19db..cb31b0a791bc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_verify.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/ops/ssm/k3_kda_verify.py @@ -23,7 +23,7 @@ class K3KdaVerify(OpWrapper): """Return the gated-norm core output bf16 [N (1 + num_spec), H, 128] of the verify tokens whose fused projection rows are ``proj``; updates the slots' conv caches, state and per-draft states in place.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS = ( Cell( why="one TP16 request, seven drafts and six rounds of accepted-draft replay", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/k3_situ_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/k3_situ_mul.py index e69d2f9281d2..8a0bd80419c7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/k3_situ_mul.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/k3_situ_mul.py @@ -11,7 +11,7 @@ class K3SituMul(OpWrapper): """Return `SituAndMul(beta, linear_beta)(gu)`, gate half first, in one k3_situ_mul call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"{name} caps at M={m}: TP16 dense MLP row-count independence", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/situ_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/situ_and_mul.py index b8c982647257..54a09600dde0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/situ_and_mul.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/activation/situ_and_mul.py @@ -17,7 +17,7 @@ class SituAndMul(OpWrapper): A new `[M, d]` tensor in `x.dtype`. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn.py index 7c643238e483..a43f43db512a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn.py @@ -18,7 +18,7 @@ class K3DrafterAttn(OpWrapper): """Write ``out`` [R T, num_heads * 64]: request r's T rows of ``qkv`` attend to its ``ctx_len[r]`` cached rows (pages: row r of ``page_table``) and to its block's own k / v in ``qkv``. Returns None.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn_qknorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn_qknorm.py index 39d7cc603a27..e54d267da1a0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn_qknorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_drafter_attn_qknorm.py @@ -15,7 +15,7 @@ class K3DrafterAttnQknorm(OpWrapper): """``k3_drafter_attn`` after RMSNorm (``q_w`` / ``k_w``, ``eps``) and NeoX RoPE (``rope_base``, one position per row of ``qkv``) of every q and k head, computed in the kernel; ``qkv`` is not modified. Returns None.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"R={r}, T={t}: " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py index 3a4291d3343c..06f8a56c4d71 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_attn_vb_out.py @@ -44,7 +44,7 @@ class K3MlaAttnVbOut(OpWrapper): ``gate``) for the decode attention o of R <= 8 requests of T <= 8 tokens over their pages of ``pool``. Uses ``workspace``'s partial slots and, past 7 clusters, advances its counters.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py index fe65a88a2b8e..d7704c6d6ba3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/attention/k3_mla_qkv.py @@ -21,7 +21,7 @@ class K3MlaQkv(OpWrapper): """Return ``fused_q`` [M, heads * 576] bf16 of the M <= 64 rows of ``ag``, and store each token's cache row ``[bf16(rmsnorm(ag[t, 1536:2048]) * w_kv) | ag[t, 2048:2112]]`` into ``pool`` at its request's position.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_latent_reduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_latent_reduce.py index 4403f3c35029..f3a4c530c6f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_latent_reduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_latent_reduce.py @@ -23,7 +23,7 @@ class K3LatentReduce(OpWrapper): every rank pushed into it since the previous reduce. Empties the words it read and advances ``exchange`` by one call: every rank pushes, then reduces, the same token count on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"{tokens} decode tokens, {kind} partials: {reason}", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_oproj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_oproj.py index 6fc4d5a2347d..277194825b84 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_oproj.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_oproj.py @@ -140,7 +140,7 @@ class K3SandwichOproj(OpWrapper): (the sum alone without ``prefix``), ``normed`` = RMSNorm(attn_res(block_residual..., updated)). Advances ``workspace`` by one call: every rank of the group makes the same calls on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"{tokens} decode tokens, {snapshots} snapshots, prefix={prefix}: attention projection", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_plain.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_plain.py index ffca9c948c1d..6283cd579ab0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_plain.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_plain.py @@ -30,7 +30,7 @@ class K3SandwichPlain(OpWrapper): (with ``swiglu``, of ``silu_and_mul(x) @ weight^T``), ``normed = RMSNorm(updated) * norm_weight``. Advances ``workspace`` by one call: every rank of the group makes the same calls on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"{tokens} decode tokens, swiglu={swiglu}, draw {draw}: drafter " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_tail.py index 45de1fe88ab7..6096e2147e40 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/k3_sandwich_tail.py @@ -37,7 +37,7 @@ class K3SandwichTail(OpWrapper): (``updated`` with ``tap_updated``). ``updated_out``: store ``updated`` there; it is then the returned ``updated``. Advances ``workspace`` by one call: every rank of the group makes the same calls on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"{tokens} decode tokens, {snapshots} snapshots, prefix={prefix}: latent slice shift {shift}", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py index 056c0fec7647..75fb1b2f6178 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allgather_split.py @@ -32,7 +32,7 @@ class MnnvlAllgatherSplit(OpWrapper): Takes one turn of the workspace's Lamport rotation, as an all-reduce on it does: every rank of the group makes the same MNNVL calls on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=f"{label}: bf16 rounding and fp32 words at every decode and wide token count", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py index feff2441b0aa..9fd00844d824 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_allreduce_attn_res.py @@ -20,7 +20,7 @@ class MnnvlAllreduceAttnRes(OpWrapper): sum alone without ``prefix_sum``), ``normed`` = RMSNorm(attn_res(block_residual..., updated)). Advances ``workspace`` by one call: every rank of the group makes the same calls on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py index f65bdaad67cd..2639b6b8850a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/comm/mnnvl_fusion_allreduce.py @@ -38,7 +38,7 @@ class MnnvlFusionAllreduce(OpWrapper): must hold the call (:func:`required_buffer_bytes`). Advances ``workspace`` by one call: every rank of the group makes the same MNNVL calls on it in the same order.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv.py index 3492e2d425b4..c20fb35854ee 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv.py @@ -16,7 +16,7 @@ class K3CtmGemv(OpWrapper): """Return `bf16(x @ weight.T)`, fp32-accumulated, in one k3_ctm_gemv call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="MLA o_proj, all K resident in one CTA", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_long.py index 0262010ca34f..f98e5d9fc293 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_long.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_long.py @@ -16,7 +16,7 @@ class K3CtmGemvLong(OpWrapper): """Return `bf16(x @ weight.T)`, columns >= `sig_col0` (if >= 0) as `bf16(sigmoid(.))`, in one call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="MLA fused W_a/W_g, sigmoid gate columns and split 6", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_swiglu.py index 7bd40322423e..365b78cae354 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_swiglu.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_swiglu.py @@ -17,7 +17,7 @@ class K3CtmGemvSwiglu(OpWrapper): """Return `bf16(silu_and_mul(gu) @ weight.T)`, gate columns first, in one k3_ctm_gemv_swiglu call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="drafter down, seven k-tiles split unevenly as four and three", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_wide.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_wide.py index 5cfd29a372e0..fef16d453434 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_wide.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_ctm_gemv_wide.py @@ -16,7 +16,7 @@ class K3CtmGemvWide(OpWrapper): """Return `x @ weight.T` in bf16 (columns >= `sig_col0` as sigmoid) or fp32, in one call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="KDA q/k/v/g/f_a/b, final output block has eight rows", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_decode_gemv.py index e8f2a6c71fe4..21dc01f0cd5e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_decode_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_decode_gemv.py @@ -18,7 +18,7 @@ class K3DecodeGemv(OpWrapper): """Return ``x @ weight^T`` as a new bf16 ``[M, N]`` tensor (fp32 accumulation, one bf16 rounding) in one k3_decode_gemv call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="KDA o_proj, short K resident in one CTA per weight tile", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_head_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_head_gemv.py index 8b5687ddbdf7..e5e547b2a9f2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_head_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/gemm/k3_head_gemv.py @@ -29,7 +29,7 @@ class K3HeadGemv(OpWrapper): ``workspace``, whose schedule and ``chunk_tiles`` the call takes. The calls that share a workspace must be ordered on one stream.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="TP16 LM-head shard, stream-K workspace across all decode row counts", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe.py index 77b4d94ec223..96bef3fa0784 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe.py @@ -80,7 +80,7 @@ class K3Moe(OpWrapper): counters before its grid-dependency wait. Between the two, some kernel must wait for its predecessor before it triggers its dependents (``griddepcontrol.wait``, then ``launch_dependents``), or the stream is synchronized.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="M<=8 random routing over the TP4 x EP4 expert window", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe_front.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe_front.py index e693b766999e..ec50a666558e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe_front.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_moe_front.py @@ -39,7 +39,7 @@ class K3MoeFront(OpWrapper): ``publish``: also release the workspace's per-token ready words, which the next ``moe/k3_moe`` call on a head_flags state (``head=workspace``) acquires; every publishing call is followed by exactly one such call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="four-rank gathered head and TP16-width shared activation, every decode row count", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_route_quant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_route_quant.py index ac528382a74d..542f68a35945 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_route_quant.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/k3_route_quant.py @@ -27,7 +27,7 @@ class K3RouteQuant(OpWrapper): grid-dependency wait, before the outputs are written; for a dependent that waits for this whole grid before reading them, as ``trtllm::k3_moe`` does.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( @@ -94,7 +94,7 @@ class K3RouteQuant(OpWrapper): with interleaved eager calls and retain the isolated calls' bits. Invalid shapes, dtypes, layouts and token counts raise ValueError; the tests retain those refusals and both cold-capture refusals. - Certification is SM100 only; the stock comparator requires SM10.x. + The stock comparator requires SM10.x. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/kimi_k3_noaux_tc_mxfp8_quant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/kimi_k3_noaux_tc_mxfp8_quant.py index a0550a28fd72..43fbf55381f1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/kimi_k3_noaux_tc_mxfp8_quant.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/moe/kimi_k3_noaux_tc_mxfp8_quant.py @@ -16,7 +16,7 @@ class KimiK3NoauxTcMxfp8Quant(OpWrapper): to noaux_tc_op), linear UE8M0 scales, one per 32 elements. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_fwd.py index 3fa360685256..3fd325bf1865 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_fwd.py @@ -19,7 +19,7 @@ class AttnResFwd(OpWrapper): per-candidate rsqrt(mean square + eps), softmax probabilities and logits, all newly allocated. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py index 4b990f21abfb..c2ea442743b7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/attn_res_rmsnorm_fwd.py @@ -19,7 +19,7 @@ class AttnResRmsnormFwd(OpWrapper): it writes the output, so a dependent PDL kernel must wait on the grid dependency before reading it. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/k3_embed_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/k3_embed_norm.py index 51ab00a7fb91..4e698d428a1a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/k3_embed_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/k3_embed_norm.py @@ -13,7 +13,7 @@ class K3EmbedNorm(OpWrapper): """Write ``table[ids]`` into ``raw`` (zero rows for ids outside ``[0, V)``) and return ``raw * rsqrt(mean(raw^2) + eps) * weight`` as a new bf16 tensor, in one k3_embed_norm call.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/rms_norm_gated_token_major.py index 55ee4c52b3ac..052607c14ed2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/norm/rms_norm_gated_token_major.py @@ -17,7 +17,7 @@ class RmsNormGatedTokenMajor(OpWrapper): the result is quantized to float8_e4m3fn in the same kernel. A new tensor. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/causal_conv1d_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/causal_conv1d_fwd.py index 609c1f0eba79..1087c67dbb15 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/causal_conv1d_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/causal_conv1d_fwd.py @@ -17,7 +17,7 @@ class CausalConv1dFwd(OpWrapper): inputs into `conv_states[cache_indices[b]]`. Returns None. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/chunk_kda.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/chunk_kda.py index cad5b759e452..26eeb35d472f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/chunk_kda.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/chunk_kda.py @@ -16,7 +16,7 @@ class ChunkKda(OpWrapper): it as the keyword arguments it reads them from. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( Cell( why="TP16: two ragged sequences mixing fresh and continuing states below four chunks", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/k3_kda_decode_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/k3_kda_decode_attn.py index 4e4ac43e9973..50dd8b465e09 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/k3_kda_decode_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/k3_kda_decode_attn.py @@ -26,7 +26,7 @@ class K3KdaDecodeAttn(OpWrapper): [R, 7168]: the projection ``x w^T`` and the plain decode of ``ssm/kda_decode`` on it, in one launch. Updates the slots' conv and state pools in place; advances ``buffers`` by one launch.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS = tuple( Cell( why=f"{requests} active TP16 decode rows and {8 - requests} zero-padded rows on three manager layers", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_decode.py index 016c4f3a17cf..764ec5c1cf60 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_decode.py @@ -15,7 +15,7 @@ class KdaDecode(OpWrapper): """Writes the decode output into ``output`` [B, 1, HV, 128] bf16; updates ``state`` at ``ssm_state_indices`` and, with ``update_conv_cache``, the conv pools there. Returns None.""" - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = tuple( Cell( why=( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_post_conv.py index 85cff8774aaf..a3fcb20cc147 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_post_conv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_post_conv.py @@ -15,7 +15,7 @@ class KdaPostConv(OpWrapper): L2-normalized along the head dim, v is copied as is. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_prefill.py index 6f573059be60..f89beeb352fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_prefill.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/ops/ssm/kda_prefill.py @@ -17,7 +17,7 @@ class KdaPrefill(OpWrapper): V-first `[slots, H, V, K]`). Q/K normalization and gate/beta activation are fused. Returns new `o`. """ - ARCHS = frozenset({Arch.SM_100}) + ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) CELLS: tuple[Cell, ...] = ( *( Cell( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index cb250f4d4948..440bbbec8aeb 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -162,11 +162,7 @@ l0_b200: - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py - unittest/_torch/modules/kimi_kda/test_kimi_kda_decode_state_parity.py - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py - # Kimi K3's generic path (prefill, mixed and wide steps) in the modeling_v2 catalog: its own entries - # and the sm_100 cells of the shared entries it calls. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py + # Shared catalog entries used by K3's generic path. - unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py::test_kimi_k3_config - unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py::test_kimi_k3_routed_hidden - unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py @@ -179,41 +175,8 @@ l0_b200: - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py - # Kimi K3 decode kernels (CuTe DSL): the GEMVs, the head GEMV and the - # embedding, and the cluster-wait and tcgen05-fence lints over them. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py - # modeling_v2 catalog entries with sm_100 receipts (the K3 ones skip on - # other architectures). - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py - unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py - # The Kimi K3 target's decode GEMVs, LM head and embedding on those entries. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py - # The Kimi K3 target's DSpark drafter on the drafter entries, against the stock block forward. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py - # The Kimi K3 targets' construction checks and the drift check of tp16_moetp16ep1's copied modules. - # l0_b300's modeling_v2 entry, the only other one that collects them, is waived (nvbugs/6853741). - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py - # tp16_moetp16ep1's MoE decode path: the layers and steps it takes (its kernels: l0_gb200_multi_gpus.yml). - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py - # The Kimi K3 targets' generic-path MoE routing, outside the TRTLLM-Gen kernel, against the kernel's own. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py # KDA runtime state-pool and mixed-batch core coverage. - unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py - unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py @@ -343,16 +306,8 @@ l0_b200: # step that runs them (SM 100; they skip on other GPUs). - unittest/_torch/kernels/test_spec_step_copies.py - unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k "test_spec_decode_graph_step_gather_kernel" - # Kimi K3 DSpark drafter attention (CuTe DSL, SM 100) and its modeling_v2 catalog entries, plus the sm_100 cells of - # fused_qk_norm_rope, which the drafter calls. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py + # Shared fused_qk_norm_rope cells used by the K3 drafter. - unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py - # Kimi K3 DSpark decode kernels (CuTe DSL, SM 100), one GPU: the acceptance against vocabulary-sharded target - # logits and the drafter's context K / V. Their multi-rank tests are in l0_gb200_multi_gpus.yml. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py # Vocab-parallel token selection kernels on a single shard. - unittest/_torch/speculative/test_vocab_parallel.py - unittest/_torch/moe/kernels/parallel @@ -422,17 +377,6 @@ l0_b200: - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_field_completeness # ------------- Prefix-aware scheduling E2E tests --------------- - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke - # ------------- Kimi K3 decode MoE kernels and their catalog entries (sm_100) --------------- - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py # ------------- Visual Gen tests --------------- - unittest/_torch/visual_gen/kernels/test_nvfp4_conv3d.py - unittest/_torch/visual_gen/test_media_decode.py @@ -527,23 +471,6 @@ l0_b200: - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[t2va] TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[fl2va] TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[ref2va] TIMEOUT (30) - # ------------- Kimi K3 KDA / MLA decode kernels (sm_100) and their modeling_v2 entries --------------- - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 4c421dd99214..0f6d0f16a38e 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -142,45 +142,6 @@ l0_b300: # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm --ignore=unittest/_torch/modeling_v2/models/kimi_k3_vl TIMEOUT (30) - # K3 single-GPU and host tests; collective launchers have dedicated multi-GPU entries. - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_step.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index ca9023a1faac..5ff0d96256f0 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -77,28 +77,7 @@ l0_gb200_multi_gpus: - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_preserves_cuda_graph_addresses - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_rejects_wrong_group_membership - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_growth_keeps_captured_graphs - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py - # Kimi K3 collective kernels (sm_100) and the catalog entries over caller-owned state, 4 ranks - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py - # Kimi K3 DSpark decode kernels' cross-rank exchanges (sm_100): the sharded acceptance and the Markov chain, 4 ranks - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py - unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py - # The Kimi K3 target's decode-path collectives over those entries, 4 ranks - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py - # tp16_moetp16ep1's MoE decode path: the engines pushing into the latent exchange, 4 ranks - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_failure_is_collective_and_bounded diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index ce21a01c9ddc..981dd10b3037 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -6,7 +6,7 @@ l0_gb300: - condition: ranges: system_gpu_count: - gte: 1 + gte: 4 lte: 4 wildcards: gpu: @@ -83,3 +83,78 @@ l0_gb300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-tp4] TIMEOUT (120) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-dp4] TIMEOUT (120) - accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_mtp[tp_size=4-ep_size=4] TIMEOUT (60) + +# K3 kernel and wrapper checks that require one GB300 GPU. +- condition: + ranges: + system_gpu_count: + gte: 1 + lte: 1 + wildcards: + gpu: + - '*gb110*' + - '*gb300*' + linux_distribution_name: ubuntu* + cpu: aarch64 + terms: + stage: pre_merge + backend: pytorch + tests: + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 97e375284bc5..5c9f71c687c5 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -1,5 +1,6 @@ version: 0.0.1 l0_gb300_multi_gpus: +# K3 collectives and decode-path exchanges, four ranks in one NVLink domain. - condition: ranges: system_gpu_count: @@ -12,12 +13,18 @@ l0_gb300_multi_gpus: linux_distribution_name: ubuntu* cpu: aarch64 terms: - stage: post_merge + stage: pre_merge backend: pytorch tests: - # One ordered sequence in one job: the cases share module-global rank state. - - unittest/_torch/modeling_v2/comm + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py @@ -26,6 +33,24 @@ l0_gb300_multi_gpus: - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py + +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + wildcards: + gpu: + - '*gb110*' + - '*gb300*' + linux_distribution_name: ubuntu* + cpu: aarch64 + terms: + stage: post_merge + backend: pytorch + tests: + # One ordered sequence in one job: the cases share module-global rank state. + - unittest/_torch/modeling_v2/comm - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_identity_vs_mtp3[dep4] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_mtp3_acceptance[modeling_v2-dep4] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_mtp3_acceptance[stock-dep4] diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_engines.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_engines.py index 6fa7812394ea..eab2c09f8799 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_engines.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_engines.py @@ -26,7 +26,7 @@ def is_sm100() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) def bits(t: torch.Tensor) -> torch.Tensor: diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py index 803aaed333e6..d498c19ac56d 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_k3_moe_front_op_matrix.py @@ -938,7 +938,9 @@ def _run_one_rank(args) -> int: global R, OPS, WS_A, WS_B, PLAIN, FLAGS, LAYER_WEIGHTS, BIAS, EXPERT_BUFFERS, OFFSET, WL, WE global EXPERT_RAW R = ls.Rank(args) - assert torch.cuda.get_device_capability() == (10, 0), "k3_moe_front and k3_moe need sm_100" + assert torch.cuda.get_device_capability() in ((10, 0), (10, 3)), ( + "k3_moe_front and k3_moe need sm_100" + ) import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock path's ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.moe import ( k3_moe as moe, diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_cells.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_cells.py index aed027319f2b..534407f33bae 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_cells.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/_kda_cells.py @@ -32,8 +32,8 @@ def sm100() -> bool: - """sm_100 exactly: the architecture these entries are certified on (a missing receipt reads as unknown).""" - return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + """The SM100 and SM103 architectures covered by these entries.""" + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) def load_ops(): diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py index 5ee7df26d5c6..9d204c68b345 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py @@ -34,10 +34,10 @@ def _is_sm100() -> bool: if not torch.cuda.is_available(): return False - return torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_fused_moe needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_fused_moe needs sm_100 or sm_103") H, TOP_K, NUM_EXPERTS, SV = 3584, 16, 896, 32 I_TP, E_LOCAL, MOE_TP, TP_RANK, EP_RANK = 768, 224, 4, 1, 1 # one rank of experts TP4 x EP4 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py index e7df09103004..9bc374dbc906 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py @@ -95,7 +95,7 @@ def namespace(): def _supported() -> bool: if MPI is None or not torch.cuda.is_available() or torch.cuda.device_count() < WORLD: return False - return torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.get_device_capability() in ((10, 0), (10, 3)) pytestmark = [ diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py index 36a8ffd8f777..32ee49cc5584 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py @@ -40,10 +40,10 @@ def _is_sm100() -> bool: if not torch.cuda.is_available(): return False - return torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe_m1 needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe_m1 needs sm_100 or sm_103") H, NUM_EXPERTS, TOP_K, SV = 3584, 896, 16, 32 GATE_CAP, LINEAR_CAP = ( diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py index 9bac14aeb286..ac3d853d5000 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py @@ -41,10 +41,10 @@ def _is_sm100() -> bool: if not torch.cuda.is_available(): return False - return torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe_m2 needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe_m2 needs sm_100 or sm_103") H, NUM_EXPERTS, TOP_K, SV = 3584, 896, 16, 32 GATE_CAP, LINEAR_CAP = ( diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py index 280d906bc68a..31da0607172d 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py @@ -97,7 +97,7 @@ def namespace(): def _supported() -> bool: if MPI is None or not torch.cuda.is_available() or torch.cuda.device_count() < WORLD: return False - return torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.get_device_capability() in ((10, 0), (10, 3)) pytestmark = [ diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py index f2e7998ff3e7..b747b41a7d54 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py @@ -41,10 +41,10 @@ def _is_sm100() -> bool: if not torch.cuda.is_available(): return False - return torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe needs sm_100 or sm_103") H, TOP_K, NUM_EXPERTS, SV = 3584, 16, 896, 32 I_TP, E_LOCAL, MOE_TP, TP_RANK, EP_RANK = 768, 224, 4, 1, 1 # one rank of experts TP4 x EP4 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py index 2f9082d93928..ad8808ce424b 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py @@ -27,7 +27,8 @@ assert torch.cuda.is_available(), "attn_res_fwd requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) HIDDEN_SIZE = 7168 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py index 9e371d4c7a60..9516c9967e71 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py @@ -28,7 +28,8 @@ assert torch.cuda.is_available(), "attn_res_rmsnorm_fwd requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) HIDDEN_SIZE = 7168 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py index c36118083ab4..49d2f8bb045b 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py @@ -17,7 +17,8 @@ assert torch.cuda.is_available(), "causal_conv1d_fwd requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) WIDTH = 4 # Kimi K3's short_conv_kernel_size diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py index 4404276a2a22..e9aa56891a19 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py @@ -22,7 +22,8 @@ assert torch.cuda.is_available(), "chunk_kda requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) HEAD_DIM = kr.HEAD_DIM diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py index 82ea0637863d..cb5a5360e699 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py @@ -17,7 +17,8 @@ assert torch.cuda.is_available(), "k3_ctm_gemv requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py index 20d74d4c1309..de593e97b90d 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py @@ -17,7 +17,8 @@ assert torch.cuda.is_available(), "k3_ctm_gemv_long requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py index 32569113e62d..09a8bc67a1fd 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py @@ -17,7 +17,8 @@ assert torch.cuda.is_available(), "k3_ctm_gemv_swiglu requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py index 78f5e5404dda..af3786e29f38 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py @@ -20,7 +20,8 @@ assert torch.cuda.is_available(), "k3_ctm_gemv_wide requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) TOL = 8e-3 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py index 537c3278dfb9..9f2b52a5e505 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py @@ -18,7 +18,8 @@ assert torch.cuda.is_available(), "k3_decode_gemv requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) TOL = 8e-3 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py index d98232ab5731..c6c9f6794e3b 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py @@ -10,9 +10,9 @@ assert torch.cuda.is_available(), "k3_drafter_attn requires a CUDA device" -if torch.cuda.get_device_capability() != (10, 0): - # The CuTe DSL kernel uses tcgen05, TMA and clusters; it is certified on sm_100 (B200 / GB200) only. - pytest.skip("k3_drafter_attn is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + # The CuTe DSL kernel uses tcgen05, TMA and clusters; it is certified on SM100 and SM103. + pytest.skip("k3_drafter_attn is certified on sm_100 and sm_103", allow_module_level=True) from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.attention.k3_drafter_attn import ( # noqa: E402 k3_drafter_attn, diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py index 02ff7d504612..4e1cb5da5ac8 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py @@ -9,9 +9,9 @@ assert torch.cuda.is_available(), "k3_drafter_attn requires a CUDA device" -if torch.cuda.get_device_capability() != (10, 0): - # The CuTe DSL kernel uses tcgen05, TMA and clusters; it is certified on sm_100 (B200 / GB200) only. - pytest.skip("k3_drafter_attn is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + # The CuTe DSL kernel uses tcgen05, TMA and clusters; it is certified on SM100 and SM103. + pytest.skip("k3_drafter_attn is certified on sm_100 and sm_103", allow_module_level=True) from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.ops.attention.k3_drafter_attn_qknorm import ( # noqa: E402 k3_drafter_attn_qknorm, diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py index b0034836d60d..d5456114093c 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py @@ -26,7 +26,8 @@ assert torch.cuda.is_available(), "k3_embed_norm requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) VOCAB, HIDDEN = 163840, 7168 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py index e02c60f6ccbb..33939e0bec9d 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py @@ -19,7 +19,8 @@ assert torch.cuda.is_available(), "k3_head_gemv requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) TOL = 8e-3 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py index 924f4e96982a..c5e893cc78fb 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py @@ -15,8 +15,8 @@ assert torch.cuda.is_available(), "k3_latent_reduce requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("k3_latent_reduce is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("k3_latent_reduce is certified on sm_100 and sm_103", allow_module_level=True) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py index c97a4188f849..e26501d59ecb 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py @@ -62,11 +62,11 @@ def _is_sm100() -> bool: if not torch.cuda.is_available(): return False major, minor = torch.cuda.get_device_capability() - # sm_100 exactly: the receipts' architecture (a missing receipt reads as unknown). - return (major, minor) == (10, 0) + # SM100 and SM103 are covered by the kernel certification tests. + return (major, minor) in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="certified on SM100 only") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="certified on SM100 and SM103") H, NOPE, LAT, PE, QL, V = 6, 128, 512, 64, 1536, 128 DQK, PAGE = LAT + PE, 64 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py index 99018d2c45c7..8454aaf89496 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py @@ -53,11 +53,11 @@ def _is_sm100() -> bool: if not torch.cuda.is_available(): return False major, minor = torch.cuda.get_device_capability() - # sm_100 exactly: the receipts' architecture (a missing receipt reads as unknown). - return (major, minor) == (10, 0) + # SM100 and SM103 are covered by the kernel certification tests. + return (major, minor) in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="certified on SM100 only") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="certified on SM100 and SM103") H, NOPE, PE, QK, LAT, QL, V = 6, 128, 64, 192, 512, 1536, 128 DQK, PAGE = LAT + PE, 64 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py index 870d3e27ec62..ba5bc4e46663 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py @@ -62,10 +62,10 @@ def _is_sm100() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe needs sm_100 or sm_103") @pytest.fixture(autouse=True) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py index f32afe2e80c4..b0d58b06990d 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py @@ -18,10 +18,10 @@ def _is_sm100() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe_front needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_moe_front needs sm_100 or sm_103") # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py index 7c88da4eb934..69edac097771 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py @@ -18,7 +18,7 @@ "kimi_k3_mxfp4__sm_100__tp16_moetp16ep1.ops.moe.k3_moe_m1" ).k3_moe_m1 -pytestmark = pytest.mark.skipif(not engines.is_sm100(), reason="k3_moe_m1 needs sm_100") +pytestmark = pytest.mark.skipif(not engines.is_sm100(), reason="k3_moe_m1 needs sm_100 or sm_103") ENGINES = [engines.Engine("k3_moe_m1", cell.spec["m"]) for cell in k3_moe_m1.CELLS] IDS = [f"m{cell.spec['m']}" for cell in k3_moe_m1.CELLS] diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py index 4b6afdc98838..e76c9a5ab2f2 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py @@ -18,7 +18,7 @@ "kimi_k3_mxfp4__sm_100__tp16_moetp16ep1.ops.moe.k3_moe_m2" ).k3_moe_m2 -pytestmark = pytest.mark.skipif(not engines.is_sm100(), reason="k3_moe_m2 needs sm_100") +pytestmark = pytest.mark.skipif(not engines.is_sm100(), reason="k3_moe_m2 needs sm_100 or sm_103") ENGINE = engines.Engine("k3_moe_m2", k3_moe_m2.CELLS[0].spec["m"]) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py index e0a750bdd4fc..abec03490ed7 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py @@ -36,10 +36,10 @@ def _is_sm100() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) -pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_route_quant needs sm_100") +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="k3_route_quant needs sm_100 or sm_103") @pytest.fixture(autouse=True) diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py index 3224e6b470f6..299e378da39d 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py @@ -15,8 +15,8 @@ assert torch.cuda.is_available(), "k3_sandwich_oproj requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("k3_sandwich_oproj is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("k3_sandwich_oproj is certified on sm_100 and sm_103", allow_module_level=True) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py index 1c7d379532fc..4f3d1943625e 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py @@ -15,8 +15,8 @@ assert torch.cuda.is_available(), "k3_sandwich_plain requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("k3_sandwich_plain is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("k3_sandwich_plain is certified on sm_100 and sm_103", allow_module_level=True) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py index c1411bd39d6b..6bf7a7051e6b 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py @@ -15,8 +15,8 @@ assert torch.cuda.is_available(), "k3_sandwich_tail requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("k3_sandwich_tail is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("k3_sandwich_tail is certified on sm_100 and sm_103", allow_module_level=True) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py index fb85cd558e0b..d8a9048364a5 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py @@ -23,7 +23,8 @@ assert torch.cuda.is_available(), "k3_situ_mul requires a CUDA device" # The receipts are sm_100 ones; CI's other architectures skip this file. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) TOL = 8e-3 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py index ec9e1d6012ce..ab3e3749b57c 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py @@ -23,7 +23,7 @@ __extra_import_path__ = ["../.."] from _validating import validating -pytestmark = pytest.mark.skipif(not kc.sm100(), reason="certified on SM100 only") +pytestmark = pytest.mark.skipif(not kc.sm100(), reason="certified on SM100 and SM103") LAYERS = 3 STEPS = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py index ba4c82a5b600..5beefec93e3b 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py @@ -16,7 +16,8 @@ assert torch.cuda.is_available(), "kda_post_conv requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) HEAD_DIM = 128 # Kimi K3's KDA head_dim diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py index 6001c86d0fa1..5fd7c6b196f9 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py @@ -19,7 +19,8 @@ assert torch.cuda.is_available(), "kda_prefill requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) HEAD_DIM, LOWER_BOUND, SCALE, CHUNK = kr.HEAD_DIM, kr.LOWER_BOUND, kr.SCALE, kr.CHUNK diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py index b52427943efc..d0e6d8150414 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py @@ -16,9 +16,9 @@ assert torch.cuda.is_available(), "the Kimi K3 target's decode collectives require CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): # The target is certified on sm_100 (GB200) only, as are the catalog entries it calls. - pytest.skip("the Kimi K3 target runs on sm_100 only", allow_module_level=True) + pytest.skip("K3 decode-path kernels require SM100 or SM103", allow_module_level=True) if torch.cuda.device_count() < _rank_job.WORLD_SIZE: # One rank per device: fewer visible devices cannot host the check's world size. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py index 8f2f83983d71..0b6b21cae090 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py @@ -50,8 +50,8 @@ from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), - reason="the K3 decode kernels run on sm_100 only", + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="the K3 decode kernels run on sm_100 and sm_103", ) TOL = 8e-3 # max |y - ref| / max |ref|, ref a float64 product diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py index dbaf003daf05..e05ece1e9b23 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py @@ -49,8 +49,8 @@ from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), - reason="the K3 drafter entries run on sm_100 only", + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="the K3 drafter entries run on sm_100 and sm_103", ) HIDDEN, INTER, HEADS, KV_HEADS, HEAD_DIM, LAYERS, PAGE = 7168, 896, 6, 1, 64, 2, 64 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py index aef529d49750..7ac19a8a8c6e 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py @@ -66,7 +66,7 @@ def test_gate_routes_outside_the_kernel(target): def _sm_100(): - return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) def _gate(target): @@ -172,7 +172,7 @@ def _clear_rows(logits, bias, routed_scaling_factor): @pytest.mark.skipif( - not _sm_100(), reason="the TRTLLM-Gen MXFP4 cubins and the fused route + quant run on sm_100" + not _sm_100(), reason="the TRTLLM-Gen MXFP4 cubins and fused route + quant need SM100 or SM103" ) @TARGETS def test_routing_outside_the_kernel_matches_inside(target, monkeypatch): diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py index d91f34f0b578..92e19461bc2f 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py @@ -20,7 +20,8 @@ assert torch.cuda.is_available(), "kimi_k3_noaux_tc_mxfp8_quant requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) NUM_EXPERTS = 896 diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py index d53299b07c8d..cd4c39d7fc4b 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py @@ -16,9 +16,9 @@ assert torch.cuda.is_available(), "the Kimi K3 target's MoE decode path requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): # The target is certified on sm_100 (GB200) only, as are the catalog entries it calls. - pytest.skip("the Kimi K3 target runs on sm_100 only", allow_module_level=True) + pytest.skip("K3 decode-path kernels require SM100 or SM103", allow_module_level=True) if torch.cuda.device_count() < _rank_job.WORLD_SIZE: # One rank per device: fewer visible devices cannot host the check's world size. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py index 4c9aa757517c..163349e0c960 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py @@ -15,8 +15,8 @@ assert torch.cuda.is_available(), "mnnvl_allgather_split requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("mnnvl_allgather_split is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("mnnvl_allgather_split is certified on sm_100 and sm_103", allow_module_level=True) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py index 1b2d1676611f..d96716bb3068 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py @@ -15,8 +15,10 @@ assert torch.cuda.is_available(), "mnnvl_allreduce_attn_res requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("mnnvl_allreduce_attn_res is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip( + "mnnvl_allreduce_attn_res is certified on sm_100 and sm_103", allow_module_level=True + ) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py index f7444af490ba..987968de2638 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py @@ -15,8 +15,8 @@ assert torch.cuda.is_available(), "mnnvl_fusion_allreduce requires CUDA devices" -if torch.cuda.get_device_capability() != (10, 0): - pytest.skip("mnnvl_fusion_allreduce is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("mnnvl_fusion_allreduce is certified on sm_100 and sm_103", allow_module_level=True) # Each case starts its own W-rank mpirun over the visible devices; under xdist several workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py index c91d8a1a6559..468f7fb60f6f 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py @@ -16,7 +16,8 @@ assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) # Kimi K3's KDA: head_dim 128, rms_norm_eps 1e-5, 96 heads (6 / 12 / 24 per rank at TP16 / TP8 / TP4). diff --git a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py index 62123daf1e24..23376a7807ca 100644 --- a/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py @@ -17,7 +17,8 @@ assert torch.cuda.is_available(), "situ_and_mul requires a CUDA device" pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="certified on sm_100 and sm_103", ) # Kimi K3's activation_situ_beta and activation_situ_linear_beta. diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index b6eaa3f1ee81..0880039618ba 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -17,10 +17,12 @@ from __future__ import annotations +from pathlib import Path + import pytest __extra_import_path__ = ["~/jenkins/scripts/cbts"] -from cbts.blocks import Block, Stage, YAMLIndex +from cbts.blocks import Block, Stage, YAMLIndex, block_matches_stage, parse_stages_from_groovy from cbts.rules.base import PRInputs from cbts.rules.modeling_v2 import ModelingV2Rule, _is_mv2_claim, _kept_entries @@ -200,3 +202,62 @@ def test_k3_family_paths_select_kernel_collective_and_host_tests(rule, relative) assert result.scope == "modelingv2only" other = _selected(_apply(rule, f"{_PY}/models/gpt_oss/csrc/kernel.cu")) assert not (selected - {_CLAIMS}) & other + + +def test_k3_ci_runs_once_on_the_required_gb300_topology(): + repo = Path(__file__).resolve().parents[4] + prefix = "unittest/_torch/modeling_v2/models/kimi_k3_vl/" + expected = {p.name for p in (repo / "tests" / prefix).glob("test_*.py")} + collective_tests = { + "test_k3_latent_reduce.py", + "test_k3_markov.py", + "test_k3_mnnvl_comm.py", + "test_k3_moe_front.py", + "test_k3_moe_push.py", + "test_k3_sandwich.py", + "test_k3_spec_accept_sharded.py", + "test_modeling_v2_k3_latent_reduce_op_matrix.py", + "test_modeling_v2_k3_moe_front_op_matrix.py", + "test_modeling_v2_k3_sandwich_oproj_op_matrix.py", + "test_modeling_v2_k3_sandwich_plain_op_matrix.py", + "test_modeling_v2_k3_sandwich_tail_op_matrix.py", + "test_modeling_v2_kimi_k3_decode_comm_op_matrix.py", + "test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py", + "test_modeling_v2_mnnvl_allgather_split_op_matrix.py", + "test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py", + "test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py", + } + index = YAMLIndex.load(repo / "tests/integration/test_lists/test-db") + stages = parse_stages_from_groovy(repo / "jenkins/L0_Test.groovy") + seen = set() + for block in index.blocks: + for entry in block.tests: + target = entry.split()[0].split("::")[0] + if not target.startswith(prefix): + # Broad directory entries must exclude this family outside its + # dedicated stages, or the GPU coverage would be duplicated. + if prefix.startswith(target.rstrip("/") + "/"): + assert any( + prefix.rstrip("/") == arg.removeprefix("--ignore=").rstrip("/") + for arg in entry.split() + if arg.startswith("--ignore=") + ), (block.yaml_stem, entry) + continue + name = target.removeprefix(prefix) + assert name not in seen, (block.yaml_stem, name) + seen.add(name) + if block.yaml_stem == "l0_cpu": + assert name not in collective_tests + continue + multi = name in collective_tests + assert block.yaml_stem == ("l0_gb300_multi_gpus" if multi else "l0_gb300") + count = 4 if multi else 1 + assert block.condition["ranges"]["system_gpu_count"] == {"gte": count, "lte": count} + assert block.condition["terms"]["stage"] == "pre_merge" + matched = { + stage.name + for stage in stages.values() + if stage.yaml_stem == block.yaml_stem and block_matches_stage(block, stage) + } + assert matched == {"GB300-4_GPUs-PyTorch-2" if multi else "GB300-PyTorch-1"} + assert seen == expected From 32dca99aa7bfcf2602c5c7b2eb652929053243f7 Mon Sep 17 00:00:00 2001 From: Tao Li Date: Thu, 29 Oct 2026 07:47:35 -0700 Subject: [PATCH 11/11] Use existing B300 and GB300 stages for K3 tests Signed-off-by: Tao Li --- jenkins/L0_Test.groovy | 2 - .../test_lists/test-db/l0_b300.yml | 59 +++++++++++++ .../test_lists/test-db/l0_gb300.yml | 83 +++++-------------- .../test-db/l0_gb300_multi_gpus.yml | 34 -------- .../scripts/cbts/test_modeling_v2_rule.py | 8 +- 5 files changed, 85 insertions(+), 101 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index d9075e450ab4..0fded65aa33e 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -6675,9 +6675,7 @@ def launchTestJobs(pipeline, testFilter, globalVars) "GB200-4_GPUs-PyTorch-Post-Merge-1": ["auto:gb200-x4", "l0_gb200_multi_gpus", 1, 2, 4, 1, false, true], "GB200-4_GPUs-PyTorch-Post-Merge-2": ["auto:gb200-x4", "l0_gb200_multi_gpus", 2, 2, 4, 1, false, true], "GB10-PyTorch-Post-Merge-1": ["gb10x-single", "l0_gb10", 1, 1], - "GB300-PyTorch-1": ["auto:gb300-flex", "l0_gb300", 1, 1, 1, 1, true, false], "GB300-4_GPUs-PyTorch-1": ["auto:gb300-x4", "l0_gb300", 1, 1, 4, 1, true, false], - "GB300-4_GPUs-PyTorch-2": ["auto:gb300-x4", "l0_gb300_multi_gpus", 1, 1, 4, 1, true, false], "GB300-4_GPUs-PyTorch-Post-Merge-1": ["auto:gb300-x4", "l0_gb300_multi_gpus", 1, 3, 4, 1, true, false], "GB300-4_GPUs-PyTorch-Post-Merge-2": ["auto:gb300-x4", "l0_gb300_multi_gpus", 2, 3, 4, 1, true, false], "GB300-4_GPUs-PyTorch-Post-Merge-3": ["auto:gb300-x4", "l0_gb300_multi_gpus", 3, 3, 4, 1, true, false], diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 0f6d0f16a38e..8f1ccee6744e 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -142,6 +142,65 @@ l0_b300: # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm --ignore=unittest/_torch/modeling_v2/models/kimi_k3_vl TIMEOUT (30) + # K3 single-GPU kernel and wrapper checks. + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index 981dd10b3037..73def12b07d1 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -6,7 +6,7 @@ l0_gb300: - condition: ranges: system_gpu_count: - gte: 4 + gte: 1 lte: 4 wildcards: gpu: @@ -84,12 +84,12 @@ l0_gb300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-dp4] TIMEOUT (120) - accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_mtp[tp_size=4-ep_size=4] TIMEOUT (60) -# K3 kernel and wrapper checks that require one GB300 GPU. +# K3 collectives and decode-path exchanges, four ranks in one NVLink domain. - condition: ranges: system_gpu_count: - gte: 1 - lte: 1 + gte: 4 + lte: 4 wildcards: gpu: - '*gb110*' @@ -100,61 +100,20 @@ l0_gb300: stage: pre_merge backend: pytorch tests: - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctx_kv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_drafter_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_embed.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_fused_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_head_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_decode_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_pools_past_2g.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_kda_verify.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mla_q.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front_geometry.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m1.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_m2.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_route_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_tcgen05_fences.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_attn_res_rmsnorm_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_causal_conv1d_fwd.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_chunk_kda.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_long.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_swiglu.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_ctm_gemv_wide.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_drafter_attn_qknorm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_embed_norm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_head_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_decode_attn.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_kda_verify.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_attn_vb_out.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_mla_qkv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m1.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_m2.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_native.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_route_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_situ_mul.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_decode.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_post_conv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kda_prefill.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_construction.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_gemv.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drafter.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_drift.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_moe_routing.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_noaux_tc_mxfp8_quant.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_moe.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py + - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 5c9f71c687c5..58d14ea21ba6 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -1,39 +1,5 @@ version: 0.0.1 l0_gb300_multi_gpus: -# K3 collectives and decode-path exchanges, four ranks in one NVLink domain. -- condition: - ranges: - system_gpu_count: - gte: 4 - lte: 4 - wildcards: - gpu: - - '*gb110*' - - '*gb300*' - linux_distribution_name: ubuntu* - cpu: aarch64 - terms: - stage: pre_merge - backend: pytorch - tests: - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_latent_reduce.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_markov.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_mnnvl_comm.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_front.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_moe_push.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_sandwich.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_spec_accept_sharded.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_latent_reduce_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_moe_front_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_oproj_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_plain_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_k3_sandwich_tail_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_decode_comm_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_route_b_decode_moe_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allgather_split_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_allreduce_attn_res_op_matrix.py - - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_mnnvl_fusion_allreduce_op_matrix.py - - condition: ranges: system_gpu_count: diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py index 0880039618ba..370d102a25b0 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -204,7 +204,7 @@ def test_k3_family_paths_select_kernel_collective_and_host_tests(rule, relative) assert not (selected - {_CLAIMS}) & other -def test_k3_ci_runs_once_on_the_required_gb300_topology(): +def test_k3_ci_runs_once_on_b300_or_gb300_by_topology(): repo = Path(__file__).resolve().parents[4] prefix = "unittest/_torch/modeling_v2/models/kimi_k3_vl/" expected = {p.name for p in (repo / "tests" / prefix).glob("test_*.py")} @@ -250,7 +250,7 @@ def test_k3_ci_runs_once_on_the_required_gb300_topology(): assert name not in collective_tests continue multi = name in collective_tests - assert block.yaml_stem == ("l0_gb300_multi_gpus" if multi else "l0_gb300") + assert block.yaml_stem == ("l0_gb300" if multi else "l0_b300") count = 4 if multi else 1 assert block.condition["ranges"]["system_gpu_count"] == {"gte": count, "lte": count} assert block.condition["terms"]["stage"] == "pre_merge" @@ -259,5 +259,7 @@ def test_k3_ci_runs_once_on_the_required_gb300_topology(): for stage in stages.values() if stage.yaml_stem == block.yaml_stem and block_matches_stage(block, stage) } - assert matched == {"GB300-4_GPUs-PyTorch-2" if multi else "GB300-PyTorch-1"} + assert matched == ( + {"GB300-4_GPUs-PyTorch-1"} if multi else {"B300-PyTorch-1", "B300-PyTorch-2"} + ) assert seen == expected