diff --git a/.cmake-format.json b/.cmake-format.json index e7eb73ae25fb..21bf59072b21 100644 --- a/.cmake-format.json +++ b/.cmake-format.json @@ -25,6 +25,27 @@ "NAMESPACE": "+", "ARCHS": "+" } + }, + "add_modeling_v2_family": { + "pargs": 1 + }, + "add_modeling_v2_shared": { + "pargs": 1, + "kwargs": { + "ARCHS": "+", + "SOURCES": "+", + "DEPS": "+", + "REQUIRES": "+" + } + }, + "add_modeling_v2_target": { + "pargs": 2, + "kwargs": { + "ARCHS": "+", + "SOURCES": "+", + "DEPS": "+", + "REQUIRES": "+" + } } }, "_help_vartags": [ diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2d06dffa3268..3ebf14e3b741 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -387,6 +387,11 @@ # Individual handles rather than a team, like SCAFFOLDING: this is one bounded # experiment with named owners, not a standing domain. /tensorrt_llm/_torch/_experimental/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju +# What every modeling_v2 family library shares (the discovery and the common/ +# macros): the modeling_v2 owners, and the kernels owners because it is CUDA and +# links against the stock libraries. +/cpp/tensorrt_llm/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @NVIDIA/trt-llm-kernels-devs +/cpp/cmake/modules/modeling_v2.cmake @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @NVIDIA/trt-llm-kernels-devs /tests/unittest/_torch/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju /tests/integration/defs/accuracy/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju # The two catalog categories whose contracts state kernel behaviour the @@ -399,6 +404,10 @@ /tests/unittest/_torch/modeling_v2/models/glm_moe_dsa @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @dongfengy @Tabrizian @xwang233 /cpp/tensorrt_llm/kernels/glmMoeDsa @NVIDIA/trt-llm-kernels-devs @dongfengy @Tabrizian @xwang233 /cpp/tensorrt_llm/thop/glmMoeDsa/ @NVIDIA/trt-llm-kernels-devs @dongfengy @Tabrizian @xwang233 +# A family's native code lives in its csrc/ directories: the kernels owners +# co-own it. Last, so it applies under the family rules above; a family with its +# own owners adds them here. +/tensorrt_llm/_torch/_experimental/modeling_v2/models/**/csrc/ @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @NVIDIA/trt-llm-kernels-devs ## TensorRT-LLM LLM Disaggregated /examples/disaggregated @NVIDIA/trt-llm-disagg-devs @NVIDIA/trt-llm-doc-owners diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 31c58f03a236..9c68ecd60f6c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -49,6 +49,12 @@ option(BUILD_DEEP_EP "Build the Deep EP module" ON) option(BUILD_NCCL_EXTENSIONS "Build the nccl-extensions Python wheel" ON) option(BUILD_DEEP_GEMM "Build the DeepGEMM module" ON) option(BUILD_FLASH_MLA "Build the FlashMLA module" ON) +set(TRTLLM_MODELING_V2_TARGETS + "all" + CACHE + STRING + "modeling_v2 native targets to build: all, none, or a list of / globs" +) option(BUILD_MICRO_BENCHMARKS "Build C++ micro benchmarks" OFF) option(NVTX_DISABLE "Disable all NVTX features" ON) option(WARNING_IS_ERROR "Treat all warnings as errors" OFF) diff --git a/cpp/cmake/modules/modeling_v2.cmake b/cpp/cmake/modules/modeling_v2.cmake new file mode 100644 index 000000000000..8cb64739272a --- /dev/null +++ b/cpp/cmake/modules/modeling_v2.cmake @@ -0,0 +1,685 @@ +# +# 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. +# + +#[[ +Native (C++/CUDA) code of modeling_v2 targets. + +A model family's native code lives with its Python code, under +tensorrt_llm/_torch/_experimental/modeling_v2/models//: + + csrc/CMakeLists.txt declares the family, its shared code and its targets + csrc/*.cu|*.h|*.cpp code several targets of the family share + /csrc/... one target's own kernels (.cu) and torch bindings (.cpp) + +Each family builds into one torch-op library, libth_modeling_v2_.so, +installed next to libth_common.so. Nothing in it is linked into +libtensorrt_llm.so or libth_common.so; the Python loader +(tensorrt_llm/_torch/_experimental/modeling_v2/_native.py) opens it only when +routing picks one of the family's targets. The "Native code" section of the +modeling_v2 README is the guide for authors; cpp/tensorrt_llm/modeling_v2/ +holds only what every family shares (this discovery, the common/ macros). + +A family's csrc/CMakeLists.txt uses three functions; names are the Python +package names: + + add_modeling_v2_family() + add_modeling_v2_shared( ARCHS ... SOURCES ... [DEPS ...] [REQUIRES ...]) + add_modeling_v2_target( ARCHS ... [SOURCES ...] [DEPS ...] [REQUIRES ...]) + +ARCHS NN, NNa or NNf. Each is intersected with the build's requested + architectures, the way flash_mla and deep_ep choose theirs; a target + with an empty intersection is not built and its family library + does not carry it. +SOURCES Paths or globs relative to the unit's csrc/ directory: the + family's for add_modeling_v2_shared, /csrc/ for a target. + .cu files are compiled for ARCHS without torch; .cpp files are the + torch bindings. Includes of the family's own code are written from + the models/ directory: "/csrc/x.h", + "//csrc/x.h". A target without SOURCES has no csrc/ + and no code of its own: it runs its family's shared code, and + declaring it is what makes the library carry it. + Source properties the family's CMakeLists.txt sets with + set_source_files_properties (e.g. COMPILE_OPTIONS) apply to shared + and target sources alike. +DEPS From a closed list: cublas, cublaslt, cuda_driver, internal_cutlass, + nccl, trtllm_gen_fmha_interface. +REQUIRES Build features: MULTI_DEVICE, NVSHMEM. Unmet requirements skip the + target. + +Like deep_ep and flash_mla, each family writes cuda_architectures.txt into its +build directory: the architectures its library was built for, empty when it +has none. scripts/build_wheel.py installs the library only when that list is +not empty. +#]] + +include_guard(GLOBAL) + +set(_MODELING_V2_KNOWN_DEPS cublas cublaslt cuda_driver internal_cutlass nccl + trtllm_gen_fmha_interface) +set(_MODELING_V2_KNOWN_REQUIRES MULTI_DEVICE NVSHMEM) + +# --- Naming ------------------------------------------------------------------ + +# A name usable as a C++ or torch identifier: "__" collapsed to "_", since C++ +# reserves identifiers that contain a double underscore. +function(_modeling_v2_ident out name) + set(result "${name}") + while(result MATCHES "__") + string(REPLACE "__" "_" result "${result}") + endwhile() + set(${out} + "${result}" + PARENT_SCOPE) +endfunction() + +function(_modeling_v2_check_family_name family) + if(NOT family MATCHES "^[a-z][a-z0-9_]*[a-z0-9]$" OR family MATCHES "__") + message( + FATAL_ERROR + "modeling_v2: family '${family}' is not a snake_case package name") + endif() +endfunction() + +# ____, the identity the Python target directory +# carries (see tensorrt_llm/_torch/_experimental/modeling_v2/README.md). +function(_modeling_v2_check_target_name target) + string(REPLACE "__" ";" segments "${target}") + list(LENGTH segments count) + if(NOT count EQUAL 3 OR NOT target MATCHES "^[a-z][a-z0-9_]*$") + message( + FATAL_ERROR + "modeling_v2: target '${target}' is not a ____ name" + ) + endif() + list(GET segments 1 sm) + if(NOT sm MATCHES "^sm_[0-9]+[a-z0-9]*$") + message( + FATAL_ERROR + "modeling_v2: target '${target}' has no sm_ segment (got '${sm}')" + ) + endif() +endfunction() + +# --- Selection, requirements, architectures ----------------------------------- + +# Whether TRTLLM_MODELING_V2_TARGETS selects /. +function(_modeling_v2_selected out family target) + set(filter "${TRTLLM_MODELING_V2_TARGETS}") + if(filter STREQUAL "" OR filter STREQUAL "all") + set(${out} + TRUE + PARENT_SCOPE) + return() + endif() + if(filter STREQUAL "none") + set(${out} + FALSE + PARENT_SCOPE) + return() + endif() + set(selected FALSE) + foreach(pattern IN LISTS filter) + # Glob to regex: escape the dots, then * -> [^/]*, ? -> [^/]. + string(REPLACE "." "\\." regex "${pattern}") + string(REPLACE "*" "[^/]*" regex "${regex}") + string(REPLACE "?" "[^/]" regex "${regex}") + if("${family}/${target}" MATCHES "^${regex}$") + set(selected TRUE) + endif() + endforeach() + set(${out} + ${selected} + PARENT_SCOPE) +endfunction() + +# Empty when every requirement holds, otherwise the reason. +function(_modeling_v2_unmet_requirements out) + set(reasons "") + foreach(requirement IN LISTS ARGN) + if(NOT requirement IN_LIST _MODELING_V2_KNOWN_REQUIRES) + message( + FATAL_ERROR + "modeling_v2: unknown REQUIRES '${requirement}' (known: ${_MODELING_V2_KNOWN_REQUIRES})" + ) + endif() + if(requirement STREQUAL "MULTI_DEVICE" AND NOT ENABLE_MULTI_DEVICE) + list(APPEND reasons "built with ENABLE_MULTI_DEVICE=OFF") + elseif(requirement STREQUAL "NVSHMEM" AND NOT ENABLE_NVSHMEM) + list(APPEND reasons "built with ENABLE_NVSHMEM=OFF") + endif() + endforeach() + list(JOIN reasons ", " reasons) + set(${out} + "${reasons}" + PARENT_SCOPE) +endfunction() + +#[[ +Intersect declared architectures with the requested ones. + + NN -> NN-real if NN is requested + NNa -> NNa-real if NN is requested + NNf -> NNf-real if a requested architecture is in family NN: the same + major version and a minor version >= NN's + +`requested` defaults to CMAKE_CUDA_ARCHITECTURES_ORIG, the plain numbers the +build was asked for -- the list flash_mla intersects with too. This does not +reuse set_cuda_architectures(), which drops NNa silently and falls back to +placeholder kernels. +#]] +function(_modeling_v2_resolve_archs out) + cmake_parse_arguments(PARSE_ARGV 1 arg "" "" "DECLARED;REQUESTED") + if(NOT DEFINED arg_REQUESTED) + set(arg_REQUESTED ${CMAKE_CUDA_ARCHITECTURES_ORIG}) + endif() + set(resolved "") + foreach(arch IN LISTS arg_DECLARED) + if(NOT arch MATCHES "^([0-9]+)([af]?)$") + message( + FATAL_ERROR "modeling_v2: ARCHS entry '${arch}' is not NN, NNa or NNf") + endif() + set(number "${CMAKE_MATCH_1}") + set(suffix "${CMAKE_MATCH_2}") + if(suffix STREQUAL "f") + math(EXPR family_major "${number} / 10") + set(in_family FALSE) + foreach(requested IN LISTS arg_REQUESTED) + math(EXPR requested_major "${requested} / 10") + if(requested_major EQUAL family_major AND requested GREATER_EQUAL + number) + set(in_family TRUE) + endif() + endforeach() + if(in_family) + list(APPEND resolved "${number}f-real") + endif() + elseif(number IN_LIST arg_REQUESTED) + list(APPEND resolved "${number}${suffix}-real") + endif() + endforeach() + list(REMOVE_DUPLICATES resolved) + set(${out} + "${resolved}" + PARENT_SCOPE) +endfunction() + +# --- Families, shared code and targets ---------------------------------------- + +function(add_modeling_v2_family family) + _modeling_v2_check_family_name("${family}") + get_filename_component(dir_name "${CMAKE_CURRENT_SOURCE_DIR}" NAME) + get_filename_component(package_dir "${CMAKE_CURRENT_SOURCE_DIR}" DIRECTORY) + get_filename_component(package_name "${package_dir}" NAME) + if(NOT dir_name STREQUAL "csrc" OR NOT package_name STREQUAL family) + message( + FATAL_ERROR + "modeling_v2: family '${family}' is declared in models/${family}/csrc/CMakeLists.txt, " + "not in ${CMAKE_CURRENT_SOURCE_DIR}") + endif() + get_property(families GLOBAL PROPERTY MODELING_V2_FAMILIES) + if(family IN_LIST families) + message(FATAL_ERROR "modeling_v2: family '${family}' is declared twice") + endif() + set_property(GLOBAL APPEND PROPERTY MODELING_V2_FAMILIES "${family}") + set_property(GLOBAL PROPERTY MODELING_V2_${family}_SOURCE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}") + set_property(GLOBAL PROPERTY MODELING_V2_${family}_PACKAGE_DIR + "${package_dir}") + set_property(GLOBAL PROPERTY MODELING_V2_${family}_BINARY_DIR + "${CMAKE_CURRENT_BINARY_DIR}") +endfunction() + +function(_modeling_v2_require_family family) + get_property(families GLOBAL PROPERTY MODELING_V2_FAMILIES) + if(NOT family IN_LIST families) + message( + FATAL_ERROR + "modeling_v2: call add_modeling_v2_family(${family}) before declaring its code" + ) + endif() + get_property(family_dir GLOBAL PROPERTY MODELING_V2_${family}_SOURCE_DIR) + if(NOT family_dir STREQUAL CMAKE_CURRENT_SOURCE_DIR) + message( + FATAL_ERROR + "modeling_v2: declare family '${family}' code in ${family_dir}/CMakeLists.txt" + ) + endif() +endfunction() + +# Expand SOURCES (paths or globs) relative to `base`. +function(_modeling_v2_expand_sources out base) + set(sources "") + foreach(item IN LISTS ARGN) + if(IS_ABSOLUTE "${item}") + message( + FATAL_ERROR + "modeling_v2: SOURCES must be relative to ${base}, got '${item}'") + endif() + if(item MATCHES "[*?]") + file(GLOB matched CONFIGURE_DEPENDS "${base}/${item}") + if(NOT matched) + message(FATAL_ERROR "modeling_v2: '${base}/${item}' matches no file") + endif() + list(SORT matched) + list(APPEND sources ${matched}) + else() + if(NOT EXISTS "${base}/${item}") + message(FATAL_ERROR "modeling_v2: '${base}/${item}' does not exist") + endif() + list(APPEND sources "${base}/${item}") + endif() + endforeach() + set(${out} + "${sources}" + PARENT_SCOPE) +endfunction() + +function(_modeling_v2_check_deps) + foreach(dep IN LISTS ARGN) + if(NOT dep IN_LIST _MODELING_V2_KNOWN_DEPS) + message( + FATAL_ERROR + "modeling_v2: unknown DEPS '${dep}' (known: ${_MODELING_V2_KNOWN_DEPS})" + ) + endif() + endforeach() +endfunction() + +# Compile-time usage requirements of DEPS on one OBJECT library. +function(_modeling_v2_apply_compile_deps object_lib) + foreach(dep IN LISTS ARGN) + if(dep STREQUAL "internal_cutlass") + target_include_directories( + ${object_lib} + PRIVATE + $ + ) + elseif(dep STREQUAL "trtllm_gen_fmha_interface") + target_link_libraries(${object_lib} PRIVATE trtllm_gen_fmha_interface) + elseif(dep STREQUAL "nccl" AND ENABLE_MULTI_DEVICE) + target_include_directories(${object_lib} PRIVATE ${MPI_C_INCLUDE_DIRS}) + endif() + endforeach() +endfunction() + +# Link-time libraries of DEPS, for the family library. +function(_modeling_v2_link_libraries out) + set(libraries "") + foreach(dep IN LISTS ARGN) + if(dep STREQUAL "cublas") + list(APPEND libraries ${CUBLAS_LIB}) + elseif(dep STREQUAL "cublaslt") + list(APPEND libraries ${CUBLASLT_LIB}) + elseif(dep STREQUAL "cuda_driver") + list(APPEND libraries ${CUDA_DRV_LIB}) + elseif(dep STREQUAL "nccl" AND ENABLE_MULTI_DEVICE) + list(APPEND libraries ${NCCL_LIB} ${MPI_C_LIBRARIES}) + endif() + endforeach() + set(${out} + "${libraries}" + PARENT_SCOPE) +endfunction() + +#[[ +Create the OBJECT libraries of one unit (a target, or a family's shared code): +`_kernels` for the .cu sources, compiled for `archs` without torch, and +`_ops` for the .cpp torch bindings. Both get the unit's namespaces as +compile definitions (see common/torchLibrary.h and common/namespace.h). +#]] +function( + _modeling_v2_add_objects + name + family + torch_ns + scope + archs + sources + deps) + set(cu_sources "") + set(cpp_sources "") + foreach(source IN LISTS sources) + if(source MATCHES "\\.cu$") + list(APPEND cu_sources "${source}") + elseif(source MATCHES "\\.cpp$") + list(APPEND cpp_sources "${source}") + elseif(NOT source MATCHES "\\.(h|cuh|hpp|inl)$") + message( + FATAL_ERROR + "modeling_v2: '${source}' is neither a .cu kernel source nor a .cpp binding" + ) + endif() + endforeach() + + # Shared code is assembled from the modeling_v2 directory, where source + # properties the family set in its own directory are not visible. + get_property(family_dir GLOBAL PROPERTY MODELING_V2_${family}_SOURCE_DIR) + foreach(source IN LISTS cu_sources cpp_sources) + foreach(property COMPILE_OPTIONS COMPILE_DEFINITIONS) + get_source_file_property(value "${source}" DIRECTORY "${family_dir}" + ${property}) + if(value) + set_source_files_properties("${source}" PROPERTIES ${property} + "${value}") + endif() + endforeach() + endforeach() + + set(definitions + TRTLLM_MV2_TORCH_NS=${torch_ns} TRTLLM_MV2_FAMILY_TORCH_NS=mv2_${family} + TRTLLM_MV2_SCOPE=${scope}) + # Family code includes its own headers from the models/ directory. + get_filename_component(models_dir "${family_dir}" DIRECTORY) + get_filename_component(models_dir "${models_dir}" DIRECTORY) + + set(objects "") + if(cu_sources) + add_library(${name}_kernels OBJECT ${cu_sources}) + set_target_properties( + ${name}_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON CUDA_ARCHITECTURES + "${archs}") + target_compile_definitions(${name}_kernels PRIVATE ${definitions}) + target_include_directories(${name}_kernels PRIVATE "${models_dir}") + _modeling_v2_apply_compile_deps(${name}_kernels ${deps}) + list(APPEND objects ${name}_kernels) + endif() + if(cpp_sources) + add_library(${name}_ops OBJECT ${cpp_sources}) + set_target_properties(${name}_ops PROPERTIES POSITION_INDEPENDENT_CODE ON) + target_compile_definitions(${name}_ops PRIVATE ${definitions}) + target_include_directories(${name}_ops PRIVATE "${models_dir}") + target_link_libraries(${name}_ops PRIVATE ${TORCH_LIBRARIES}) + _modeling_v2_apply_compile_deps(${name}_ops ${deps}) + list(APPEND objects ${name}_ops) + endif() + set_property(GLOBAL APPEND PROPERTY MODELING_V2_${family}_OBJECTS ${objects}) + set_property(GLOBAL APPEND PROPERTY MODELING_V2_${family}_DEPS ${deps}) + set_property(GLOBAL APPEND PROPERTY MODELING_V2_${family}_ARCHS ${archs}) +endfunction() + +# Decide whether a unit builds: "" when it does, otherwise why not. +function(_modeling_v2_skip_reason out family target declared requires) + _modeling_v2_selected(selected "${family}" "${target}") + _modeling_v2_unmet_requirements(unmet ${requires}) + _modeling_v2_resolve_archs(archs DECLARED ${declared}) + if(NOT selected) + set(reason + "not selected by TRTLLM_MODELING_V2_TARGETS=${TRTLLM_MODELING_V2_TARGETS}" + ) + elseif(NOT unmet STREQUAL "") + set(reason "${unmet}") + elseif(NOT archs) + set(reason + "declares ${declared}; this build requests ${CMAKE_CUDA_ARCHITECTURES_ORIG}" + ) + else() + set(reason "") + endif() + set(${out} + "${reason}" + PARENT_SCOPE) +endfunction() + +function(add_modeling_v2_shared family) + cmake_parse_arguments(PARSE_ARGV 1 arg "" "" "ARCHS;SOURCES;DEPS;REQUIRES") + _modeling_v2_require_family("${family}") + get_property(already GLOBAL PROPERTY MODELING_V2_${family}_HAS_SHARED) + if(already) + message( + FATAL_ERROR + "modeling_v2: family '${family}' declares add_modeling_v2_shared twice") + endif() + set_property(GLOBAL PROPERTY MODELING_V2_${family}_HAS_SHARED TRUE) + if(NOT arg_ARCHS OR NOT arg_SOURCES) + message( + FATAL_ERROR + "modeling_v2: add_modeling_v2_shared(${family}) needs ARCHS and SOURCES" + ) + endif() + _modeling_v2_check_deps(${arg_DEPS}) + # Recorded now, built by _modeling_v2_add_family_library only when a target of + # the family builds: shared code with no target has nothing to serve. + foreach(field ARCHS SOURCES DEPS REQUIRES) + set_property(GLOBAL PROPERTY MODELING_V2_${family}_SHARED_${field} + "${arg_${field}}") + endforeach() +endfunction() + +function(add_modeling_v2_target family target) + cmake_parse_arguments(PARSE_ARGV 2 arg "" "" "ARCHS;SOURCES;DEPS;REQUIRES") + _modeling_v2_require_family("${family}") + _modeling_v2_check_target_name("${target}") + if(NOT arg_ARCHS) + message( + FATAL_ERROR + "modeling_v2: add_modeling_v2_target(${family} ${target}) needs ARCHS") + endif() + _modeling_v2_check_deps(${arg_DEPS}) + get_property(package_dir GLOBAL PROPERTY MODELING_V2_${family}_PACKAGE_DIR) + set(target_dir "${package_dir}/${target}/csrc") + if(arg_SOURCES AND NOT IS_DIRECTORY "${target_dir}") + message( + FATAL_ERROR + "modeling_v2: target '${target}' keeps its code in ${target_dir}") + endif() + if(NOT arg_SOURCES) + get_property(has_shared GLOBAL PROPERTY MODELING_V2_${family}_HAS_SHARED) + if(NOT has_shared) + message( + FATAL_ERROR + "modeling_v2: target '${family}/${target}' has no SOURCES, so it runs " + "its family's shared code; declare add_modeling_v2_shared(${family} ...) " + "before it") + endif() + endif() + + get_property(targets GLOBAL PROPERTY MODELING_V2_${family}_TARGETS) + if(target IN_LIST targets) + message( + FATAL_ERROR "modeling_v2: target '${family}/${target}' is declared twice") + endif() + set_property(GLOBAL APPEND PROPERTY MODELING_V2_${family}_TARGETS "${target}") + + _modeling_v2_ident(target_ident "${target}") + set(torch_ns "mv2_${family}_${target_ident}") + get_property(namespaces GLOBAL PROPERTY MODELING_V2_NAMESPACES) + if(torch_ns IN_LIST namespaces) + message( + FATAL_ERROR + "modeling_v2: torch namespace '${torch_ns}' of ${family}/${target} is already taken" + ) + endif() + set_property(GLOBAL APPEND PROPERTY MODELING_V2_NAMESPACES "${torch_ns}") + + _modeling_v2_skip_reason(reason "${family}" "${target}" "${arg_ARCHS}" + "${arg_REQUIRES}") + if(reason STREQUAL "" AND NOT arg_SOURCES) + # A target without code of its own runs only where its shared code builds. + get_property(shared_archs_declared GLOBAL + PROPERTY MODELING_V2_${family}_SHARED_ARCHS) + get_property(shared_requires GLOBAL + PROPERTY MODELING_V2_${family}_SHARED_REQUIRES) + _modeling_v2_unmet_requirements(unmet ${shared_requires}) + _modeling_v2_resolve_archs(shared_archs DECLARED ${shared_archs_declared}) + if(NOT unmet STREQUAL "") + set(reason "its family's shared code is not built: ${unmet}") + elseif(NOT shared_archs) + set(reason "its family's shared code declares ${shared_archs_declared}") + endif() + endif() + if(NOT reason STREQUAL "") + message(STATUS "modeling_v2: ${family}/${target} not built: ${reason}") + return() + endif() + _modeling_v2_resolve_archs(archs DECLARED ${arg_ARCHS}) + if(arg_SOURCES) + message(STATUS "modeling_v2: ${family}/${target} built for ${archs}") + _modeling_v2_expand_sources(sources "${target_dir}" ${arg_SOURCES}) + _modeling_v2_add_objects( + th_modeling_v2_${family}_${target_ident} + "${family}" + "${torch_ns}" + "${family}::${target_ident}" + "${archs}" + "${sources}" + "${arg_DEPS}") + else() + message( + STATUS + "modeling_v2: ${family}/${target} carried for ${archs} (family shared code only)" + ) + set_property(GLOBAL APPEND PROPERTY MODELING_V2_${family}_ARCHS ${archs}) + endif() + set_property(GLOBAL APPEND PROPERTY MODELING_V2_${family}_BUILT "${target}") +endfunction() + +# --- Family libraries --------------------------------------------------------- + +# The registration unit of a family library: it claims the family's torch +# namespace and every built target's, so a second library claiming any of them +# fails when it is loaded, and it registers built_targets() -- which targets +# this library carries -- the way th_common answers feature queries such as +# mhc_fused_hc_mma_enabled(). +function(_modeling_v2_registration_source out family binary_dir) + get_property(built GLOBAL PROPERTY MODELING_V2_${family}_BUILT) + list(JOIN built "," built_csv) + set(source + "// Generated by cpp/cmake/modules/modeling_v2.cmake for family ${family}. Do not edit.\n" + ) + string(APPEND source "#include \n#include \n\n") + string(APPEND source "namespace\n{\n") + string(APPEND source + "//! The targets of this family compiled into this library.\n") + string(APPEND source "std::string builtTargets()\n{\n") + string(APPEND source " return \"${built_csv}\";\n}\n} // namespace\n\n") + string(APPEND source "TORCH_LIBRARY(mv2_${family}, m)\n{\n") + string(APPEND source " m.def(\"built_targets() -> str\");\n}\n\n") + string(APPEND source + "TORCH_LIBRARY_IMPL(mv2_${family}, CompositeExplicitAutograd, m)\n{\n") + string(APPEND source " m.impl(\"built_targets\", &builtTargets);\n}\n") + foreach(target IN LISTS built) + _modeling_v2_ident(target_ident "${target}") + string( + APPEND source + "\nTORCH_LIBRARY(mv2_${family}_${target_ident}, m)\n{\n (void) m;\n}\n" + ) + endforeach() + set(path "${binary_dir}/registration.cpp") + file( + GENERATE + OUTPUT "${path}" + CONTENT "${source}") + set_source_files_properties("${path}" PROPERTIES GENERATED TRUE) + set(${out} + "${path}" + PARENT_SCOPE) +endfunction() + +function(_modeling_v2_add_family_library family) + get_property(binary_dir GLOBAL PROPERTY MODELING_V2_${family}_BINARY_DIR) + get_property(built GLOBAL PROPERTY MODELING_V2_${family}_BUILT) + if(NOT built) + # Like deep_ep and flash_mla without a matching architecture: an empty + # cuda_architectures.txt, and nothing for build_wheel.py to install. + file(WRITE "${binary_dir}/cuda_architectures.txt" "") + return() + endif() + + get_property(has_shared GLOBAL PROPERTY MODELING_V2_${family}_HAS_SHARED) + if(has_shared) + foreach(field ARCHS SOURCES DEPS REQUIRES) + get_property(shared_${field} GLOBAL + PROPERTY MODELING_V2_${family}_SHARED_${field}) + endforeach() + # The target filter does not apply: a built target needs its shared code. + _modeling_v2_unmet_requirements(unmet ${shared_REQUIRES}) + _modeling_v2_resolve_archs(shared_archs DECLARED ${shared_ARCHS}) + if(NOT unmet STREQUAL "") + message(STATUS "modeling_v2: ${family}/shared not built: ${unmet}") + elseif(NOT shared_archs) + message( + STATUS + "modeling_v2: ${family}/shared not built: declares ${shared_ARCHS}; " + "this build requests ${CMAKE_CUDA_ARCHITECTURES_ORIG}") + else() + get_property(source_dir GLOBAL PROPERTY MODELING_V2_${family}_SOURCE_DIR) + _modeling_v2_expand_sources(sources "${source_dir}" ${shared_SOURCES}) + message(STATUS "modeling_v2: ${family}/shared built for ${shared_archs}") + _modeling_v2_add_objects( + th_modeling_v2_${family}_shared + "${family}" + "mv2_${family}" + "${family}::shared" + "${shared_archs}" + "${sources}" + "${shared_DEPS}") + endif() + endif() + + get_property(objects GLOBAL PROPERTY MODELING_V2_${family}_OBJECTS) + get_property(deps GLOBAL PROPERTY MODELING_V2_${family}_DEPS) + get_property(archs GLOBAL PROPERTY MODELING_V2_${family}_ARCHS) + list(REMOVE_DUPLICATES deps) + list(REMOVE_DUPLICATES archs) + + set(library th_modeling_v2_${family}) + _modeling_v2_registration_source(registration "${family}" "${binary_dir}") + set(object_sources "") + foreach(object IN LISTS objects) + list(APPEND object_sources $) + endforeach() + add_library(${library} MODULE "${registration}" ${object_sources}) + _modeling_v2_link_libraries(dep_libraries ${deps}) + # Linked the way th_common is: torch, th_utils, Python (which + # , included by stock headers such as thop/thUtils.h, + # references) and libtensorrt_llm.so, whose singletons (logger, mcast + # registry, driver wrapper) stay single. + target_link_libraries( + ${library} + PRIVATE ${TORCH_LIBRARIES} + th_utils + ${Python3_LIBRARIES} + ${SHARED_TARGET} + ${CUDA_RT_LIB} + ${CUDA_DRV_LIB} + ${dep_libraries}) + set(version_script + "${PROJECT_SOURCE_DIR}/tensorrt_llm/modeling_v2/common/familyLibrary.version" + ) + set_target_properties( + ${library} + PROPERTIES PREFIX "lib" + LIBRARY_OUTPUT_DIRECTORY "${binary_dir}" + POSITION_INDEPENDENT_CODE ON + LINKER_LANGUAGE CXX + LINK_DEPENDS "${version_script}" + BUILD_RPATH "$ORIGIN;$ORIGIN/../../nvidia/nccl/lib") + target_link_options( + ${library} PRIVATE "-Wl,--version-script=${version_script}" + "-Wl,--no-undefined") + add_dependencies(modeling_v2_libs ${library}) + file(WRITE "${binary_dir}/cuda_architectures.txt" "${archs}") +endfunction() + +# Build every declared family's library. Called by +# cpp/tensorrt_llm/modeling_v2/CMakeLists.txt after it added the families. +function(add_modeling_v2_family_libraries) + get_property(families GLOBAL PROPERTY MODELING_V2_FAMILIES) + foreach(family IN LISTS families) + _modeling_v2_add_family_library("${family}") + endforeach() +endfunction() diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 8fa6c33c8eee..fc84a4fc1942 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -341,6 +341,8 @@ endif() if(BUILD_PYT) add_subdirectory(thop) + # modeling_v2 family libraries: torch-op libraries loaded on demand. + add_subdirectory(modeling_v2) endif() add_subdirectory(nanobind) diff --git a/cpp/tensorrt_llm/modeling_v2/CMakeLists.txt b/cpp/tensorrt_llm/modeling_v2/CMakeLists.txt new file mode 100644 index 000000000000..3785811be8cf --- /dev/null +++ b/cpp/tensorrt_llm/modeling_v2/CMakeLists.txt @@ -0,0 +1,67 @@ +# +# 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. +# + +# Native code of modeling_v2 targets: one torch-op library per model family, +# loaded only when routing picks one of its targets. A family keeps its native +# code with its Python code, in +# tensorrt_llm/_torch/_experimental/modeling_v2/models//csrc/ (declared +# by the CMakeLists.txt there) and //csrc/; this directory holds +# only what every family shares. See cpp/cmake/modules/modeling_v2.cmake. + +include(modeling_v2) + +add_custom_target(modeling_v2_libs) + +get_filename_component( + modeling_v2_models_dir + "${PROJECT_SOURCE_DIR}/../tensorrt_llm/_torch/_experimental/modeling_v2/models" + ABSOLUTE) +file( + GLOB family_lists + LIST_DIRECTORIES false + CONFIGURE_DEPENDS "${modeling_v2_models_dir}/*/csrc/CMakeLists.txt") +list(SORT family_lists) + +# A family removed or renamed since an earlier configure leaves its build +# directory, and a library in it, behind: empty its cuda_architectures.txt so +# that build_wheel.py neither installs that library nor keeps it installed. +set(families "") +foreach(family_list IN LISTS family_lists) + get_filename_component(csrc_dir "${family_list}" DIRECTORY) + get_filename_component(family_dir "${csrc_dir}" DIRECTORY) + get_filename_component(family "${family_dir}" NAME) + list(APPEND families "${family}") +endforeach() +file( + GLOB previous_arch_files + LIST_DIRECTORIES false + "${CMAKE_CURRENT_BINARY_DIR}/*/cuda_architectures.txt") +foreach(arch_file IN LISTS previous_arch_files) + get_filename_component(binary_dir "${arch_file}" DIRECTORY) + get_filename_component(name "${binary_dir}" NAME) + if(NOT name IN_LIST families) + file(WRITE "${arch_file}" "") + endif() +endforeach() + +foreach(family_list IN LISTS family_lists) + get_filename_component(csrc_dir "${family_list}" DIRECTORY) + get_filename_component(family_dir "${csrc_dir}" DIRECTORY) + get_filename_component(family "${family_dir}" NAME) + add_subdirectory("${csrc_dir}" "${CMAKE_CURRENT_BINARY_DIR}/${family}") +endforeach() +add_modeling_v2_family_libraries() diff --git a/cpp/tensorrt_llm/modeling_v2/README.md b/cpp/tensorrt_llm/modeling_v2/README.md new file mode 100644 index 000000000000..301cb1167af2 --- /dev/null +++ b/cpp/tensorrt_llm/modeling_v2/README.md @@ -0,0 +1,19 @@ + + + +# modeling_v2 family libraries: what they share + +Each modeling_v2 model family keeps its C++/CUDA with its Python code, in +`tensorrt_llm/_torch/_experimental/modeling_v2/models//csrc/` and +`//csrc/`, and builds into its own torch-op library, +`libth_modeling_v2_.so`. The guide for authors is the "Native code" +section of `tensorrt_llm/_torch/_experimental/modeling_v2/README.md`. + +This directory holds only what every family library shares; adding a model +changes nothing here: + +| Path | Role | +|---|---| +| `CMakeLists.txt` | Finds every `models//csrc/CMakeLists.txt` and builds the family libraries (`cpp/cmake/modules/modeling_v2.cmake`). | +| `common/namespace.h`, `common/torchLibrary.h` | The per-unit C++ and torch namespace macros native code registers through. | +| `common/familyLibrary.version` | The version script: a family library exports no symbol. | diff --git a/cpp/tensorrt_llm/modeling_v2/common/familyLibrary.version b/cpp/tensorrt_llm/modeling_v2/common/familyLibrary.version new file mode 100644 index 000000000000..59b171adfea3 --- /dev/null +++ b/cpp/tensorrt_llm/modeling_v2/common/familyLibrary.version @@ -0,0 +1,7 @@ +/* A modeling_v2 family library exports nothing: its torch ops register from + * static initializers when the library is loaded, and keeping every symbol + * local lets two family libraries define the same C++ names in one process. + * See cpp/cmake/modules/modeling_v2.cmake. */ +{ + local: *; +}; diff --git a/cpp/tensorrt_llm/modeling_v2/common/namespace.h b/cpp/tensorrt_llm/modeling_v2/common/namespace.h new file mode 100644 index 000000000000..0bad346d62a4 --- /dev/null +++ b/cpp/tensorrt_llm/modeling_v2/common/namespace.h @@ -0,0 +1,41 @@ +/* + * 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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +// add_modeling_v2_target() and add_modeling_v2_shared() define TRTLLM_MV2_SCOPE +// as `::` or `::shared` (cpp/cmake/modules/modeling_v2.cmake). +#ifndef TRTLLM_MV2_SCOPE +#error "modeling_v2 native code is compiled through add_modeling_v2_target() or add_modeling_v2_shared()" +#endif + +//! Opens tensorrt_llm::modeling_v2:::: (or ::shared) for the unit +//! being compiled. Every target has its own C++ namespace, so sibling targets can +//! carry identically named code in one family library. +#define TRTLLM_MV2_NAMESPACE_BEGIN \ + TRTLLM_NAMESPACE_BEGIN \ + namespace modeling_v2::TRTLLM_MV2_SCOPE \ + { + +#define TRTLLM_MV2_NAMESPACE_END \ + } \ + TRTLLM_NAMESPACE_END + +//! The fully qualified namespace TRTLLM_MV2_NAMESPACE_BEGIN opens. +#define TRTLLM_MV2_NS ::tensorrt_llm::modeling_v2::TRTLLM_MV2_SCOPE diff --git a/cpp/tensorrt_llm/modeling_v2/common/torchLibrary.h b/cpp/tensorrt_llm/modeling_v2/common/torchLibrary.h new file mode 100644 index 000000000000..2b5de20fcd2b --- /dev/null +++ b/cpp/tensorrt_llm/modeling_v2/common/torchLibrary.h @@ -0,0 +1,39 @@ +/* + * 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. + */ + +#pragma once + +#include + +// add_modeling_v2_target() defines TRTLLM_MV2_TORCH_NS as mv2__ +// and add_modeling_v2_shared() as mv2_ (cpp/cmake/modules/modeling_v2.cmake). The +// family library claims each namespace once with TORCH_LIBRARY, so op files only +// ever add fragments to it. +#ifndef TRTLLM_MV2_TORCH_NS +#error "modeling_v2 torch bindings are compiled through add_modeling_v2_target() or add_modeling_v2_shared()" +#endif + +// One level of indirection so that TORCH_LIBRARY_*'s own # and ## operators see +// the expanded namespace rather than the macro name. +#define TRTLLM_MV2_DETAIL_LIBRARY_FRAGMENT(ns, m) TORCH_LIBRARY_FRAGMENT(ns, m) +#define TRTLLM_MV2_DETAIL_LIBRARY_IMPL(ns, k, m) TORCH_LIBRARY_IMPL(ns, k, m) + +//! TORCH_LIBRARY_FRAGMENT on this unit's own namespace: declare op schemas. +#define TRTLLM_MV2_LIBRARY_FRAGMENT(m) TRTLLM_MV2_DETAIL_LIBRARY_FRAGMENT(TRTLLM_MV2_TORCH_NS, m) + +//! TORCH_LIBRARY_IMPL on this unit's own namespace: register kernels for dispatch key `k`. +#define TRTLLM_MV2_LIBRARY_IMPL(k, m) TRTLLM_MV2_DETAIL_LIBRARY_IMPL(TRTLLM_MV2_TORCH_NS, k, m) diff --git a/jenkins/scripts/cbts/README.md b/jenkins/scripts/cbts/README.md index 047fa7692588..20453820c954 100644 --- a/jenkins/scripts/cbts/README.md +++ b/jenkins/scripts/cbts/README.md @@ -50,7 +50,7 @@ Nine rules, registered in `cbts/command/main.py::RULE_CLASSES`: | `TestListRule` | `testlistonly` | `tests/integration/test_lists/test-db/*.yml` | | `VisualGenRule` | `visualgenonly` | `examples/visual_gen/**`, `scripts/visualgen_eval/**`, `tensorrt_llm/_torch/visual_gen/**`, `tensorrt_llm/media/**`, `tensorrt_llm/visual_gen/**` (excl. docs; reference images such as `cat_piano.png` ARE test fixtures and stay claimed; outward-facing files force fallback) | | `SpecDecRule` | `specdeconly` | `tensorrt_llm/_torch/speculative/**`, `tensorrt_llm/models/{eagle,medusa,redrafter}/**`, `examples/{eagle,medusa,redrafter,draft_target_model,ngram}/**`, `examples/llm-api/llm_speculative_decoding.py` (excl. docs; other suffixes incl. images kept as potential test fixtures) | -| `ModelingV2Rule` | `modelingv2only` | `tensorrt_llm/_torch/_experimental/modeling_v2/**` (excl. docs) | +| `ModelingV2Rule` | `modelingv2only` | `tensorrt_llm/_torch/_experimental/modeling_v2/**`, `cpp/tensorrt_llm/modeling_v2/**`, `cpp/cmake/modules/modeling_v2.cmake` (excl. docs); per model node for changes under `models/` | | `AgentFlowRule` | `agentflowonly` | `agent-flow/**` (excl. docs and out-of-scope metadata) | | `OpenEngineRule` | `openengineonly` | `tensorrt_llm/grpc/openengine/**` (excl. docs) | | `DocsRule` | `docsonly` | Outside noop paths: `docs/**` and verified Python docstring edits → `CPU-Build_Docs` plus complete `l0_cpu`; other `**/*.md` / `**/*.rst` → `CPU-Build_Docs` only | @@ -67,7 +67,7 @@ See `cbts/rules/README.md` for per-rule logic. | `testlistonly` | `TestListRule` fired solo: PR only adds entries under `tests/integration/test_lists/test-db/*.yml`. | | `visualgenonly` | `VisualGenRule` fired solo: PR only touches VisualGen internal source paths (`examples/visual_gen/**`, `scripts/visualgen_eval/**`, `tensorrt_llm/_torch/visual_gen/**`; excl. docs; image fixtures like `cat_piano.png` are claimed). Narrows to blocks containing VG test entries. Outward-facing files under `tensorrt_llm/visual_gen/**` and `tensorrt_llm/media/**` (eagerly imported by `trtllm-serve`) force `null` fallback. | | `specdeconly` | `SpecDecRule` fired solo: PR only touches speculative-decoding source paths (`tensorrt_llm/_torch/speculative/**`, `tensorrt_llm/models/{eagle,medusa,redrafter}/**`, `examples/{eagle,medusa,redrafter,draft_target_model,ngram}/**`, `examples/llm-api/llm_speculative_decoding.py`; excl. docs). Narrows to blocks containing spec-dec test entries (eagle / medusa / redrafter / ngram / draft-target-model / MTP). | -| `modelingv2only` | `ModelingV2Rule` fired solo: PR only touches the modeling_v2 subtree (`tensorrt_llm/_torch/_experimental/modeling_v2/**`; excl. docs, which are a fifth of the subtree — every catalog entry carries a contract document). Narrows to blocks containing modeling_v2 test entries (`unittest/_torch/modeling_v2/`, `test_modeling_v2_*`). No outward-facing fallback is needed: nothing beyond the resolver module is imported unless `LLM(modeling_v2=...)` is on, and its callers outside the subtree (`_torch/pyexecutor/model_loader.py`, `_torch/models/modeling_auto.py`) are left unclaimed, so touching the shared resolver falls back to baseline. | +| `modelingv2only` | `ModelingV2Rule` fired solo: PR only touches modeling_v2 code -- the subtree (`tensorrt_llm/_torch/_experimental/modeling_v2/**`, each family's native code in its `csrc/` included: one torch-op library per model family) or what the family libraries share (`cpp/tensorrt_llm/modeling_v2/**`, `cpp/cmake/modules/modeling_v2.cmake`); excl. docs, which are a fifth of the Python subtree — every catalog entry carries a contract document. A change under `models/` selects its model node's unit tests and target accuracy gates, plus whole-tree checks; shared code or an unreadable routing tree selects every modeling_v2 entry. No outward-facing fallback is needed: nothing beyond the resolver module is imported, and no family library is loaded, unless `LLM(modeling_v2=...)` is on; the CPU checks that keep this true run in every scope. The callers outside the subtree (`_torch/pyexecutor/model_loader.py`, `_torch/models/modeling_auto.py`) are left unclaimed, so touching them falls back to baseline. | | `agentflowonly` | `AgentFlowRule` fired solo: PR only touches `agent-flow/**` source or test files (excl. docs). Runs `CPU-AgentFlow-UnitTest`. | | `openengineonly` | `OpenEngineRule` fired solo: PR only touches `tensorrt_llm/grpc/openengine/**` source files (excl. docs). Narrows to the registered OpenEngine unit tests: the stub-based ones on the always-run `CPU-Generic-*` stages, plus `test_capability_conformance.py` on `A10-PyTorch-*`, which needs a GPU. | | `docsonly` | `DocsRule` fired solo: all documentation changes run `CPU-Build_Docs`; changes under `docs/**` also run the complete `l0_cpu` suite on matching `CPU-Generic-*` stages. Docs-only orchestration skips the SBSA track, so its resolved ARM stage is not launched. | diff --git a/jenkins/scripts/cbts/cbts/rules/README.md b/jenkins/scripts/cbts/cbts/rules/README.md index f8f722363eb5..e3605f493dec 100644 --- a/jenkins/scripts/cbts/cbts/rules/README.md +++ b/jenkins/scripts/cbts/cbts/rules/README.md @@ -16,7 +16,7 @@ for the overall CBTS architecture. | `test_list.py` | `TestListRule` | `testlistonly` | `tests/integration/test_lists/test-db/*.yml` | | `visual_gen.py` | `VisualGenRule` | `visualgenonly` | `examples/visual_gen/**`, `scripts/visualgen_eval/**`, `tensorrt_llm/_torch/visual_gen/**`, `tensorrt_llm/media/**`, `tensorrt_llm/visual_gen/**` (each excl. docs) | | `spec_dec.py` | `SpecDecRule` | `specdeconly` | `tensorrt_llm/_torch/speculative/**`, `tensorrt_llm/models/{eagle,medusa,redrafter}/**`, `examples/{eagle,medusa,redrafter,draft_target_model,ngram}/**`, `examples/llm-api/llm_speculative_decoding.py` (each excl. docs) | -| `modeling_v2.py` | `ModelingV2Rule` | `modelingv2only` | `tensorrt_llm/_torch/_experimental/modeling_v2/**` (excl. docs) | +| `modeling_v2.py` | `ModelingV2Rule` | `modelingv2only` | `tensorrt_llm/_torch/_experimental/modeling_v2/**`, `cpp/tensorrt_llm/modeling_v2/**`, `cpp/cmake/modules/modeling_v2.cmake` (excl. docs); per model node for changes under `models/` | | `agent_flow.py` | `AgentFlowRule` | `agentflowonly` | `agent-flow/**` (excl. docs and out-of-scope metadata) → the single `CPU-AgentFlow-UnitTest` stage; not test-db-driven | | `openengine.py` | `OpenEngineRule` | `openengineonly` | `tensorrt_llm/grpc/openengine/**` (excl. docs) → the `l0_cpu` block containing `unittest/grpc/openengine/` | | `docs.py` | `DocsRule` | `docsonly` | Outside noop paths: `docs/**` and verified Python docstring edits → `CPU-Build_Docs` plus complete `l0_cpu`; other `**/*.md` / `**/*.rst` → `CPU-Build_Docs` only | @@ -303,7 +303,20 @@ Outcomes: Path-only rule. Claims non-documentation source changes under `tensorrt_llm/_torch/_experimental/modeling_v2/`, the self-contained second modeling path (one flat forward per checkpoint/arch/parallel triple, assembled -from a catalog of op wrappers). +from a catalog of op wrappers), including each family's native code +(`models//csrc/`, `models///csrc/`: one torch-op library +per model family, loaded only when routing picks one of the family's targets), +plus what every family library shares: `cpp/tensorrt_llm/modeling_v2/` and the +native build module `cpp/cmake/modules/modeling_v2.cmake`. + +Scope — per model node, else all of modeling_v2: source changes under +`models/` (native `csrc/` included) select the nearest model node's unit tests +and the accuracy gates of its targets. Family-shared code selects the whole +family. A covering directory entry is retained when no more specific entry +exists, unless it ignores that node. Whole-tree unit checks are always retained. +Shared Python or native infrastructure selects every modeling_v2 entry, as does +an unreadable routing tree. If no modeling_v2 entry matches, the rule falls back +to baseline. `.md` exclusion carries more weight here than elsewhere: every catalog entry ships a contract document, so roughly a fifth of the subtree is @@ -318,9 +331,10 @@ in `backend: pytorch` blocks beside everything else. A block belongs to modeling_v2 iff one of its `tests:` entries matches `_MV2_ENTRY_PATTERNS`: -- `unittest/_torch/modeling_v2/` — the op-level catalog matrix, carried - as whole-directory entries (one on `l0_gb300`, one on - `l0_gb300_multi_gpus` for the 4-rank collectives). +- `unittest/_torch/modeling_v2/` — the unit tests: one directory entry + (`--ignore` of `comm/`) that collects the catalog matrix and every + family's tests under `models/`, plus `comm/` on + `l0_gb300_multi_gpus` for the 4-rank collectives. - `test_modeling_v2_` — the accuracy gates, and any future unit file. Both markers are exact by construction rather than by luck: every test @@ -331,9 +345,9 @@ carve-out is needed. Outward fallback: not needed, and by design rather than by accident. Nothing beyond the resolver module is imported unless `LLM(modeling_v2=...)` -is on — the model loader calls `modeling_v2_resolve`, which returns -immediately when the switch is off, and the routing modules are imported -lazily behind it. The callers outside the subtree, +is on — the model loader calls `modeling_v2_resolve_on_all_ranks`, which +returns immediately when the switch is off, and the routing modules are +imported lazily behind it. The callers outside the subtree, `tensorrt_llm/_torch/pyexecutor/model_loader.py` (decides) and `tensorrt_llm/_torch/models/modeling_auto.py` (reads the decision off `ModelConfig`), are deliberately left unclaimed: a change to the shared diff --git a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py index 791827e35e28..fee9b67e260b 100644 --- a/jenkins/scripts/cbts/cbts/rules/modeling_v2.py +++ b/jenkins/scripts/cbts/cbts/rules/modeling_v2.py @@ -16,7 +16,9 @@ modeling_v2 is a second modeling path living entirely under `tensorrt_llm/_torch/_experimental/modeling_v2/`: one self-contained forward per (checkpoint, GPU arch, parallel topology), assembled from a catalog of -op wrappers. +op wrappers. Native code lives beside each model in `csrc/`; shared build +infrastructure lives under `cpp/tensorrt_llm/modeling_v2/` and +`cpp/cmake/modules/modeling_v2.cmake` and selects every modeling_v2 entry. Block selection: a block belongs to modeling_v2 iff one of its `tests:` entries matches a marker in `_MV2_ENTRY_PATTERNS`. Those markers are exact @@ -95,7 +97,12 @@ # Source-path prefixes the rule may claim. Tests under tests/** are left # to TestsDefRule; the two scopes combine via _TESTSONLY_FAMILY. -_MV2_SRC_PREFIXES: tuple[str, ...] = ("tensorrt_llm/_torch/_experimental/modeling_v2/",) +_MV2_SRC_PREFIXES: tuple[str, ...] = ( + "tensorrt_llm/_torch/_experimental/modeling_v2/", + "cpp/tensorrt_llm/modeling_v2/", +) +_CPP_BUILD_FILES = frozenset({"cpp/cmake/modules/modeling_v2.cmake"}) +_IGNORE = re.compile(r"--ignore=(\S+)") _UNIT_ROOT = "unittest/_torch/modeling_v2" _ACCURACY_ROOT = "accuracy/modeling_v2" @@ -126,7 +133,7 @@ def _is_mv2_claim(path: str) -> bool: fixture, so the rule keeps claiming it and re-runs the stages (safe over-run). """ - if not path.startswith(_MV2_SRC_PREFIXES): + if not (path.startswith(_MV2_SRC_PREFIXES) or path in _CPP_BUILD_FILES): return False if is_docs_path(path): return False @@ -134,7 +141,9 @@ def _is_mv2_claim(path: str) -> bool: def _entry_is_mv2(entry: str) -> bool: - return any(p in entry for p in _MV2_ENTRY_PATTERNS) + return entry.split()[0].rstrip("/") == _UNIT_ROOT or any( + p in entry for p in _MV2_ENTRY_PATTERNS + ) def _mv2_entries(block) -> list[str]: @@ -174,12 +183,20 @@ def _kept_entries(entries: list[str], prefixes: set[str]) -> list[str]: file in a directory above a prefix is kept as well; it does not count as sitting under the prefix, so it never displaces the covering entry. """ - targets = {entry: _entry_target(entry) for entry in entries} + targets = {entry: _entry_target(entry).split()[0] for entry in entries} kept = {e for e, t in targets.items() if _is_whole_tree_test(t)} for prefix in prefixes: under = {e for e, t in targets.items() if _target_in_filter_subtree(t, prefix)} if not under: - covering = [e for e, t in targets.items() if _target_in_filter_subtree(prefix, t)] + covering = [ + e + for e, t in targets.items() + if _target_in_filter_subtree(prefix, t) + and not any( + _target_in_filter_subtree(prefix, ignored.rstrip("/")) + for ignored in _IGNORE.findall(e) + ) + ] under = {max(covering, key=lambda e: len(targets[e]))} if covering else set() kept |= under kept |= {e for e, t in targets.items() if _is_ancestor_file(t, prefix)} @@ -226,6 +243,8 @@ def _gates_under(self, node: Path) -> Optional[set[str]]: def _prefixes_for(self, path: str) -> Optional[set[str]]: """Filter prefixes a changed file selects; None selects every modeling_v2 entry.""" + if not path.startswith(_MV2_SRC_PREFIXES[0]): + return None rel = Path(path[len(_MV2_SRC_PREFIXES[0]) :]) if rel.parts[0] != "models" or len(rel.parts) < 3: return None diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 76f3c1cdd261..b858578c77ac 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -792,6 +792,31 @@ def has_sm90_or_newer(cuda_architectures: str) -> bool: return False +def install_modeling_v2_libs(mv2_build_dir: Path, lib_dir: Path, + install_file) -> None: + """Install the modeling_v2 family libraries the way deep_ep and flash_mla are. + + Each family's build directory holds cuda_architectures.txt, written at + configure time: the architectures its library was built for, empty when + there is none (cpp/cmake/modules/modeling_v2.cmake). As deep_ep's package + is, a family library installed by an earlier build that this build did not + produce is removed, so it is neither loaded nor shipped. + """ + installed = set() + for arch_file in sorted(mv2_build_dir.glob("*/cuda_architectures.txt")): + if not arch_file.read_text().strip().strip(";"): + continue + libraries = list(arch_file.parent.glob("libth_modeling_v2_*.so")) + assert len(libraries) == 1, ( + "Exactly one modeling_v2 library should be present in " + f"{arch_file.parent}: {libraries}") + install_file(libraries[0], lib_dir / libraries[0].name) + installed.add(libraries[0].name) + for stale in lib_dir.glob("libth_modeling_v2_*.so"): + if stale.name not in installed: + stale.unlink() + + def stage_nccl_extensions_package(wheel: Path, staging_dir: Path) -> None: """Stage NCCL-EP namespace packages from the intermediate wheel.""" package_root = staging_dir / "3rdparty" / "nccl_extensions" @@ -1018,7 +1043,7 @@ def main(*, targets = [ "tensorrt_llm", "th_common", "bindings", "deep_ep", "deep_gemm", - "pg_utils", "flash_mla" + "pg_utils", "flash_mla", "modeling_v2_libs" ] build_nccl_extensions_enabled = has_sm90_or_newer(cuda_architectures) if build_nccl_extensions_enabled: @@ -1341,6 +1366,8 @@ def copy_resolving_symlink(src_path, dst_path): lib_dir / "libtensorrt_llm.so") install_file(build_dir / "tensorrt_llm/thop/libth_common.so", lib_dir / "libth_common.so") + install_modeling_v2_libs(build_dir / "tensorrt_llm/modeling_v2", + lib_dir, install_file) if os.path.exists( build_dir / "tensorrt_llm/executor/cache_transmission/ucx_utils/libtensorrt_llm_ucx_wrapper.so" diff --git a/setup.py b/setup.py index 6ff03926441a..2b49833f7e47 100644 --- a/setup.py +++ b/setup.py @@ -179,6 +179,7 @@ def has_ext_modules(self): package_data = [ 'libs/libtensorrt_llm.so', 'libs/libth_common.so', + 'libs/libth_modeling_v2_*.so', 'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so', 'libs/libtensorrt_llm_nixl_wrapper.so', @@ -269,8 +270,13 @@ def warn_on_build_skew(precompiled_location: str) -> None: import subprocess NATIVE_BUILD_INPUTS = [ - "cpp/", "3rdparty/", "setup.py", "scripts/build_wheel.py", - "requirements.txt" + "cpp/", + "3rdparty/", + "setup.py", + "scripts/build_wheel.py", + "requirements.txt", + # modeling_v2 families keep their native code with their Python code. + ":(glob)tensorrt_llm/_torch/_experimental/modeling_v2/models/**/csrc/**" ] def head_of(checkout: str) -> str | None: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index f9c8568348f3..fe0fa24e2b0b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -111,11 +111,16 @@ checkpoint of the same shape is ungated. ``` _router_index.py architectures[0] -> routing module. Small, stable, test-guarded. +_native.py loads a target's native code (see "Native code" below) models// routing.py one forward-reading decision tree per architecture family + fake_ops.py fakes of the family's shared native ops, if it has any + csrc/ the family's CMakeLists.txt and the native code its targets share ____/ modeling.py weights.py ops// entries only this target calls + fake_ops.py fakes of the target's native ops, if it has any + csrc/ the target's own kernels and torch bindings catalog// entries more than one target calls ``` @@ -164,6 +169,64 @@ read as a hierarchy that was never one -- every level had exactly one child, and each needed an `__init__.py` whose only job was to exist. The class name carries the same triple; `test_modeling_v2_claims.py` asserts they agree. +### Native code + +A target's C++/CUDA lives with its Python code: `/csrc/`, and the family's +`csrc/` for code several targets share. Each family builds into one torch-op library, +`libth_modeling_v2_.so`, installed next to `libth_common.so`; nothing of it is +linked into the stock libraries, and it loads only when routing picks one of the +family's targets. Adding a model with kernels edits nothing under `cpp/`. +`models/gpt_oss/` is the example: one kernel, used by its weight loader. + +Adding native code to target `foo_nvfp4__sm_103__tp4` of family `foo`: + +1. **Declare it** in `models/foo/csrc/CMakeLists.txt`: + + ```cmake + add_modeling_v2_family(foo) + add_modeling_v2_target(foo foo_nvfp4__sm_103__tp4 ARCHS 103a SOURCES *.cu *.cpp) + ``` + +2. **Write the kernels and their torch bindings** in `foo_nvfp4__sm_103__tp4/csrc/`, in + the target's namespaces: + + ```cpp + TRTLLM_MV2_NAMESPACE_BEGIN // tensorrt_llm::modeling_v2::foo::foo_nvfp4_sm_103_tp4 + ... + TRTLLM_MV2_NAMESPACE_END + TRTLLM_MV2_LIBRARY_FRAGMENT(m) { m.def("my_op(Tensor x) -> Tensor"); } + TRTLLM_MV2_LIBRARY_IMPL(CUDA, m) { m.impl("my_op", &TRTLLM_MV2_NS::myOp); } + ``` + + Include the family's own headers from `models/` (`"foo/foo_nvfp4__sm_103__tp4/csrc/x.h"`); + any stock header is fine, another target's or family's code is not. + +3. **Load it** from the target's `__init__.py`, `native_ops = load_target_ops(__name__)` + (from `_native`); the target's modules use `from . import native_ops`. + +4. **Fake every op** in the target's `fake_ops.py`, which the loader imports right after + the library opens: `@torch.library.register_fake(f"{namespace_of(__name__)}::my_op")`. + +5. **Test it** under `tests/unittest/_torch/modeling_v2/models/foo/`. The modeling_v2 entry + of `l0_b300.yml` collects it, so a target that runs on GB300 changes no test list; a + test for another GPU, or for several GPUs, is listed in that GPU's test list. + +Code several targets call goes once in the family's `csrc/`, declared with +`add_modeling_v2_shared(foo ARCHS ... SOURCES ...)`; its ops are in `torch.ops.mv2_foo`, +reached with `load_family_ops(__name__)` and faked in `models/foo/fake_ops.py`. A target +with no code of its own declares no `SOURCES`. + +- **`ARCHS`** (`NN`, `NNa`, `NNf`) is intersected with the build's architectures; a target + with none is not built, and is a `NoTarget` at run time. `REQUIRES` (`MULTI_DEVICE`, + `NVSHMEM`) and `DEPS` (a closed list) are the other keywords; a per-file compile option + goes on the file with `set_source_files_properties`. +- **Loading**: a library this installation does not carry, or a target it did not build, + is a `NoTarget` (`auto` keeps the built-in model); one that is installed but fails to + load, or lacks an op a `fake_ops.py` declares, is a `NativeLoadError`. +- **Building only your family**: `build_wheel.py -D "TRTLLM_MODELING_V2_TARGETS=foo/*"`, and + after an edit `cmake --build cpp/build --target th_modeling_v2_foo`. +- `test_modeling_v2_native_claims.py` checks these rules on CPU; run it before sending the PR. + ### Default and specialized targets The parallel segment of a target's name is either one topology (for diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index 96e5209c99dc..16d312d3e7b1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -31,6 +31,7 @@ ModelingV2Mode, NoTarget, modeling_v2_resolve, + modeling_v2_resolve_on_all_ranks, step_contract_enabled, ) @@ -39,5 +40,6 @@ "ModelingV2Mode", "NoTarget", "modeling_v2_resolve", + "modeling_v2_resolve_on_all_ranks", "step_contract_enabled", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_native.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_native.py new file mode 100644 index 000000000000..a96e8f52edda --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_native.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A target's native (C++/CUDA) code: where it is built, and how the target loads it. + +Native code of a model family lives with its Python code, in ``models//csrc/`` +(the family's ``CMakeLists.txt`` and shared code) and ``models///csrc/`` +(a target's own), and builds into one torch-op library, +``libth_modeling_v2_.so``, installed next to ``libth_common.so``. Nothing of +it is in the stock libraries, and nothing loads it until routing picks a target of +the family. The API is four names: + +* :func:`load_target_ops` -- a target package with native code starts with + ``native_ops = load_target_ops(__name__)``; +* :func:`load_family_ops` -- the family's shared ops, ``load_family_ops(__name__)`` + from any module of the family; +* :func:`namespace_of` -- the torch namespace a ``fake_ops.py`` registers its fakes + in, ``namespace_of(__name__)``; +* :class:`NativeLoadError` -- a library that is installed but broken. + +``route`` imports a target only on the branch that returns it, so loading is part of +routing, and it follows how the stock code treats its optional native modules +(``flash_mla``, ``deep_ep``): a library this installation does not have, or a target +it did not build (it answers ``torch.ops.mv2_.built_targets()``), is a +:class:`NoTarget`, so ``auto`` keeps the built-in model; a library that fails to +load, or lacks an op a ``fake_ops.py`` declares, is a :class:`NativeLoadError`, +never a fallback. The family's ``fake_ops.py`` registers when the library opens, a +target's when the target is found in it. + +The "Native code" section of the modeling_v2 README is the guide for authors. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import re +from pathlib import Path +from typing import Any, Dict, FrozenSet, Optional, Tuple + +from ._router_index import NoTarget + +_PACKAGE = __name__.rpartition(".")[0] +_MODELS_PACKAGE = f"{_PACKAGE}.models" + +# tensorrt_llm/libs, where build_wheel.py installs libth_common.so and the family libraries. +_LIBS_DIR = Path(__file__).resolve().parents[3] / "libs" + +_built_targets: Dict[str, FrozenSet[str]] = {} +_family_ops: Dict[str, Any] = {} + + +class NativeLoadError(RuntimeError): + """A family library is installed but cannot be used. + + Raised when the library fails to open or an op a ``fake_ops.py`` declares is + missing from it. It is deliberately not a :class:`NoTarget`, so ``auto`` does + not hide a broken installation behind the built-in model. + """ + + +def namespace_of(module: str) -> str: + """The torch namespace of the ops a module's ``fake_ops.py`` fakes. + + ``mv2__`` for a module of a target package, ``mv2_`` for + a module of the family itself (its shared ops). Target directories are named + ``____``, so the module path tells the two apart. + """ + return _namespace(*_split(module)) + + +# --- Naming (mirrors cpp/cmake/modules/modeling_v2.cmake) -------------------- + + +def _identifier(name: str) -> str: + """``name`` with ``__`` collapsed, usable as a C++ or torch identifier.""" + return re.sub("_{2,}", "_", name) + + +def _namespace(family: str, target: Optional[str] = None) -> str: + return f"mv2_{family}" if target is None else f"mv2_{family}_{_identifier(target)}" + + +def _library_name(family: str) -> str: + return f"libth_modeling_v2_{family}.so" + + +def _split(module: str) -> Tuple[str, Optional[str]]: + """``(family, target)`` of a module under ``models.``; ``target`` is None + for a module of the family itself.""" + prefix = f"{_MODELS_PACKAGE}." + parts = module[len(prefix) :].split(".") if module.startswith(prefix) else [] + if not parts or not parts[0]: + raise ValueError(f"{module!r} is not a module of a modeling_v2 family") + target = parts[1] if len(parts) > 1 and "__" in parts[1] else None + return parts[0], target + + +# --- Loading ------------------------------------------------------------------ + + +def _import_fakes(module: str, family: str) -> None: + """Import a ``fake_ops`` module, if there is one; it registers fakes of ops the + family library must define.""" + if importlib.util.find_spec(module) is None: + return + try: + importlib.import_module(module) + except (AttributeError, RuntimeError, ValueError) as error: + raise NativeLoadError( + f"{module} does not match the ops in {_library_name(family)}: {error}" + ) from None + + +def _open_library(path: Path, family: str) -> FrozenSet[str]: + """dlopen a family library and return the targets it carries.""" + import torch + + torch.ops.load_library(str(path)) + answer = getattr(torch.ops, _namespace(family)).built_targets() + return frozenset(t for t in answer.split(",") if t) + + +def _load_family_library(family: str) -> FrozenSet[str]: + """Open ``libth_modeling_v2_.so`` once and return the targets it carries.""" + if family in _built_targets: + return _built_targets[family] + path = _LIBS_DIR / _library_name(family) + if not path.is_file(): + raise NoTarget( + f"this installation has no {path.name}: the build compiles a family's library " + "only when its CUDA architectures include one of the family's targets" + ) + try: + built = _open_library(path, family) + except (AttributeError, OSError, RuntimeError) as error: + raise NativeLoadError(f"{path} is installed but cannot be loaded: {error}") from None + _import_fakes(f"{_MODELS_PACKAGE}.{family}.fake_ops", family) + _built_targets[family] = built + return built + + +def load_target_ops(package: str) -> Any: + """Load the native code of the target package ``package`` and return its ops + namespace, ``torch.ops.mv2__``. + + Call it from the target's ``__init__.py`` with ``__name__``. Raises + :class:`NoTarget` when this installation does not carry the target + and :class:`NativeLoadError` when its library is broken. + """ + family, target = _split(package) + if target is None or package != f"{_MODELS_PACKAGE}.{family}.{target}": + raise ValueError(f"{package!r} is not a models.. package") + built = _load_family_library(family) + if target not in built: + carried = ", ".join(sorted(built)) or "none" + raise NoTarget( + f"{_library_name(family)} does not carry {target}: this build did not compile " + f"it (its ARCHS miss the build's, a REQUIRES is unmet, or " + f"TRTLLM_MODELING_V2_TARGETS excluded it); it carries {carried}" + ) + _import_fakes(f"{package}.fake_ops", family) + import torch + + return getattr(torch.ops, _namespace(family, target)) + + +def load_family_ops(module: str) -> Any: + """Load the native code of the family ``module`` belongs to and return the ops of + its shared code, ``torch.ops.mv2_``. + + Call it with ``__name__`` from any module of the family, where the op is used; + after the first call it is a dictionary lookup. Raises like + :func:`load_target_ops`, whose target check the target's ``__init__.py`` runs + first. + """ + ops = _family_ops.get(module) + if ops is None: + family = _split(module)[0] + _load_family_library(family) + import torch + + ops = _family_ops[module] = getattr(torch.ops, _namespace(family)) + return ops diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index b9cd4f09fac1..3814683f7fd7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -205,3 +205,57 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation" ) return None + + +def _qualified(target: Optional[type]) -> Optional[str]: + return None if target is None else f"{target.__module__}.{target.__qualname__}" + + +def modeling_v2_resolve_on_all_ranks( + config: "ModelConfig", llm_args: "TorchLlmArgs" +) -> Optional[type]: + """:func:`modeling_v2_resolve`, agreed by every rank of the engine. + + Each rank routes on its own facts -- its device, and whether this + installation built the target's native code -- so one rank can land on the + built-in model while its peers build the target, and the first collective + then hangs instead of failing. Every rank therefore reports its outcome, and + all of them raise unless the outcomes agree. A rank's error is reported + after the exchange, not before, so its peers are never left waiting in it. + + ``off`` returns before any exchange: nothing in this package runs then. + """ + if ModelingV2Mode.of(llm_args) is ModelingV2Mode.OFF: + return None + + target = error = None + try: + target = modeling_v2_resolve(config, llm_args) + except Exception as caught: # noqa: BLE001 -- re-raised below, after every rank has reported + error = caught + outcome = (_qualified(target), None if error is None else f"{type(error).__name__}: {error}") + + mapping = config.mapping + outcomes = [outcome] + if mapping.world_size > 1: + from tensorrt_llm._torch.distributed.communicator import Distributed + + outcomes = Distributed.get(mapping).allgather(outcome) + + if error is not None: + raise error + failed = {rank: why for rank, (_, why) in enumerate(outcomes) if why is not None} + if failed: + detail = "; ".join(f"rank {rank}: {why}" for rank, why in sorted(failed.items())) + raise RuntimeError(f"modeling_v2: other ranks failed to resolve a target ({detail})") + if len({built for built, _ in outcomes}) > 1: + detail = "; ".join( + f"rank {rank}: {built or 'built-in'}" for rank, (built, _) in enumerate(outcomes) + ) + raise RuntimeError( + f"modeling_v2: ranks disagree on the implementation to build ({detail}); " + "every rank has to run the same code" + ) + if target is not None and mapping.rank == 0: + logger.warning(f"modeling_v2: building {_qualified(target)} on all {len(outcomes)} rank(s)") + return target diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/csrc/CMakeLists.txt b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/csrc/CMakeLists.txt new file mode 100644 index 000000000000..a700ed5efa4e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/csrc/CMakeLists.txt @@ -0,0 +1,25 @@ +# +# 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. +# + +# gpt_oss native code. Each target keeps its own in /csrc/; see the +# modeling_v2 README. +add_modeling_v2_family(gpt_oss) + +add_modeling_v2_target( + gpt_oss gpt_oss_120b__sm_10x__default + ARCHS 100f + SOURCES *.cu *.cpp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/__init__.py index 0b5adcfe48f8..81d9013e3e5e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/__init__.py @@ -1,3 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """gpt-oss-120b / sm_10x / default: every parallel layout the routing table admits.""" + +from tensorrt_llm._torch._experimental.modeling_v2._native import load_target_ops + +native_ops = load_target_ops(__name__) # torch.ops.mv2_gpt_oss_gpt_oss_120b_sm_10x_default diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.cu new file mode 100644 index 000000000000..4eaed982506a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.cu @@ -0,0 +1,67 @@ +/* + * 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. + */ + +#include "gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.h" + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +// One 32-bit word per thread: the four scales (m, c..c+3) of a 4-column group stay together in the swizzled +// layout, so each word moves whole. In words, one 128x4 tile is 128 words. +__global__ void swizzleBlockScalesKernel( + uint32_t const* __restrict__ in, uint32_t* __restrict__ out, int64_t rows, int64_t words, int64_t total) +{ + for (int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) + { + int64_t const word = i % words; + int64_t const m = (i / words) % rows; + int64_t const expert = i / (words * rows); + int64_t const dst + = expert * rows * words + (m / 128) * 128 * words + word * 128 + (m % 32) * 4 + (m % 128) / 32; + out[dst] = in[i]; + } +} + +} // namespace + +cudaError_t invokeSwizzleBlockScales( + uint8_t const* in, uint8_t* out, int64_t experts, int64_t rows, int64_t cols, cudaStream_t stream) +{ + int64_t const words = cols / 4; + int64_t const total = experts * rows * words; + if (total == 0) + { + return cudaSuccess; + } + constexpr int kThreads = 256; + int64_t const blocks = std::min((total + kThreads - 1) / kThreads, 65535); + swizzleBlockScalesKernel<<(blocks), kThreads, 0, stream>>>( + reinterpret_cast(in), reinterpret_cast(out), rows, words, total); + return cudaGetLastError(); +} + +} // namespace kernels + +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.h b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.h new file mode 100644 index 000000000000..1d4cd171e16a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.h @@ -0,0 +1,39 @@ +/* + * 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. + */ + +#pragma once + +#include "tensorrt_llm/modeling_v2/common/namespace.h" + +#include +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Rearranges MXFP4 block scales, `[experts, rows, cols]` E8M0 bytes, into the 128x4 swizzled layout the MoE +//! kernel reads: within an expert, the byte at `(m, c)` moves to offset +//! `(m / 128) * 512 * (cols / 4) + (c / 4) * 512 + (m % 32) * 16 + ((m % 128) / 32) * 4 + c % 4`. +//! Requires `rows % 128 == 0` and `cols % 4 == 0`. +cudaError_t invokeSwizzleBlockScales( + uint8_t const* in, uint8_t* out, int64_t experts, int64_t rows, int64_t cols, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_MV2_NAMESPACE_END diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScalesOp.cpp b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScalesOp.cpp new file mode 100644 index 000000000000..112198a5a529 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScalesOp.cpp @@ -0,0 +1,68 @@ +/* + * 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. + */ + +#include "gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/modeling_v2/common/namespace.h" +#include "tensorrt_llm/modeling_v2/common/torchLibrary.h" + +#include +#include +#include + +#include + +TRTLLM_MV2_NAMESPACE_BEGIN + +namespace +{ + +at::Tensor swizzleBlockScales(at::Tensor const& scales) +{ + TORCH_CHECK( + scales.is_cuda() && scales.is_contiguous(), "swizzle_block_scales: scales must be a contiguous CUDA tensor"); + TORCH_CHECK(scales.scalar_type() == at::kByte, "swizzle_block_scales: scales must be uint8 (E8M0), got ", + scales.scalar_type()); + TORCH_CHECK( + scales.dim() == 3, "swizzle_block_scales: scales must be [experts, rows, cols], got ", scales.dim(), " dims"); + int64_t const rows = scales.size(1); + int64_t const cols = scales.size(2); + TORCH_CHECK(rows % 128 == 0 && cols % 4 == 0, + "swizzle_block_scales: needs rows % 128 == 0 and cols % 4 == 0, got rows ", rows, ", cols ", cols); + TORCH_CHECK(reinterpret_cast(scales.data_ptr()) % 4 == 0, + "swizzle_block_scales: scales must start on a 4-byte boundary"); + c10::cuda::CUDAGuard const deviceGuard{scales.device()}; + at::Tensor out = at::empty_like(scales); + cudaError_t const status = kernels::invokeSwizzleBlockScales(scales.data_ptr(), out.data_ptr(), + scales.size(0), rows, cols, at::cuda::getCurrentCUDAStream(scales.get_device())); + TLLM_CHECK_WITH_INFO(status == cudaSuccess, "swizzle_block_scales launch failed: %s", cudaGetErrorString(status)); + return out; +} + +} // namespace + +TRTLLM_MV2_NAMESPACE_END + +TRTLLM_MV2_LIBRARY_FRAGMENT(m) +{ + m.def("swizzle_block_scales(Tensor scales) -> Tensor"); +} + +TRTLLM_MV2_LIBRARY_IMPL(CUDA, m) +{ + m.impl("swizzle_block_scales", &TRTLLM_MV2_NS::swizzleBlockScales); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/fake_ops.py new file mode 100644 index 000000000000..a71f9e242756 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/fake_ops.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fake (meta) implementations of this target's native ops, built from its csrc/. + +The loader imports this module right after the target is found in the family +library, so a fake for an op the library does not define fails the load. +""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2._native import namespace_of + +_NS = namespace_of(__name__) + + +@torch.library.register_fake(f"{_NS}::swizzle_block_scales") +def _(scales): + return torch.empty_like(scales) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/weights.py index 190cd8faad3c..48ed6f1a1813 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/weights.py @@ -49,6 +49,8 @@ import torch import torch.nn as nn +from . import native_ops + #: `layers=ALL_LAYERS` is the common case and the default. A sentinel rather #: than a lambda so the table reads as data, and rather than `True` so the #: field has one type of answer: which layers. @@ -184,7 +186,11 @@ def _fc1_perm(rows: int, device) -> torch.Tensor: def _swizzle_scales(s: torch.Tensor) -> torch.Tensor: """Per-expert 128x4 block-scale swizzle, viewed back as `[E, M, C]`: the byte at `(m, c)` moves to flat offset `(m//128)*512*(C//4) + (c//4)*512 - + (m%32)*16 + ((m%128)//32)*4 + (c%4)`.""" + + (m%32)*16 + ((m%128)//32)*4 + (c%4)`. On the device this target's own + kernel (csrc/swizzleBlockScales.cu) moves it in one pass; on the host, + a torch permutation does.""" + if s.is_cuda: + return native_ops.swizzle_block_scales(s) e, m, c = s.shape assert m % 128 == 0 and c % 4 == 0, (m, c) v = s.reshape(e, m // 128, 4, 32, c // 4, 4) # (e, m/128, (m%128)/32, m%32, c/4, c%4) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 67a3fbba4a6f..f967e92d7069 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -43,7 +43,7 @@ from tensorrt_llm.quantization.utils.fp4_utils import float4_e2m1x2 from ...llmapi.llm_args import LoadFormat -from .._experimental.modeling_v2 import modeling_v2_resolve +from .._experimental.modeling_v2 import modeling_v2_resolve_on_all_ranks from ..locality_domain.policy import LocalityDomainPolicy from ..model_config import ModelConfig from ..models import AutoModelForCausalLM @@ -2096,7 +2096,7 @@ def _load_and_validate_config( # modeling_v2 target reads the arguments as the engine will run with # them and never overrides those defaults. config = replace(config, - modeling_v2_target=modeling_v2_resolve( + modeling_v2_target=modeling_v2_resolve_on_all_ranks( config, self.llm_args)) if uses_mtp_head_checkpoint(self.spec_config): diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 8ec539b91268..c832c16e75af 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -135,34 +135,9 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - # modeling_v2 unit tests, one entry per catalog op test and one per model - # family that has unit tests, so CBTS can select the ones a change reaches. - # The whole-tree tests that need no GPU are in l0_cpu; the comm tests in - # l0_gb300_multi_gpus. - - unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py - - unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py - - unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py - - unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py - - unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py - - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py - - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py - - unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py - - unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py - - unittest/_torch/modeling_v2/models/deepseek_v3 - - unittest/_torch/modeling_v2/models/qwen4_exp - - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_weight_table.py - - unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py - - unittest/_torch/modeling_v2/test_modeling_v2_no_stale_claims.py - - unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py - - unittest/_torch/modeling_v2/test_modeling_v2_weight_table.py + # 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) - 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 b14442eab318..5b3ec00ab910 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -69,6 +69,9 @@ l0_cpu: # derivation and attention-call arguments, its routed-expert scalars and its weight table. # The rest of unittest/_torch/modeling_v2 needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py + - unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py + - unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py + - unittest/_torch/modeling_v2/test_modeling_v2_rank_agreement.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.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/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/test_modeling_v2_gpt_oss_swizzle_block_scales.py b/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/test_modeling_v2_gpt_oss_swizzle_block_scales.py new file mode 100644 index 000000000000..ab41f587f048 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/test_modeling_v2_gpt_oss_swizzle_block_scales.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""gpt-oss's own kernel: the MXFP4 block-scale swizzle of its weight loader. + +Importing the target package loads ``libth_modeling_v2_gpt_oss.so`` through +``_native.load_target_ops``, as routing does: a pass here means the family library +was built for this GPU, packaged, opened, carries the target, and registered the +fake. The kernel only moves bytes, so it must match the torch swizzle the loader +runs on the host bit for bit. +""" + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss.gpt_oss_120b__sm_10x__default import ( # noqa: E501 + native_ops, + weights, +) + +assert torch.cuda.is_available(), "the gpt_oss family library needs a CUDA device" +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability()[0] != 10, reason="built for sm_100f only" +) + + +def test_swizzle_block_scales(): + g = torch.Generator(device="cuda").manual_seed(0) + # gpt-oss-120b's FC1 and FC2 scales (FC2's 90 scales per row padded to 96), and no experts: + # the kernel, on the device, against the torch permutation the loader runs on the host. + for shape in ((4, 5888, 92), (3, 384, 96), (0, 128, 4)): + s = torch.randint(0, 256, shape, generator=g, device="cuda", dtype=torch.uint8) + assert torch.equal(weights._swizzle_scales(s).cpu(), weights._swizzle_scales(s.cpu())), ( + shape + ) + + def scales(rows, cols): + return torch.zeros(2, rows, cols, device="cuda", dtype=torch.uint8) + + for bad, match in ( + (scales(127, 4), "got rows 127, cols 4"), + (scales(128, 6), "got rows 128, cols 6"), + (scales(128, 8)[:, :, :4], "contiguous"), + ): + with pytest.raises(RuntimeError, match=match): + native_ops.swizzle_block_scales(bad) + + with FakeTensorMode(): + out = native_ops.swizzle_block_scales(scales(256, 8)) + assert out.shape == (2, 256, 8) and out.dtype == torch.uint8 diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index b0f0d86d4c06..5f79f65e7115 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -121,9 +121,11 @@ def test_every_target_ships_both_products(): def test_targets_do_not_share_files(): """Isolation is the property being demonstrated; measure it. - Targets are allowed to import the catalog and nothing else of each - other's. A shared helper between two targets is the first step back to - the abstraction this package exists to avoid. + Targets are allowed to import the catalog, and the loader of their own + native code (``_native``), and nothing else of each other's. A shared + helper between two targets is the first step back to the abstraction this + package exists to avoid. The native half of the same rule is + ``test_modeling_v2_native_claims.py``. """ for arch in _ARCHS: for target_dir, name in _routed_targets(arch): @@ -133,10 +135,10 @@ def test_targets_do_not_share_files(): dots, tail = match.group(1), match.group(2) if len(dots) == 1: continue # sibling within the target - assert tail.startswith("catalog"), ( + assert tail.startswith(("catalog", "_native")), ( f"{name}: {source_file.name} reaches outside its own " f"directory for {tail!r}; targets may import the " - f"catalog and nothing else" + f"catalog and _native and nothing else" ) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py index 226dbc8232b2..0529d8b775ef 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py @@ -31,16 +31,22 @@ import pytest import torch +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss import routing -from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss.gpt_oss_120b__sm_10x__default import ( # noqa: E501 - weights as D, -) from tensorrt_llm._torch.moe.fused_moe.quantization import ( _get_weight_alignment, _requires_rank_local_tp_padding, ) from tensorrt_llm.mapping import Mapping +# Importing the target loads its family library, as routing does. +try: + from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss.gpt_oss_120b__sm_10x__default import ( # noqa: E501 + weights as D, + ) +except NoTarget as no_target: + pytest.skip(str(no_target), allow_module_level=True) + pytestmark = pytest.mark.cpu_only _INTER = 2880 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 new file mode 100644 index 000000000000..86c8f5ee16c2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py @@ -0,0 +1,390 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The boundary between stock code and modeling_v2 native code, read off disk. + +modeling_v2 native code lives with the Python code of its family, in +``models//csrc/`` and ``models///csrc/``, and builds one +library per family that loads only when routing picks one of the family's +targets (``_native``). Three things rest on that staying true: ``off`` loading +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; 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. +""" + +from __future__ import annotations + +import ast +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2 import _native + +_REPO = Path(__file__).resolve().parents[4] +# What every family shares: the discovery and the common/ macros. +_INFRA = _REPO / "cpp/tensorrt_llm/modeling_v2" +_PY = _REPO / "tensorrt_llm/_torch/_experimental/modeling_v2" +_MODELS = _PY / "models" +_CMAKE_MODULE = _REPO / "cpp/cmake/modules/modeling_v2.cmake" +_TEST_DB = _REPO / "tests/integration/test_lists/test-db" + +pytestmark = [ + pytest.mark.cpu_only, + pytest.mark.skipif(not _INFRA.is_dir(), reason="needs the source checkout"), +] + +_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))} + + +def arch_covers(arch: str, capability: tuple[int, int]) -> bool: + """Whether code built for ``arch`` (``NN``, ``NNa`` or ``NNf``) runs on ``capability``. + + ``NNa`` runs on exactly NN; ``NN`` on NN and later minor versions of the same + major version; ``NNf`` on every member of family NN, the same way. + """ + suffix = arch[-1] if arch[-1] in "af" else "" + major, minor = divmod(int(arch[:-1] if suffix else arch), 10) + if suffix == "a": + return capability == (major, minor) + return capability[0] == major and capability[1] >= minor + + +# --- Reading the CMake declarations ------------------------------------------ + + +def _calls(text: str, function: str) -> list[list[str]]: + """Argument tokens of every ``function(...)`` call in a CMakeLists.txt.""" + text = re.sub(r"#.*", "", text) + return [m.group(1).split() for m in re.finditer(rf"\b{function}\s*\(([^)]*)\)", text)] + + +def _keyword_values(tokens: list[str], keyword: str) -> list[str]: + keywords = {"ARCHS", "SOURCES", "DEPS", "REQUIRES"} + values, active = [], False + for token in tokens: + if token in keywords: + active = token == keyword + elif active: + values.append(token) + return values + + +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")): + 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" + ) + families[calls[0][0]] = cmake.parent + return families + + +def _targets() -> list[tuple[str, str, list[str]]]: + """(family, target, tokens) of every add_modeling_v2_target call.""" + out = [] + for csrc in _families().values(): + for tokens in _calls((csrc / "CMakeLists.txt").read_text(), "add_modeling_v2_target"): + out.append((tokens[0], tokens[1], tokens[2:])) + return out + + +def _units() -> list[tuple[str, str | None, Path]]: + """(family, target or None for shared code, csrc directory) of every unit with code.""" + units = [(family, None, csrc) for family, csrc in _families().items()] + for family, target, _ in _targets(): + csrc = _MODELS / family / target / "csrc" + if csrc.is_dir(): + units.append((family, target, csrc)) + return units + + +def _native_files(root: Path) -> list[Path]: + return sorted(p for p in root.rglob("*") if p.is_file() and p.suffix in _NATIVE_SUFFIXES) + + +# --- Layout ------------------------------------------------------------------ + + +def test_native_code_is_where_the_build_and_the_loader_find_it(): + """Native code and CMake live only in a family's or target's csrc/, every + native family and target is a Python one, a target with native code loads it, + and its ARCHS cover the architecture its name promises.""" + for family, csrc in _families().items(): + assert csrc.parent.name == family and (csrc.parent / "__init__.py").is_file(), ( + f"{csrc.relative_to(_REPO)}/CMakeLists.txt declares family {family!r}; " + f"it belongs in models/{family}/csrc/ of the Python family {family}" + ) + + misplaced = [] + for path in sorted(_MODELS.rglob("*")): + if not path.is_file() or not ( + path.suffix in _NATIVE_SUFFIXES or path.name == "CMakeLists.txt" + ): + continue + parts = path.relative_to(_MODELS).parts + if path.name == "CMakeLists.txt": + ok = len(parts) == 3 and parts[1] == "csrc" + else: + ok = (len(parts) >= 3 and parts[1] == "csrc") or ( + len(parts) >= 4 and parts[2] == "csrc" + ) + if not ok: + misplaced.append(str(path.relative_to(_REPO))) + assert not misplaced, ( + "native code outside models//csrc/ or models///csrc/ " + f"(CMakeLists.txt only in models//csrc/): {misplaced}" + ) + + declared, with_code = set(), set() + for family, target, tokens in _targets(): + declared.add((family, target)) + package = _MODELS / family / target / "__init__.py" + assert package.is_file() and "load_target_ops(__name__)" in package.read_text(), ( + f"native target {family}/{target}: its Python package must start with " + "native_ops = load_target_ops(__name__)" + ) + if _keyword_values(tokens, "SOURCES"): + with_code.add((family, target)) + archs = _keyword_values(tokens, "ARCHS") + sm = target.split("__")[1][len("sm_") :] + assert sm in _SM_GROUPS or sm.isdigit(), f"{family}/{target}: unknown sm segment sm_{sm}" + for capability in _SM_GROUPS[sm] if sm in _SM_GROUPS else (divmod(int(sm), 10),): + assert any(arch_covers(a, capability) for a in archs), ( + f"{family}/{target} declares ARCHS {archs}, none of which runs on " + f"sm_{capability[0]}{capability[1]}, the architecture in its name" + ) + for csrc in sorted(_MODELS.glob("*/*/csrc")): + unit = (csrc.parent.parent.name, csrc.parent.name) + assert unit in with_code, f"{csrc.relative_to(_REPO)} holds code no target's SOURCES builds" + for package in sorted(_MODELS.glob("*/*/__init__.py")): + unit = (package.parent.parent.name, package.parent.name) + assert "load_target_ops(" not in package.read_text() or unit in declared, ( + f"{package.relative_to(_REPO)} loads native code no add_modeling_v2_target declares" + ) + + +# --- What native code may reach ----------------------------------------------- + +_INCLUDE = re.compile(r'^\s*#\s*include\s*"([^"]+)"', re.MULTILINE) +_RAW_REGISTRATION = re.compile(r"\bTORCH_LIBRARY(?:_FRAGMENT|_IMPL)?\s*\(") +_OP_DEF = re.compile(r'\bm\.def\(\s*"(\w+)\(') + + +def test_native_code_stays_in_its_own_unit(): + """Native code may include stock headers, but of modeling_v2 only its own + target's csrc/, its family's csrc/ and common/; it registers through the + TRTLLM_MV2_ macros, into its own namespace only; every op it defines has a + fake in the unit's fake_ops.py (which the loader imports); and every file + carries NVIDIA's Apache-2.0 header, which the C++ license checker only checks + under cpp/.""" + python_families = {p.name for p in _MODELS.iterdir() if (p / "__init__.py").is_file()} + problems = [] + for family, target, csrc in _units(): + allowed = [f"{family}/csrc/", "tensorrt_llm/modeling_v2/common/"] + if target is not None: + allowed.append(f"{family}/{target}/csrc/") + ops = set() + for path in [*_native_files(csrc), *csrc.glob("CMakeLists.txt")]: + text = path.read_text(errors="ignore") + where = str(path.relative_to(_REPO)) + head = text[:1200] + apache = ( + "SPDX-License-Identifier: Apache-2.0" in head + or "Apache License, Version 2.0" in head + ) + if "NVIDIA CORPORATION" not in head or not apache: + problems.append(f"{where}: no NVIDIA Apache-2.0 license header") + if path.name == "CMakeLists.txt": + continue + for include in _INCLUDE.findall(text): + head = include.split("/", 1)[0] + if head in python_families or include.startswith("tensorrt_llm/modeling_v2/"): + if not include.startswith(tuple(allowed)): + problems.append(f'{where}: "{include}" reaches another target or family') + elif ".." in include or "/" not in include and (csrc / include).exists(): + problems.append( + f'{where}: "{include}": include it from models/, "{family}/..."' + ) + if _RAW_REGISTRATION.search(text): + problems.append(f"{where}: use the TRTLLM_MV2_LIBRARY_* macros, not TORCH_LIBRARY*") + ops.update(_OP_DEF.findall(text)) + if ops: + python_dir = _MODELS / family if target is None else _MODELS / family / target + fake_ops = python_dir / "fake_ops.py" + text = fake_ops.read_text() if fake_ops.is_file() else "" + missing = sorted(op for op in ops if f'::{op}"' not in text) + if missing: + problems.append(f"{fake_ops.relative_to(_REPO)}: no fake for {missing}") + assert not problems, "\n".join(problems) + + +# --- What stock code may reach ------------------------------------------------ + +_STOCK_PY_REFERENCE = re.compile( + r"modeling_v2[./]models\b|torch\.ops\.mv2_|register_fake\(\s*f?[\"']mv2_" +) +_STOCK_CPP_REFERENCE = re.compile(r'#\s*include\s*"tensorrt_llm/modeling_v2/') +# A path into one family's directory; the discovery's models/*/csrc glob names none. +_STOCK_FAMILY_PATH = re.compile(r"_experimental/modeling_v2/models/[a-z]") + + +def test_stock_does_not_reach_modeling_v2(): + """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(path.read_text(errors="ignore")) + ] + for root in ("cpp/tensorrt_llm", "cpp/include", "cpp/tests", "cpp/cmake"): + for path in sorted((_REPO / root).rglob("*")): + if not path.is_file() or path.is_relative_to(_INFRA): + continue + # Skip generated cubin sources; they include nothing of interest. + if path.suffix in _NATIVE_SUFFIXES and path.stat().st_size < (1 << 20): + pattern = _STOCK_CPP_REFERENCE + elif path.name == "CMakeLists.txt" or path.suffix == ".cmake": + pattern = _STOCK_FAMILY_PATH + else: + continue + 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((_REPO / name).read_text()): + offenders.append(name) + assert not offenders, f"stock code reaches modeling_v2 targets or families: {offenders}" + + +# --- Targets load only when routed ---------------------------------------------- + + +def _is_docstring_or_future(node: ast.stmt) -> bool: + docstring = isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant) + future = isinstance(node, ast.ImportFrom) and node.module == "__future__" + return docstring or future + + +def test_native_code_loads_only_when_routing_picks_its_target(): + """``models/__init__.py`` imports only the families' routing modules, a + family package holds a docstring only, and a routing module imports a target + only inside ``route``; anything else would load native code on every resolve. + A target calls exactly the ops it declares: choosing a path by whether an op + exists (hasattr, or getattr with a default) would mean ``require`` no longer + promises the gated code ran.""" + models_init = _MODELS / "__init__.py" + for node in ast.parse(models_init.read_text()).body: + routing_import = ( + isinstance(node, ast.ImportFrom) + and node.level == 1 + and node.module is not None + and "." not in node.module + and [alias.name for alias in node.names] == ["routing"] + ) + assert _is_docstring_or_future(node) or routing_import, ( + f"{models_init.relative_to(_REPO)}:{node.lineno} -- imports something besides " + "the families' routing modules" + ) + for init in sorted(_MODELS.glob("*/__init__.py")): + for node in ast.parse(init.read_text()).body: + assert _is_docstring_or_future(node), ( + f"{init.relative_to(_REPO)}:{node.lineno} -- a family package holds a docstring only" + ) + 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), ( + f"{routing.relative_to(_REPO)}:{node.lineno} -- imports .{node.module} at " + "module level; a routing module imports a target only inside route" + ) + probe = re.compile(r"\bhasattr\(\s*torch\.ops\b|\bgetattr\(\s*torch\.ops[\w.]*\s*,[^,()]+,") + probing = [ + str(p.relative_to(_REPO)) for p in sorted(_PY.rglob("*.py")) if probe.search(p.read_text()) + ] + assert not probing, f"op probing in {probing}" + + +# --- Test lists ---------------------------------------------------------------- + +_MV2_ENTRY = re.compile(r"^\s*-\s+((?:unittest/_torch/modeling_v2|accuracy/modeling_v2)\S*)(.*)$") + + +def _runs(target: str, options: str, unit: str) -> bool: + """Whether a test-db entry runs tests under the directory ``unit``: it names + the directory or something in it, or a directory above it that it does not + ``--ignore``.""" + if target == unit or target.startswith(unit + "/"): + return True + if not unit.startswith(target.rstrip("/") + "/"): + return False + ignored = (path.rstrip("/") for path in re.findall(r"--ignore=(\S+)", options)) + return not any(unit == path or unit.startswith(path + "/") for path in ignored) + + +def test_test_lists_run_every_native_family(): + """Some test-db entry runs each native family's tests, and accuracy files are + named after a family, since CBTS selects them by family name.""" + entries = [ + (yml.name, *match.groups()) + for yml in sorted(_TEST_DB.glob("*.yml")) + for match in map(_MV2_ENTRY.match, yml.read_text().splitlines()) + if match + ] + for family in _families(): + unit = f"unittest/_torch/modeling_v2/models/{family}" + accuracy = f"accuracy/modeling_v2/test_{family}.py" + has_unit_tests = any((_REPO / "tests" / unit).rglob("test_*.py")) + assert any( + target.startswith(accuracy) or (has_unit_tests and _runs(target, options, unit)) + for _, target, options in entries + ), f"family {family!r} has native code but no test-db entry runs {unit} or {accuracy}" + families = {p.name for p in _MODELS.iterdir() if (p / "__init__.py").is_file()} + for yml, target, _ in entries: + match = re.match(r"accuracy/modeling_v2/test_(\w+)\.py", target) + assert not match or match.group(1) in families, ( + f"{yml}: {target} is named after no modeling_v2 family" + ) + + +# --- CMake and Python agree ------------------------------------------------------ + + +@pytest.mark.skipif(shutil.which("cmake") is None, reason="needs cmake") +def test_cmake_and_python_agree_on_names_and_architectures(tmp_path): + """The loader and these checks compute identifiers, namespaces and + architecture coverage in Python; the build computes them in modeling_v2.cmake.""" + names = ["gpt_oss", "minimax_m3", "r1_0528_nvfp4__sm_10x__default"] + script = tmp_path / "probe.cmake" + lines = [f'include("{_CMAKE_MODULE}")'] + for name in names: + lines += [ + f'_modeling_v2_ident(i "{name}")', + f'message(STATUS "name {name} ${{i}}")', + ] + lines += [ + "_modeling_v2_resolve_archs(r DECLARED 80 90a 100f 103a 120f REQUESTED 86 90 103 107)", + 'message(STATUS "archs ${r}")', + ] + script.write_text("\n".join(lines) + "\n") + output = subprocess.run( + ["cmake", "-P", str(script)], capture_output=True, text=True, check=True + ).stdout + for name in names: + assert f"name {name} {_native._identifier(name)}" in output, output + assert "archs 90a-real;100f-real;103a-real" in output, output + # The family namespace the build registers is the loader's. + assert _native._namespace("gpt_oss", "gpt_oss_120b__sm_10x__default") == ( + "mv2_gpt_oss_gpt_oss_120b_sm_10x_default" + ) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py new file mode 100644 index 000000000000..eb969fa92536 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What ``_native`` decides, and that the installed family libraries load. + +The first tests drive the loader's decisions without opening anything: which +outcomes are ``NoTarget`` (this installation does not carry the target, so +``auto`` keeps the built-in model) and which are ``NativeLoadError`` (an +installed library is broken, which is never a fallback). + +The last is the load smoke test of the installed wheel: every family library in +``tensorrt_llm/libs`` opens in a fresh process, names the targets it carries, and +defines no op the target's fake_ops.py does not fake. A library aborts its +process from a static initializer if two libraries claim one torch namespace, so +each one is opened in a child process. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +import textwrap +import types +from pathlib import Path + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2 import _native +from tensorrt_llm._torch._experimental.modeling_v2._native import ( + NativeLoadError, + load_family_ops, + load_target_ops, + namespace_of, +) +from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + +pytestmark = pytest.mark.cpu_only + +_FAMILY = "tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss" +_TARGET = f"{_FAMILY}.gpt_oss_120b__sm_10x__default" +_TARGET_NS = "mv2_gpt_oss_gpt_oss_120b_sm_10x_default" + + +@pytest.fixture +def libs(monkeypatch, tmp_path): + """An empty libs directory, a cold library cache, and ``libs.install(opens)``, + which puts a gpt_oss library there; ``opens`` stands in for dlopen, and no + fake_ops module is imported, since the stand-in defines no op.""" + monkeypatch.setattr(_native, "_LIBS_DIR", tmp_path) + monkeypatch.setattr(_native, "_built_targets", {}) + monkeypatch.setattr(_native, "_family_ops", {}) + + def install(opens): + (tmp_path / "libth_modeling_v2_gpt_oss.so").write_bytes(b"") + monkeypatch.setattr(_native, "_open_library", lambda path, family: opens()) + monkeypatch.setattr(_native, "_import_fakes", lambda module, family: None) + + return types.SimpleNamespace(install=install) + + +def test_names(): + assert namespace_of(_TARGET) == _TARGET_NS + assert namespace_of(f"{_TARGET}.ops.moe.entry") == _TARGET_NS + assert namespace_of(f"{_FAMILY}.fake_ops") == "mv2_gpt_oss" + with pytest.raises(ValueError): + namespace_of("tensorrt_llm._torch.models") + with pytest.raises(ValueError, match="models.."): + load_target_ops(f"{_TARGET}.weights") + + +def test_what_is_not_built_and_what_is_broken(libs): + with pytest.raises(NoTarget, match="has no libth_modeling_v2_gpt_oss.so"): + load_target_ops(_TARGET) + + def fail(): + raise OSError("undefined symbol: _ZN12tensorrt_llm6common") + + libs.install(fail) + with pytest.raises(NativeLoadError, match="undefined symbol"): + load_target_ops(_TARGET) + + libs.install(lambda: frozenset({"other__sm_90__default"})) + with pytest.raises(NoTarget, match="does not carry gpt_oss_120b__sm_10x__default"): + load_target_ops(_TARGET) + + +def test_a_library_opens_once_and_registers_the_family_fakes_first(libs, monkeypatch): + ops = types.SimpleNamespace(mv2_gpt_oss=object(), **{_TARGET_NS: object()}) + monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(ops=ops)) + opened, imported = [], [] + + def opens(): + opened.append(1) + return frozenset({"gpt_oss_120b__sm_10x__default"}) + + libs.install(opens) + monkeypatch.setattr(_native, "_import_fakes", lambda module, family: imported.append(module)) + assert load_target_ops(_TARGET) is getattr(ops, _TARGET_NS) + assert load_family_ops(f"{_TARGET}.weights") is ops.mv2_gpt_oss + assert load_family_ops(f"{_TARGET}.weights") is ops.mv2_gpt_oss + assert len(opened) == 1 + assert imported == [f"{_FAMILY}.fake_ops", f"{_TARGET}.fake_ops"] + + +_LOAD_ONE_FAMILY = textwrap.dedent( + """ + import json, sys + import torch + import tensorrt_llm # noqa: F401 -- loads libth_common.so + from tensorrt_llm._torch._experimental.modeling_v2 import _native + + built = sorted(_native._load_family_library(sys.argv[1])) + schemas = sorted(s.name for s in torch._C._jit_get_all_schemas() if s.name.startswith("mv2_")) + print(json.dumps({"built": built, "ops": schemas})) + """ +) + + +def test_every_installed_family_library_loads(): + libraries = sorted(_native._LIBS_DIR.glob("libth_modeling_v2_*.so")) + if not libraries: + pytest.skip("this installation has no modeling_v2 family library") + models = Path(_native.__file__).resolve().parent / "models" + for library in libraries: + family = library.name[len("libth_modeling_v2_") : -len(".so")] + result = subprocess.run( + [sys.executable, "-c", _LOAD_ONE_FAMILY, family], + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"{library.name} failed to load:\n{result.stdout}\n{result.stderr}" + ) + report = json.loads(result.stdout.strip().splitlines()[-1]) + assert report["built"], f"{library.name} carries no target" + for target in report["built"]: + assert (models / family / target / "__init__.py").is_file(), ( + f"{library.name} carries {target}, which has no Python target package" + ) + prefix = f"{_native._namespace(family, target)}::" + target_ops = [op[len(prefix) :] for op in report["ops"] if op.startswith(prefix)] + fake_ops = models / family / target / "fake_ops.py" + text = fake_ops.read_text() if fake_ops.is_file() else "" + unfaked = [op for op in target_ops if f'::{op}"' not in text] + assert not unfaked, f"{family}/{target} defines {unfaked} with no fake in fake_ops.py" diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_rank_agreement.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_rank_agreement.py new file mode 100644 index 000000000000..290ea9a8e1ab --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_rank_agreement.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``modeling_v2_resolve_on_all_ranks``: every rank builds the same code, or every rank fails. + +Each rank routes on its own facts -- its device, whether this installation built +the target's native code -- so ranks can disagree, and a disagreement surfaces +as a hang in the first collective unless the resolver catches it. These tests +stand in for the other ranks with a fake ``Distributed`` whose ``allgather`` +returns their outcomes. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2 import _router_index +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + modeling_v2_resolve_on_all_ranks, +) +from tensorrt_llm._torch.distributed import communicator + +pytestmark = pytest.mark.cpu_only + + +class Target: + """Stands in for a target class.""" + + +_TARGET = f"{Target.__module__}.{Target.__qualname__}" + + +@pytest.fixture +def ranks(monkeypatch): + """``ranks(mode, mine, others, world_size=2)``: resolve with this rank's own + decision ``mine`` (a target, None or an exception) and the other ranks' + outcomes ``others``; returns the result and the number of exchanges.""" + calls = [] + + def resolve_on_all_ranks(mode, mine, others, world_size=2): + def resolve(config, llm_args): + if isinstance(mine, Exception): + raise mine + return mine + + def allgather(outcome): + calls.append(outcome) + return [outcome, *others] + + monkeypatch.setattr(_router_index, "modeling_v2_resolve", resolve) + monkeypatch.setattr( + communicator.Distributed, + "get", + staticmethod(lambda mapping: SimpleNamespace(allgather=allgather)), + ) + calls.clear() + config = SimpleNamespace(mapping=SimpleNamespace(rank=0, world_size=world_size)) + return modeling_v2_resolve_on_all_ranks(config, SimpleNamespace(modeling_v2=mode)) + + resolve_on_all_ranks.calls = calls + return resolve_on_all_ranks + + +def test_off_and_one_rank_exchange_nothing(ranks): + assert ranks("off", AssertionError("routing ran under off"), []) is None + assert ranks.calls == [] + assert ranks("auto", Target, [], world_size=1) is Target + assert ranks.calls == [] + + +def test_ranks_build_the_same_code_or_all_fail(ranks): + assert ranks("require", Target, [(_TARGET, None)]) is Target + assert ranks("auto", None, [(None, None)]) is None + with pytest.raises(RuntimeError, match="ranks disagree.*rank 1: built-in"): + ranks("auto", Target, [(None, None)]) + with pytest.raises(RuntimeError, match="rank 1: NativeLoadError"): + ranks("auto", Target, [(None, "NativeLoadError: libth_modeling_v2_x.so cannot be loaded")]) + # This rank's own error is raised after the exchange, not before it, which + # would leave the other ranks waiting in it. + with pytest.raises(ValueError, match="no target fits"): + ranks("require", ValueError("modeling_v2='require', but no target fits"), [(_TARGET, None)]) + assert len(ranks.calls) == 1 diff --git a/tests/unittest/scripts/cbts/test_modeling_v2_rule.py b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py new file mode 100644 index 000000000000..88026448641b --- /dev/null +++ b/tests/unittest/scripts/cbts/test_modeling_v2_rule.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +# 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. +"""Tests for CBTS ModelingV2Rule: per-family selection, native code included.""" + +from __future__ import annotations + +import pytest + +__extra_import_path__ = ["~/jenkins/scripts/cbts"] +from cbts.blocks import Block, Stage, YAMLIndex +from cbts.rules.base import PRInputs +from cbts.rules.modeling_v2 import ModelingV2Rule, _is_mv2_claim, _kept_entries + +pytestmark = pytest.mark.cpu_only + +_PY = "tensorrt_llm/_torch/_experimental/modeling_v2" +_CPP = "cpp/tensorrt_llm/modeling_v2" +_UNIT = "unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30)" +_COMM = "unittest/_torch/modeling_v2/comm" +_KIMI_UNIT = "unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_kimi_k3_native.py" +_DEEPSEEK_ACCURACY = "accuracy/modeling_v2/test_deepseek_v3.py::TestDeepseek::test_gsm8k[tp4]" +_GPT_OSS_ACCURACY = "accuracy/modeling_v2/test_gpt_oss.py::TestGptOss::test_gsm8k[tp1]" +_CLAIMS = "unittest/_torch/modeling_v2/test_modeling_v2_claims.py" + + +@pytest.fixture +def rule(tmp_path) -> ModelingV2Rule: + yaml_index = YAMLIndex() + yaml_index.blocks = [ + Block( + yaml_stem="l0_b300", + block_index=0, + condition={}, + tests=[ + "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4", + _UNIT, + _GPT_OSS_ACCURACY, + ], + ), + Block( + yaml_stem="l0_gb300_multi_gpus", + block_index=0, + condition={}, + tests=[_COMM, _DEEPSEEK_ACCURACY], + ), + Block(yaml_stem="l0_b200", block_index=0, condition={}, tests=[_KIMI_UNIT]), + Block(yaml_stem="l0_cpu", block_index=0, condition={}, tests=["unittest/tools", _CLAIMS]), + ] + stages = { + "B300-PyTorch-1": Stage("B300-PyTorch-1", "l0_b300", "aarch64", 1, 2), + "GB300-4_GPUs-PyTorch-1": Stage( + "GB300-4_GPUs-PyTorch-1", "l0_gb300_multi_gpus", "aarch64", 1, 1 + ), + "B200-PyTorch-1": Stage("B200-PyTorch-1", "l0_b200", "x86_64", 1, 1), + "CPU-Generic-x86-1": Stage("CPU-Generic-x86-1", "l0_cpu", "x86_64", 1, 1), + } + for family, target, cls in ( + ("gpt_oss", "gpt_oss_120b__sm_10x__default", "GptOss"), + ("deepseek_v3", "target", "Deepseek"), + ("kimi_k3_vl", "target", "Kimi"), + ): + root = tmp_path / _PY / "models" / family + (root / target).mkdir(parents=True) + (root / target / "modeling.py").touch() + (root / "routing.py").write_text(f"from .{target}.modeling import {cls}\n") + return ModelingV2Rule(yaml_index, stages, repo_root=tmp_path) + + +def _apply(rule: ModelingV2Rule, *files: str): + return rule.apply(PRInputs(changed_files=list(files), diffs={})) + + +def _selected(result) -> set[str]: + return { + entry + for prefixes in result.block_filters.values() + for entries in prefixes.values() + for entry in entries + } + + +def test_claims_and_families(): + """Python and native modeling_v2 sources are claimed; docs and stock are not. + + A path under models// belongs to that family; anything else claimed is + shared by every family. + """ + for path in ( + f"{_PY}/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/swizzleBlockScales.cu", + f"{_PY}/models/gpt_oss/csrc/CMakeLists.txt", + f"{_CPP}/common/namespace.h", + "cpp/cmake/modules/modeling_v2.cmake", + ): + assert _is_mv2_claim(path), path + for path in ( + f"{_CPP}/README.md", + "cpp/tensorrt_llm/thop/attentionOp.cpp", + "tensorrt_llm/_torch/pyexecutor/model_loader.py", + ): + assert not _is_mv2_claim(path), path + + +def test_entry_runs_family(): + for entry, runs in ( + ("unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm", True), + ("unittest/_torch/modeling_v2/models/gpt_oss/test_x.py::test_y", True), + ("accuracy/modeling_v2/test_gpt_oss.py::TestGptOss", True), + ("unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/models", False), + ("unittest/_torch/modeling_v2/models/gpt_oss_mini", False), + ("accuracy/modeling_v2/test_deepseek_v3.py", False), + ): + prefixes = { + "unittest/_torch/modeling_v2/models/gpt_oss", + "accuracy/modeling_v2/test_gpt_oss.py", + } + assert bool(_kept_entries([entry], prefixes)) is runs, entry + + +def test_scope(rule): + """A family change selects what runs that family's tests, plus the CPU checks. + + Families union; shared code selects every modeling_v2 entry; an unreadable family + selects all modeling_v2 entries. + """ + result = _apply(rule, f"{_PY}/models/gpt_oss/gpt_oss_120b__sm_10x__default/csrc/k.cu") + assert result.scope == "modelingv2only" and result.sanity_relevant is False + assert _selected(result) == {_UNIT, _GPT_OSS_ACCURACY, _CLAIMS} + result = _apply( + rule, f"{_PY}/models/deepseek_v3/routing.py", f"{_PY}/models/kimi_k3_vl/csrc/k.cu" + ) + assert _selected(result) == {_UNIT, _DEEPSEEK_ACCURACY, _KIMI_UNIT, _CLAIMS} + result = _apply(rule, f"{_CPP}/common/torchLibrary.h") + assert _selected(result) == { + _UNIT, + _COMM, + _KIMI_UNIT, + _DEEPSEEK_ACCURACY, + _GPT_OSS_ACCURACY, + _CLAIMS, + } + assert _apply(rule, "tensorrt_llm/_torch/pyexecutor/model_loader.py") is None + rule.yaml_index.blocks[0].tests.remove(_UNIT) + result = _apply(rule, f"{_PY}/models/minimax_m3/csrc/attnMega.cu") + assert result.scope == "modelingv2only" + assert _COMM in _selected(result) # unreadable family falls back to all modeling_v2 diff --git a/tests/unittest/scripts/test_build_wheel_modeling_v2.py b/tests/unittest/scripts/test_build_wheel_modeling_v2.py new file mode 100644 index 000000000000..9f416d1cfba5 --- /dev/null +++ b/tests/unittest/scripts/test_build_wheel_modeling_v2.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only check of how build_wheel.py installs the modeling_v2 family libraries.""" + +import importlib.util +import shutil +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + + +def _family(build_dir: Path, family: str, archs: str, library: bool) -> None: + family_dir = build_dir / family + family_dir.mkdir(parents=True) + (family_dir / "cuda_architectures.txt").write_text(archs) + if library: + (family_dir / f"libth_modeling_v2_{family}.so").write_bytes(b"\x7fELF") + + +def test_install_modeling_v2_libs(tmp_path: Path) -> None: + """A family library built for some architecture is installed. + + One this build did not produce is removed. + """ + script = Path(__file__).resolve().parents[3] / "scripts" / "build_wheel.py" + spec = importlib.util.spec_from_file_location("build_wheel", script) + build_wheel = importlib.util.module_from_spec(spec) + spec.loader.exec_module(build_wheel) + + build_dir, lib_dir = tmp_path / "build", tmp_path / "libs" + lib_dir.mkdir() + _family(build_dir, "gpt_oss", "100f-real", library=True) + _family(build_dir, "glm_moe_dsa", "", library=True) # left by an earlier build + (lib_dir / "libth_common.so").write_bytes(b"stock") + (lib_dir / "libth_modeling_v2_glm_moe_dsa.so").write_bytes(b"earlier build") + (lib_dir / "libth_modeling_v2_removed.so").symlink_to(lib_dir / "gone.so") + + build_wheel.install_modeling_v2_libs(build_dir, lib_dir, shutil.copy) + assert sorted(p.name for p in lib_dir.iterdir()) == [ + "libth_common.so", + "libth_modeling_v2_gpt_oss.so", + ] + + _family(build_dir, "qwen4_exp", "103a-real", library=False) + with pytest.raises(AssertionError, match="Exactly one modeling_v2 library"): + build_wheel.install_modeling_v2_libs(build_dir, lib_dir, shutil.copy)