diff --git a/CMakeLists.txt b/CMakeLists.txt index 1098879680..36208a3ea4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -39,6 +39,8 @@ find_library(TRTMC_TRT_LIBRARY REQUIRED ) +include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/EdgeLLM.cmake") + option(TRTMC_ENABLE_BYOK "Enable the optional TVM-FFI BYOK bridge" ON) set(TRTMC_HAS_TVM_FFI OFF) if(TRTMC_ENABLE_BYOK) diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index dbb5c93bb8..8a27e4292e 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -8,5 +8,12 @@ #include int main(int argc, char** argv) { - return trtmc::cli::run(argc, argv, std::cout, std::cerr); + // The executable owns the console: keep result output machine-readable even + // when loaded libraries write C++ diagnostics to std::cout. Do not change + // library logger levels or the output behavior of embedded runtime APIs. + std::ostream result(std::cout.rdbuf()); + std::cout.rdbuf(std::cerr.rdbuf()); + const int status = trtmc::cli::run(argc, argv, result, std::cerr); + result.flush(); + return status; } diff --git a/cmake/EdgeLLM.cmake b/cmake/EdgeLLM.cmake new file mode 100644 index 0000000000..678e33dced --- /dev/null +++ b/cmake/EdgeLLM.cmake @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Optional native dependency provisioning. Model builds never acquire dependencies. +option(TRTMC_ENABLE_EDGELLM "Install the pinned native Edge-LLM SDK and builder" OFF) +if(NOT TRTMC_ENABLE_EDGELLM) + return() +endif() +if(CMAKE_CROSSCOMPILING) + message(FATAL_ERROR "Edge-LLM cross compilation is not supported") +endif() +include("${CMAKE_CURRENT_LIST_DIR}/edgellm/CheckNative.cmake") +option(TRTMC_EDGELLM_ALL_KERNELS "Build all pinned Edge operator groups supported by the native GPU" OFF) +option(TRTMC_EDGELLM_ONNX "Install the pinned ONNX exporter and native engine builder" OFF) +set(_edge_cute_groups "fmha|gdn") +set(_edge_cute_cli_groups "fmha,gdn") +if(TRTMC_EDGELLM_ALL_KERNELS) + set(_edge_cute_groups ALL) + set(_edge_cute_cli_groups ALL) +endif() +set(_edge_build_targets edgellmCore NvInfer_edgellm_plugin) +set(_edge_onnx_byproducts "") +if(TRTMC_EDGELLM_ONNX) + list(APPEND _edge_build_targets llm_build) + list(APPEND _edge_onnx_byproducts "${CMAKE_BINARY_DIR}/_deps/edgellm/install/bin/edgellm-onnx-build") +endif() +set(_edge_version "0.10.1") +set(_edge_revision "e8b29522938901f6df19ebeedd4b69bc8edbcd97") +set(_edge_root "${CMAKE_BINARY_DIR}/_deps/edgellm") +set(_edge_prefix "${_edge_root}/install") +find_package(EdgeLLM ${_edge_version} EXACT CONFIG QUIET) +if(EdgeLLM_FOUND AND NOT EdgeLLM_PREFIX STREQUAL _edge_prefix) + _edgellm_json_include(_edge_json_include) + _edgellm_check_json_headers("${_edge_json_include}" "${EdgeLLM_PREFIX}/include/edgellm/3rdParty/nlohmannJson") + if(NOT EdgeLLM_REVISION STREQUAL _edge_revision) + message(FATAL_ERROR "EdgeLLM package does not match the pinned GitHub revision") + endif() + if(TRTMC_EDGELLM_ALL_KERNELS AND NOT EdgeLLM_ALL_KERNELS) + message(FATAL_ERROR "EdgeLLM package lacks requested full native operator coverage; rebuild with TRTMC_EDGELLM_ALL_KERNELS=ON") + endif() + if(TRTMC_EDGELLM_ONNX AND (NOT EdgeLLM_ONNX OR NOT EXISTS "${EdgeLLM_ONNX_BUILDER}")) + message(FATAL_ERROR "EdgeLLM package lacks requested ONNX tools; rebuild with TRTMC_EDGELLM_ONNX=ON") + endif() + install(FILES "$" DESTINATION "${CMAKE_INSTALL_LIBDIR}" COMPONENT EdgeLLM) + return() +endif() + +include(ExternalProject) +include(CMakePackageConfigHelpers) +find_package(Python3 3.10 REQUIRED COMPONENTS Interpreter) +find_package(Threads REQUIRED) +set(TRTMC_EDGELLM_TRT_ROOT "$ENV{TRT_ROOT}" CACHE PATH "Native TensorRT SDK, including its Python wheel") +set(TRTMC_EDGELLM_CUDA_ARCHITECTURE "${CMAKE_CUDA_ARCHITECTURES}" CACHE STRING "One local GPU architecture for Edge-LLM") +set(TRTMC_EDGELLM_JOBS 2 CACHE STRING "Parallel Edge-LLM native and AOT compilation jobs") +set(TRTMC_EDGELLM_WHEELHOUSE "" CACHE PATH "Optional complete offline Python wheelhouse") +set(TRTMC_EDGELLM_GIT_MIRROR "" CACHE PATH "Optional local mirror of the pinned upstream Git repository") +if(NOT TRTMC_EDGELLM_CUDA_ARCHITECTURE MATCHES "^[0-9]+$") + message(FATAL_ERROR "Set TRTMC_EDGELLM_CUDA_ARCHITECTURE to one local GPU architecture, e.g. 80") +endif() +if(NOT EXISTS "${TRTMC_EDGELLM_TRT_ROOT}/include/NvInfer.h") + message(FATAL_ERROR "TRTMC_EDGELLM_TRT_ROOT must contain the native TensorRT SDK") +endif() +_edgellm_check_gpu("${TRTMC_EDGELLM_CUDA_ARCHITECTURE}") +_edgellm_trt_version("${TRTMC_EDGELLM_TRT_ROOT}/include" _edge_trt_version) +set(_edge_source "${_edge_root}/source") +set(_edge_build "${_edge_root}/build") +set(_edge_python "${_edge_prefix}/libexec/trtmc-edge-llm/bin/python") +set(_edge_repository "https://github.com/NVIDIA/TensorRT-Edge-LLM.git") +if(TRTMC_EDGELLM_GIT_MIRROR) + set(_edge_repository "${TRTMC_EDGELLM_GIT_MIRROR}") +endif() +set(_edge_template_dir "${CMAKE_CURRENT_LIST_DIR}/edgellm") +_edgellm_json_include(_edge_json_include) +file(MAKE_DIRECTORY "${_edge_prefix}/lib/cmake/EdgeLLM" "${_edge_prefix}/include/edgellm/cpp" + "${_edge_prefix}/include/edgellm/3rdParty/nlohmannJson/include" + "${_edge_prefix}/include/edgellm/3rdParty/stb" "${_edge_prefix}/include/edgellm/3rdParty/miniaudio") +configure_file("${_edge_template_dir}/CheckNative.cmake" "${_edge_prefix}/lib/cmake/EdgeLLM/CheckNative.cmake" COPYONLY) +foreach(_script IN ITEMS Prepare Install) + configure_file("${_edge_template_dir}/${_script}.cmake.in" "${_edge_root}/${_script}.cmake" @ONLY) +endforeach() +configure_file("${_edge_template_dir}/EdgeLLMConfig.cmake.in" + "${_edge_prefix}/lib/cmake/EdgeLLM/EdgeLLMConfig.cmake" @ONLY) +write_basic_package_version_file("${_edge_prefix}/lib/cmake/EdgeLLM/EdgeLLMConfigVersion.cmake" + VERSION "${_edge_version}" COMPATIBILITY ExactVersion) +ExternalProject_Add(trtmc_edgellm_dependency + PREFIX "${_edge_root}/ep" SOURCE_DIR "${_edge_source}" BINARY_DIR "${_edge_build}" + GIT_REPOSITORY "${_edge_repository}" GIT_TAG "${_edge_revision}" + GIT_SUBMODULES_RECURSE TRUE UPDATE_DISCONNECTED TRUE + LIST_SEPARATOR | + # Preparation installs tools; it does not patch upstream sources. Keep it in + # the configure step so template changes invalidate disconnected builds too. + CONFIGURE_COMMAND "${CMAKE_COMMAND}" -P "${_edge_root}/Prepare.cmake" + COMMAND "${_edge_prefix}/libexec/trtmc-edge-llm/bin/cmake" + -S -B -DCMAKE_BUILD_TYPE=Release -DCMAKE_POSITION_INDEPENDENT_CODE=ON + "-DCMAKE_CUDA_COMPILER=${CMAKE_CUDA_COMPILER}" + "-DCMAKE_CUDA_ARCHITECTURES=${TRTMC_EDGELLM_CUDA_ARCHITECTURE}" + "-DCUDA_DIR=${CUDAToolkit_LIBRARY_ROOT}" "-DCUDAToolkit_ROOT=${CUDAToolkit_LIBRARY_ROOT}" + "-DCUDA_CTK_VERSION=${CUDAToolkit_VERSION_MAJOR}.${CUDAToolkit_VERSION_MINOR}" + "-DTRT_PACKAGE_DIR=${TRTMC_EDGELLM_TRT_ROOT}" "-DPython3_EXECUTABLE=${_edge_python}" + -DEDGELLM_WHEEL_PAYLOAD_DIR=unused "-DENABLE_CUTE_DSL=${_edge_cute_groups}" + "-DCUTE_DSL_ARTIFACT_TAG=sm_${TRTMC_EDGELLM_CUDA_ARCHITECTURE}" + BUILD_COMMAND "${CMAKE_COMMAND}" --build --target ${_edge_build_targets} + --parallel "${TRTMC_EDGELLM_JOBS}" + INSTALL_COMMAND "${CMAKE_COMMAND}" -P "${_edge_root}/Install.cmake" + BUILD_BYPRODUCTS "${_edge_prefix}/lib/libedgellmCore.a" + "${_edge_prefix}/lib/libNvInfer_edgellm_plugin.so" + "${_edge_prefix}/lib/libcutedsl.a" ${_edge_onnx_byproducts} + LOG_DOWNLOAD ON LOG_CONFIGURE ON LOG_BUILD ON LOG_INSTALL ON LOG_OUTPUT_ON_FAILURE ON) +ExternalProject_Add_StepDependencies(trtmc_edgellm_dependency configure "${_edge_root}/Prepare.cmake") +ExternalProject_Add_StepDependencies(trtmc_edgellm_dependency install "${_edge_root}/Install.cmake") +# Generated package targets refer to declared future byproducts; their build dependency +# prevents consumers from compiling or linking until installation completes. +find_package(EdgeLLM ${_edge_version} EXACT CONFIG REQUIRED + PATHS "${_edge_prefix}/lib/cmake/EdgeLLM" NO_DEFAULT_PATH) +add_dependencies(EdgeLLM::Core trtmc_edgellm_dependency) +add_dependencies(EdgeLLM::Plugin trtmc_edgellm_dependency) +install(DIRECTORY "${_edge_prefix}/" DESTINATION . USE_SOURCE_PERMISSIONS COMPONENT EdgeLLM) +# Family DSOs may use lib64; their dynamically loaded plugin must remain adjacent. +install(FILES "$" DESTINATION "${CMAKE_INSTALL_LIBDIR}" COMPONENT EdgeLLM) diff --git a/cmake/edgellm/CheckNative.cmake b/cmake/edgellm/CheckNative.cmake new file mode 100644 index 0000000000..03dd575b39 --- /dev/null +++ b/cmake/edgellm/CheckNative.cmake @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Read the complete TensorRT SDK version using the compiler, including aliased macros. +# include_dir: native SDK include directory; output: caller variable receiving x.y.z.build. +function(_edgellm_trt_version include_dir output) + set(_version) + set(_probe "${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/edgellm-version.cpp") + file(WRITE "${_probe}" "#include \n") + foreach(_part IN ITEMS MAJOR MINOR PATCH BUILD) + file(APPEND "${_probe}" "TRTMC_EDGE_${_part}=NV_TENSORRT_${_part}\n") + endforeach() + execute_process(COMMAND "${CMAKE_CXX_COMPILER}" -E -P -I "${include_dir}" "${_probe}" + OUTPUT_VARIABLE _expanded COMMAND_ERROR_IS_FATAL ANY) + foreach(_part IN ITEMS MAJOR MINOR PATCH BUILD) + if(NOT _expanded MATCHES "TRTMC_EDGE_${_part}=[ \t]*([0-9]+)") + message(FATAL_ERROR "Cannot determine TensorRT ${_part} from ${include_dir}") + endif() + list(APPEND _version "${CMAKE_MATCH_1}") + endforeach() + list(JOIN _version "." _version) + set(${output} "${_version}" PARENT_SCOPE) +endfunction() + +# Require the installed package GPU architecture to be present on this build host. +function(_edgellm_check_gpu architecture) + execute_process(COMMAND nvidia-smi --query-gpu=compute_cap --format=csv,noheader + OUTPUT_VARIABLE _sms RESULT_VARIABLE _result OUTPUT_STRIP_TRAILING_WHITESPACE) + string(REPLACE "." "" _sms "${_sms}") + string(REPLACE "\n" ";" _sms "${_sms}") + if(NOT _result EQUAL 0 OR NOT architecture IN_LIST _sms) + message(FATAL_ERROR "EdgeLLM requires a local GPU with architecture ${architecture}") + endif() +endfunction() + +# A version label alone is not an ABI guarantee: development headers can retain +# 3.12.0 while changing parser layouts inside the same C++ ABI namespace. +function(_edgellm_json_include output) + get_target_property(_includes nlohmann_json::nlohmann_json INTERFACE_INCLUDE_DIRECTORIES) + foreach(_include IN LISTS _includes) + string(REGEX REPLACE "^\\$$" "\\1" _include "${_include}") + if(EXISTS "${_include}/nlohmann/json.hpp") + set(${output} "${_include}" PARENT_SCOPE) + return() + endif() + endforeach() + message(FATAL_ERROR "Cannot locate nlohmann_json headers for EdgeLLM ABI validation") +endfunction() + +function(_edgellm_check_json_headers include_dir vendor_dir) + set(_header "${include_dir}/nlohmann/json.hpp") + set(_single "${vendor_dir}/single_include/nlohmann/json.hpp") + set(_multiple "${vendor_dir}/include/nlohmann/json.hpp") + if(NOT EXISTS "${_header}" OR NOT EXISTS "${_single}" OR NOT EXISTS "${_multiple}") + message(FATAL_ERROR "Missing nlohmann_json headers for EdgeLLM ABI validation") + endif() + file(SHA256 "${_header}" _actual) + file(SHA256 "${_single}" _expected_single) + if(_actual STREQUAL _expected_single) + return() + endif() + file(GLOB_RECURSE _headers RELATIVE "${vendor_dir}/include" "${vendor_dir}/include/nlohmann/*.hpp") + foreach(_relative IN LISTS _headers) + if(EXISTS "${include_dir}/${_relative}") + file(SHA256 "${include_dir}/${_relative}" _actual) + file(SHA256 "${vendor_dir}/include/${_relative}" _expected) + if(_actual STREQUAL _expected) + continue() + endif() + endif() + message(FATAL_ERROR "EdgeLLM requires the pinned nlohmann_json headers, not only the same version label. Set nlohmann_json_DIR to an installation of the pinned Edge 3rdParty/nlohmannJson dependency. Mismatch: ${_relative}") + endforeach() +endfunction() diff --git a/cmake/edgellm/EdgeLLMConfig.cmake.in b/cmake/edgellm/EdgeLLMConfig.cmake.in new file mode 100644 index 0000000000..bade6c9dba --- /dev/null +++ b/cmake/edgellm/EdgeLLMConfig.cmake.in @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +include(CMakeFindDependencyMacro) +find_dependency(CUDAToolkit) +find_dependency(Threads) +find_dependency(nlohmann_json 3.12.0 EXACT) +include("${CMAKE_CURRENT_LIST_DIR}/CheckNative.cmake") +get_filename_component(EdgeLLM_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../.." ABSOLUTE) +_edgellm_json_include(_edge_json_include) +# During first provisioning these future headers do not exist yet; Prepare +# performs the same check after checkout and before installing or building tools. +if(EXISTS "${EdgeLLM_PREFIX}/include/edgellm/3rdParty/nlohmannJson/include/nlohmann/json.hpp") + _edgellm_check_json_headers("${_edge_json_include}" "${EdgeLLM_PREFIX}/include/edgellm/3rdParty/nlohmannJson") +endif() +set(EdgeLLM_VERSION "@_edge_version@") +set(EdgeLLM_REVISION "@_edge_revision@") +# Tool availability only; model admission and orchestration remain family-owned. +set(EdgeLLM_ALL_KERNELS "@TRTMC_EDGELLM_ALL_KERNELS@") +set(EdgeLLM_ONNX "@TRTMC_EDGELLM_ONNX@") +set(EdgeLLM_ONNX_BUILDER "${EdgeLLM_PREFIX}/bin/edgellm-onnx-build") +set(EdgeLLM_CUDA_VERSION "@CUDAToolkit_VERSION@") +set(EdgeLLM_TENSORRT_VERSION "@_edge_trt_version@") +set(EdgeLLM_ARCH "@CMAKE_SYSTEM_PROCESSOR@") +set(EdgeLLM_CUDA_ARCHITECTURE "@TRTMC_EDGELLM_CUDA_ARCHITECTURE@") +if(CMAKE_CROSSCOMPILING OR NOT CMAKE_SYSTEM_PROCESSOR STREQUAL EdgeLLM_ARCH) + message(FATAL_ERROR "EdgeLLM is a native-only package for ${EdgeLLM_ARCH}") +endif() +if(NOT CUDAToolkit_VERSION_MAJOR EQUAL @CUDAToolkit_VERSION_MAJOR@ OR + NOT CUDAToolkit_VERSION_MINOR EQUAL @CUDAToolkit_VERSION_MINOR@) + message(FATAL_ERROR "EdgeLLM requires the CUDA SDK it was built with: ${EdgeLLM_CUDA_VERSION}") +endif() +set(EdgeLLM_PYTHON_EXECUTABLE "${EdgeLLM_PREFIX}/libexec/trtmc-edge-llm/bin/python") +set(EdgeLLM_BUILDER_LAUNCHER "${EdgeLLM_PREFIX}/bin/edgellm-builder") +find_path(EdgeLLM_TRT_INCLUDE_DIR NvInfer.h HINTS "$ENV{TRT_ROOT}" "@TRTMC_EDGELLM_TRT_ROOT@" PATH_SUFFIXES include REQUIRED) +find_library(EdgeLLM_TRT_LIBRARY nvinfer HINTS "$ENV{TRT_ROOT}" "@TRTMC_EDGELLM_TRT_ROOT@" PATH_SUFFIXES lib lib64 REQUIRED) +find_library(EdgeLLM_PARSER_LIBRARY nvonnxparser HINTS "$ENV{TRT_ROOT}" "@TRTMC_EDGELLM_TRT_ROOT@" PATH_SUFFIXES lib lib64 REQUIRED) +_edgellm_trt_version("${EdgeLLM_TRT_INCLUDE_DIR}" _edge_current_trt) +if(NOT _edge_current_trt STREQUAL EdgeLLM_TENSORRT_VERSION) + message(FATAL_ERROR "EdgeLLM requires TensorRT ${EdgeLLM_TENSORRT_VERSION}; found ${_edge_current_trt}") +endif() +_edgellm_check_gpu("${EdgeLLM_CUDA_ARCHITECTURE}") +if(NOT TARGET EdgeLLM::Core) + add_library(EdgeLLM::Core STATIC IMPORTED) + set_target_properties(EdgeLLM::Core PROPERTIES + IMPORTED_LOCATION "${EdgeLLM_PREFIX}/lib/libedgellmCore.a" + INTERFACE_INCLUDE_DIRECTORIES "${EdgeLLM_PREFIX}/include;${EdgeLLM_PREFIX}/include/edgellm/cpp;${EdgeLLM_PREFIX}/include/edgellm/3rdParty/nlohmannJson/include;${EdgeLLM_PREFIX}/include/edgellm/3rdParty/stb;${EdgeLLM_PREFIX}/include/edgellm/3rdParty/miniaudio;${EdgeLLM_TRT_INCLUDE_DIR}" + INTERFACE_LINK_LIBRARIES "${EdgeLLM_PREFIX}/lib/libcutedsl.a;${EdgeLLM_TRT_LIBRARY};${EdgeLLM_PARSER_LIBRARY};CUDA::cudart;CUDA::cuda_driver;Threads::Threads;${CMAKE_DL_LIBS}") + add_library(EdgeLLM::Plugin SHARED IMPORTED) + set_target_properties(EdgeLLM::Plugin PROPERTIES IMPORTED_LOCATION "${EdgeLLM_PREFIX}/lib/libNvInfer_edgellm_plugin.so") +endif() diff --git a/cmake/edgellm/Install.cmake.in b/cmake/edgellm/Install.cmake.in new file mode 100644 index 0000000000..562c720b40 --- /dev/null +++ b/cmake/edgellm/Install.cmake.in @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +cmake_minimum_required(VERSION 3.20) +file(INSTALL "@_edge_build@/cpp/libedgellmCore.a" DESTINATION "@_edge_prefix@/lib") +file(INSTALL "@_edge_build@/libNvInfer_edgellm_plugin.so" DESTINATION "@_edge_prefix@/lib" FOLLOW_SYMLINK_CHAIN) +file(INSTALL "@_edge_source@/cpp/kernels/cuteDSLArtifact/@CMAKE_SYSTEM_PROCESSOR@/sm_@TRTMC_EDGELLM_CUDA_ARCHITECTURE@/libcutedsl_@CMAKE_SYSTEM_PROCESSOR@.a" + DESTINATION "@_edge_prefix@/lib" RENAME libcutedsl.a) +file(INSTALL "@_edge_source@/cpp" DESTINATION "@_edge_prefix@/include/edgellm" FILES_MATCHING PATTERN "*.h" PATTERN "*.cuh") +foreach(_third_party IN ITEMS nlohmannJson stb miniaudio) + file(INSTALL "@_edge_source@/3rdParty/${_third_party}" DESTINATION "@_edge_prefix@/include/edgellm/3rdParty" + FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp") +endforeach() +file(MAKE_DIRECTORY "@_edge_prefix@/bin" "@_edge_prefix@/share/trtmc") +file(WRITE "@_edge_prefix@/bin/edgellm-builder" [=[#!/bin/sh +set -eu +prefix=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd) +exec "$prefix/libexec/trtmc-edge-llm/bin/python" -I -c 'from experimental.builder.cli import main; main()' "$@" +]=]) +file(CHMOD "@_edge_prefix@/bin/edgellm-builder" PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE) +if("@TRTMC_EDGELLM_ONNX@") + file(INSTALL "@_edge_build@/examples/llm/llm_build" DESTINATION "@_edge_prefix@/bin" + TYPE PROGRAM RENAME edgellm-onnx-build) +endif() +set(_all_kernels false) +if("@TRTMC_EDGELLM_ALL_KERNELS@") + set(_all_kernels true) +endif() +set(_onnx false) +if("@TRTMC_EDGELLM_ONNX@") + set(_onnx true) +endif() +execute_process(COMMAND "@_edge_python@" -I -c "import tensorrt; print(tensorrt.__version__)" + OUTPUT_VARIABLE _trt_version OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) +# Paths are relative to the installation prefix, preserving relocatability. +file(WRITE "@_edge_prefix@/share/trtmc/edge-llm.json" "{\n \"schema_version\": 1,\n \"version\": \"@_edge_version@\",\n \"revision\": \"@_edge_revision@\",\n \"arch\": \"@CMAKE_SYSTEM_PROCESSOR@\",\n \"architectures\": [@TRTMC_EDGELLM_CUDA_ARCHITECTURE@],\n \"cuda_version\": \"@CUDAToolkit_VERSION_MAJOR@.@CUDAToolkit_VERSION_MINOR@\",\n \"tensorrt_version\": \"${_trt_version}\",\n \"python\": \"libexec/trtmc-edge-llm/bin/python\",\n \"builder\": \"bin/edgellm-builder\",\n \"all_native_kernels\": ${_all_kernels},\n \"onnx\": ${_onnx},\n \"onnx_builder\": \"bin/edgellm-onnx-build\",\n \"plugin\": \"lib/libNvInfer_edgellm_plugin.so\"\n}\n") diff --git a/cmake/edgellm/Prepare.cmake.in b/cmake/edgellm/Prepare.cmake.in new file mode 100644 index 0000000000..bf19447034 --- /dev/null +++ b/cmake/edgellm/Prepare.cmake.in @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Executed only by the explicit CMake dependency build, never by model dispatch. +cmake_minimum_required(VERSION 3.20) +include("@_edge_template_dir@/CheckNative.cmake") +_edgellm_check_json_headers("@_edge_json_include@" "@_edge_source@/3rdParty/nlohmannJson") +function(run) + execute_process(COMMAND ${ARGV} COMMAND_ERROR_IS_FATAL ANY) +endfunction() +execute_process(COMMAND "@Python3_EXECUTABLE@" -I -c "import ensurepip" + RESULT_VARIABLE _has_ensurepip OUTPUT_QUIET ERROR_QUIET) +if(_has_ensurepip EQUAL 0) + run("@Python3_EXECUTABLE@" -I -m venv --copies "@_edge_prefix@/libexec/trtmc-edge-llm") +else() + # Debian minimal Python may omit ensurepip; use an already installed bootstrapper. + run("@Python3_EXECUTABLE@" -I -m virtualenv --copies --no-download --no-periodic-update + "@_edge_prefix@/libexec/trtmc-edge-llm") +endif() +set(_pip_options --isolated install --no-user) +if(NOT "@TRTMC_EDGELLM_WHEELHOUSE@" STREQUAL "") + list(APPEND _pip_options --no-index --find-links "@TRTMC_EDGELLM_WHEELHOUSE@") +endif() +file(GLOB _trt_wheels "@TRTMC_EDGELLM_TRT_ROOT@/python/tensorrt-*-cp@Python3_VERSION_MAJOR@@Python3_VERSION_MINOR@-none-linux_@CMAKE_SYSTEM_PROCESSOR@.whl") +list(LENGTH _trt_wheels _wheel_count) +if(NOT _wheel_count EQUAL 1) + message(FATAL_ERROR "Expected exactly one TensorRT SDK wheel matching the native Python ABI") +endif() +run("@_edge_python@" -I -m pip ${_pip_options} --report "@_edge_prefix@/pip-report.json" + ${_trt_wheels} numpy==2.2.6 transformers==5.14.1 jinja2==3.1.6 + scikit-build-core==0.11.6 wheel==0.45.1 cmake==3.31.10 ninja==1.13.0 + "cuda-python>=@CUDAToolkit_VERSION_MAJOR@.@CUDAToolkit_VERSION_MINOR@,<@CUDAToolkit_VERSION_MAJOR@.@CUDAToolkit_VERSION_MINOR@.999" + "nvidia-cutlass-dsl[cu@CUDAToolkit_VERSION_MAJOR@]==4.7.0" "cupy-cuda@CUDAToolkit_VERSION_MAJOR@x==13.6.0") +run("@_edge_python@" -I -m pip ${_pip_options} --no-deps --no-build-isolation "@_edge_source@") +if("@TRTMC_EDGELLM_ONNX@") + # Original exporter dependencies, isolated from the caller environment. Export + # is CPU-side; native TensorRT compilation still runs on the inference GPU. + set(_torch_options ${_pip_options}) + if("@TRTMC_EDGELLM_WHEELHOUSE@" STREQUAL "") + list(APPEND _torch_options --index-url https://download.pytorch.org/whl/cpu) + endif() + run("@_edge_python@" -I -m pip ${_torch_options} "torch==2.13.0") + run("@_edge_python@" -I -m pip ${_pip_options} "tensorrt-edgellm[export]==@_edge_version@") +endif() +run("@_edge_python@" -I -m pip --isolated check) +run("@_edge_python@" -I -c "print(__import__('tensorrt').__version__)") +set(ENV{PATH} "@CUDAToolkit_BIN_DIR@:$ENV{PATH}") +run("@_edge_python@" -I "@_edge_source@/kernelSrcs/build_cutedsl.py" + --gpu_arch "sm_@TRTMC_EDGELLM_CUDA_ARCHITECTURE@" --arch "@CMAKE_SYSTEM_PROCESSOR@" + --kernels "@_edge_cute_cli_groups@" --cuda-version "@CUDAToolkit_VERSION_MAJOR@.@CUDAToolkit_VERSION_MINOR@" --jobs "@TRTMC_EDGELLM_JOBS@") diff --git a/cmake/edgellm/README.md b/cmake/edgellm/README.md new file mode 100644 index 0000000000..89ae872efc --- /dev/null +++ b/cmake/edgellm/README.md @@ -0,0 +1,72 @@ +# Pinned native Edge-LLM package + +Edge-LLM is optional. The default `TRTMC_ENABLE_EDGELLM=OFF` neither downloads +nor builds it. Enable it once while installing Model Connect; ordinary model +builds only use the installed package and never fetch or install dependencies. +Cross compilation is rejected. Configure and build on the inference GPU host. + +```bash +cmake -S . -B build \ + -DTRTMC_ENABLE_EDGELLM=ON \ + -DTRTMC_EDGELLM_CUDA_ARCHITECTURE=80 \ + -DTRTMC_EDGELLM_TRT_ROOT="$TRT_ROOT" \ + -DCUDAToolkit_ROOT="$CUDA_ROOT" \ + -DCMAKE_CUDA_COMPILER="$CUDA_ROOT/bin/nvcc" \ + -DCMAKE_INSTALL_PREFIX="$PWD/install" +cmake --build build --parallel 8 +cmake --install build +export CMAKE_PREFIX_PATH="$PWD/install${CMAKE_PREFIX_PATH:+:$CMAKE_PREFIX_PATH}" +``` + +The regular project dependencies remain required, including nlohmann_json +**3.12.0 with the exact pinned upstream headers** when Edge is enabled. A +development snapshot can retain that version label but change parser layouts; +mixing it with the static SDK causes undefined behavior. The package checks +header content against its vendored dependency (single or multiple headers). +If rejected, install `3rdParty/nlohmannJson` from the pinned Edge checkout into +a separate prefix and configure with that installation’s `nlohmann_json_DIR`. The Python +interpreter needs `ensurepip` or an already installed `virtualenv` bootstrapper. +The native CUDA SDK must include NVCC, NVRTC, cuRAND headers and driver link +libraries; the TensorRT SDK must contain its matching CPython wheel. + +The provider first uses `find_package(EdgeLLM 0.10.1 EXACT CONFIG)`. If absent, +CMake `ExternalProject` clones the public NVIDIA TensorRT-Edge-LLM repository at +`e8b29522938901f6df19ebeedd4b69bc8edbcd97` (v0.10.1), initializes the pinned +submodules, builds the native core/plugin and FMHA/GDN CuTe archives, and installs +an isolated direct-builder Python environment. It does not modify the caller +Python environment. Downloads happen only during this +explicit dependency build. `TRTMC_EDGELLM_WHEELHOUSE` selects a complete offline +Python wheelhouse; `TRTMC_EDGELLM_GIT_MIRROR` optionally supplies a local Git +mirror, still checked out at the immutable upstream commit. + +Upstream 0.10.1 does not export a CMake SDK package, so these compact templates +supply that installation boundary. `EdgeLLM::Core` exposes the installed static +core, headers, CuTe archive and native dependencies. Consumers requiring CUDA +device linking enable separable compilation and device-symbol resolution. +`EdgeLLM::Plugin` identifies the plugin DSO; adapters load it, rather than linking +it twice. `EdgeLLM_PYTHON_EXECUTABLE` and `EdgeLLM_BUILDER_LAUNCHER` expose the +isolated upstream `experimental.builder.cli.main` API. + +`share/trtmc/edge-llm.json` records the pin, native architecture, CUDA/TensorRT +versions and prefix-relative Python/plugin paths. The manifest is written only +after successful installation. Build-tree package files live under +`build/_deps/edgellm/install`; `cmake --install` copies the package into the final +prefix. Package discovery rejects mismatched native CPU/GPU and SDK versions. +Model support and routing policies belong exclusively to the model families. + +Set `TRTMC_EDGELLM_ALL_KERNELS=ON` to provision all upstream operator groups +supported by the native GPU. Set `TRTMC_EDGELLM_ONNX=ON` to additionally install +the original Python exporter (including its pinned CPU PyTorch dependencies) +and original C++ `llm_build` executable as `bin/edgellm-onnx-build`. Families invoke +the exporter using the installed Python and obtain the native builder path from +`onnx_builder` in the manifest. These options describe SDK capabilities, not +qualified model support. Reusing an installed package that lacks a requested +capability is an error; no dependency installation occurs during model builds. +The CUDA and TensorRT shared libraries must remain available to the executable. + +Run the existing runtime and family checks against this installation +(some tests require a local GPU): + +```bash +ctest --test-dir build --output-on-failure +``` diff --git a/core/builder/tensorrt_model_connect/__init__.py b/core/builder/tensorrt_model_connect/__init__.py index 101f6dcb4b..407e0c8a14 100644 --- a/core/builder/tensorrt_model_connect/__init__.py +++ b/core/builder/tensorrt_model_connect/__init__.py @@ -3,12 +3,14 @@ """TensorRT Model Connect build API.""" -from .build import BuildRequest, build +from .build import BuildExecutionInputs, BuildRequest, NamedCheckpoint, build from .bundle_writer import BundleWriter from .graph_transform import GraphTransform __all__ = [ + "BuildExecutionInputs", "BuildRequest", + "NamedCheckpoint", "BundleWriter", "GraphTransform", "build", diff --git a/core/builder/tensorrt_model_connect/build.py b/core/builder/tensorrt_model_connect/build.py index d0d5d68638..69e4722f67 100644 --- a/core/builder/tensorrt_model_connect/build.py +++ b/core/builder/tensorrt_model_connect/build.py @@ -7,6 +7,8 @@ import hashlib import importlib +import os +import platform import re import sys from dataclasses import dataclass @@ -20,6 +22,52 @@ _ID = re.compile(r"[a-z][a-z0-9_]*\Z") +@dataclass(frozen=True) +class NamedCheckpoint: + """One explicitly named local checkpoint; the family owns role semantics.""" + + role: str + model_dir: Path + + def __post_init__(self) -> None: + _validate_id("checkpoint role", self.role) + if not isinstance(self.model_dir, Path): + raise TypeError("checkpoint model_dir must be a Path") + if not self.model_dir.is_dir(): + raise ValueError(f"checkpoint must be an existing local directory: {self.model_dir}") + + +@dataclass(frozen=True) +class BuildExecutionInputs: + """Optional family-owned execution variant and immutable local companions. + + Core transports these inputs without interpreting variants, fetching models, + or inferring compatibility. A family must explicitly implement the capability. + """ + + variant: str + checkpoints: tuple[NamedCheckpoint, ...] = () + + def __post_init__(self) -> None: + _validate_id("execution variant", self.variant) + if not isinstance(self.checkpoints, tuple) or any( + not isinstance(checkpoint, NamedCheckpoint) for checkpoint in self.checkpoints + ): + raise TypeError("checkpoints must be a tuple of NamedCheckpoint values") + roles = [checkpoint.role for checkpoint in self.checkpoints] + if len(roles) != len(set(roles)): + raise ValueError("checkpoint roles must be unique") + self.validate_local() + + def validate_local(self) -> None: + """Recheck local availability before dispatch, without acquiring inputs.""" + for checkpoint in self.checkpoints: + if not checkpoint.model_dir.is_dir(): + raise ValueError( + f"checkpoint must be an existing local directory: {checkpoint.model_dir}" + ) + + @dataclass(frozen=True) class BuildRequest: """Inputs shared by the build core and one family-owned builder.""" @@ -72,6 +120,61 @@ def __post_init__(self) -> None: raise ValueError("graph_transform must be callable when provided") +def subprocess_environment(overrides: dict[str, str], *, + prepend_paths: dict[str, str] | None = None) -> dict[str, str]: + """Copy the parent environment for one child without mutating process state. + + Callers own explicit tool settings; this helper only merges values and + prepends search paths using the executing platform's path separator. + """ + environment = os.environ.copy() + environment.update(overrides) + for name, value in (prepend_paths or {}).items(): + previous = environment.get(name) + environment[name] = value + (os.pathsep + previous if previous else "") + return environment + + +def cmake_prefixes() -> list[Path]: + """Return explicit standard CMake prefixes followed by the Python prefix.""" + prefixes = [ + Path(value) for value in os.environ.get("CMAKE_PREFIX_PATH", "").split(os.pathsep) if value + ] + return [*prefixes, Path(sys.prefix)] + + +def detect_local_platform() -> dict: + """Return executing GPU and native SDK identity without selecting a model. + + Returns: + OS/release, CPU architecture, GPU SM, CUDA and TensorRT versions. + + Raises: + ImportError: Native SDK Python bindings are unavailable. + RuntimeError: CUDA cannot identify the executing device. + """ + import tensorrt as trt + from cuda.bindings import runtime + + def checked(result): + if int(result[0]) != 0: + raise RuntimeError(f"CUDA device discovery failed: {result[0]}") + return result[1] + + device = checked(runtime.cudaGetDevice()) + gpu = checked(runtime.cudaGetDeviceProperties(device)) + cuda = checked(runtime.cudaRuntimeGetVersion()) + release = platform.freedesktop_os_release() if sys.platform == "linux" else {} + return { + "os": sys.platform, + "os_version": release.get("VERSION_ID", platform.release()), + "arch": platform.machine(), + "sm": gpu.major * 10 + gpu.minor, + "cuda_version": f"{cuda // 1000}.{cuda % 1000 // 10}", + "tensorrt_version": trt.__version__, + } + + def _validate_id(field: str, value: object) -> str: if not isinstance(value, str) or _ID.fullmatch(value) is None: raise ValueError( @@ -130,16 +233,36 @@ def _select_backend(backend: str) -> None: sys.modules["tensorrt"] = rtx -def build(request: BuildRequest) -> None: - """Run one family builder and publish its bundle on success.""" +def build(request: BuildRequest, *, execution: BuildExecutionInputs | None = None) -> None: + """Run one family builder and atomically publish its bundle on success. + + Explicit execution inputs require the optional family build_with_inputs hook. + Missing support fails before constructing the writer; failures never retry a + different variant or silently invoke the ordinary builder. + """ + + if execution is not None: + if not isinstance(execution, BuildExecutionInputs): + raise TypeError("execution must be BuildExecutionInputs") + execution.validate_local() family = _resolve_family(request) _select_backend(request.backend) family_module = _load_family(family) + extended_build = None + if execution is not None: + extended_build = getattr(family_module, "build_with_inputs", None) + if not callable(extended_build): + raise NotImplementedError( + f"family {family!r} does not support explicit build execution inputs" + ) writer = BundleWriter(request.output_path) try: with graph_transform(request.graph_transform): - family_module.build(request, writer) + if execution is None: + family_module.build(request, writer) + else: + extended_build(request, writer, execution) writer.finish() except BaseException: writer.abort() diff --git a/core/builder/tensorrt_model_connect/build_cli.py b/core/builder/tensorrt_model_connect/build_cli.py index 9042a68320..0d81dded9f 100644 --- a/core/builder/tensorrt_model_connect/build_cli.py +++ b/core/builder/tensorrt_model_connect/build_cli.py @@ -10,7 +10,7 @@ from pathlib import Path from typing import Sequence -from .build import BuildRequest, _load_family, build +from .build import BuildExecutionInputs, BuildRequest, NamedCheckpoint, _load_family, build from .model_support import load_model_metadata, resolve_family @@ -35,6 +35,14 @@ def _parser() -> argparse.ArgumentParser: build_parser.add_argument("--fp32-layer", type=int, action="append", default=[]) build_parser.add_argument("--dynamic-kv-cache", action="store_true") build_parser.add_argument("--verbose", action="store_true") + build_parser.add_argument("--execution-variant", help="Explicit family-owned execution variant") + build_parser.add_argument( + "--companion", + action="append", + default=[], + metavar="ROLE=LOCAL_DIR", + help="Named existing local checkpoint; repeat for multiple distinct roles", + ) prepare_parser = commands.add_parser( "prepare-structure", help="Prepare one structure request without rebuilding its model bundle", @@ -47,8 +55,28 @@ def _parser() -> argparse.ArgumentParser: return parser +def _execution_inputs(args: argparse.Namespace) -> BuildExecutionInputs | None: + """Parse only explicit local inputs; no variant list or model acquisition.""" + if args.command != "build": + return None + if args.execution_variant is None: + if args.companion: + raise ValueError("--companion requires --execution-variant") + return None + checkpoints = [] + for value in args.companion: + role, separator, directory = value.partition("=") + if not separator or not role or not directory: + raise ValueError("--companion must be ROLE=LOCAL_DIR") + if "://" in directory: + raise ValueError("--companion requires a local directory, not a URI") + checkpoints.append(NamedCheckpoint(role, Path(directory))) + return BuildExecutionInputs(args.execution_variant, tuple(checkpoints)) + + def main(argv: Sequence[str] | None = None) -> int: args = _parser().parse_args(argv) + execution = _execution_inputs(args) model_dir = _resolve_model(args.model, args.revision) family, support = resolve_family(load_model_metadata(model_dir)) if args.command == "prepare-structure": @@ -74,27 +102,29 @@ def main(argv: Sequence[str] | None = None) -> int: f"family {family!r} does not support task {task!r}; " f"choose one of: {', '.join(support.tasks)}" ) - build( - BuildRequest( - model_dir=model_dir, - output_path=args.output, - precision=args.precision or support.default_precision, - backend=args.backend, - family=family, - task=task, - max_sequence_length=args.max_sequence_length, - image_height=args.image_height, - image_width=args.image_width, - video_num_frames=args.video_num_frames, - max_batch_size=args.max_batch_size, - tensor_parallel_size=args.tensor_parallel_size, - context_parallel_size=args.context_parallel_size, - quantization=args.quantization, - fp32_layers=tuple(args.fp32_layer), - dynamic_kv_cache=args.dynamic_kv_cache, - verbose=args.verbose, - ) + request = BuildRequest( + model_dir=model_dir, + output_path=args.output, + precision=args.precision or support.default_precision, + backend=args.backend, + family=family, + task=task, + max_sequence_length=args.max_sequence_length, + image_height=args.image_height, + image_width=args.image_width, + video_num_frames=args.video_num_frames, + max_batch_size=args.max_batch_size, + tensor_parallel_size=args.tensor_parallel_size, + context_parallel_size=args.context_parallel_size, + quantization=args.quantization, + fp32_layers=tuple(args.fp32_layer), + dynamic_kv_cache=args.dynamic_kv_cache, + verbose=args.verbose, ) + if execution is None: + build(request) + else: + build(request, execution=execution) return 0 diff --git a/core/builder/tests/test_build.py b/core/builder/tests/test_build.py index 5eff69de1b..fc94a5a99c 100644 --- a/core/builder/tests/test_build.py +++ b/core/builder/tests/test_build.py @@ -4,6 +4,7 @@ from __future__ import annotations import importlib +from contextlib import contextmanager import sys from dataclasses import FrozenInstanceError, replace from pathlib import Path @@ -11,7 +12,7 @@ import pytest -from tensorrt_model_connect import BuildRequest +from tensorrt_model_connect import BuildExecutionInputs, BuildRequest, NamedCheckpoint, build_cli build_core = importlib.import_module("tensorrt_model_connect.build") @@ -287,3 +288,295 @@ def abort(self) -> None: with pytest.raises(OSError, match="publish failed"): build_core.build(_request(tmp_path)) assert events == ["finish", "abort"] + + +@pytest.mark.parametrize("value", ["", "/one", "/one:/two", ":/one::/two:"]) +def test_cmake_prefixes_preserve_standard_search_order(monkeypatch, value): + monkeypatch.setenv("CMAKE_PREFIX_PATH", value) + monkeypatch.setattr(build_core.sys, "prefix", "/python") + expected = [Path(item) for item in value.split(build_core.os.pathsep) if item] + assert build_core.cmake_prefixes() == [*expected, Path("/python")] + + +def test_cmake_prefixes_without_environment_use_python_prefix(monkeypatch): + monkeypatch.delenv("CMAKE_PREFIX_PATH", raising=False) + monkeypatch.setattr(build_core.sys, "prefix", "/python") + assert build_core.cmake_prefixes() == [Path("/python")] + # Constructing explicit child-tool settings must not change the caller's + # package search order or mutate an inherited search path. + monkeypatch.setenv("TEST_TOOL_SEARCH_PATH", "/original") + monkeypatch.delenv("TEST_TOOL_NEW_PATH", raising=False) + child = build_core.subprocess_environment( + {"CMAKE_PREFIX_PATH": "/child"}, + prepend_paths={"TEST_TOOL_SEARCH_PATH": "/first", "TEST_TOOL_NEW_PATH": "/new"}, + ) + assert child["CMAKE_PREFIX_PATH"] == "/child" + assert child["TEST_TOOL_SEARCH_PATH"] == "/first" + build_core.os.pathsep + "/original" + assert child["TEST_TOOL_NEW_PATH"] == "/new" + assert build_core.cmake_prefixes() == [Path("/python")] + assert build_core.os.environ["TEST_TOOL_SEARCH_PATH"] == "/original" + assert "TEST_TOOL_NEW_PATH" not in build_core.os.environ + + +@pytest.fixture +def native_platform_bindings(monkeypatch): + from unittest.mock import Mock + + runtime = SimpleNamespace( + cudaGetDevice=Mock(return_value=(0, 3)), + cudaGetDeviceProperties=Mock(return_value=(0, SimpleNamespace(major=8, minor=6))), + cudaRuntimeGetVersion=Mock(return_value=(0, 13030)), + ) + monkeypatch.setitem(sys.modules, "tensorrt", SimpleNamespace(__version__="11.1.0.106")) + monkeypatch.setitem(sys.modules, "cuda.bindings", SimpleNamespace(runtime=runtime)) + monkeypatch.setattr(build_core.sys, "platform", "linux") + monkeypatch.setattr(build_core.platform, "machine", lambda: "x86_64") + monkeypatch.setattr( + build_core.platform, "freedesktop_os_release", lambda: {"VERSION_ID": "24.04"} + ) + monkeypatch.setattr(build_core.platform, "release", lambda: "fallback-release") + return runtime + + +def test_native_platform_uses_executing_cuda_device_and_full_sdk(native_platform_bindings): + assert build_core.detect_local_platform() == { + "os": "linux", + "os_version": "24.04", + "arch": "x86_64", + "sm": 86, + "cuda_version": "13.3", + "tensorrt_version": "11.1.0.106", + } + native_platform_bindings.cudaGetDevice.assert_called_once_with() + native_platform_bindings.cudaGetDeviceProperties.assert_called_once_with(3) + native_platform_bindings.cudaRuntimeGetVersion.assert_called_once_with() + + +@pytest.mark.parametrize( + "failing", ["cudaGetDevice", "cudaGetDeviceProperties", "cudaRuntimeGetVersion"] +) +def test_native_platform_propagates_cuda_discovery_failure(native_platform_bindings, failing): + getattr(native_platform_bindings, failing).return_value = (35,) + with pytest.raises(RuntimeError, match="CUDA device discovery failed: 35"): + build_core.detect_local_platform() + + +def test_native_platform_retains_nonlinux_identity(native_platform_bindings, monkeypatch): + monkeypatch.setattr(build_core.sys, "platform", "win32") + result = build_core.detect_local_platform() + assert result["os"] == "win32" + assert result["os_version"] == "fallback-release" + + +def execution_request(root: Path) -> BuildRequest: + return BuildRequest(root, root / "model.bundle", "example", "text_generation", "fp16") + + +def inputs(root: Path) -> BuildExecutionInputs: + return BuildExecutionInputs("paired", (NamedCheckpoint("draft", root),)) + + +def test_execution_inputs_are_immutable(tmp_path): + execution = inputs(tmp_path) + with pytest.raises(FrozenInstanceError): + execution.variant = "other" + with pytest.raises(FrozenInstanceError): + execution.checkpoints[0].role = "other" + assert execution.checkpoints[0].model_dir is tmp_path + + +@pytest.mark.parametrize("value", ["", "../bad", "UPPER", "a-b", "a.b"]) +def test_invalid_role_and_variant(tmp_path, value): + with pytest.raises(ValueError, match="lowercase identifier"): + NamedCheckpoint(value, tmp_path) + with pytest.raises(ValueError, match="lowercase identifier"): + BuildExecutionInputs(value) + + +def test_execution_requires_immutable_typed_companions(tmp_path): + checkpoint = NamedCheckpoint("draft", tmp_path) + with pytest.raises(TypeError, match="tuple"): + BuildExecutionInputs("paired", [checkpoint]) + with pytest.raises(TypeError, match="NamedCheckpoint"): + BuildExecutionInputs("paired", (object(),)) + with pytest.raises(ValueError, match="unique"): + BuildExecutionInputs("paired", (checkpoint, checkpoint)) + with pytest.raises(TypeError, match="Path"): + NamedCheckpoint("draft", str(tmp_path)) + + +def test_local_checkpoint_required_and_rechecked(tmp_path, monkeypatch): + with pytest.raises(ValueError, match="existing local directory"): + NamedCheckpoint("draft", tmp_path / "missing") + file = tmp_path / "file" + file.write_text("not a directory") + with pytest.raises(ValueError, match="existing local directory"): + NamedCheckpoint("draft", file) + directory = tmp_path / "companion" + directory.mkdir() + execution = inputs(directory) + directory.rmdir() + monkeypatch.setattr(build_core, "_select_backend", lambda _: pytest.fail("backend touched")) + with pytest.raises(ValueError, match="existing local directory"): + build_core.build(execution_request(tmp_path), execution=execution) + + +def test_untyped_execution_fails_before_side_effects(tmp_path, monkeypatch): + monkeypatch.setattr(build_core, "_select_backend", lambda _: pytest.fail("backend touched")) + with pytest.raises(TypeError, match="BuildExecutionInputs"): + build_core.build(execution_request(tmp_path), execution={"variant": "paired"}) + + +@pytest.mark.parametrize("hook", [None, 17]) +def test_missing_capability_fails_before_writer(tmp_path, monkeypatch, hook): + monkeypatch.setattr( + build_core, + "_load_family", + lambda _: SimpleNamespace( + build=lambda *_: pytest.fail("ordinary fallback invoked"), build_with_inputs=hook + ), + ) + monkeypatch.setattr(build_core, "BundleWriter", lambda _: pytest.fail("writer created")) + with pytest.raises(NotImplementedError, match="does not support explicit"): + build_core.build(execution_request(tmp_path), execution=inputs(tmp_path)) + + +def test_exact_envelope_and_existing_transaction_are_preserved(tmp_path, monkeypatch): + events = [] + original_request, execution = execution_request(tmp_path), inputs(tmp_path) + + @contextmanager + def transform(value): + assert value is original_request.graph_transform + events.append("enter") + yield + events.append("exit") + + class Writer: + def __init__(self, path): + assert path == original_request.output_path + events.append("writer") + + def finish(self): + events.append("finish") + + def abort(self): + pytest.fail("unexpected abort") + + def extended(actual_request, writer, actual_execution): + assert actual_request is original_request and actual_execution is execution + assert isinstance(writer, Writer) + events.append("extended") + + monkeypatch.setattr(build_core, "graph_transform", transform) + monkeypatch.setattr(build_core, "BundleWriter", Writer) + monkeypatch.setattr( + build_core, + "_load_family", + lambda _: SimpleNamespace( + build=lambda *_: pytest.fail("ordinary fallback invoked"), build_with_inputs=extended + ), + ) + build_core.build(original_request, execution=execution) + assert events == ["writer", "enter", "extended", "exit", "finish"] + + +@pytest.mark.parametrize("failure", [RuntimeError("failed"), KeyboardInterrupt()]) +def test_explicit_failure_aborts_real_writer_without_replacing_bundle( + tmp_path, monkeypatch, failure +): + build_request = execution_request(tmp_path) + build_request.output_path.write_bytes(b"previous valid publication") + + def extended(actual, writer, execution): + writer.set_header(family=actual.family, task=actual.task, backend=actual.backend) + writer.add_json("test.json", {"variant": execution.variant}) + raise failure + + monkeypatch.setattr( + build_core, + "_load_family", + lambda _: SimpleNamespace( + build=lambda *_: pytest.fail("ordinary fallback invoked"), build_with_inputs=extended + ), + ) + with pytest.raises(type(failure)) as caught: + build_core.build(build_request, execution=inputs(tmp_path)) + assert caught.value is failure + assert build_request.output_path.read_bytes() == b"previous valid publication" + assert sorted(path.name for path in tmp_path.iterdir()) == ["model.bundle"] + + +def test_variant_without_companions_is_explicit_and_supported(tmp_path, monkeypatch): + seen = [] + + def extended(actual, writer, execution): + seen.append(execution) + writer.set_header(family=actual.family, task=actual.task, backend=actual.backend) + writer.add_json("test.json", {"variant": execution.variant}) + + monkeypatch.setattr( + build_core, "_load_family", lambda _: SimpleNamespace(build_with_inputs=extended) + ) + execution = BuildExecutionInputs("embedded") + build_core.build(execution_request(tmp_path), execution=execution) + assert seen == [execution] and (tmp_path / "model.bundle").is_file() + + +def test_ordinary_build_ignores_available_optional_hook(tmp_path, monkeypatch): + def ordinary(actual, writer): + writer.set_header(family=actual.family, task=actual.task, backend=actual.backend) + writer.add_json("test.json", {"ordinary": True}) + + monkeypatch.setattr( + build_core, + "_load_family", + lambda _: SimpleNamespace( + build=ordinary, build_with_inputs=lambda *_: pytest.fail("optional hook invoked") + ), + ) + build_core.build(execution_request(tmp_path)) + + +@pytest.mark.parametrize( + "options", + [ + ["--companion", "draft=/missing"], + ["--execution-variant", "paired", "--companion", "missing_separator"], + ["--execution-variant", "paired", "--companion", "=path"], + ["--execution-variant", "paired", "--companion", "draft="], + ["--execution-variant", "paired", "--companion", "draft=https://example.com/model"], + ["--execution-variant", ""], + ], +) +def test_bad_cli_execution_rejected_before_primary_model_acquisition(monkeypatch, options): + monkeypatch.setattr(build_cli, "_resolve_model", lambda *_: pytest.fail("model acquisition")) + with pytest.raises(ValueError): + build_cli.main(["build", "model-id", "-o", "/tmp/example.bundle", *options]) + + +def test_cli_forwards_exact_variant_and_named_local_paths(tmp_path, monkeypatch): + (tmp_path / "config.json").write_text('{"model_type":"gpt2"}') + companion = tmp_path / "checkpoint=local" + companion.mkdir() + seen = [] + monkeypatch.setattr( + build_cli, "build", lambda request, **kwargs: seen.append((request, kwargs)) + ) + build_cli.main( + [ + "build", + str(tmp_path), + "-o", + str(tmp_path / "out.bundle"), + "--execution-variant", + "paired", + "--companion", + f"draft={companion}", + ] + ) + actual_request, kwargs = seen[0] + assert actual_request.model_dir == tmp_path + assert kwargs == { + "execution": BuildExecutionInputs("paired", (NamedCheckpoint("draft", companion),)) + } diff --git a/core/runtime/bundle/bundle_format.cpp b/core/runtime/bundle/bundle_format.cpp index eb11f20674..2c1661f088 100644 --- a/core/runtime/bundle/bundle_format.cpp +++ b/core/runtime/bundle/bundle_format.cpp @@ -6,6 +6,7 @@ #include "runtime/bundle/bundle_format.h" #include +#include #include #include #include @@ -215,6 +216,32 @@ std::vector BundleReader::read_section(std::string_view name) const { return data; } +void BundleReader::copy_section(std::string_view name, std::ostream& output) const { + const auto* section = find_section(name); + if (section == nullptr) + throw std::runtime_error("Bundle section not found: " + std::string(name)); + const auto offset = checked_section_file_offset(*section, data_offset_, file_size_, path_); + if (offset > static_cast(std::numeric_limits::max())) + throw std::runtime_error("Bundle section has an unsupported file offset: " + path_); + std::ifstream input(path_, std::ios::binary); + input.seekg(static_cast(offset)); + if (!input || !output) + throw std::runtime_error("Cannot copy bundle section: " + std::string(name)); + std::array buffer; + auto remaining = section->length; + while (remaining != 0) { + const auto count = + static_cast(std::min(remaining, buffer.size())); + input.read(buffer.data(), count); + if (!input) + throw std::runtime_error("Failed reading bundle section: " + std::string(name)); + output.write(buffer.data(), count); + if (!output) + throw std::runtime_error("Failed writing bundle section: " + std::string(name)); + remaining -= static_cast(count); + } +} + BundleInfo InspectBundle(const std::string& bundle_path) { return BundleReader(bundle_path).info(); } diff --git a/core/runtime/include/trtmc/bundle.h b/core/runtime/include/trtmc/bundle.h index 154d697ef1..2b32776d2e 100644 --- a/core/runtime/include/trtmc/bundle.h +++ b/core/runtime/include/trtmc/bundle.h @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include #include @@ -44,6 +45,12 @@ class BundleReader { const BundleSectionInfo* find_section(std::string_view name) const noexcept; std::vector read_section(std::string_view name) const; + /// Copy a named section to an output stream using bounded working memory. + /// @param name Validated bundle section name. + /// @param output Caller-owned stream; may contain partial data on failure. + /// @throws std::runtime_error If the section is absent or input/output fails. + void copy_section(std::string_view name, std::ostream& output) const; + private: std::string path_; BundleInfo info_; diff --git a/core/runtime/tests/test_bundle_format_v1.cpp b/core/runtime/tests/test_bundle_format_v1.cpp index 0da636b49b..6c5faa420f 100644 --- a/core/runtime/tests/test_bundle_format_v1.cpp +++ b/core/runtime/tests/test_bundle_format_v1.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -52,6 +53,36 @@ bool read_throws(const std::filesystem::path& path) { } } +/// Verify bounded-chunk section boundaries, empty sections and late I/O failures. +void test_copy_section(const std::filesystem::path& directory) { + const auto path = directory / "stream.bundle"; + const std::string payload(2 * 64 * 1024 + 17, 'x'); + const std::string header = + R"({"format":1,"family":"fake","task":"text","backend":"fake","sections":{"data":{"offset":3,"length":)" + + std::to_string(payload.size()) + R"(},"empty":{"offset":0,"length":0}}})"; + write_bundle(path, header, "PRE" + payload + "POST"); + const trtmc::BundleReader reader(path.string()); + std::ostringstream output; + reader.copy_section("data", output); + check(output.str() == payload, "stream copy preserves boundaries across chunks"); + reader.copy_section("empty", output); + check(output.str() == payload, "empty section appends nothing"); + auto fails = [&](const char* section, std::ostream& destination) { + try { + reader.copy_section(section, destination); + return false; + } catch (const std::runtime_error&) { + return true; + } + }; + check(fails("missing", output), "stream copy rejects missing sections"); + std::ostringstream broken; + broken.setstate(std::ios::badbit); + check(fails("data", broken), "stream copy reports output failure"); + std::filesystem::resize_file(path, 16 + header.size() + 3 + payload.size() - 1); + check(fails("data", output), "stream copy detects truncation after validation"); +} + } // namespace int main() { @@ -103,6 +134,7 @@ int main() { "PLAN"); check(read_throws(out_of_bounds), "out of bounds section rejected"); + test_copy_section(directory); std::filesystem::remove_all(directory); std::cerr << (failures == 0 ? "ALL PASSED\n" : "SOME FAILED\n"); return failures; diff --git a/families/qwen3_8/EDGE_LLM.md b/families/qwen3_8/EDGE_LLM.md new file mode 100644 index 0000000000..1160f29d8e --- /dev/null +++ b/families/qwen3_8/EDGE_LLM.md @@ -0,0 +1,60 @@ +# Qwen3.8 DSpark Edge-LLM adapter + +The Qwen3.8 family owns its configuration admission, ONNX command mapping, +bundle assets and C++ runtime orchestration. It does not reuse another Qwen +family. Standalone builds retain the original native path; this change admits +Edge only for an explicit mixed-NVFP4 target plus DSpark companion. + +## Build and inference + +Provision the optional native SDK with `TRTMC_EDGELLM_ALL_KERNELS=ON` and +`TRTMC_EDGELLM_ONNX=ON` using the +[pinned package instructions](../../cmake/edgellm/README.md). The source is +GitHub Edge-LLM0.10.1 at `e8b29522938901f6df19ebeedd4b69bc8edbcd97`. +Configure `CMAKE_PREFIX_PATH` for the installed package and compile the runtime +with `TRTMC_ENABLE_EDGELLM=ON`. Cross compilation is unsupported. + +The Python build API accepts `BuildExecutionInputs(variant="dspark", +checkpoints=(NamedCheckpoint("draft", draft_path),))` through its `execution` +argument. The CLI equivalent adds `--execution-variant dspark` and +`--companion draft=/path/to/draft` to an ordinary build invocation. +The family invokes the original Edge Python ONNX exporter and native +`edgellm-onnx-build`; it does not alter source tensors or pad safetensors headers. +The exporter resolves the draft LM head from the target checkpoint. Both +speculative engines, embedding/head sidecars and tokenizer assets are bundled; +checkpoint weights are not duplicated in the bundle. + +The runtime calls the original Edge speculative inference constructor with +proposal block7, verify8, drafting topK1/step1 and DSpark scheduling disabled. +It preserves supported sampling controls. Engine preparation errors warn and +try native once with the same requested execution variant; native currently +rejects the unmapped DSpark variant explicitly. No failure substitutes an +ordinary base-only decoder. Inference errors propagate without fallback. + +## Validated exact profile + +- Target: `RadixArk/Qwen3.8-27B-NVFP4`, revision + `319f741cce68d7914884900c138a1fbb70a42f30`. +- Draft: `RadixArk/Qwen3.8-27B-DSpark`, revision + `b9a5dbdf03bc999c6c73c426b19c2d9041cea393`. +- Native SM120, CUDA13.3, TensorRT11.1.0.106, FP16 execution with source mixed + NVFP4/FP8 metadata retained, TP1/batch1, input/KV capacity1024. +- Actual Model Connect paired build and public CLI inference passed. +- Independent greedy oracle: exact token match, NED0.0 against0.15. + The saved CPU FP32 reference was reused after checkpoint-byte and reference + function verification; it was not regenerated during this run. +- Original Edge `llm_basic` prompt and128-token sampled profile + (temperature1/topK50/topP1): ROUGE-1 **0.4246**, ROUGE-L **0.2458**, above + unchanged **0.25/0.20** gates. Chat enabled and thinking disabled. +- Publication regressions:11 existing family Python tests plus106 existing + builder/architecture tests pass; both existing native C++ tests pass. + +These results qualify only the exact text profile, not standalone NVFP4, +multimodal inputs, other checkpoints/platforms or statistical sampling parity. +The local run reused existing owning E2E helpers with an explicit companion; +this pair is not yet a registered pytest manifest case. + +A first inference attempt exposed incompatible development JSON headers sharing +Edge’s3.12.0 version label. Matching the exact pinned headers fixed the crash +without changing Edge or the engines. The generic SDK now checks header content +rather than relying only on the version label. diff --git a/families/qwen3_8/dispatch.py b/families/qwen3_8/dispatch.py new file mode 100644 index 0000000000..d833bf71ac --- /dev/null +++ b/families/qwen3_8/dispatch.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned complete-network route map; native is the default.""" + +from __future__ import annotations + +import json +import logging +import os +from pathlib import Path +import tempfile +import traceback + +from . import edge_llm + +_LOG = logging.getLogger(__name__) + +# Exact qualified native DSpark route; no unvalidated platform or ordinary offload. +EDGE_DISPATCH = { + ("linux", "x86_64", 120, "fp16"): edge_llm.prepare_dspark, +} + + +def candidate(request, raw: dict) -> bool: + """Return whether this family's model/request contract can delegate to Edge.""" + config = raw.get("text_config", raw) + source_quantization = edge_llm.checkpoint_quantization(Path(request.model_dir), raw) + return ( + isinstance(config, dict) + and raw.get("model_type") == "qwen3_5" + and ("output_gate_type" in config and "mlp_only_layers" not in config) + and config.get("linear_key_head_dim") == config.get("linear_value_head_dim") == 128 + and not config.get("num_experts") + and source_quantization == "nvfp4" + and request.backend == "trt" and request.task == "text_generation" + and request.precision.lower() == "fp16" + and request.quantization in {None, source_quantization} + and request.max_batch_size == request.tensor_parallel_size == request.context_parallel_size == 1 + and not request.dynamic_kv_cache and not request.fp32_layers and request.graph_transform is None + and all(value is None for value in (request.image_height, request.image_width, request.video_num_frames)) + ) + + +def build(request, writer, native, *, draft_dir: Path) -> None: + """Dispatch locally or warn and retry native once with the original request. + + Args: + request: Unmodified Model Connect build request. + writer: Unpublished bundle writer. + native: This family's original native builder callback. + + Raises: + Exception: Common input/publication error, or native build error with + Edge cause after a failed preparation. Cancellation never retries. + """ + raw = json.loads((Path(request.model_dir) / "config.json").read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError("checkpoint config.json must contain an object") + config = raw.get("text_config", raw) + if not isinstance(config, dict): + raise ValueError("checkpoint text_config must contain an object") + if not candidate(request, raw): + native(request, writer) + return + capacity = config.get("max_position_embeddings") + if type(capacity) is not int or capacity <= 0: + raise ValueError("checkpoint max_position_embeddings must be a positive integer") + if request.max_sequence_length and request.max_sequence_length > capacity: + raise ValueError("max_sequence_length exceeds checkpoint context capacity") + failure = None + descriptor, name = tempfile.mkstemp(prefix=f".{request.output_path.name}.edge-", suffix=".log", + dir=request.output_path.parent) + os.close(descriptor) + log_path = Path(name) + with tempfile.TemporaryDirectory(prefix="trtmc-qwen3_8-edge-") as directory: + try: + target = edge_llm.local_target() + key = (target["os"], target["arch"], target["sm"], request.precision.lower()) + adapter = EDGE_DISPATCH.get(key) + if adapter is not None: + files, marker = adapter(request, raw, target, Path(directory), log_path, draft_dir) + except Exception as error: + failure = error + with log_path.open("a", encoding="utf-8") as log: + traceback.print_exception(error, file=log) + _LOG.warning("qwen3_8 Edge build failed: %s. Diagnostics: %s. " + "Retrying native once with the unchanged request.", error, log_path, exc_info=True) + else: + # Edge preparation did not touch writer; publication cannot fallback. + if adapter is not None: + edge_llm.publish(request, writer, files, marker) + return + log_path.unlink() # A platform non-match is not an Edge failure. + try: + native(request, writer) + except Exception as error: + if failure is not None: + raise error from failure + raise diff --git a/families/qwen3_8/edge_llm.py b/families/qwen3_8/edge_llm.py new file mode 100644 index 0000000000..c086d7fc09 --- /dev/null +++ b/families/qwen3_8/edge_llm.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Family-owned paired adapter to the pinned Edge ONNX builder API.""" + +from __future__ import annotations + +import json +from pathlib import Path +import shutil +import subprocess + +from tensorrt_model_connect.build import cmake_prefixes, detect_local_platform, subprocess_environment + +EDGE_REVISION = "e8b29522938901f6df19ebeedd4b69bc8edbcd97" + + +def local_target() -> dict: + """Return the executing worker identity supplied by generic build mechanics.""" + return detect_local_platform() + + +def installed_package(target: dict) -> dict: + """Resolve CMake installation via standard prefixes; never install anything. + + Args: + target: Executing device and SDK identity. + + Returns: + Validated package metadata with absolute Python and plugin paths. + + Raises: + FileNotFoundError: No CMake installation or required artifact exists. + ValueError: Pin, architecture, SDK or contained-path contract differs. + """ + for prefix in cmake_prefixes(): + manifest = prefix / "share/trtmc/edge-llm.json" + if not manifest.is_file(): + continue + package = json.loads(manifest.read_text(encoding="utf-8")) + if package.get("schema_version") != 1 or package.get("revision") != EDGE_REVISION: + raise ValueError(f"Edge package has an unsupported revision/schema: {manifest}") + if package.get("version") != "0.10.1" or package.get("arch") != target["arch"]: + raise ValueError("Edge package version/architecture differs from executing worker") + if target["sm"] not in package.get("architectures", []): + raise ValueError("Edge package was not built for this local GPU") + cuda_version = ".".join(str(package.get("cuda_version", "")).split(".")[:2]) + if cuda_version != target["cuda_version"] or package.get("tensorrt_version") != target["tensorrt_version"]: + raise ValueError("Edge package CUDA/TensorRT differs from executing worker") + for name in ("python", "plugin") + (("onnx_builder",) if package.get("onnx") else ()): + relative = Path(package[name]) + path = (prefix / relative).resolve() + if relative.is_absolute() or not path.is_relative_to(prefix.resolve()): + raise ValueError(f"Edge package {name} must be contained in its installation") + if not path.is_file(): + raise FileNotFoundError(f"Edge package {name} is missing: {path}") + package[name] = str(path) + return package + raise FileNotFoundError("Edge-LLM is not installed; enable the optional Edge-LLM CMake dependency " + "and set CMAKE_PREFIX_PATH to its install prefix") + + + +def checkpoint_quantization(model_dir: Path, raw: dict) -> str | None: + """Admit plain weights or the documented mixed ModelOpt checkpoint. + + Edge owns per-layer decoding and FP8 KV interpretation. Require matching + embedded/sidecar metadata rather than inventing or converting a format. + """ + if any((model_dir / name).exists() for name in ("quantize_config.json", "quant_config.json")): + return None + embedded = raw.get("quantization_config") + nested = raw.get("text_config", {}).get("quantization_config") + sidecar = model_dir / "hf_quant_config.json" + if not embedded and not nested and not sidecar.exists(): + return "none" + if not isinstance(embedded, dict) or (nested and nested != embedded): + return None + if embedded.get("quant_method") != "modelopt" or embedded.get("quant_algo") != "MIXED_PRECISION": + return None + layers = embedded.get("quantized_layers") + if not isinstance(layers, dict) or not layers or any(not isinstance(v, dict) for v in layers.values()): + return None + if {v.get("quant_algo") for v in layers.values()} != {"FP8", "NVFP4"}: + return None + if not sidecar.is_file(): + return None + value = json.loads(sidecar.read_text(encoding="utf-8")) + quant = value.get("quantization") if isinstance(value, dict) else None + if not isinstance(quant, dict) or quant.get("quant_algo") != "MIXED_PRECISION": + return None + if quant.get("quantized_layers") != layers: + return None + if quant.get("kv_cache_quant_algo") != "FP8" or not embedded.get("kv_cache_scheme"): + return None + return "nvfp4" + +_PROMPT_PROGRAM = r"""import json, sys +from pathlib import Path +from transformers import AutoTokenizer +checkpoint = Path(sys.argv[sys.argv.index("--model-dir") + 1]) +engine = Path(sys.argv[sys.argv.index("--engine-dir") + 1]) +tokenizer = AutoTokenizer.from_pretrained(checkpoint, local_files_only=True, trust_remote_code=False) +slot = "Qwen38SingleUserContentSlot" +formats = {} +for thinking in (False, True): + options = dict(tokenize=False, add_generation_prompt=True, enable_thinking=thinking) + rendered = tokenizer.apply_chat_template([dict(role="user", content=slot)], **options) + if rendered.count(slot) != 1: + raise ValueError("Qwen3.8 source does not preserve a single user prompt") + prefix, suffix = rendered.split(slot) + whitespace = "".join(chr(cp) for cp in range(sys.maxunicode + 1) if chr(cp).isspace()) + for probe in ("", " leading and trailing ", "first\nsecond", "世界", whitespace + "text" + whitespace): + actual = tokenizer.apply_chat_template([dict(role="user", content=probe)], **options) + if actual != prefix + probe.strip() + suffix: + raise ValueError("Qwen3.8 source user content is not a prefix/suffix mapping") + formats[str(thinking).lower()] = dict(prefix=prefix, suffix=suffix) +(engine / "trtmc_single_user_prompts.json").write_text(json.dumps(formats, ensure_ascii=False)) +""" + + + +def prepare_dspark(request, raw: dict, target: dict, staging: Path, log_path: Path, + draft_dir: Path) -> tuple[dict, dict]: + """Map the paired request to the original exporter and native ONNX builder. + + Edge supplies the draft LM head from its target checkpoint. No safetensors + padding or weight conversion is needed, and baked weights are not bundled + twice. Both plans and their complete runtime assets must exist to publish. + """ + package = installed_package(target) + if package.get("onnx") is not True: + raise ValueError("Qwen3.8 DSpark requires an ONNX-enabled Edge SDK") + source, draft_dir = Path(request.model_dir).resolve(), draft_dir.resolve() + if checkpoint_quantization(source, raw) != "nvfp4": + raise ValueError("Qwen3.8 DSpark requires the mixed NVFP4 target") + if not list(source.glob("*.safetensors")) or not list(draft_dir.glob("*.safetensors")): + raise ValueError("Qwen3.8 DSpark requires both local safetensors checkpoints") + config = raw.get("text_config", raw) + limit = request.max_sequence_length or min(int(config["max_position_embeddings"]), 256) + if not 8 < limit <= 1024: + raise ValueError("Qwen3.8 DSpark requires capacity above verification size8 and at most1024") + checkpoint = staging / "edge_llm/checkpoint" + checkpoint.mkdir(parents=True) + for name in ("config.json", "tokenizer.json", "tokenizer_config.json", "generation_config.json", + "chat_template.jinja"): + if (source / name).is_file(): + shutil.copy2(source / name, checkpoint / name) + (checkpoint / "draft").mkdir() + shutil.copy2(draft_dir / "config.json", checkpoint / "draft/config.json") + engine, onnx = staging / "edge_llm/engine", staging / "onnx" + env = subprocess_environment( + {"EDGELLM_PLUGIN_PATH": package["plugin"]}, + prepend_paths={"LD_LIBRARY_PATH": str(Path(package["plugin"]).parent)}, + ) + for role, subdirectory, flag in (("draft", "dspark_draft", "--specDraft"), + ("base", "llm", "--specBase")): + commands = [ + [package["python"], "-I", "-m", "tensorrt_edgellm.scripts.export", + str(source), str(onnx), f"--dspark-{role}", "--dspark-draft-dir", str(draft_dir), + "--skip-visual", "--skip-audio"], + [package["onnx_builder"], "--onnxDir", str(onnx / subdirectory), + "--engineDir", str(engine), flag, "--maxInputLen", str(min(limit, 1024)), + "--maxKVCacheCapacity", str(limit), "--maxBatchSize", "1", + "--maxVerifyTreeSize", "8", "--maxDraftTreeSize", "7"], + ] + with log_path.open("a", encoding="utf-8") as log: + for command in commands: + log.write(json.dumps(command) + "\n") + log.flush() + subprocess.run(command, check=True, stdout=log, stderr=subprocess.STDOUT, + cwd=staging, env=env) + shutil.rmtree(onnx) # Only this preparation's successfully consumed intermediates. + with log_path.open("a", encoding="utf-8") as log: + subprocess.run([package["python"], "-I", "-c", _PROMPT_PROGRAM, + "--model-dir", str(checkpoint), "--engine-dir", str(engine)], + check=True, stdout=log, stderr=subprocess.STDOUT, cwd=staging, env=env) + required = ("spec_base.engine", "spec_draft.engine", "base_config.json", "draft_config.json", + "embedding.safetensors", "dspark_heads.safetensors", "dspark_heads_info.json", + "tokenizer.json", "tokenizer_config.json", + "processed_chat_template.json", "trtmc_single_user_prompts.json") + for name in required: + if not (engine / name).is_file() or (engine / name).stat().st_size == 0: + raise ValueError(f"Edge ONNX builder did not produce required artifact: {name}") + for role in ("base", "draft"): + built = json.loads((engine / f"{role}_config.json").read_text(encoding="utf-8")) + if built.get("spec_decode_type") != "dspark" or built.get("dspark_config", {}).get("block_size") != 7: + raise ValueError("Edge ONNX builder returned a different DSpark contract") + files = {} + for directory in (engine, checkpoint): + for path in sorted(directory.rglob("*")): + if path.is_symlink(): + raise ValueError(f"Edge output must not contain symlinks: {path}") + if path.is_file(): + files[path.relative_to(staging).as_posix()] = path + return files, { + "version": 1, "edge_revision": EDGE_REVISION, "target": target, "precision": "fp16", + "max_sequence_length": limit, "max_input_length": min(limit, 1024), "max_batch_size": 1, + "checkpoint_quantization": "nvfp4", "artifacts": list(files), + "execution_variant": "dspark", "builder_flow": "onnx", "dspark_block_size": 7, + } + + +def publish(request, writer, files: dict, marker: dict) -> None: + """Stream complete Edge sections; publication errors must not retry native.""" + writer.set_header(family=request.family, task=request.task, backend=request.backend) + for name, path in files.items(): + with path.open("rb") as source, writer.open_section(name) as destination: + shutil.copyfileobj(source, destination, length=1024 * 1024) + writer.add_json("edge_llm.json", marker) diff --git a/families/qwen3_8/model.py b/families/qwen3_8/model.py index bfe571aa04..70995b5935 100644 --- a/families/qwen3_8/model.py +++ b/families/qwen3_8/model.py @@ -148,3 +148,50 @@ def build(request, writer) -> None: path = model_dir / filename if path.is_file(): writer.add_bytes(filename, path.read_bytes()) + + +def build_with_inputs(request, writer, execution) -> None: + """Build the retained mixed-NVFP4 Qwen3.8 / DSpark block7 pair.""" + from . import dispatch, edge_llm + + if execution.variant != "dspark" or tuple(x.role for x in execution.checkpoints) != ("draft",): + raise ValueError("Qwen3.8 paired execution requires variant=dspark and one draft checkpoint") + draft_dir = execution.checkpoints[0].model_dir + raw = json.loads((request.model_dir / "config.json").read_text()) + draft = json.loads((draft_dir / "config.json").read_text()) + if not dispatch.candidate(request, raw) or edge_llm.checkpoint_quantization(request.model_dir, raw) != "nvfp4": + raise ValueError("The retained Qwen3.8 DSpark pair requires a matching mixed-NVFP4 base") + base = raw.get("text_config", raw) + if not isinstance(draft, dict) or draft.get("architectures") != ["DSparkDraftModel"]: + raise ValueError("Expected a DSparkDraftModel companion") + for name in ("hidden_size", "vocab_size"): + if type(draft.get(name)) is not int or draft[name] != base.get(name): + raise ValueError(f"Qwen3.8 DSpark base and draft disagree on {name}") + if draft.get("num_target_layers") != base.get("num_hidden_layers"): + raise ValueError("Qwen3.8 DSpark target layer count differs from base") + config = draft.get("dspark_config") + if not isinstance(config, dict) or config.get("block_size", draft.get("block_size")) != 7: + raise ValueError("Qwen3.8 DSpark maps the upstream block7 / verify8 profile") + layers = config.get("target_layer_ids", draft.get("target_layer_ids")) + if (not isinstance(layers, list) or not layers + or any(type(i) is not int or not 0 <= i < base["num_hidden_layers"] for i in layers) + or len(set(layers)) != len(layers)): + raise ValueError("Invalid Qwen3.8 DSpark target layer IDs") + mask = config.get("mask_token_id", draft.get("mask_token_id")) + if type(mask) is not int or not 0 <= mask < base["vocab_size"]: + raise ValueError("Invalid Qwen3.8 DSpark mask token") + limit = request.max_sequence_length or min(base["max_position_embeddings"], 256) + capacity = draft.get("max_position_embeddings") + if type(capacity) is not int or not 8 < limit <= capacity: + raise ValueError("Requested context exceeds DSpark draft capacity or block minimum") + if draft.get("quantization_config") or any( + (draft_dir / name).exists() + for name in ("hf_quant_config.json", "quantize_config.json", "quant_config.json") + ): + raise ValueError("This Qwen3.8 DSpark profile requires unquantized draft weights") + + def native_pair(original_request, original_writer): + # A failure must never replace the requested pair with base-only decoding. + raise NotImplementedError("Native Qwen3.8 does not implement the requested DSpark variant") + + dispatch.build(request, writer, native_pair, draft_dir=draft_dir) diff --git a/families/qwen3_8/runtime/CMakeLists.txt b/families/qwen3_8/runtime/CMakeLists.txt index af423a3801..b74b007112 100644 --- a/families/qwen3_8/runtime/CMakeLists.txt +++ b/families/qwen3_8/runtime/CMakeLists.txt @@ -21,7 +21,10 @@ target_link_libraries(trtmc_model_qwen3_8 PRIVATE nlohmann_json::nlohmann_json ${TRTMC_CUDART_LIBRARY} ) -target_compile_options(trtmc_model_qwen3_8 PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_model_qwen3_8 PRIVATE + "$<$:-Wall;-Wextra;-Wpedantic>" + "$<$:-Xcompiler=-Wall,-Wextra>" +) set_target_properties(trtmc_model_qwen3_8 PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" BUILD_RPATH "\$ORIGIN" @@ -71,3 +74,26 @@ if(TRTMC_BUILD_TESTS) SKIP_RETURN_CODE 77 ) endif() + +# Complete-network offload is family-owned and absent from native-only builds. +if(TARGET EdgeLLM::Core) + target_sources(trtmc_model_qwen3_8 PRIVATE + edge_llm/adapter.cpp + edge_llm/device_link.cu + ) + target_compile_definitions(trtmc_model_qwen3_8 PRIVATE TRTMC_HAS_EDGE_LLM=1) + target_link_libraries(trtmc_model_qwen3_8 PRIVATE EdgeLLM::Core) + set_target_properties(trtmc_model_qwen3_8 PROPERTIES + CUDA_ARCHITECTURES "${EdgeLLM_CUDA_ARCHITECTURE}" + CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON + ) +endif() + +if(TARGET EdgeLLM::Core) + add_custom_command(TARGET trtmc_model_qwen3_8 POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + $ $ + VERBATIM + ) +endif() diff --git a/families/qwen3_8/runtime/edge_llm/adapter.cpp b/families/qwen3_8/runtime/edge_llm/adapter.cpp new file mode 100644 index 0000000000..a23687fdb2 --- /dev/null +++ b/families/qwen3_8/runtime/edge_llm/adapter.cpp @@ -0,0 +1,273 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "families/qwen3_8/runtime/edge_llm/adapter.h" + +#include "families/qwen3_8/runtime/edge_llm/request.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc::qwen3_8::edge_llm { +namespace { +namespace fs = std::filesystem; + +/// Turn CUDA failures into caller-visible load or inference errors. +void check_cuda(cudaError_t result) { + if (result != cudaSuccess) + throw std::runtime_error(std::string("Qwen3.8 Edge CUDA error: ") + + cudaGetErrorString(result)); +} + +/// Reject an engine built for a different local GPU or CUDA/TensorRT runtime. +void validate_target(const nlohmann::json& target) { + utsname host{}; + if (uname(&host) != 0) + throw std::runtime_error("Cannot identify Qwen3.8 Edge runtime host"); + std::ifstream release("/etc/os-release"); + std::string line, os_version; + while (std::getline(release, line)) { + if (line.rfind("VERSION_ID=", 0) == 0) { + os_version = line.substr(11); + if (os_version.size() >= 2 && os_version.front() == char(34) && + os_version.back() == char(34)) + os_version = os_version.substr(1, os_version.size() - 2); + } + } + int device = 0, cuda_version = 0; + check_cuda(cudaGetDevice(&device)); + check_cuda(cudaRuntimeGetVersion(&cuda_version)); + cudaDeviceProp gpu{}; + check_cuda(cudaGetDeviceProperties(&gpu, device)); + const int trt_version = getInferLibVersion(); + const std::string trt = + std::to_string(trt_version / 10000) + "." + std::to_string((trt_version % 10000) / 100) + + "." + std::to_string(trt_version % 100) + "." + std::to_string(getInferLibBuildVersion()); + const std::string cuda = + std::to_string(cuda_version / 1000) + "." + std::to_string((cuda_version % 1000) / 10); + if (target.at("os") != "linux" || target.at("os_version") != os_version || + target.at("arch") != host.machine || target.at("sm") != gpu.major * 10 + gpu.minor || + target.at("cuda_version") != cuda || target.at("tensorrt_version") != trt) + throw std::runtime_error( + "Qwen3.8 Edge bundle requires its build GPU and CUDA/TensorRT stack"); +} + +/// Own extracted engine/checkpoint files until after the Edge runtime is destroyed. +class Artifacts { + public: + explicit Artifacts(const BundleReader& bundle, const nlohmann::json& marker) { + std::set names; + for (const auto& entry : marker.at("artifacts")) { + const auto name = entry.get(); + if (!safe_artifact_path(name) || !names.insert(name).second || + !bundle.find_section(name)) + throw std::runtime_error("Invalid Qwen3.8 Edge artifact: " + name); + } + std::vector required_files{ + "edge_llm/engine/tokenizer.json", "edge_llm/engine/tokenizer_config.json", + "edge_llm/engine/processed_chat_template.json", + "edge_llm/engine/trtmc_single_user_prompts.json", "edge_llm/checkpoint/config.json"}; + + for (const auto* name : + {"spec_base.engine", "spec_draft.engine", "base_config.json", "draft_config.json", + "embedding.safetensors", "dspark_heads.safetensors", "dspark_heads_info.json"}) + required_files.push_back(std::string("edge_llm/engine/") + name); + required_files.push_back("edge_llm/checkpoint/draft/config.json"); + for (const auto& required : required_files) + if (!names.count(required) || bundle.find_section(required)->length == 0) + throw std::runtime_error("Required Qwen3.8 Edge artifact missing: " + required); + std::string pattern = (fs::temp_directory_path() / "trtmc-qwen3_8-edge-XXXXXX").string(); + if (!mkdtemp(pattern.data())) + throw std::runtime_error("Cannot create Qwen3.8 Edge artifact directory"); + root_ = pattern; + try { + for (const auto& name : names) { + const auto destination = root_ / name; + fs::create_directories(destination.parent_path()); + std::ofstream output(destination, std::ios::binary); + bundle.copy_section(name, output); + output.close(); + if (!output) + throw std::runtime_error("Cannot extract Qwen3.8 Edge artifact: " + name); + } + } catch (...) { + cleanup(); + throw; + } + } + ~Artifacts() { cleanup(); } + Artifacts(const Artifacts&) = delete; + Artifacts& operator=(const Artifacts&) = delete; + std::string engine() const { return (root_ / "edge_llm/engine").string(); } + std::string checkpoint() const { return (root_ / "edge_llm/checkpoint").string(); } + + private: + void cleanup() noexcept { + std::error_code ignored; + fs::remove_all(root_, ignored); + } + fs::path root_; +}; + +/// Close the plugin handle after runtime destruction; registrations remain mapped. +struct CloseLibrary { + void operator()(void* handle) const noexcept { + if (handle) + dlclose(handle); + } +}; + +/// Initialize the CMake-installed adjacent plugin without process-global environment mutation. +std::unique_ptr load_plugin() { + Dl_info location{}; + if (!dladdr(reinterpret_cast(&create), &location) || !location.dli_fname) + throw std::runtime_error("Cannot locate Qwen3.8 family library"); + const auto path = + fs::absolute(location.dli_fname).parent_path() / "libNvInfer_edgellm_plugin.so"; + std::unique_ptr plugin( + dlopen(path.c_str(), RTLD_NOW | RTLD_GLOBAL | RTLD_NODELETE)); + if (!plugin) + throw std::runtime_error("Cannot load CMake-installed Edge plugin: " + + std::string(dlerror())); + using Initialize = bool (*)(void*, const char*); + auto initialize = reinterpret_cast(dlsym(plugin.get(), "initEdgellmPlugins")); + if (!initialize || !initialize(static_cast(&trt_edgellm::gLogger), "")) + throw std::runtime_error("Cannot initialize Qwen3.8 Edge plugin"); + return plugin; +} + +/// Stream ownership is independent of construction success and outlives the Edge instance. +class Stream { + public: + Stream() { check_cuda(cudaStreamCreateWithFlags(&value_, cudaStreamNonBlocking)); } + ~Stream() { cudaStreamDestroy(value_); } + Stream(const Stream&) = delete; + Stream& operator=(const Stream&) = delete; + cudaStream_t get() const { return value_; } + + private: + cudaStream_t value_{nullptr}; +}; + +/// Delegate the complete DSpark block7 algorithm to the pinned runtime. +std::unique_ptr make_runtime(const Artifacts& artifacts, + cudaStream_t stream) { + trt_edgellm::rt::SpecDecodeDraftingConfig drafting{}; + drafting.draftingTopK = 1; + drafting.draftingStep = 1; + drafting.verifySize = 8; + drafting.dflashBlockSize = 0; + drafting.dsparkSchedulerMode = trt_edgellm::rt::DSparkSchedulerMode::kOff; + drafting.dsparkConfidenceThreshold = 0.0F; + drafting.dsparkMinProposalLen = 1; + drafting.dsparkMaxProposalLen = 0; + return std::make_unique( + artifacts.engine(), "", std::unordered_map{}, drafting, stream, + trt_edgellm::rt::ContextCacheConfig{}, artifacts.checkpoint(), + (fs::path(artifacts.checkpoint()) / "draft").string()); +} + +/// Thin persistent Edge API adapter; serialization prevents concurrent use of Edge request state. +class EdgeTask final : public ITextGeneration { + public: + EdgeTask(const BundleReader& bundle, const nlohmann::json& marker) + : artifacts_(bundle, marker), plugin_(load_plugin()), + runtime_(make_runtime(artifacts_, stream_.get())), + capacity_(marker.at("max_sequence_length").get()), + input_limit_(marker.at("max_input_length").get()) { + std::ifstream input(fs::path(artifacts_.engine()) / "trtmc_single_user_prompts.json"); + source_prompts_ = nlohmann::json::parse(input); + for (const auto* mode : {"false", "true"}) + for (const auto* part : {"prefix", "suffix"}) + if (!source_prompts_.at(mode).at(part).is_string()) + throw std::runtime_error("Invalid Qwen3.8 source prompt mapping"); + } + + std::int32_t default_max_new_tokens() const override { return std::min(128, capacity_ - 1); } + + /// Drain work from failed requests before destroying the runtime and its weight buffers. + ~EdgeTask() override { cudaStreamSynchronize(stream_.get()); } + + /// Invoke Edge once; failures propagate without attempting native inference. + TextResult generate(const std::string& prompt, const TextGenerationConfig& config) override { + auto effective = config; + std::string source_prompt = prompt; + if (config.use_chat_template) { + const auto& format = source_prompts_.at(config.enable_thinking ? "true" : "false"); + source_prompt = format.at("prefix").get() + source_user_content(prompt) + + format.at("suffix").get(); + effective.use_chat_template = false; + } + auto request = make_request(source_prompt, effective, default_max_new_tokens()); + std::lock_guard lock(mutex_); + const auto counts = runtime_->countPromptTokens(request); + if (counts.size() != 1) + throw std::runtime_error("Qwen3.8 Edge returned invalid prompt counts"); + validate_capacity(counts.front(), input_limit_, capacity_, request.maxGenerateLength); + trt_edgellm::rt::LLMGenerationResponse response{}; + if (!runtime_->handleRequest(request, response, stream_.get()) || + response.outputIds.size() != 1 || response.outputTexts.size() != 1 || + response.outputIds.front().empty() || + response.outputIds.front().size() > static_cast(request.maxGenerateLength)) + throw std::runtime_error("Qwen3.8 Edge generation failed"); + if (response.finishReasons.size() != 1 || + (response.finishReasons.front() != trt_edgellm::rt::FinishReason::kEndId && + response.finishReasons.front() != trt_edgellm::rt::FinishReason::kLength)) + throw std::runtime_error("Qwen3.8 Edge generation did not complete successfully"); + // This API does not expose per-request stage times; zero means unavailable. + return {std::move(response.outputTexts.front()), std::move(response.outputIds.front())}; + } + + private: + // Reverse destruction order keeps weights, plugin and stream alive throughout Edge teardown. + Artifacts artifacts_; + std::unique_ptr plugin_; + Stream stream_; + std::unique_ptr runtime_; + int capacity_; + int input_limit_; + nlohmann::json source_prompts_; + std::mutex mutex_; +}; +} // namespace + +ITask* create(const BundleReader& bundle) { + const auto bytes = bundle.read_section("edge_llm.json"); + const auto marker = nlohmann::json::parse(bytes.begin(), bytes.end()); + if (marker.at("version") != 1 || marker.at("edge_revision") != kRevision || + marker.at("max_sequence_length").get() <= 1 || + marker.at("max_input_length").get() <= 0 || + marker.at("max_input_length").get() > marker.at("max_sequence_length").get() || + marker.at("max_batch_size") != 1 || marker.at("precision") != "fp16" || + !marker.at("artifacts").is_array()) + throw std::runtime_error("Invalid Qwen3.8 Edge bundle contract"); + const auto variant = marker.value("execution_variant", ""); + if (!valid_execution_variant(variant, marker.value("dspark_block_size", 0))) + throw std::runtime_error("Unsupported Qwen3.8 Edge execution variant"); + if (marker.value("builder_flow", "") != "onnx" || + marker.value("checkpoint_quantization", "") != "nvfp4" || + marker.at("max_sequence_length").get() <= 8 || + marker.at("max_sequence_length").get() > 1024 || + marker.at("target").at("arch") != "x86_64" || marker.at("target").at("sm") != 120) + throw std::runtime_error("Qwen3.8 DSpark requires the ONNX builder flow"); + validate_target(marker.at("target")); + return new EdgeTask(bundle, marker); +} + +} // namespace trtmc::qwen3_8::edge_llm diff --git a/families/qwen3_8/runtime/edge_llm/adapter.h b/families/qwen3_8/runtime/edge_llm/adapter.h new file mode 100644 index 0000000000..86cf5a89ce --- /dev/null +++ b/families/qwen3_8/runtime/edge_llm/adapter.h @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "trtmc/bundle.h" +#include "trtmc/task.h" + +namespace trtmc::qwen3_8::edge_llm { + +/// Create a persistent Edge task from a self-contained bundle; throws on load failure. +ITask* create(const BundleReader& bundle); + +} // namespace trtmc::qwen3_8::edge_llm diff --git a/families/qwen3_8/runtime/edge_llm/contract.h b/families/qwen3_8/runtime/edge_llm/contract.h new file mode 100644 index 0000000000..aedecb8588 --- /dev/null +++ b/families/qwen3_8/runtime/edge_llm/contract.h @@ -0,0 +1,61 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "trtmc/task.h" + +#include +#include +#include +#include + +namespace trtmc::qwen3_8::edge_llm { + +inline constexpr const char* kRevision = "e8b29522938901f6df19ebeedd4b69bc8edbcd97"; + +/// DSpark proposal length excludes the one additional verification token. +inline bool valid_execution_variant(const std::string& variant, int block_size) { + return variant == "dspark" && block_size == 7; +} + +/// Return whether an artifact is a normalized file below one of the two Edge roots. +inline bool safe_artifact_path(const std::string& name) { + if (name.find('\\') != std::string::npos || name.find('\0') != std::string::npos) + return false; + const std::filesystem::path path(name); + if (path.is_absolute() || path.filename().empty()) + return false; + for (const auto& part : path) + if (part == "." || part == "..") + return false; + return path.generic_string() == name && + (name.rfind("edge_llm/engine/", 0) == 0 || name.rfind("edge_llm/checkpoint/", 0) == 0); +} + +/// Reject invalid sampling settings and controls with no equivalent Edge request API. +inline void validate_generation(const TextGenerationConfig& c) { + if (!std::isfinite(c.temperature) || c.temperature < 0 || !std::isfinite(c.top_p) || + c.top_p <= 0 || c.top_p > 1 || c.top_k < 0) + throw std::invalid_argument("Invalid Qwen3.8 Edge sampling parameters"); + if (c.min_p != 0 || c.seed != -1 || c.eos_token_id != -1 || c.repetition_penalty != 1 || + !c.lora_adapter_id.empty() || c.stop_on_boxed_answer || + (c.text_generation_mode != "auto" && c.text_generation_mode != "autoregressive") || + c.source_language_token_id != -1 || c.forced_bos_token_id != -1 || c.guidance_scale != -1 || + c.cfg_scale != -1 || c.num_steps != -1 || c.sde_gamma != -1 || !c.initial_latents.empty() || + !c.condition_latents.empty() || !c.condition_mask.empty() || !c.sampling_steps.empty() || + !c.sde_noises.empty() || c.block_length != 0 || c.confidence_threshold != -1) + throw std::invalid_argument( + "Requested generation controls are unsupported by Qwen3.8 Edge"); +} + +/// Enforce prompt and total capacity without allowing Edge to silently clip generation. +inline void validate_capacity(int prompt_tokens, int input_limit, int capacity, + std::int64_t generated_tokens) { + if (prompt_tokens <= 0 || prompt_tokens > input_limit || generated_tokens <= 0 || + generated_tokens > static_cast(capacity) - prompt_tokens) + throw std::invalid_argument("Qwen3.8 Edge prompt and generation exceed bundle capacity"); +} + +} // namespace trtmc::qwen3_8::edge_llm diff --git a/families/qwen3_8/runtime/edge_llm/device_link.cu b/families/qwen3_8/runtime/edge_llm/device_link.cu new file mode 100644 index 0000000000..1ea4768d60 --- /dev/null +++ b/families/qwen3_8/runtime/edge_llm/device_link.cu @@ -0,0 +1,5 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +// Enable the final CUDA device-link step for Edge's static runtime dependencies. diff --git a/families/qwen3_8/runtime/edge_llm/request.h b/families/qwen3_8/runtime/edge_llm/request.h new file mode 100644 index 0000000000..5a8e8a2e75 --- /dev/null +++ b/families/qwen3_8/runtime/edge_llm/request.h @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "families/qwen3_8/runtime/edge_llm/contract.h" + +#include + +namespace trtmc::qwen3_8::edge_llm { + +/// Match the source Jinja trim filter without locale-dependent ASCII-only trimming. +inline std::string source_user_content(std::string text) { + // Python str.strip whitespace, including C0 separators and Unicode spaces. + constexpr const char* whitespace[] = { + "\t", "\n", "\v", "\f", "\r", + "\x1c", "\x1d", "\x1e", "\x1f", " ", + "\xc2\x85", "\xc2\xa0", "\xe1\x9a\x80", "\xe2\x80\x80", "\xe2\x80\x81", + "\xe2\x80\x82", "\xe2\x80\x83", "\xe2\x80\x84", "\xe2\x80\x85", "\xe2\x80\x86", + "\xe2\x80\x87", "\xe2\x80\x88", "\xe2\x80\x89", "\xe2\x80\x8a", "\xe2\x80\xa8", + "\xe2\x80\xa9", "\xe2\x80\xaf", "\xe2\x81\x9f", "\xe3\x80\x80"}; + bool changed = true; + while (changed && !text.empty()) { + changed = false; + for (const std::string space : whitespace) { + if (text.compare(0, space.size(), space) == 0) { + text.erase(0, space.size()); + changed = true; + } + if (text.size() >= space.size() && + text.compare(text.size() - space.size(), space.size(), space) == 0) { + text.resize(text.size() - space.size()); + changed = true; + } + } + } + if (text.rfind("", 0) == 0 && text.size() >= 16 && + text.compare(text.size() - 16, 16, "") == 0) + throw std::invalid_argument("Qwen3.8 single-user prompt contains no user query"); + return text; +} + +/// Map Model Connect text arguments to the pinned Edge API; rejects unmapped controls. +inline trt_edgellm::rt::LLMGenerationRequest +make_request(const std::string& prompt, const TextGenerationConfig& config, int default_length) { + validate_generation(config); + trt_edgellm::rt::LLMGenerationRequest request{}; + request.requests.resize(1); + request.requests.front().messages.push_back({"user", {{"text", prompt}}}); + request.applyChatTemplate = config.use_chat_template; + request.enableThinking = config.enable_thinking; + request.temperature = config.temperature; + request.topK = config.top_k; + request.topP = config.top_p; + request.maxGenerateLength = config.max_new_tokens > 0 ? config.max_new_tokens : default_length; + return request; +} + +} // namespace trtmc::qwen3_8::edge_llm diff --git a/families/qwen3_8/runtime/plugin.cpp b/families/qwen3_8/runtime/plugin.cpp index 49a8b4c57f..114a0f4533 100644 --- a/families/qwen3_8/runtime/plugin.cpp +++ b/families/qwen3_8/runtime/plugin.cpp @@ -9,6 +9,9 @@ #include "families/qwen3_8/runtime/plugin_helpers.h" #include "families/qwen3_8/runtime/recurrent_state.h" #include "trtmc/runtime/family_factory.h" +#ifdef TRTMC_HAS_EDGE_LLM +#include "families/qwen3_8/runtime/edge_llm/adapter.h" +#endif #include #include @@ -119,6 +122,14 @@ std::string chat_template(const BundleReader& bundle) { } // namespace ITask* create(const FamilyContext& context) { + if (context.reader.find_section("edge_llm.json")) { +#ifdef TRTMC_HAS_EDGE_LLM + return edge_llm::create(context.reader); +#else + throw std::runtime_error("Qwen3.8 Edge bundle requires a runtime configured with " + "-DTRTMC_ENABLE_EDGELLM=ON; rebuild and install Model Connect"); +#endif + } const RuntimeConfig config = parse_runtime_config(context.reader); auto decoder = load_engine(context.backend, require_section(context.reader, "engine.plan"), "qwen3_8 decoder"); diff --git a/families/qwen3_8/tests/test_e2e.py b/families/qwen3_8/tests/test_e2e.py index cd5ad94989..5e284af876 100644 --- a/families/qwen3_8/tests/test_e2e.py +++ b/families/qwen3_8/tests/test_e2e.py @@ -133,7 +133,7 @@ def _thresholds(case_name: str) -> dict[str, float]: return thresholds -def _build_bundle(manifest: dict, model_dir: Path, bundle: Path) -> None: +def _build_bundle(manifest: dict, model_dir: Path, bundle: Path, *, execution=None) -> None: quantization = manifest.get("quantization") assert quantization is None or isinstance(quantization, str) fp32_layers = tuple(manifest.get("fp32_layers", ())) @@ -148,7 +148,8 @@ def _build_bundle(manifest: dict, model_dir: Path, bundle: Path) -> None: tensor_parallel_size=manifest["tensor_parallel_size"], quantization=quantization, fp32_layers=fp32_layers, - ) + ), + execution=execution, ) assert bundle.is_file() and bundle.stat().st_size > 0, bundle @@ -227,7 +228,7 @@ def _run_native( completed = subprocess.run( command, - check=True, + check=False, capture_output=True, text=True, timeout=600, @@ -235,6 +236,9 @@ def _run_native( ) record_evidence("commands", {"argv": getattr(completed, "args", None)}) record_evidence("native", {"stdout": getattr(completed, "stdout", None), "stderr": getattr(completed, "stderr", None)}) + (tmp_path / "native.stdout.log").write_text(completed.stdout, encoding="utf-8") + (tmp_path / "native.stderr.log").write_text(completed.stderr, encoding="utf-8") + completed.check_returncode() if tp_size == 1: return json.loads(completed.stdout) @@ -353,16 +357,102 @@ def _hf_reference( "bf16": torch.bfloat16, } assert reference_precision in dtypes, reference_precision - model = ( - AutoModelForCausalLM.from_pretrained( + if case.get("reference_decode_modelopt_mixed", False): + # Preserve the declared FP32 oracle; packed bytes are not floating weights. + # This does not emulate compiled activation or KV quantization. + from modelopt.torch.export.quant_utils import QUANTIZATION_FP8, from_quantized_weight + from modelopt.torch.quantization.qtensor import NVFP4QTensor + from safetensors.torch import load_file + from transformers import AutoConfig, GenerationConfig, Qwen3_5ForCausalLM + + assert not trust_remote_code and reference_precision == "fp32" + raw = json.loads((model_dir / "config.json").read_text()) + quant = json.loads((model_dir / "hf_quant_config.json").read_text())["quantization"] + embedded = raw["quantization_config"] + assert embedded["quant_method"] == "modelopt" + assert embedded["quant_algo"] == quant["quant_algo"] == "MIXED_PRECISION" + layers = quant["quantized_layers"] + assert layers and embedded["quantized_layers"] == layers + assert all( + policy["quant_algo"] == "FP8" + or (policy["quant_algo"] == "NVFP4" and policy["group_size"] == 16) + for policy in layers.values() + ) + state = {} + for shard in sorted(model_dir.glob("*.safetensors")): + tensors = load_file(str(shard), device="cpu") + assert not state.keys() & tensors.keys(), "duplicate checkpoint tensors" + state.update(tensors) + packed = {key for key, value in state.items() if value.dtype == torch.uint8} + fp8 = { + key for key, value in state.items() + if key.endswith(".weight") and value.dtype == torch.float8_e4m3fn + } + quantized = packed | fp8 + assert packed and fp8 and quantized == {key + ".weight" for key in layers} + assert all(layers[key.removesuffix(".weight")]["quant_algo"] == "NVFP4" for key in packed) + assert all(layers[key.removesuffix(".weight")]["quant_algo"] == "FP8" for key in fp8) + for key in packed: + weight = state[key] + assert key.endswith(".weight") and weight.ndim == 2, key + assert weight.shape[-1] % 8 == 0, key + prefix = key.removesuffix("weight") + scale = state[prefix + "weight_scale"] + double_scale = state[prefix + "weight_scale_2"] + shape = (weight.shape[0], weight.shape[1] * 2) + assert scale.dtype == torch.float8_e4m3fn + assert scale.shape == (shape[0], shape[1] // 16), key + assert torch.isfinite(scale.float()).all() and (scale.float() >= 0).all(), key + assert double_scale.numel() == 1 and torch.isfinite(double_scale).all(), key + assert (double_scale > 0).all(), key + state[key] = NVFP4QTensor(shape, torch.float32, weight).dequantize( + dtype=torch.float32, scale=scale, double_scale=double_scale, + block_sizes={-1: 16}, fast=False, + ) + assert state[key].shape == shape and torch.isfinite(state[key]).all(), key + for key in fp8: + weight = state[key] + scale = state[key.removesuffix("weight") + "weight_scale"] + assert weight.ndim == 2 and scale.numel() == 1, key + assert torch.isfinite(scale).all() and (scale > 0).all(), key + state[key] = from_quantized_weight( + weight, scale, QUANTIZATION_FP8, torch.float32, + ) + assert state[key].shape == weight.shape and torch.isfinite(state[key]).all(), key + for key in list(state): + if key.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): + assert key.rsplit(".", 1)[0] + ".weight" in quantized, key + scale = state.pop(key) + assert torch.isfinite(scale.float()).all() and (scale.float() >= 0).all(), key + # Use the original official text-only class and checkpoint prefix conversion. + # HF owns its declared ignoring of unused visual/MTP keys; do not filter weights. + config = AutoConfig.from_pretrained( + model_dir, local_files_only=True, trust_remote_code=False, + ).get_text_config(decoder=True) + assert config.model_type == "qwen3_5_text" + assert not getattr(config, "quantization_config", None) + model, loading = Qwen3_5ForCausalLM.from_pretrained( + None, config=config, state_dict=state, dtype=torch.float32, + output_loading_info=True, + ) + assert all(not loading.get(key) for key in ( + "missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs", + )), loading + if (model_dir / "generation_config.json").is_file(): + model.generation_config = GenerationConfig.from_pretrained( + model_dir, local_files_only=True, + ) + del state, tensors + else: + model = AutoModelForCausalLM.from_pretrained( model_dir, local_files_only=True, trust_remote_code=trust_remote_code, dtype=dtypes[reference_precision], ) - .eval() - .to("cuda") - ) + reference_device = case.get("reference_device", "cuda") + assert reference_device in {"cpu", "cuda"}, reference_device + model = model.eval().to(reference_device) inputs = _render_prompt(tokenizer, prompt, case).to(model.device) prompt_ids = inputs["input_ids"][0].tolist() if "expected_prompt_token_ids" in case: diff --git a/tools/tests/test_architecture.py b/tools/tests/test_architecture.py index ef10edad44..ad1bf66425 100644 --- a/tools/tests/test_architecture.py +++ b/tools/tests/test_architecture.py @@ -356,7 +356,15 @@ def test_shared_python_and_native_trees_are_closed_minimal_sets() -> None: "tools/tests/test_pr_metadata.py", "tools/tests/test_public_source_hygiene.py", } - expected_cmake = {"cmake/trtmcConfig.cmake.in"} + expected_cmake = { + "cmake/trtmcConfig.cmake.in", + "cmake/EdgeLLM.cmake", + "cmake/edgellm/CheckNative.cmake", + "cmake/edgellm/EdgeLLMConfig.cmake.in", + "cmake/edgellm/Install.cmake.in", + "cmake/edgellm/Prepare.cmake.in", + "cmake/edgellm/README.md", + } expected_third_party = { "third_party/stb/stb_image.h", "third_party/stb/stb_image_resize2.h", diff --git a/website/docs/api/python-builder.md b/website/docs/api/python-builder.md index 980b803521..643776e18f 100644 --- a/website/docs/api/python-builder.md +++ b/website/docs/api/python-builder.md @@ -29,6 +29,36 @@ resolved API directly. decides whether that directory is a Hugging Face snapshot or a prepared checkpoint; `BuildRequest` does not perform another discovery pass. +## Optional execution inputs + +`build(request, execution=...)` accepts an optional, frozen +`BuildExecutionInputs` descriptor. It contains a family-owned `variant` string +and a tuple of `NamedCheckpoint(role, model_dir)` descriptors. Import these +public types from `tensorrt_model_connect`. Companion directories must already +exist locally; the core does not download them or infer compatible model pairs. +Roles must be unique. Variant and role names are lowercase identifiers. + +Providing execution inputs requires the selected family to implement +`build_with_inputs(request, writer, execution)`. The family validates the +variant, checkpoint roles, compatibility and execution semantics. A missing +hook fails before bundle creation; the core never substitutes ordinary +base-only generation or another variant. With no execution inputs, the +existing `build(request, writer)` family call is unchanged. + +The build CLI exposes the same optional contract: + +```text +trtmc build LOCAL_TARGET -o model.bundle \ + --execution-variant FAMILY_VARIANT \ + --companion ROLE=LOCAL_COMPANION_DIR +``` + +Replace the uppercase placeholders with values from the selected family's +recipe; they are not literal supported identifiers. Repeat `--companion` only +for distinct roles. A companion requires `--execution-variant`; URLs and +implicit companion downloads are unsupported. The generic API does not itself +qualify any speculative algorithm or checkpoint pair. + ## Optional graph transform `BuildRequest.graph_transform` is an in-place callback invoked on the completed diff --git a/website/docs/architecture/build-pipeline.md b/website/docs/architecture/build-pipeline.md index 36dc14df30..f3351401d3 100644 --- a/website/docs/architecture/build-pipeline.md +++ b/website/docs/architecture/build-pipeline.md @@ -11,6 +11,7 @@ model ID/local snapshot -> choose family default task or validate --task -> import the selected families..model -> call build(BuildRequest, BundleWriter) + or the explicitly requested family build_with_inputs hook -> atomically publish format-1 bundle ``` @@ -33,12 +34,20 @@ sizes, family-owned quantization selection, FP32 layer overrides, direct dynamic-KV opt-in, and optional graph transform. Each family must implement or explicitly reject every non-default request it receives. +Optional `BuildExecutionInputs` travel separately from `BuildRequest`. Core +checks descriptor types, unique roles and existing local directories; only +the selected family interprets variant names and companion compatibility. +See the [Python Build API](../api/python-builder.md#optional-execution-inputs). + ## Family build `families//model.py` exposes a plain `build(request, writer)` function. It reads model config and weights, constructs the TensorRT network and engines, and writes family-owned named sections. Builder inheritance and shared model -topology helpers are forbidden. +topology helpers are forbidden. A family may instead delegate a complete +network to an installed optimized runtime through a family-owned adapter. +Model-specific admission, builder mapping, runtime orchestration and validation +remain in that family; shared dependency provisioning contains no model policy. The graph-transform callback, when present, receives the live TensorRT network immediately before serialization. This is the build-time half of the explicit diff --git a/website/docs/features/model-families.md b/website/docs/features/model-families.md index cb3a6e18f6..39c683d602 100644 --- a/website/docs/features/model-families.md +++ b/website/docs/features/model-families.md @@ -66,6 +66,22 @@ long-context qualification remain outside this contract. The committed-token-per-forward receipt is an algorithmic diagnostic, not a wall-clock speedup claim. +### Qwen3.8 paired ONNX execution + +The Qwen3.8 family owns explicit mixed-NVFP4 target plus DSpark block7 execution +through the optional pinned native Edge-LLM SDK. The qualified profile uses +`RadixArk/Qwen3.8-27B-NVFP4` and `RadixArk/Qwen3.8-27B-DSpark`, text-only +FP16 execution with the source mixed NVFP4/FP8 metadata, TP1/batch1 on SM120. +Standalone builds retain the original native path; ordinary experimental Edge +offload and other platform routes are not enabled by this change. + +Provision the [native SDK](../user-guides/configure-runtime.md#optional-native-edge-llm-sdk), +then add `--execution-variant dspark --companion draft=/path/to/draft` to the build +CLI. See the [owning Qwen3.8 recipe](https://github.com/NVIDIA/TensorRT-Model-Connect/blob/main/families/qwen3_8/EDGE_LLM.md) +for exact revisions, capacities, sampling controls and quality evidence. +The local paired qualification is not a registered manifest case and does not +imply CI coverage of that pair or statistical sampling equivalence. + ## Runtime and validation The directory name is also the runtime DSO identity: diff --git a/website/docs/user-guides/configure-runtime.md b/website/docs/user-guides/configure-runtime.md index 82f3f62dca..fbbc07a82f 100644 --- a/website/docs/user-guides/configure-runtime.md +++ b/website/docs/user-guides/configure-runtime.md @@ -31,3 +31,26 @@ Unsupported values fail; they are not silently ignored. See [Configuration and Backends](../features/config-and-backends.md), [Quantization](../features/quantization.md), and [Multi-Device Execution](../features/multi-device.md). + +## Optional native Edge-LLM SDK + +Provision Edge-LLM explicitly when building Model Connect, not during model +builds or inference. `TRTMC_ENABLE_EDGELLM=ON` selects the public Edge-LLM +0.10.1 snapshot at `e8b29522938901f6df19ebeedd4b69bc8edbcd97`. The default is +`OFF`. Configure and build on the inference GPU host; cross compilation is +rejected. The package must match the native CPU/GPU and CUDA/TensorRT stack. +Runtime compilation must also use the exact pinned JSON dependency headers; +the SDK rejects same-version development headers with an incompatible C++ ABI. + +- `TRTMC_EDGELLM_ALL_KERNELS=ON` requests all upstream operator groups supported + by the local GPU. +- `TRTMC_EDGELLM_ONNX=ON` also installs the original Python exporter and native + C++ ONNX engine builder. It does not select a model's build flow. +- `CMAKE_PREFIX_PATH` points builders and runtime compilation to the installed + SDK. Reuse fails explicitly if a requested capability is absent. + +Follow the repository's [pinned SDK installation instructions](https://github.com/NVIDIA/TensorRT-Model-Connect/blob/main/cmake/edgellm/README.md) +for dependencies, native architecture selection and offline provisioning. +Each family decides whether and how to use the package. Installing the SDK +is not evidence that a model, precision, input modality or execution variant +has passed validation. Ordinary model builds never install missing SDK tools.