diff --git a/CMakeLists.txt b/CMakeLists.txt index 2d653e3281..ab80d774d1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,11 +10,20 @@ set(CMAKE_CXX_EXTENSIONS OFF) set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) -set(CMAKE_BUILD_RPATH_USE_ORIGIN TRUE) +if(WIN32) + set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON) +else() + set(CMAKE_BUILD_RPATH_USE_ORIGIN TRUE) +endif() + +set(TRTMC_CXX_WARNINGS + "$<$:-Wall;-Wextra;-Wpedantic>" +) include(GNUInstallDirs) find_package(CUDAToolkit REQUIRED) find_package(nlohmann_json 3.11 REQUIRED) +option(TRTMC_BUILD_BACKEND_TRT "Build the standard TensorRT backend DSO" ON) set(_trtmc_dependency_roots /usr @@ -22,26 +31,34 @@ set(_trtmc_dependency_roots /usr/local/cuda /opt/tensorrt ) +if(DEFINED ENV{CUDA_PATH}) + list(PREPEND _trtmc_dependency_roots "$ENV{CUDA_PATH}") +endif() if(DEFINED ENV{TRT_ROOT}) list(PREPEND _trtmc_dependency_roots "$ENV{TRT_ROOT}") endif() -find_path(TRTMC_TRT_INCLUDE_DIR - NAMES NvInferRuntime.h - HINTS ${_trtmc_dependency_roots} - PATH_SUFFIXES include include/zapped_headers - REQUIRED -) -find_library(TRTMC_TRT_LIBRARY - NAMES nvinfer libnvinfer.so.11 - HINTS ${_trtmc_dependency_roots} - PATH_SUFFIXES lib lib64 lib/aarch64-linux-gnu lib/x86_64-linux-gnu - REQUIRED -) +if(TRTMC_BUILD_BACKEND_TRT) + find_path(TRTMC_TRT_INCLUDE_DIR + NAMES NvInferRuntime.h + HINTS ${_trtmc_dependency_roots} + PATH_SUFFIXES include include/zapped_headers + REQUIRED + ) + find_library(TRTMC_TRT_LIBRARY + NAMES nvinfer libnvinfer.so.11 nvinfer_11 + HINTS ${_trtmc_dependency_roots} + PATH_SUFFIXES lib lib64 lib/x64 lib/aarch64-linux-gnu lib/x86_64-linux-gnu + REQUIRED + ) +endif() option(TRTMC_ENABLE_BYOK "Enable the optional TVM-FFI BYOK bridge" ON) set(TRTMC_HAS_TVM_FFI OFF) if(TRTMC_ENABLE_BYOK) + if(NOT TRTMC_BUILD_BACKEND_TRT) + message(FATAL_ERROR "TRTMC_ENABLE_BYOK=ON requires TRTMC_BUILD_BACKEND_TRT=ON") + endif() find_package(Python3 COMPONENTS Interpreter QUIET) if(Python3_Interpreter_FOUND) execute_process( @@ -77,7 +94,11 @@ endif() # Family CMake files use these two variables to link the shared Engine API. set(TRTMC_CUDA_INCLUDE_DIR ${CUDAToolkit_INCLUDE_DIRS}) -set(TRTMC_CUDART_LIBRARY CUDA::cudart) +if(WIN32 AND CMAKE_CUDA_RUNTIME_LIBRARY STREQUAL "Static") + set(TRTMC_CUDART_LIBRARY CUDA::cudart_static) +else() + set(TRTMC_CUDART_LIBRARY CUDA::cudart) +endif() add_library(trtmc_core SHARED core/runtime/bundle/bundle_format.cpp @@ -95,18 +116,22 @@ target_include_directories(trtmc_core target_include_directories(trtmc_core SYSTEM PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) target_link_libraries(trtmc_core PUBLIC - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} PRIVATE nlohmann_json::nlohmann_json ) -target_compile_options(trtmc_core PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_core PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_core PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) add_library(trtmc_runtime SHARED core/runtime/loader/family_loader.cpp + core/runtime/platform/dynamic_library.cpp ) target_include_directories(trtmc_runtime PUBLIC @@ -120,12 +145,16 @@ target_link_libraries(trtmc_runtime trtmc_core ${CMAKE_DL_LIBS} ) -target_compile_options(trtmc_runtime PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_runtime PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_runtime PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) +if(TRTMC_BUILD_BACKEND_TRT) add_library(trtmc_backend_trt SHARED core/runtime/tensorrt/trt_backend.cpp core/runtime/tensorrt/trt_logger.cpp @@ -144,16 +173,21 @@ target_link_libraries(trtmc_backend_trt PRIVATE trtmc_core ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${CMAKE_DL_LIBS} ) -target_compile_options(trtmc_backend_trt PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_backend_trt PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_backend_trt PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) +endif() option(TRTMC_BUILD_BACKEND_RTX "Build TensorRT-RTX backend DSO" OFF) +set(TRTMC_RTX_RUNTIME_DIR "" CACHE PATH "TensorRT-RTX runtime DLL directory on Windows") if(TRTMC_BUILD_BACKEND_RTX) if(NOT IS_DIRECTORY "${TRTMC_RTX_INCLUDE_DIR}" OR NOT EXISTS "${TRTMC_RTX_INCLUDE_DIR}/NvInfer.h") @@ -167,7 +201,7 @@ if(TRTMC_BUILD_BACKEND_RTX) ) endif() find_library(TRTMC_RTX_LIBRARY - NAMES tensorrt_rtx + NAMES tensorrt_rtx tensorrt_rtx_1_6 PATHS "${TRTMC_RTX_LIBRARY_DIR}" NO_DEFAULT_PATH ) @@ -177,8 +211,22 @@ if(TRTMC_BUILD_BACKEND_RTX) ) endif() + if(WIN32) + find_file(TRTMC_RTX_RUNTIME_DLL + NAMES tensorrt_rtx_1_6.dll tensorrt_rtx.dll + PATHS "${TRTMC_RTX_RUNTIME_DIR}" + NO_DEFAULT_PATH + ) + if(NOT TRTMC_RTX_RUNTIME_DLL) + message(FATAL_ERROR + "TRTMC_BUILD_BACKEND_RTX=ON requires TRTMC_RTX_RUNTIME_DIR with the runtime DLL" + ) + endif() + endif() + add_library(trtmc_backend_rtx SHARED core/runtime/tensorrt/rtx_backend.cpp + core/runtime/tensorrt/runtime_cache_persistence.cpp core/runtime/tensorrt/trt_logger.cpp core/runtime/tensorrt/trt_module_impl.cpp ) @@ -195,14 +243,30 @@ if(TRTMC_BUILD_BACKEND_RTX) PRIVATE trtmc_core ${TRTMC_RTX_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} + ) + target_compile_options(trtmc_backend_rtx PRIVATE ${TRTMC_CXX_WARNINGS}) + target_compile_definitions(trtmc_backend_rtx PRIVATE + TRTMC_TRT_DELETE_OBJECTS=1 + TRTMC_TRT_HAS_TENSOR_ALIAS_API=1 + TRTMC_TRT_HAS_COMMUNICATOR_API=1 ) - target_compile_options(trtmc_backend_rtx PRIVATE -Wall -Wextra -Wpedantic) set_target_properties(trtmc_backend_rtx PROPERTIES OUTPUT_NAME trtmc_backend_trt_rtx + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) + if(WIN32) + add_custom_command(TARGET trtmc_backend_rtx POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${TRTMC_RTX_RUNTIME_DLL}" + "$" + COMMENT "Copying TensorRT-RTX runtime DLL next to the backend" + ) + endif() endif() if(TRTMC_HAS_TVM_FFI) @@ -226,7 +290,7 @@ if(TRTMC_HAS_TVM_FFI) ) target_link_libraries(trtmc_byok_tvm_ffi PRIVATE ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} nlohmann_json::nlohmann_json ) @@ -234,7 +298,7 @@ if(TRTMC_HAS_TVM_FFI) TRTMC_HAS_TRT=1 TRTMC_HAS_TVM_FFI=1 ) - target_compile_options(trtmc_byok_tvm_ffi PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(trtmc_byok_tvm_ffi PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_byok_tvm_ffi PROPERTIES BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN;${TRTMC_TVM_FFI_LIBRARY_DIR}" @@ -250,6 +314,9 @@ endif() # A family owns its target, sources, dependencies, warnings, and installation. # The root knows only the directory convention; adding a family never changes # a central source or target list. +set(TRTMC_RUNTIME_MODELS "" CACHE STRING + "Semicolon-separated family runtimes to build; empty builds every family" +) file(GLOB _trtmc_family_runtime_cmake CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/families/*/runtime/CMakeLists.txt" ) @@ -257,6 +324,9 @@ foreach(_trtmc_runtime_cmake IN LISTS _trtmc_family_runtime_cmake) get_filename_component(_trtmc_runtime_dir "${_trtmc_runtime_cmake}" DIRECTORY) get_filename_component(_trtmc_family_dir "${_trtmc_runtime_dir}" DIRECTORY) get_filename_component(_trtmc_family "${_trtmc_family_dir}" NAME) + if(TRTMC_RUNTIME_MODELS AND NOT _trtmc_family IN_LIST TRTMC_RUNTIME_MODELS) + continue() + endif() add_subdirectory( "${_trtmc_runtime_dir}" "${CMAKE_BINARY_DIR}/families/${_trtmc_family}" @@ -266,10 +336,12 @@ endforeach() add_library(trtmc_cli STATIC apps/cli/cli.cpp apps/cli/io.cpp + apps/cli/windows_media.cpp ) target_include_directories(trtmc_cli PUBLIC ${PROJECT_SOURCE_DIR}/apps PRIVATE + ${PROJECT_SOURCE_DIR}/core ${PROJECT_SOURCE_DIR}/core/runtime/include ${PROJECT_SOURCE_DIR}/third_party/stb ) @@ -280,17 +352,29 @@ target_link_libraries(trtmc_cli nlohmann_json::nlohmann_json ${CMAKE_DL_LIBS} ) +if(WIN32) + target_link_libraries(trtmc_cli PRIVATE mfplat mfreadwrite mfuuid ole32) +endif() target_compile_definitions(trtmc_cli PRIVATE TRTMC_VERSION_STRING="${PROJECT_VERSION}") -target_compile_options(trtmc_cli PRIVATE -Wall -Wextra -Wpedantic) -set_source_files_properties(apps/cli/io.cpp PROPERTIES - COMPILE_OPTIONS "-Wno-missing-field-initializers;-Wno-pedantic" -) +target_compile_options(trtmc_cli PRIVATE ${TRTMC_CXX_WARNINGS}) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + set_source_files_properties(apps/cli/io.cpp PROPERTIES + COMPILE_OPTIONS "-Wno-missing-field-initializers;-Wno-pedantic" + ) +endif() add_executable(trtmc apps/cli/main.cpp) +if(WIN32) + target_sources(trtmc PRIVATE + apps/cli/windows_utf8_argv.cpp + apps/cli/trtmc_windows_utf8.manifest + ) +endif() target_include_directories(trtmc PRIVATE ${PROJECT_SOURCE_DIR}/apps) target_link_libraries(trtmc PRIVATE trtmc_cli) -target_compile_options(trtmc PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}" ) @@ -305,7 +389,7 @@ if(TRTMC_BUILD_EXAMPLES) trtmc_runtime nlohmann_json::nlohmann_json ) - target_compile_options(trtmc_benchmark_worker PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(trtmc_benchmark_worker PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_benchmark_worker PROPERTIES BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}" @@ -320,7 +404,7 @@ if(TRTMC_BUILD_EXAMPLES) trtmc_runtime nlohmann_json::nlohmann_json ) - target_compile_options(trtmc_dataset_benchmark PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(trtmc_dataset_benchmark PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_dataset_benchmark PROPERTIES BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}" @@ -331,12 +415,12 @@ if(TRTMC_BUILD_TESTS) add_executable(test_bundle_format_v1 core/runtime/tests/test_bundle_format_v1.cpp) target_include_directories(test_bundle_format_v1 PRIVATE ${PROJECT_SOURCE_DIR}/core) target_link_libraries(test_bundle_format_v1 PRIVATE trtmc_core) - target_compile_options(test_bundle_format_v1 PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_bundle_format_v1 PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME bundle_format_v1 COMMAND test_bundle_format_v1) add_executable(test_task_api core/runtime/tests/test_task_api.cpp) target_include_directories(test_task_api PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include) - target_compile_options(test_task_api PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_task_api PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME task_api COMMAND test_task_api) add_executable(test_cli apps/cli/tests/test_cli.cpp) @@ -345,29 +429,33 @@ if(TRTMC_BUILD_TESTS) ${PROJECT_SOURCE_DIR}/apps ) target_link_libraries(test_cli PRIVATE trtmc_cli) - target_compile_options(test_cli PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_cli PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME cli COMMAND test_cli) set(_trtmc_test_runtime_root "${CMAKE_BINARY_DIR}/tests/runtime") add_library(trtmc_test_backend_fake SHARED core/runtime/tests/fake_backend.cpp) target_include_directories(trtmc_test_backend_fake PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include) - target_link_libraries(trtmc_test_backend_fake PRIVATE CUDA::cudart) + target_link_libraries(trtmc_test_backend_fake PRIVATE ${TRTMC_CUDART_LIBRARY}) set_target_properties(trtmc_test_backend_fake PROPERTIES OUTPUT_NAME trtmc_backend_fake LIBRARY_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}" + RUNTIME_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" ) add_library(trtmc_test_backend_fake_rtx SHARED core/runtime/tests/fake_backend.cpp) target_include_directories(trtmc_test_backend_fake_rtx PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include ) - target_link_libraries(trtmc_test_backend_fake_rtx PRIVATE CUDA::cudart) + target_link_libraries(trtmc_test_backend_fake_rtx PRIVATE ${TRTMC_CUDART_LIBRARY}) target_compile_definitions(trtmc_test_backend_fake_rtx PRIVATE TRTMC_FAKE_BACKEND_NAME="trt_rtx" ) set_target_properties(trtmc_test_backend_fake_rtx PROPERTIES OUTPUT_NAME trtmc_backend_trt_rtx LIBRARY_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}" + RUNTIME_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" ) add_library(trtmc_test_family_fake SHARED core/runtime/tests/fake_family.cpp) @@ -376,6 +464,8 @@ if(TRTMC_BUILD_TESTS) set_target_properties(trtmc_test_family_fake PROPERTIES OUTPUT_NAME trtmc_model_fake LIBRARY_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}" + RUNTIME_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" BUILD_RPATH "\$ORIGIN/../.." ) @@ -385,7 +475,7 @@ if(TRTMC_BUILD_TESTS) ${PROJECT_SOURCE_DIR}/core ) target_link_libraries(test_family_loader PRIVATE trtmc_runtime) - target_compile_options(test_family_loader PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_family_loader PRIVATE ${TRTMC_CXX_WARNINGS}) add_dependencies(test_family_loader trtmc_test_backend_fake trtmc_test_backend_fake_rtx @@ -393,6 +483,52 @@ if(TRTMC_BUILD_TESTS) ) add_test(NAME family_loader COMMAND test_family_loader "${_trtmc_test_runtime_root}") + add_executable(test_dynamic_library core/runtime/tests/test_dynamic_library.cpp) + target_include_directories(test_dynamic_library PRIVATE ${PROJECT_SOURCE_DIR}/core) + target_link_libraries(test_dynamic_library PRIVATE trtmc_runtime) + target_compile_definitions(test_dynamic_library PRIVATE + TRTMC_TEST_LIBRARY="$" + ) + target_compile_options(test_dynamic_library PRIVATE ${TRTMC_CXX_WARNINGS}) + add_dependencies(test_dynamic_library trtmc_test_backend_fake) + add_test(NAME dynamic_library COMMAND test_dynamic_library) + + add_executable(test_runtime_cache_persistence + core/runtime/tests/test_runtime_cache_persistence.cpp + core/runtime/tensorrt/runtime_cache_persistence.cpp + ) + target_include_directories(test_runtime_cache_persistence PRIVATE ${PROJECT_SOURCE_DIR}/core) + target_compile_options(test_runtime_cache_persistence PRIVATE ${TRTMC_CXX_WARNINGS}) + add_test(NAME runtime_cache_persistence COMMAND test_runtime_cache_persistence) + + if(WIN32) + add_executable(test_windows_utf8_argv + apps/cli/tests/test_windows_utf8_argv.cpp + apps/cli/windows_utf8_argv.cpp + apps/cli/trtmc_windows_utf8.manifest + ) + target_include_directories(test_windows_utf8_argv PRIVATE ${PROJECT_SOURCE_DIR}/apps) + add_test(NAME windows_utf8_argv COMMAND test_windows_utf8_argv) + + add_executable(test_windows_media apps/cli/tests/test_windows_media.cpp) + target_include_directories(test_windows_media PRIVATE + ${PROJECT_SOURCE_DIR}/apps + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_link_libraries(test_windows_media PRIVATE trtmc_cli) + add_test(NAME windows_media COMMAND test_windows_media) + endif() + + if(TRTMC_BUILD_BACKEND_TRT OR TRTMC_BUILD_BACKEND_RTX) + if(TRTMC_BUILD_BACKEND_TRT) + set(_trtmc_dynamic_test_backend trtmc_backend_trt) + set(_trtmc_dynamic_test_include ${TRTMC_TRT_INCLUDE_DIR}) + set(_trtmc_dynamic_test_library ${TRTMC_TRT_LIBRARY}) + else() + set(_trtmc_dynamic_test_backend trtmc_backend_rtx) + set(_trtmc_dynamic_test_include ${TRTMC_RTX_INCLUDE_DIR}) + set(_trtmc_dynamic_test_library ${TRTMC_RTX_LIBRARY}) + endif() add_executable(test_trt_module_dynamic_input core/runtime/tests/test_trt_module_dynamic_input.cpp ) @@ -401,32 +537,33 @@ if(TRTMC_BUILD_TESTS) ${PROJECT_SOURCE_DIR}/core ) target_include_directories(test_trt_module_dynamic_input SYSTEM PRIVATE - ${TRTMC_TRT_INCLUDE_DIR} + ${_trtmc_dynamic_test_include} ${CUDAToolkit_INCLUDE_DIRS} ) target_link_libraries(test_trt_module_dynamic_input PRIVATE - trtmc_backend_trt + ${_trtmc_dynamic_test_backend} trtmc_core - ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${_trtmc_dynamic_test_library} + ${TRTMC_CUDART_LIBRARY} ) - target_compile_options(test_trt_module_dynamic_input PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_trt_module_dynamic_input PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME trt_module_dynamic_input COMMAND test_trt_module_dynamic_input) set_tests_properties(trt_module_dynamic_input PROPERTIES SKIP_RETURN_CODE 77 LABELS gpu) + endif() if(TRTMC_BUILD_EXAMPLES) add_executable(test_dataset_answer apps/benchmark/tests/native/test_dataset_answer.cpp) target_include_directories(test_dataset_answer PRIVATE ${PROJECT_SOURCE_DIR}/apps/benchmark/native ) - target_compile_options(test_dataset_answer PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_dataset_answer PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME dataset_answer COMMAND test_dataset_answer) add_executable(test_benchmark_worker_e2e apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp ) target_link_libraries(test_benchmark_worker_e2e PRIVATE nlohmann_json::nlohmann_json) - target_compile_options(test_benchmark_worker_e2e PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_benchmark_worker_e2e PRIVATE ${TRTMC_CXX_WARNINGS}) add_dependencies(test_benchmark_worker_e2e trtmc_benchmark_worker trtmc_test_backend_fake @@ -453,7 +590,7 @@ if(TRTMC_BUILD_TESTS AND TRTMC_HAS_TVM_FFI) ${TRTMC_TRT_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} ) - target_compile_options(test_byok_shape_spec PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_byok_shape_spec PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME byok_shape_spec COMMAND test_byok_shape_spec) endif() @@ -466,7 +603,7 @@ if(TRTMC_BUILD_EXAMPLES AND TRTMC_HAS_TVM_FFI) ${TRTMC_TVM_FFI_INCLUDE_DIR} ) target_link_libraries(trtmc_byok_identity_copy PRIVATE - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} ) set_target_properties(trtmc_byok_identity_copy PROPERTIES @@ -488,14 +625,14 @@ if(TRTMC_BUILD_EXAMPLES AND TRTMC_HAS_TVM_FFI) target_link_libraries(test_byok_tvm_ffi PRIVATE trtmc_byok_tvm_ffi ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} ) target_compile_definitions(test_byok_tvm_ffi PRIVATE TRTMC_HAS_TRT=1 TRTMC_HAS_TVM_FFI=1 ) - target_compile_options(test_byok_tvm_ffi PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_byok_tvm_ffi PRIVATE ${TRTMC_CXX_WARNINGS}) add_dependencies(test_byok_tvm_ffi trtmc_byok_identity_copy) add_test(NAME byok_tvm_ffi COMMAND test_byok_tvm_ffi $ @@ -510,15 +647,25 @@ install(TARGETS trtmc_core trtmc_runtime EXPORT trtmcTargets ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} ) -install(TARGETS trtmc_backend_trt - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} -) +if(TARGET trtmc_backend_trt) + install(TARGETS trtmc_backend_trt + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) +endif() if(TARGET trtmc_backend_rtx) install(TARGETS trtmc_backend_rtx + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) + if(WIN32) + install(FILES "${TRTMC_RTX_RUNTIME_DLL}" DESTINATION ${CMAKE_INSTALL_BINDIR}) + endif() endif() if(TRTMC_HAS_TVM_FFI) install(TARGETS trtmc_byok_tvm_ffi diff --git a/apps/benchmark/native/benchmark_worker.cpp b/apps/benchmark/native/benchmark_worker.cpp index 0abbf93e7a..533ccf2c31 100644 --- a/apps/benchmark/native/benchmark_worker.cpp +++ b/apps/benchmark/native/benchmark_worker.cpp @@ -421,10 +421,10 @@ Json run_generate_audio(trtmc::ITask& task, const Json& request, const Timing& t return measure( timing, [&]() { return interface.generate_audio(prompt, config); }, [](const trtmc::AudioResult& result) { - const double seconds = - result.sample_rate > 0 - ? static_cast(result.samples.size()) / result.sample_rate - : 0.0; + const double seconds = result.sample_rate > 0 && result.channels > 0 + ? static_cast(result.samples.size()) / + result.channels / result.sample_rate + : 0.0; return Json{{"output_samples", result.samples.size()}, {"num_samples", result.samples.size()}, {"output_audio_seconds", seconds}, @@ -458,10 +458,10 @@ Json run_speak(trtmc::ITask& task, const Json& request, const Timing& timing) { [](const auto& value) { const auto& result = value.first; return Json{{"input_audio_seconds", value.second}, - {"output_audio_seconds", - result.sample_rate > 0 - ? static_cast(result.samples.size()) / result.sample_rate - : 0.0}, + {"output_audio_seconds", result.sample_rate > 0 && result.channels > 0 + ? static_cast(result.samples.size()) / + result.channels / result.sample_rate + : 0.0}, {"output_samples", result.samples.size()}, {"num_samples", result.samples.size()}, {"sample_rate", result.sample_rate}}; diff --git a/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp b/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp index 0313a15f51..50c5231867 100644 --- a/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp +++ b/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp @@ -11,6 +11,10 @@ #include #include +#if defined(_WIN32) +#include +#endif + namespace { using Json = nlohmann::json; @@ -39,12 +43,14 @@ void write_bundle(const std::filesystem::path& path) { throw std::runtime_error("failed to write fake bundle"); } +#if !defined(_WIN32) std::string shell_quote(const std::string& value) { std::string result{"'"}; for (const char character : value) result += character == '\'' ? "'\\''" : std::string(1, character); return result + "'"; } +#endif } // namespace @@ -81,10 +87,19 @@ int main(int argc, char** argv) { throw std::runtime_error("failed to write worker request"); } +#if defined(_WIN32) + const std::string worker_argument = '"' + std::string(argv[1]) + '"'; + const std::string request_argument = '"' + request_path.string() + '"'; + const std::string output_argument = '"' + output_path.string() + '"'; + check(_spawnl(_P_WAIT, argv[1], worker_argument.c_str(), "--request", + request_argument.c_str(), "--output", output_argument.c_str(), nullptr) == 0, + "worker process completed"); +#else const std::string command = shell_quote(argv[1]) + " --request " + shell_quote(request_path.string()) + " --output " + shell_quote(output_path.string()); check(std::system(command.c_str()) == 0, "worker process completed"); +#endif std::ifstream output_file(output_path); Json result; diff --git a/apps/cli/cli.cpp b/apps/cli/cli.cpp index 7cf58f10d7..8785ee174d 100644 --- a/apps/cli/cli.cpp +++ b/apps/cli/cli.cpp @@ -6,6 +6,8 @@ #include "cli/cli.h" #include "cli/io.h" +#include "cli/windows_media.h" +#include "runtime/platform/dynamic_library.h" #include "trtmc/runtime/family_loader.h" #include @@ -14,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -123,8 +124,10 @@ const std::unordered_map& command_specs() { "--num-steps", "--guidance-scale", "--cfg-scale"}}}, {"generate-video", {CommandKind::kGenerateVideo, - {"--prompt", "--image", "--output", "--negative-prompt", "--height", "--width", - "--num-steps", "--seed", "--guidance-scale", "--cfg-scale", "--initial-latents-raw"}}}, + {"--prompt", "--output", "--negative-prompt", "--height", "--width", "--num-frames", + "--num-steps", "--seed", "--guidance-scale", "--cfg-scale", "--initial-latents-raw", + "--first-frame", "--last-frame", "--reference-image", "--reference-video", + "--reference-audio"}}}, {"solve", {CommandKind::kSolve, {"--branch", "--trunk"}}}, {"forecast", {CommandKind::kForecast, {"--input", "--mask", "--frequency"}}}, {"control", {CommandKind::kControl, {"--image", "--state", "--output"}}}, @@ -142,21 +145,24 @@ bool is_byok_option(const std::string& option) { void load_byok_extension(const Command& command) { using LoadKernelFn = const char* (*)(const char*, const char*, const char*) noexcept; - const fs::path extension = fs::path(command.runtime_root) / "libtrtmc_byok_tvm_ffi.so"; - dlerror(); - void* handle = dlopen(extension.c_str(), RTLD_NOW | RTLD_LOCAL); + const fs::path extension = + fs::path(command.runtime_root) / internal::dynamic_library_filename("trtmc_byok_tvm_ffi"); + std::string loader_error; + auto handle = internal::open_dynamic_library( + extension, internal::DynamicLibraryVisibility::local, &loader_error); if (handle == nullptr) { - const char* error = dlerror(); throw std::runtime_error("unable to load BYOK extension '" + extension.string() + - "': " + (error != nullptr ? error : "unknown dlopen error")); + "': " + loader_error); } - static auto* handles = new std::vector; - handles->push_back(handle); - dlerror(); - auto load = reinterpret_cast(dlsym(handle, "trtmc_load_byok_kernel")); - if (const char* error = dlerror(); error != nullptr || load == nullptr) { - throw std::runtime_error("BYOK extension is missing trtmc_load_byok_kernel"); + auto load = reinterpret_cast( + internal::dynamic_library_symbol(handle, "trtmc_load_byok_kernel", &loader_error)); + if (load == nullptr) { + (void)internal::close_dynamic_library(handle); + throw std::runtime_error("BYOK extension is missing trtmc_load_byok_kernel: " + + loader_error); } + static auto* handles = new std::vector; + handles->push_back(handle); if (const char* error = load(command.options.at("--byok-library").c_str(), command.options.at("--byok-function").c_str(), command.options.at("--byok-name").c_str())) { @@ -464,6 +470,7 @@ ImageGenerationConfig image_config(const Command& command) { config.negative_prompt = command.options.at("--negative-prompt"); config.height = int_option(command, "--height", 0, 1); config.width = int_option(command, "--width", 0, 1); + config.video_num_frames = int_option(command, "--num-frames", 0, 1); config.num_steps = int_option(command, "--num-steps", -1, 1); config.seed = int_option(command, "--seed", -1); config.guidance_scale = float_option(command, "--guidance-scale", -1.0F); @@ -538,6 +545,102 @@ ImageResult generate_image(const Command& command, ITask& task) { return require_interface(task).generate_image(prompt, config); } +VideoImageInput load_video_image(const std::string& path) { + auto decoded = read_image(path); + VideoImageInput result; + result.pixels = std::move(decoded.pixels); + result.height = decoded.height; + result.width = decoded.width; + result.channels = 3; + return result; +} + +VideoGenerationRequest video_request(const Command& command, IVideoGeneration& generator) { + VideoGenerationRequest request; + request.prompt = require_option(command, "--prompt"); + request.config = image_config(command); + + const bool has_first = has_option(command, "--first-frame"); + const bool has_last = has_option(command, "--last-frame"); + if ((has_first || has_last) && !command.video_references.empty()) + throw std::invalid_argument("key frames cannot be combined with media references"); + if (has_first || has_last) { + request.mode = VideoGenerationMode::kFirstLastFrameToVideoAudio; + if (has_first) + request.first_frame = load_video_image(command.options.at("--first-frame")); + if (has_last) + request.last_frame = load_video_image(command.options.at("--last-frame")); + return request; + } + if (command.video_references.empty()) + return request; + + request.mode = VideoGenerationMode::kReferenceToVideoAudio; + const auto policy = generator.reference_media_decode_policy(); + request.references.reserve(command.video_references.size()); + for (const auto& argument : command.video_references) { + VideoReferenceInput reference; + reference.kind = argument.kind; + switch (argument.kind) { + case VideoReferenceKind::kImage: + reference.image = load_video_image(argument.path); + break; + case VideoReferenceKind::kVideo: + if (!policy) + throw std::runtime_error( + "loaded family does not provide a reference-media decode policy"); + reference.video = read_video_file(argument.path, *policy); + break; + case VideoReferenceKind::kAudio: + if (!policy) + throw std::runtime_error( + "loaded family does not provide a reference-media decode policy"); + reference.audio = read_audio_file(argument.path, *policy); + break; + } + request.references.push_back(std::move(reference)); + } + return request; +} + +nlohmann::json write_generated_video(const VideoResult& result, const std::string& path) { + if (result.frames.pixels.empty() && result.frames.height == 0 && result.frames.width == 0 && + result.frames.channels == 3 && result.frames.num_frames == 0 && + result.audio.samples.empty() && result.fps == 0) { + return {{"worker", true}}; + } + validate_image_result(result.frames); + if (result.fps <= 0) + throw std::runtime_error("video result has a non-positive frame rate"); + if (!result.audio.samples.empty()) { + if (result.audio.sample_rate <= 0 || + (result.audio.channels != 1 && result.audio.channels != 2) || + result.audio.samples.size() % static_cast(result.audio.channels) != 0) { + throw std::runtime_error("video result has invalid interleaved audio metadata"); + } + require_finite(result.audio.samples, "video audio result"); + } + if (is_mp4_path(path)) { + write_mp4(result, path); + return {{"output", path}, + {"frames", result.frames.num_frames}, + {"fps", result.fps}, + {"height", result.frames.height}, + {"width", result.frames.width}, + {"audio_channels", result.audio.samples.empty() ? 0 : result.audio.channels}, + {"audio_sample_rate", result.audio.samples.empty() ? 0 : result.audio.sample_rate}}; + } + + auto payload = write_video(result.frames, path); + payload["fps"] = result.fps; + if (!result.audio.samples.empty()) { + const fs::path audio_path = fs::path(path) / "audio.wav"; + io::write_wav(result.audio, audio_path.string()); + payload["audio"] = audio_path.string(); + } + return payload; +} + const char* event_kind_name(SpeechSessionEventKind kind) { switch (kind) { case SpeechSessionEventKind::kAgentAudio: @@ -665,17 +768,17 @@ Command parse_args(int argc, char** argv) { if (name == "help") { if (argc != 2) throw std::invalid_argument("help does not accept arguments"); - return {CommandKind::kHelp, name, {}, {}, {}, {}, {}, 0, {}, false}; + return {CommandKind::kHelp, name, {}, {}, {}, {}, {}, {}, 0, {}, false}; } if (name == "version") { if (argc != 2) throw std::invalid_argument("version does not accept arguments"); - return {CommandKind::kVersion, name, {}, {}, {}, {}, {}, 0, {}, false}; + return {CommandKind::kVersion, name, {}, {}, {}, {}, {}, {}, 0, {}, false}; } if (name == "inspect") { if (argc != 3 || std::string(argv[2]).empty()) throw std::invalid_argument("inspect requires exactly one BUNDLE path"); - return {CommandKind::kInspect, name, argv[2], {}, {}, {}, {}, 0, {}, false}; + return {CommandKind::kInspect, name, argv[2], {}, {}, {}, {}, {}, 0, {}, false}; } const auto spec = command_specs().find(name); @@ -684,7 +787,7 @@ Command parse_args(int argc, char** argv) { if (argc < 3 || std::string(argv[2]).empty()) throw std::invalid_argument(name + " requires a BUNDLE path"); - Command command{spec->second.kind, name, argv[2], {}, {}, {}, {}, 0, {}, false}; + Command command{spec->second.kind, name, argv[2], {}, {}, {}, {}, {}, 0, {}, false}; for (int index = 3; index < argc; ++index) { const std::string option = argv[index]; if (option == "--runtime-root") { @@ -721,6 +824,16 @@ Command parse_args(int argc, char** argv) { command.inputs.push_back(take_value(argc, argv, index, option)); continue; } + if (option == "--reference-image" || option == "--reference-video" || + option == "--reference-audio") { + VideoReferenceKind kind = VideoReferenceKind::kImage; + if (option == "--reference-video") + kind = VideoReferenceKind::kVideo; + else if (option == "--reference-audio") + kind = VideoReferenceKind::kAudio; + command.video_references.push_back({kind, take_value(argc, argv, index, option)}); + continue; + } if (command.options.count(option) != 0) throw std::invalid_argument(option + " may be specified only once"); command.options.emplace(option, take_value(argc, argv, index, option)); @@ -1195,8 +1308,10 @@ int dispatch(const Command& command, ITask& task, std::ostream& output) { return EXIT_SUCCESS; } case CommandKind::kGenerateVideo: { - const std::string directory = require_option(command, "--output"); - write_json(output, write_video(generate_image(command, task), directory)); + const std::string path = require_option(command, "--output"); + auto& generator = require_interface(task); + auto request = video_request(command, generator); + write_json(output, write_generated_video(generator.generate_video(request), path)); return EXIT_SUCCESS; } case CommandKind::kSolve: { @@ -1305,6 +1420,10 @@ void print_usage(std::ostream& output) { " [--kv-cache-size BYTES|GB|GiB]\n\n" "TensorRT-RTX runtime options:\n" " [--runtime-cache PATH] [--cuda-graphs]\n\n" + "Video generation options:\n" + " --prompt TEXT --output OUTPUT.mp4 [--num-frames N] [--height N] [--width N]\n" + " [--first-frame IMAGE] [--last-frame IMAGE]\n" + " [--reference-image IMAGE|--reference-video VIDEO|--reference-audio AUDIO]...\n\n" "Execution never searches for runtimes; --runtime-root is always required.\n"; } diff --git a/apps/cli/cli.h b/apps/cli/cli.h index 9fe114557f..13bb778da2 100644 --- a/apps/cli/cli.h +++ b/apps/cli/cli.h @@ -48,6 +48,11 @@ enum class CommandKind { kGenerateWorld, }; +struct VideoReferenceArgument { + VideoReferenceKind kind{VideoReferenceKind::kImage}; + std::string path; +}; + struct Command { CommandKind kind{CommandKind::kHelp}; std::string name; @@ -56,6 +61,7 @@ struct Command { std::unordered_map options; std::vector frames; std::vector inputs; + std::vector video_references; std::uint64_t kv_cache_size_bytes{0}; std::string runtime_cache_path; bool cuda_graphs{false}; diff --git a/apps/cli/io.cpp b/apps/cli/io.cpp index 472c86e721..d061039846 100644 --- a/apps/cli/io.cpp +++ b/apps/cli/io.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -24,18 +25,34 @@ namespace trtmc::cli::io { void write_wav(const AudioResult& audio, const std::string& path) { if (audio.samples.empty()) throw std::runtime_error("write_wav: empty audio"); + if (audio.sample_rate <= 0) + throw std::runtime_error("write_wav: sample rate must be positive"); + if (audio.channels != 1 && audio.channels != 2) + throw std::runtime_error("write_wav: channel count must be mono or stereo"); + if (audio.samples.size() % static_cast(audio.channels) != 0) + throw std::runtime_error("write_wav: interleaved sample count is not channel-aligned"); + if (audio.samples.size() > + static_cast(std::numeric_limits::max() / sizeof(float))) + throw std::runtime_error("write_wav: sample buffer is too large for RIFF/WAVE"); + if (audio.num_samples != 0 && + audio.num_samples != static_cast(audio.samples.size())) + throw std::runtime_error("write_wav: num_samples does not match the sample buffer"); std::ofstream output(path, std::ios::binary); if (!output) throw std::runtime_error("write_wav: cannot open " + path); - const auto num_samples = static_cast(audio.samples.size()); const std::int32_t sample_rate = audio.sample_rate; - const std::int16_t num_channels = 1; + const auto num_channels = static_cast(audio.channels); const std::int16_t bits_per_sample = 32; - const std::int32_t byte_rate = sample_rate * num_channels * (bits_per_sample / 8); const auto block_align = static_cast(num_channels * (bits_per_sample / 8)); - const std::int32_t data_size = num_samples * block_align; + const auto byte_rate_64 = static_cast(sample_rate) * block_align; + if (byte_rate_64 > std::numeric_limits::max()) + throw std::runtime_error("write_wav: byte rate exceeds the RIFF/WAVE range"); + const auto byte_rate = static_cast(byte_rate_64); + const auto data_size = static_cast(audio.samples.size() * sizeof(float)); + if (data_size > std::numeric_limits::max() - 36) + throw std::runtime_error("write_wav: output exceeds the RIFF/WAVE range"); const std::int32_t chunk_size = 36 + data_size; const std::int32_t format_size = 16; const std::int16_t audio_format = 3; diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index dbb5c93bb8..c853030487 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -4,9 +4,35 @@ */ #include "cli/cli.h" +#if defined(_WIN32) +#include "cli/windows_utf8_argv.h" +#endif +#include +#include #include -int main(int argc, char** argv) { +namespace { + +int run_cli(int argc, char** argv) { return trtmc::cli::run(argc, argv, std::cout, std::cerr); } + +} // namespace + +#if defined(_WIN32) +int wmain(int argc, wchar_t** argv) { + try { + trtmc::cli::Utf8CommandLine command_line(argc, argv); + return run_cli(command_line.argc(), command_line.argv()); + } catch (const std::exception& error) { + std::cerr << "Error: unable to decode the Windows command line as UTF-8: " << error.what() + << '\n'; + return EXIT_FAILURE; + } +} +#else +int main(int argc, char** argv) { + return run_cli(argc, argv); +} +#endif diff --git a/apps/cli/tests/test_cli.cpp b/apps/cli/tests/test_cli.cpp index 0824f5c86c..605f0529c7 100644 --- a/apps/cli/tests/test_cli.cpp +++ b/apps/cli/tests/test_cli.cpp @@ -5,7 +5,9 @@ #include "cli/cli.h" #include "cli/io.h" +#include "cli/windows_media.h" +#include #include #include #include @@ -21,6 +23,31 @@ namespace { int failures = 0; +class TestTempDirectory { + public: + TestTempDirectory() { + const auto nonce = std::chrono::steady_clock::now().time_since_epoch().count(); + path_ = + std::filesystem::temp_directory_path() / ("trtmc-cli-test-" + std::to_string(nonce)); + std::filesystem::create_directory(path_); + } + + ~TestTempDirectory() { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + + std::filesystem::path file(const char* name) const { return path_ / name; } + + private: + std::filesystem::path path_; +}; + +const TestTempDirectory& test_temp_directory() { + static const TestTempDirectory directory; + return directory; +} + void check(bool condition, const char* name) { if (!condition) { std::cerr << "FAIL: " << name << '\n'; @@ -89,7 +116,7 @@ class BareTextTask final : public trtmc::ITask { const char* task() const noexcept override { return trtmc::ITextGeneration::kTask; } }; -class EmptyImageWorker final : public trtmc::IImageGeneration { +class EmptyImageWorker final : public trtmc::IImageGeneration, public trtmc::IVideoGeneration { public: trtmc::ImageGenerationConfig seen; @@ -100,6 +127,18 @@ class EmptyImageWorker final : public trtmc::IImageGeneration { result.num_frames = 0; return result; } + + trtmc::VideoResult generate_video(const std::string&, + const trtmc::ImageGenerationConfig& config) override { + seen = config; + trtmc::VideoResult result; + result.frames.num_frames = 0; + return result; + } + + trtmc::VideoResult generate_video(const trtmc::VideoGenerationRequest& request) override { + return generate_video(request.prompt, request.config); + } }; class FakeStreamingAudio final : public trtmc::IAudioGeneration, @@ -388,9 +427,25 @@ int main() { "image replay latents option is retained"); const auto video_generation = parse({"trtmc", "generate-video", "model.bundle", "--runtime-root", "lib", "--prompt", - "cat", "--output", "cat.mp4", "--initial-latents-raw", "initial.f32"}); - check(video_generation.options.at("--initial-latents-raw") == "initial.f32", - "video replay latents option is retained"); + "cat", "--output", "cat.mp4", "--initial-latents-raw", "initial.f32", "--num-frames", + "124", "--first-frame", "first.png", "--last-frame", "last.png"}); + check(video_generation.options.at("--initial-latents-raw") == "initial.f32" && + video_generation.options.at("--num-frames") == "124" && + video_generation.options.at("--first-frame") == "first.png" && + video_generation.options.at("--last-frame") == "last.png", + "structured first/last-frame video options are retained"); + const auto reference_generation = + parse({"trtmc", "generate-video", "model.bundle", "--runtime-root", "lib", "--prompt", + "cat", "--output", "cat.mp4", "--reference-video", "clip.mp4", "--reference-image", + "still.png", "--reference-audio", "sound.wav"}); + check(reference_generation.video_references.size() == 3 && + reference_generation.video_references[0].kind == trtmc::VideoReferenceKind::kVideo && + reference_generation.video_references[0].path == "clip.mp4" && + reference_generation.video_references[1].kind == trtmc::VideoReferenceKind::kImage && + reference_generation.video_references[1].path == "still.png" && + reference_generation.video_references[2].kind == trtmc::VideoReferenceKind::kAudio && + reference_generation.video_references[2].path == "sound.wav", + "mixed video references preserve CLI order and media kind"); const auto batch = parse({"trtmc", "generate-image-batch", "model.bundle", "--runtime-root", "lib", "--prompts", "prompts.txt", "--seeds", "1,2", "--output", "images"}); @@ -441,9 +496,10 @@ int main() { run_command.options.emplace("--threshold", "0.75"); run_command.options.emplace("--use-chat-template", "true"); run_command.options.emplace("--enable-thinking", "false"); - run_command.options.emplace("--lora-adapter", "/tmp/adapter"); + const auto adapter_path = test_temp_directory().file("adapter").string(); + run_command.options.emplace("--lora-adapter", adapter_path); run_command.options.emplace("--lora-adapter-id", "demo"); - const std::filesystem::path replay_values_path = "/tmp/trtmc-cli-replay-values.f32"; + const auto replay_values_path = test_temp_directory().file("replay-values.f32"); const float replay_values[] = {0.25F, -0.5F}; { std::ofstream replay_file(replay_values_path, std::ios::binary); @@ -478,7 +534,7 @@ int main() { text.seen.condition_mask == text.seen.initial_latents && text.seen.sampling_steps == text.seen.initial_latents && text.seen.sde_noises == text.seen.initial_latents && - text.loaded_adapter_id == "demo" && text.loaded_adapter_path == "/tmp/adapter", + text.loaded_adapter_id == "demo" && text.loaded_adapter_path == adapter_path, "text sampling options reach the Task API"); auto invalid_block = run_command; invalid_block.options["--block-length"] = "invalid"; @@ -500,7 +556,7 @@ int main() { text.seen.text_generation_mode == "auto" && text.seen.block_length == 0 && text.seen.confidence_threshold == -1.0F && text.seen.temperature == 1.0F, "text diffusion options preserve Task API defaults"); - const std::filesystem::path unsupported_image_path = "/tmp/trtmc-cli-unsupported.ppm"; + const auto unsupported_image_path = test_temp_directory().file("unsupported.ppm"); { std::ofstream image_file(unsupported_image_path, std::ios::binary); image_file << "P6\n1 1\n255\n"; @@ -520,7 +576,7 @@ int main() { edit_command.name = "generate-image"; edit_command.options.emplace("--prompt", "make it night"); edit_command.options.emplace("--image", unsupported_image_path.string()); - edit_command.options.emplace("--output", "/tmp/trtmc-cli-unused.png"); + edit_command.options.emplace("--output", test_temp_directory().file("unused.png").string()); EmptyImageWorker image_only; check(dispatch_throws(edit_command, image_only), "image-generation task rejects an edit image instead of dropping it"); @@ -568,7 +624,7 @@ int main() { trtmc::cli::Command audio_command; audio_command.kind = trtmc::cli::CommandKind::kGenerateAudio; audio_command.name = "generate-audio"; - const std::filesystem::path audio_path = "/tmp/trtmc-cli-stream-test.raw"; + const auto audio_path = test_temp_directory().file("stream-test.raw"); audio_command.options.emplace("--prompt", "hello"); audio_command.options.emplace("--output", audio_path.string()); audio_command.options.emplace("--max-new-tokens", "9"); @@ -587,7 +643,7 @@ int main() { "streaming audio output format is explicit"); std::filesystem::remove(audio_path); - const std::filesystem::path transcription_path = "/tmp/trtmc-cli-transcription-stream.wav"; + const auto transcription_path = test_temp_directory().file("transcription-stream.wav"); trtmc::AudioResult transcription_audio; transcription_audio.samples = {0.25F, -0.5F, 0.75F}; transcription_audio.num_samples = 3; @@ -722,7 +778,7 @@ int main() { usage.str().find("--cuda-graphs") != std::string::npos, "help lists direct Task and backend options"); - const std::filesystem::path video_path = "/tmp/trtmc-cli-video-frame.ppm"; + const auto video_path = test_temp_directory().file("video-frame.ppm"); { std::ofstream image_file(video_path, std::ios::binary); image_file << "P6\n1 1\n255\n"; @@ -738,8 +794,8 @@ int main() { check(video_task.prompt.empty(), "missing video prompt reaches the family as empty"); std::filesystem::remove(video_path); - const std::filesystem::path forecast_values_path = "/tmp/trtmc-cli-forecast-values.f32"; - const std::filesystem::path forecast_mask_path = "/tmp/trtmc-cli-forecast-mask.f32"; + const auto forecast_values_path = test_temp_directory().file("forecast-values.f32"); + const auto forecast_mask_path = test_temp_directory().file("forecast-mask.f32"); const float forecast_values[] = {1.0F, 2.0F}; const float forecast_mask[] = {1.0F, 1.0F}; { @@ -762,9 +818,9 @@ int main() { std::filesystem::remove(forecast_values_path); std::filesystem::remove(forecast_mask_path); - const std::filesystem::path control_image_path = "/tmp/trtmc-cli-control.ppm"; - const std::filesystem::path control_state_path = "/tmp/trtmc-cli-control-state.f32"; - const std::filesystem::path control_output_path = "/tmp/trtmc-cli-control-actions.f32"; + const auto control_image_path = test_temp_directory().file("control.ppm"); + const auto control_state_path = test_temp_directory().file("control-state.f32"); + const auto control_output_path = test_temp_directory().file("control-actions.f32"); { std::ofstream image_file(control_image_path, std::ios::binary); image_file << "P6\n1 1\n255\n"; @@ -801,7 +857,16 @@ int main() { return true; } }; - const std::filesystem::path wav_path = "/tmp/trtmc-cli-io.wav"; + const auto wav_path = test_temp_directory().file("io.wav"); + trtmc::ReferenceMediaDecodePolicy canvas_policy{15, 24, 240, 1, 1, 32, 0.25, 4.0}; + check(trtmc::cli::detail::reference_video_decode_size(canvas_policy, 1, 1) == + std::make_pair(32U, 32U), + "canvas axes clamp to one alignment unit on every platform"); + canvas_policy.canvas_short_edge = 80; + canvas_policy.canvas_max_pixels = 80 * 80; + check(trtmc::cli::detail::reference_video_decode_size(canvas_policy, 1, 1) == + std::make_pair(64U, 64U), + "canvas rounding uses ties-to-even on every platform"); trtmc::AudioResult wav; wav.samples = {-1.0F, 0.25F, 1.0F}; wav.num_samples = 3; @@ -816,7 +881,7 @@ int main() { check(throws_runtime([&] { trtmc::cli::io::write_wav({}, wav_path.string()); }), "empty WAV is rejected"); - const std::filesystem::path png_path = "/tmp/trtmc-cli-io.png"; + const auto png_path = test_temp_directory().file("io.png"); const std::vector rgb{-0.5F, 0.5F, 1.5F}; trtmc::cli::io::save_png(png_path.string(), rgb, 1, 1); const auto loaded_image = trtmc::cli::io::read_image(png_path.string()); diff --git a/apps/cli/tests/test_windows_media.cpp b/apps/cli/tests/test_windows_media.cpp new file mode 100644 index 0000000000..715e928d34 --- /dev/null +++ b/apps/cli/tests/test_windows_media.cpp @@ -0,0 +1,559 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/io.h" +#include "cli/windows_media.h" + +#define WIN32_LEAN_AND_MEAN +#define NOMINMAX +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Microsoft::WRL::ComPtr; + +class TempDirectory { + public: + TempDirectory() { + path_ = std::filesystem::temp_directory_path() / + ("trtmc-windows-media-test-" + std::to_string(GetCurrentProcessId()) + "-" + + std::to_string(GetTickCount64())); + std::filesystem::create_directory(path_); + } + + ~TempDirectory() { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + + const std::filesystem::path& path() const { return path_; } + + private: + std::filesystem::path path_; +}; + +void require(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(message); +} + +void require_success(HRESULT result) { + require(SUCCEEDED(result), "Media Foundation call failed"); +} + +template +void require_throws_with(Fn&& fn, const char* needle, const char* message) { + try { + fn(); + } catch (const std::exception& error) { + require(std::string(error.what()).find(needle) != std::string::npos, message); + return; + } + throw std::runtime_error(message); +} + +} // namespace + +int run_test() { + constexpr trtmc::ReferenceMediaDecodePolicy policy{ + 15, 24, 240, 768, 768ULL * 1344, 32, 0.25, 4.0, + }; + constexpr trtmc::ReferenceMediaDecodePolicy compact_policy{ + 2, 12, 60, 256, 256ULL * 448, 16, 0.5, 2.0, + }; + require(trtmc::cli::detail::reference_video_frame_ceiling(compact_policy, 30'000, 1'001) == 60, + "frame ceiling must follow a non-default duration policy"); + require(trtmc::cli::detail::reference_video_decode_size(compact_policy, 3840, 2160) == + std::pair{448, 256}, + "decode canvas must follow non-default size and alignment policy values"); + require_throws_with( + [&] { (void)trtmc::cli::detail::reference_video_decode_size(compact_policy, 2560, 1080); }, + "configured range", "decode aspect bounds must follow the supplied policy"); + require(trtmc::cli::detail::reference_video_frame_ceiling(policy, 30'000, 1'001) == 450, + "configured video ceiling must round up fractional frame rates"); + require(trtmc::cli::detail::reference_video_decode_size(policy, 3840, 2160) == + std::pair{1344, 768}, + "4K reference video must decode onto the policy's bounded canvas"); + require(trtmc::cli::detail::reference_timeline_within_limit(policy, 0, 150'000'000), + "an exact configured presentation timeline must be accepted"); + require(!trtmc::cli::detail::reference_timeline_within_limit(policy, 150'000'000, 1), + "a non-empty sample starting at the duration limit must be rejected"); + require(!trtmc::cli::detail::reference_timeline_within_limit(policy, 149'000'000, 2'000'000), + "a sparse sample crossing the duration limit must be rejected"); + + std::vector timeline; + const std::vector first_audio{1.0F, -1.0F, 2.0F, -2.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, first_audio.data(), first_audio.size() / 2, 0, first_audio.size() / 2, 2, 20'000, + 1'000, 64); + require(timeline == std::vector({0.0F, 0.0F, 0.0F, 0.0F, 1.0F, -1.0F, 2.0F, -2.0F}), + "a positive first audio PTS must preserve leading silence"); + const std::vector gapped_audio{3.0F, -3.0F, 4.0F, -4.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, gapped_audio.data(), gapped_audio.size() / 2, 0, gapped_audio.size() / 2, 2, + 60'000, 1'000, 64); + require(timeline == std::vector({0.0F, 0.0F, 0.0F, 0.0F, 1.0F, -1.0F, 2.0F, -2.0F, 0.0F, + 0.0F, 0.0F, 0.0F, 3.0F, -3.0F, 4.0F, -4.0F}), + "an audio PTS discontinuity must preserve inter-sample silence"); + const std::vector overlapped_audio{9.0F, -9.0F, 10.0F, -10.0F, + 11.0F, -11.0F, 12.0F, -12.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, overlapped_audio.data(), overlapped_audio.size() / 2, 0, + overlapped_audio.size() / 2, 2, 70'000, 1'000, 64); + require(timeline == std::vector({0.0F, 0.0F, 0.0F, 0.0F, 1.0F, -1.0F, 2.0F, -2.0F, + 0.0F, 0.0F, 0.0F, 0.0F, 3.0F, -3.0F, 4.0F, -4.0F, + 10.0F, -10.0F, 11.0F, -11.0F, 12.0F, -12.0F}), + "a later overlapping audio sample must deterministically retain only its new tail"); + const auto timeline_before_contained_overlap = timeline; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, first_audio.data(), first_audio.size() / 2, 0, first_audio.size() / 2, 2, 0, + 1'000, 64); + require(timeline == timeline_before_contained_overlap, + "a wholly overlapped later audio sample must not alter the existing timeline"); + + std::vector trimmed_codec_timeline; + const std::vector primed_audio{99.0F, -99.0F, 1.0F, -1.0F, 2.0F, -2.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + trimmed_codec_timeline, primed_audio.data(), primed_audio.size() / 2, 1, 3, 2, -10'000, + 1'000, 16); + require(trimmed_codec_timeline == std::vector({1.0F, -1.0F, 2.0F, -2.0F}), + "negative-PTS codec priming must start at zero after using the retained source offset"); + require_throws_with( + [&] { + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + trimmed_codec_timeline, primed_audio.data(), 2, 0, 3, 2, 0, 1'000, 16); + }, + "invalid timeline layout", + "a retained audio window outside its source buffer must fail closed"); + + trtmc::cli::detail::validate_reference_video_soundtrack_format(32'000, 1); + trtmc::cli::detail::validate_reference_video_soundtrack_format(48'000, 2); + require_throws_with( + [] { trtmc::cli::detail::validate_reference_video_soundtrack_format(0, 2); }, + "invalid sample rate", "a video soundtrack with an invalid sample rate must fail closed"); + require_throws_with( + [] { trtmc::cli::detail::validate_reference_video_soundtrack_format(32'000, 3); }, + "mono or stereo", "a multichannel video soundtrack must fail closed"); + + constexpr std::uint32_t synthetic_codec_rate = 32'000; + constexpr std::uint64_t synthetic_mp3_access_unit_frames = 1'152; + constexpr std::uint64_t synthetic_mp3_padding_frames = 3 * synthetic_mp3_access_unit_frames; + constexpr std::uint64_t synthetic_mp3_padding_ticks = + (synthetic_mp3_padding_frames * 10'000'000 + synthetic_codec_rate - 1) / + synthetic_codec_rate; + constexpr std::uint64_t synthetic_public_audio_frames = 15 * synthetic_codec_rate; + require(trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, -static_cast(synthetic_mp3_padding_ticks), + synthetic_mp3_padding_ticks) && + trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, 150'000'000 + static_cast(synthetic_mp3_padding_ticks), + synthetic_mp3_padding_ticks), + "compressed empty-event timestamps must accept the exact codec-padding window"); + require(!trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, -36'000'000'000, synthetic_mp3_padding_ticks) && + !trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, 36'000'000'000, synthetic_mp3_padding_ticks), + "compressed empty-event timestamps at negative or positive hours must be rejected"); + trtmc::cli::detail::ReferenceAudioDecodeState boundary_padding_state; + require(trtmc::cli::detail::account_reference_audio_decode( + policy, 0, synthetic_public_audio_frames + synthetic_mp3_padding_frames, + synthetic_codec_rate, synthetic_mp3_padding_frames, boundary_padding_state), + "an exact three-access-unit MP3 decoded tail must be accepted"); + require(boundary_padding_state.decoded_padding_frames == synthetic_mp3_padding_frames, + "synthetic MP3 tail accounting must reach the codec-padding boundary"); + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 150'000'001, 1, synthetic_codec_rate, synthetic_mp3_padding_frames, + boundary_padding_state), + "one decoded frame beyond the MP3 padding boundary must be rejected"); + + trtmc::cli::detail::ReferenceAudioDecodeState falsified_duration_state; + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 150'000'000, synthetic_mp3_padding_frames + 1, synthetic_codec_rate, + synthetic_mp3_padding_frames, falsified_duration_state), + "a short or falsified presentation duration must not hide a decoded tail over three " + "MP3 access units"); + + trtmc::cli::detail::ReferenceAudioDecodeState future_timestamp_state; + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 36'000'000'000, 1, synthetic_codec_rate, synthetic_mp3_padding_frames, + future_timestamp_state), + "a tiny decoded sample with a far-future timestamp must not spend only one padding " + "frame"); + trtmc::cli::detail::ReferenceAudioDecodeState excessive_leading_state; + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, -130'000'000, synthetic_public_audio_frames, synthetic_codec_rate, + synthetic_mp3_padding_frames, excessive_leading_state), + "excessive negative-timestamp decoded PCM must not be hidden as leading padding"); + + constexpr std::uint64_t synthetic_aac_padding_frames = 3 * 1'024; + trtmc::cli::detail::ReferenceAudioDecodeState aac_padding_state; + require(trtmc::cli::detail::account_reference_audio_decode( + policy, 0, synthetic_public_audio_frames + synthetic_aac_padding_frames, + synthetic_codec_rate, synthetic_aac_padding_frames, aac_padding_state), + "an exact three-access-unit AAC decoded tail must be accepted"); + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 150'000'001, 1, synthetic_codec_rate, synthetic_aac_padding_frames, + aac_padding_state), + "one decoded frame beyond the AAC padding boundary must be rejected"); + + trtmc::VideoResult result; + result.frames.width = 64; + result.frames.height = 64; + result.frames.channels = 3; + result.frames.num_frames = 24; + result.fps = 24; + const std::size_t pixel_count = static_cast(result.frames.width) * + result.frames.height * result.frames.channels * + result.frames.num_frames; + result.frames.pixels.resize(pixel_count); + for (int frame = 0; frame < result.frames.num_frames; ++frame) { + for (int row = 0; row < result.frames.height; ++row) { + for (int column = 0; column < result.frames.width; ++column) { + const auto offset = + ((static_cast(frame) * result.frames.height + row) * + result.frames.width + + column) * + 3; + result.frames.pixels[offset] = static_cast(column) / 63.0F; + result.frames.pixels[offset + 1] = static_cast(row) / 63.0F; + result.frames.pixels[offset + 2] = static_cast(frame) / 23.0F; + } + } + } + + result.audio.sample_rate = 32000; + result.audio.channels = 2; + result.audio.samples.resize(static_cast(result.audio.sample_rate) * 2); + for (int frame = 0; frame < result.audio.sample_rate; ++frame) { + const float value = 0.1F * std::sin(2.0F * 3.14159265358979323846F * 440.0F * + static_cast(frame) / result.audio.sample_rate); + result.audio.samples[static_cast(frame) * 2] = value; + result.audio.samples[static_cast(frame) * 2 + 1] = value; + } + result.audio.num_samples = static_cast(result.audio.samples.size()); + + TempDirectory temporary; + const auto& temporary_root = temporary.path(); + const auto path = temporary_root / "video.mp4"; + const auto wav_path = temporary_root / "audio.wav"; + const auto mp3_path = temporary_root / "audio.mp3"; + const auto aac_path = temporary_root / "audio.m4a"; + const auto mp3_44k_path = temporary_root / "audio-44k.mp3"; + const auto aac_44k_path = temporary_root / "audio-44k.m4a"; + const auto over_limit_mp3_path = temporary_root / "over-limit.mp3"; + const auto boundary_wav_path = temporary_root / "boundary.wav"; + const auto over_limit_wav_path = temporary_root / "over-limit.wav"; + const auto over_limit_video_path = temporary_root / "over-limit.mp4"; + trtmc::cli::write_mp4(result, path.string()); + require(std::filesystem::is_regular_file(path), "MP4 output file is missing"); + require(std::filesystem::file_size(path) > 1024, "MP4 output file is empty"); + + const auto decoded = trtmc::cli::read_video_file(path.string(), policy); + require(decoded.width == 768, "decoded MP4 width must use the configured resolver"); + require(decoded.height == 768, "decoded MP4 height must use the configured resolver"); + require(decoded.channels == 3, "decoded MP4 channel mismatch"); + require(decoded.num_frames == result.frames.num_frames, "decoded MP4 frame-count mismatch"); + require(decoded.fps_numerator == result.fps, "decoded MP4 frame-rate mismatch"); + require(decoded.fps_denominator == 1, "decoded MP4 frame-rate denominator mismatch"); + require(decoded.pixels.size() == static_cast(decoded.width) * decoded.height * + decoded.num_frames * decoded.channels, + "decoded MP4 pixel-count mismatch"); + require(decoded.soundtrack.sample_rate == result.audio.sample_rate, + "decoded MP4 audio-rate mismatch"); + require(decoded.soundtrack.channels == result.audio.channels, + "decoded MP4 audio-channel mismatch"); + require(!decoded.soundtrack.samples.empty(), "decoded MP4 has no soundtrack samples"); + + const auto compact_decoded = trtmc::cli::read_video_file(path.string(), compact_policy); + require(compact_decoded.width == 256 && compact_decoded.height == 256, + "decoded MP4 canvas must follow a second policy"); + require(compact_decoded.num_frames == 12 && compact_decoded.fps_numerator == 12 && + compact_decoded.fps_denominator == 1, + "decoded MP4 sampling rate must follow a second policy"); + + const auto extracted_audio = trtmc::cli::read_audio_file(path.string(), policy); + require(extracted_audio.sample_rate == result.audio.sample_rate, + "standalone media audio-rate mismatch"); + require(extracted_audio.channels == result.audio.channels, + "standalone media audio-channel mismatch"); + require(!extracted_audio.samples.empty(), "standalone media audio decode is empty"); + + trtmc::cli::io::write_wav(result.audio, wav_path.string()); + const auto decoded_wav = trtmc::cli::read_audio_file(wav_path.string(), policy); + require(decoded_wav.sample_rate == result.audio.sample_rate, + "Media Foundation WAV audio-rate mismatch"); + require(decoded_wav.channels == result.audio.channels, + "Media Foundation WAV audio-channel mismatch"); + require(decoded_wav.num_samples == static_cast(decoded_wav.samples.size()) && + !decoded_wav.samples.empty(), + "Media Foundation WAV decode returned invalid samples"); + + trtmc::AudioResult boundary_audio; + boundary_audio.sample_rate = 8000; + boundary_audio.channels = 1; + boundary_audio.samples.resize(static_cast(boundary_audio.sample_rate) * 15); + boundary_audio.num_samples = static_cast(boundary_audio.samples.size()); + trtmc::cli::io::write_wav(boundary_audio, boundary_wav_path.string()); + const auto decoded_boundary_audio = + trtmc::cli::read_audio_file(boundary_wav_path.string(), policy); + require(decoded_boundary_audio.samples.size() == boundary_audio.samples.size(), + "configured reference audio boundary must decode successfully"); + + boundary_audio.samples.push_back(0.0F); + boundary_audio.num_samples = static_cast(boundary_audio.samples.size()); + trtmc::cli::io::write_wav(boundary_audio, over_limit_wav_path.string()); + require_throws_with( + [&] { (void)trtmc::cli::read_audio_file(over_limit_wav_path.string(), policy); }, + "configured duration limit", + "reference audio over the policy limit must fail during Media Foundation decode"); + + trtmc::VideoResult over_limit_video; + over_limit_video.frames.width = 64; + over_limit_video.frames.height = 64; + over_limit_video.frames.channels = 3; + over_limit_video.frames.num_frames = 361; + over_limit_video.fps = 24; + over_limit_video.frames.pixels.resize( + static_cast(over_limit_video.frames.width) * over_limit_video.frames.height * + over_limit_video.frames.channels * over_limit_video.frames.num_frames); + trtmc::cli::write_mp4(over_limit_video, over_limit_video_path.string()); + require_throws_with( + [&] { (void)trtmc::cli::read_video_file(over_limit_video_path.string(), policy); }, + "configured duration limit", + "reference video over the policy limit must fail during Media Foundation decode"); + + const auto decoded_pixel = [&](int frame, int row, int column, int channel) { + return decoded + .pixels[((static_cast(frame) * decoded.height + row) * decoded.width + + column) * + 3 + + channel]; + }; + require(decoded_pixel(0, 384, 672, 0) > decoded_pixel(0, 384, 96, 0), + "decoded MP4 red axis is reversed"); + require(decoded_pixel(0, 672, 384, 1) > decoded_pixel(0, 96, 384, 1), + "decoded MP4 green axis is reversed"); + require(decoded_pixel(20, 384, 384, 2) > decoded_pixel(3, 384, 384, 2), + "decoded MP4 frame order is reversed"); + + const HRESULT com_result = CoInitializeEx(nullptr, COINIT_MULTITHREADED); + const bool owns_com = SUCCEEDED(com_result); + require(owns_com || com_result == RPC_E_CHANGED_MODE, "CoInitializeEx failed"); + require_success(MFStartup(MF_VERSION, MFSTARTUP_FULL)); + + ComPtr empty_audio_buffer; + require_success(MFCreateMemoryBuffer(sizeof(float), &empty_audio_buffer)); + DWORD empty_audio_length = 1; + require_success(empty_audio_buffer->GetCurrentLength(&empty_audio_length)); + require(empty_audio_length == 0, + "synthetic Media Foundation audio buffer must start with zero current length"); + ComPtr empty_audio_sample; + require_success(MFCreateSample(&empty_audio_sample)); + require_success(empty_audio_sample->AddBuffer(empty_audio_buffer.Get())); + DWORD empty_audio_sample_length = 1; + require_success(empty_audio_sample->GetTotalLength(&empty_audio_sample_length)); + require(empty_audio_sample_length == 0, + "synthetic non-null audio sample must contain no payload"); + require_throws_with( + [&] { + trtmc::cli::detail::require_nonempty_decoded_buffer(empty_audio_sample_length, "audio"); + }, + "empty decoded audio buffer", + "a non-null sample with a zero-length audio buffer must fail closed"); + + ComPtr empty_video_sample; + require_success(MFCreateSample(&empty_video_sample)); + ComPtr empty_video_buffer; + require_success(MFCreateMemoryBuffer(4, &empty_video_buffer)); + require_success(empty_video_sample->AddBuffer(empty_video_buffer.Get())); + DWORD empty_video_length = 1; + require_success(empty_video_sample->GetTotalLength(&empty_video_length)); + require(empty_video_length == 0, + "synthetic Media Foundation video sample must have zero total length"); + require_throws_with( + [&] { trtmc::cli::detail::require_nonempty_decoded_buffer(empty_video_length, "video"); }, + "empty decoded video buffer", + "a non-null sample with a zero-length video buffer must fail closed"); + trtmc::cli::detail::require_nonempty_decoded_buffer(1, "audio"); + + ComPtr reader; + require_success(MFCreateSourceReaderFromURL(path.wstring().c_str(), nullptr, &reader)); + + ComPtr video_type; + require_success( + reader->GetNativeMediaType(MF_SOURCE_READER_FIRST_VIDEO_STREAM, 0, &video_type)); + GUID video_subtype{}; + require_success(video_type->GetGUID(MF_MT_SUBTYPE, &video_subtype)); + require(video_subtype == MFVideoFormat_H264, "MP4 video track is not H.264"); + + ComPtr audio_type; + require_success( + reader->GetNativeMediaType(MF_SOURCE_READER_FIRST_AUDIO_STREAM, 0, &audio_type)); + GUID audio_subtype{}; + require_success(audio_type->GetGUID(MF_MT_SUBTYPE, &audio_subtype)); + require(audio_subtype == MFAudioFormat_AAC, "MP4 audio track is not AAC"); + + constexpr int kCompressedBoundarySeconds = 15; + const std::size_t boundary_frames = + static_cast(result.audio.sample_rate) * kCompressedBoundarySeconds; + std::vector pcm(boundary_frames * result.audio.channels); + for (std::size_t frame = 0; frame < boundary_frames; ++frame) { + const float value = 0.1F * std::sin(2.0F * 3.14159265358979323846F * 440.0F * + static_cast(frame) / result.audio.sample_rate); + const auto quantized = static_cast(std::lround(value * 32767.0F)); + pcm[frame * 2] = quantized; + pcm[frame * 2 + 1] = quantized; + } + constexpr std::uint32_t boundary_44k_rate = 44'100; + const std::size_t boundary_44k_frames = + static_cast(boundary_44k_rate) * kCompressedBoundarySeconds; + std::vector pcm_44k(boundary_44k_frames * result.audio.channels); + for (std::size_t frame = 0; frame < boundary_44k_frames; ++frame) { + const float value = 0.1F * std::sin(2.0F * 3.14159265358979323846F * 440.0F * + static_cast(frame) / boundary_44k_rate); + const auto quantized = static_cast(std::lround(value * 32767.0F)); + pcm_44k[frame * 2] = quantized; + pcm_44k[frame * 2 + 1] = quantized; + } + const auto write_compressed_boundary = [&](const std::filesystem::path& output_path, + const GUID& subtype, + const std::vector& input_pcm, + std::uint32_t sample_rate, LONGLONG duration) { + ComPtr writer; + require_success( + MFCreateSinkWriterFromURL(output_path.wstring().c_str(), nullptr, nullptr, &writer)); + ComPtr output; + require_success(MFCreateMediaType(&output)); + require_success(output->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio)); + require_success(output->SetGUID(MF_MT_SUBTYPE, subtype)); + require_success(output->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, result.audio.channels)); + require_success(output->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate)); + require_success(output->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, + subtype == MFAudioFormat_AAC ? 24'000 : 16'000)); + if (subtype == MFAudioFormat_AAC) { + require_success(output->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16)); + require_success(output->SetUINT32(MF_MT_AAC_PAYLOAD_TYPE, 0)); + require_success(output->SetUINT32(MF_MT_AAC_AUDIO_PROFILE_LEVEL_INDICATION, 0x29)); + } + DWORD stream = 0; + require_success(writer->AddStream(output.Get(), &stream)); + ComPtr input; + require_success(MFCreateMediaType(&input)); + require_success(input->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio)); + require_success(input->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_PCM)); + require_success(input->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, result.audio.channels)); + require_success(input->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate)); + require_success(input->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16)); + const auto block_alignment = static_cast(result.audio.channels * 2); + require_success(input->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_alignment)); + require_success( + input->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, sample_rate * block_alignment)); + require_success(writer->SetInputMediaType(stream, input.Get(), nullptr)); + require_success(writer->BeginWriting()); + ComPtr buffer; + require_success(MFCreateMemoryBuffer( + static_cast(input_pcm.size() * sizeof(std::int16_t)), &buffer)); + BYTE* bytes = nullptr; + DWORD capacity = 0; + require_success(buffer->Lock(&bytes, &capacity, nullptr)); + require(capacity >= input_pcm.size() * sizeof(std::int16_t), + "Media Foundation compressed-audio input buffer is undersized"); + std::memcpy(bytes, input_pcm.data(), input_pcm.size() * sizeof(std::int16_t)); + require_success(buffer->Unlock()); + require_success( + buffer->SetCurrentLength(static_cast(input_pcm.size() * sizeof(std::int16_t)))); + ComPtr sample; + require_success(MFCreateSample(&sample)); + require_success(sample->AddBuffer(buffer.Get())); + require_success(sample->SetSampleTime(0)); + require_success(sample->SetSampleDuration(duration)); + require_success(writer->WriteSample(stream, sample.Get())); + require_success(writer->Finalize()); + }; + write_compressed_boundary(mp3_path, MFAudioFormat_MP3, pcm, result.audio.sample_rate, + 150'000'000); + write_compressed_boundary(aac_path, MFAudioFormat_AAC, pcm, result.audio.sample_rate, + 150'000'000); + write_compressed_boundary(mp3_44k_path, MFAudioFormat_MP3, pcm_44k, boundary_44k_rate, + 150'000'000); + write_compressed_boundary(aac_44k_path, MFAudioFormat_AAC, pcm_44k, boundary_44k_rate, + 150'000'000); + auto over_limit_pcm = pcm; + over_limit_pcm.resize(static_cast(result.audio.sample_rate) * + result.audio.channels * 16); + write_compressed_boundary(over_limit_mp3_path, MFAudioFormat_MP3, over_limit_pcm, + result.audio.sample_rate, 160'000'000); + + reader.Reset(); + require_success(MFShutdown()); + if (owns_com) + CoUninitialize(); + + require(std::filesystem::is_regular_file(mp3_path), "MP3 output file is missing"); + require(std::filesystem::file_size(mp3_path) > 1024, "MP3 output file is empty"); + const auto decoded_mp3 = trtmc::cli::read_audio_file(mp3_path.string(), policy); + require(decoded_mp3.sample_rate == result.audio.sample_rate, + "Media Foundation MP3 audio-rate mismatch"); + require(decoded_mp3.channels == result.audio.channels, + "Media Foundation MP3 audio-channel mismatch"); + require(decoded_mp3.num_samples == static_cast(decoded_mp3.samples.size()) && + !decoded_mp3.samples.empty(), + "Media Foundation MP3 decode returned invalid samples"); + require(decoded_mp3.samples.size() <= pcm.size() && + decoded_mp3.samples.size() >= pcm.size() - 4096, + "configured MP3 boundary must trim only codec padding"); + + require(std::filesystem::is_regular_file(aac_path), "AAC output file is missing"); + require(std::filesystem::file_size(aac_path) > 1024, "AAC output file is empty"); + const auto decoded_aac = trtmc::cli::read_audio_file(aac_path.string(), policy); + require(decoded_aac.sample_rate == result.audio.sample_rate, + "Media Foundation AAC audio-rate mismatch"); + require(decoded_aac.channels == result.audio.channels, + "Media Foundation AAC audio-channel mismatch"); + require(decoded_aac.samples.size() <= pcm.size() && + decoded_aac.samples.size() >= pcm.size() - 4096, + "configured AAC boundary must trim only codec padding"); + + const auto decoded_mp3_44k = trtmc::cli::read_audio_file(mp3_44k_path.string(), policy); + require(decoded_mp3_44k.sample_rate == static_cast(boundary_44k_rate), + "44.1 kHz MP3 boundary audio-rate mismatch"); + require(decoded_mp3_44k.samples.size() <= pcm_44k.size() && + decoded_mp3_44k.samples.size() >= pcm_44k.size() - 4096, + "configured 44.1 kHz MP3 boundary must trim only codec padding"); + const auto decoded_aac_44k = trtmc::cli::read_audio_file(aac_44k_path.string(), policy); + require(decoded_aac_44k.sample_rate == static_cast(boundary_44k_rate), + "44.1 kHz AAC boundary audio-rate mismatch"); + require(decoded_aac_44k.samples.size() <= pcm_44k.size() && + decoded_aac_44k.samples.size() >= pcm_44k.size() - 4096, + "configured 44.1 kHz AAC boundary must trim only codec padding"); + require_throws_with( + [&] { (void)trtmc::cli::read_audio_file(over_limit_mp3_path.string(), policy); }, + "configured duration", "a true 16-second MP3 must not be accepted as codec padding"); + + return 0; +} + +int main() { + try { + return run_test(); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/apps/cli/tests/test_windows_utf8_argv.cpp b/apps/cli/tests/test_windows_utf8_argv.cpp new file mode 100644 index 0000000000..3d0c524497 --- /dev/null +++ b/apps/cli/tests/test_windows_utf8_argv.cpp @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/windows_utf8_argv.h" + +#define NOMINMAX +#define WIN32_LEAN_AND_MEAN +#include +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +void test_multilingual_utf16_converts_to_utf8() { + const std::wstring multilingual = + L"Arabic \u0627\u0644\u0639\u0631\u0628\u064a\u0629 | Chinese \u4e2d\u6587 | English | " + L"French fran\u00e7ais | German Stra\u00dfe | Italian italiano | Japanese " + L"\u65e5\u672c\u8a9e | " + L"Korean \ud55c\uad6d\uc5b4 | Portuguese portugu\u00eas | Russian " + L"\u0440\u0443\u0441\u0441\u043a\u0438\u0439 | Spanish espa\u00f1ol"; + const std::string expected = + u8"Arabic \u0627\u0644\u0639\u0631\u0628\u064a\u0629 | Chinese \u4e2d\u6587 | English | " + u8"French fran\u00e7ais | German Stra\u00dfe | Italian italiano | Japanese " + u8"\u65e5\u672c\u8a9e | " + u8"Korean \ud55c\uad6d\uc5b4 | Portuguese portugu\u00eas | Russian " + u8"\u0440\u0443\u0441\u0441\u043a\u0438\u0439 | Spanish espa\u00f1ol"; + check(trtmc::cli::utf8_from_utf16(multilingual) == expected, + "11-language UTF-16 prompt converts exactly to UTF-8"); +} + +void test_command_line_preserves_prompt_and_path() { + wchar_t program[] = L"trtmc.exe"; + wchar_t prompt_flag[] = L"--prompt"; + wchar_t prompt[] = L"\u4e2d\u6587\u63d0\u793a"; + wchar_t output_flag[] = L"--output"; + wchar_t path[] = L"C:\\video\\\u6a21\u578b\\\u8f93\u5165 \u56fe\u7247.mp4"; + wchar_t* wide_argv[] = {program, prompt_flag, prompt, output_flag, path}; + + trtmc::cli::Utf8CommandLine command_line(5, wide_argv); + check(command_line.argc() == 5, "UTF-8 command line preserves argc"); + check(command_line.argv()[5] == nullptr, "UTF-8 command line is null terminated"); + check(std::string(command_line.argv()[2]) == u8"\u4e2d\u6587\u63d0\u793a", + "UTF-8 command line preserves Chinese prompt"); + check(std::string(command_line.argv()[4]) == + u8"C:\\video\\\u6a21\u578b\\\u8f93\u5165 \u56fe\u7247.mp4", + "UTF-8 command line preserves non-ASCII media path"); +} + +void test_invalid_utf16_is_rejected() { + const wchar_t invalid[] = {static_cast(0xD800), L'\0'}; + bool rejected = false; + try { + (void)trtmc::cli::utf8_from_utf16(std::wstring_view(invalid, 1)); + } catch (const std::runtime_error&) { + rejected = true; + } + check(rejected, "malformed UTF-16 command-line text is rejected"); +} + +void test_utf8_process_manifest_and_narrow_filesystem_path() { + HRSRC resource = FindResourceW(nullptr, MAKEINTRESOURCEW(1), MAKEINTRESOURCEW(24)); + check(resource != nullptr, "test executable contains an embedded process manifest"); + if (resource != nullptr) { + HGLOBAL loaded = LoadResource(nullptr, resource); + const DWORD size = SizeofResource(nullptr, resource); + const void* data = loaded == nullptr ? nullptr : LockResource(loaded); + const std::string manifest = + data == nullptr ? std::string{} : std::string(static_cast(data), size); + check(manifest.find("activeCodePage") != std::string::npos && + manifest.find(">UTF-8") != std::string::npos, + "PE manifest declares UTF-8 as the active process code page"); + } + check(GetACP() == CP_UTF8, "embedded manifest activates the UTF-8 process code page"); + + const auto wide_directory = + std::filesystem::temp_directory_path() / + (std::wstring(L"trtmc_utf8_path_\u6a21\u578b_") + std::to_wstring(GetCurrentProcessId())); + const auto wide_file = wide_directory / L"\u8f93\u5165 \u56fe\u7247.txt"; + std::error_code error; + (void)std::filesystem::remove(wide_file, error); + error.clear(); + (void)std::filesystem::remove(wide_directory, error); + error.clear(); + check(std::filesystem::create_directory(wide_directory, error) && !error, + "wide setup creates a non-ASCII test directory"); + + const std::string utf8_directory = trtmc::cli::utf8_from_utf16(wide_directory.native()); + const std::filesystem::path narrow_file = + std::filesystem::path(utf8_directory) / u8"\u8f93\u5165 \u56fe\u7247.txt"; + { + std::ofstream output(narrow_file, std::ios::binary | std::ios::trunc); + output << "utf8-path"; + check(static_cast(output), + "UTF-8 narrow std::filesystem path opens a non-ASCII output file"); + } + check(std::filesystem::exists(wide_file), + "UTF-8 narrow path resolves to the intended UTF-16 filesystem path"); + { + std::ifstream input(narrow_file, std::ios::binary); + std::string value; + input >> value; + check(value == "utf8-path", "UTF-8 narrow path reads the non-ASCII file back"); + } + + error.clear(); + check(std::filesystem::remove(wide_file, error) && !error, + "non-ASCII path test removes its file"); + error.clear(); + check(std::filesystem::remove(wide_directory, error) && !error, + "non-ASCII path test removes its directory"); +} + +} // namespace + +int main() { + test_multilingual_utf16_converts_to_utf8(); + test_command_line_preserves_prompt_and_path(); + test_invalid_utf16_is_rejected(); + test_utf8_process_manifest_and_narrow_filesystem_path(); + if (failures == 0) + std::cout << "All Windows UTF-8 argv tests passed\n"; + return failures == 0 ? 0 : 1; +} diff --git a/apps/cli/trtmc_windows_utf8.manifest b/apps/cli/trtmc_windows_utf8.manifest new file mode 100644 index 0000000000..afccd0d9ca --- /dev/null +++ b/apps/cli/trtmc_windows_utf8.manifest @@ -0,0 +1,13 @@ + + + + + + UTF-8 + true + + + diff --git a/apps/cli/windows_media.cpp b/apps/cli/windows_media.cpp new file mode 100644 index 0000000000..1166e8f040 --- /dev/null +++ b/apps/cli/windows_media.cpp @@ -0,0 +1,1434 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/windows_media.h" + +#include "cli/io.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#define NOMINMAX +#include +#include +#include +#include +#include +#include +#include +#endif + +namespace trtmc::cli { +namespace { + +std::string lowercase_extension(std::string_view path) { + const auto extension = std::filesystem::path(std::string(path)).extension().string(); + std::string lowered(extension); + std::transform(lowered.begin(), lowered.end(), lowered.begin(), + [](unsigned char value) { return static_cast(std::tolower(value)); }); + return lowered; +} + +constexpr std::uint64_t kAudioMediaTicksPerSecond = 10'000'000; + +bool reference_limit_ticks(const ReferenceMediaDecodePolicy& policy, + std::uint64_t& ticks) noexcept { + if (policy.maximum_duration_seconds == 0 || + policy.maximum_duration_seconds > + std::numeric_limits::max() / kAudioMediaTicksPerSecond) { + return false; + } + ticks = static_cast(policy.maximum_duration_seconds) * kAudioMediaTicksPerSecond; + return true; +} + +void validate_reference_media_decode_policy(const ReferenceMediaDecodePolicy& policy) { + std::uint64_t ignored_ticks = 0; + if (!reference_limit_ticks(policy, ignored_ticks) || policy.target_video_fps == 0 || + policy.maximum_source_video_fps == 0 || policy.canvas_short_edge == 0 || + policy.canvas_max_pixels == 0 || policy.canvas_multiple == 0 || + !std::isfinite(policy.minimum_aspect_ratio) || + !std::isfinite(policy.maximum_aspect_ratio) || policy.minimum_aspect_ratio <= 0.0 || + policy.maximum_aspect_ratio < policy.minimum_aspect_ratio) { + throw std::invalid_argument("reference-media decode policy values must be positive"); + } +} + +struct ReferenceVideoSize { + std::uint32_t width{0}; + std::uint32_t height{0}; +}; + +std::uint32_t round_canvas_axis(double value, const ReferenceMediaDecodePolicy& policy) { + const double scaled = value / policy.canvas_multiple; + const double lower = std::floor(scaled); + const double fraction = scaled - lower; + double rounded = lower; + if (fraction > 0.5 || (fraction == 0.5 && static_cast(lower) % 2 != 0)) + rounded += 1.0; + if (rounded > std::numeric_limits::max() / policy.canvas_multiple) + throw std::runtime_error("reference video canvas rounding overflow"); + return std::max(policy.canvas_multiple, + static_cast(rounded) * policy.canvas_multiple); +} + +ReferenceVideoSize resolve_reference_video_size(std::uint32_t source_width, + std::uint32_t source_height, + const ReferenceMediaDecodePolicy& policy) { + if (source_width == 0 || source_height == 0) + throw std::runtime_error("reference video has invalid source dimensions"); + const double ratio = static_cast(source_width) / source_height; + if (!std::isfinite(ratio) || ratio < policy.minimum_aspect_ratio || + ratio > policy.maximum_aspect_ratio) { + throw std::runtime_error("reference video aspect is outside the configured range"); + } + double width = ratio >= 1.0 ? policy.canvas_short_edge * ratio + : static_cast(policy.canvas_short_edge); + double height = ratio >= 1.0 ? static_cast(policy.canvas_short_edge) + : policy.canvas_short_edge / ratio; + const double area = width * height; + if (area > static_cast(policy.canvas_max_pixels)) { + const double scale = std::sqrt(static_cast(policy.canvas_max_pixels) / area); + width *= scale; + height *= scale; + } + return {round_canvas_axis(width, policy), round_canvas_axis(height, policy)}; +} + +struct ReferenceAudioFrameWindow { + std::uint64_t first_frame{0}; + std::uint64_t end_frame{0}; +}; + +std::uint64_t ticks_to_frames(std::uint64_t ticks, std::uint32_t sample_rate, + bool round_up) noexcept { + const std::uint64_t whole_seconds = ticks / kAudioMediaTicksPerSecond; + const std::uint64_t partial_ticks = ticks % kAudioMediaTicksPerSecond; + if (whole_seconds > std::numeric_limits::max() / sample_rate) + return std::numeric_limits::max(); + std::uint64_t frames = whole_seconds * sample_rate; + const std::uint64_t partial_product = partial_ticks * sample_rate; + const std::uint64_t partial_frames = + round_up ? (partial_product + kAudioMediaTicksPerSecond - 1) / kAudioMediaTicksPerSecond + : partial_product / kAudioMediaTicksPerSecond; + if (partial_frames > std::numeric_limits::max() - frames) + return std::numeric_limits::max(); + return frames + partial_frames; +} + +std::uint64_t ticks_to_frames_nearest(std::uint64_t ticks, std::uint32_t sample_rate) noexcept { + const std::uint64_t whole_seconds = ticks / kAudioMediaTicksPerSecond; + const std::uint64_t partial_ticks = ticks % kAudioMediaTicksPerSecond; + if (whole_seconds > std::numeric_limits::max() / sample_rate) + return std::numeric_limits::max(); + std::uint64_t frames = whole_seconds * sample_rate; + const std::uint64_t partial_product = partial_ticks * sample_rate; + const std::uint64_t partial_frames = + (partial_product + kAudioMediaTicksPerSecond / 2) / kAudioMediaTicksPerSecond; + if (partial_frames > std::numeric_limits::max() - frames) + return std::numeric_limits::max(); + return frames + partial_frames; +} + +bool frames_to_ticks_rounded_up(std::uint64_t frames, std::uint32_t sample_rate, + std::uint64_t& ticks) noexcept { + const std::uint64_t whole_seconds = frames / sample_rate; + const std::uint64_t partial_frames = frames % sample_rate; + if (whole_seconds > std::numeric_limits::max() / kAudioMediaTicksPerSecond) { + return false; + } + ticks = whole_seconds * kAudioMediaTicksPerSecond; + const std::uint64_t partial_product = partial_frames * kAudioMediaTicksPerSecond; + const std::uint64_t partial_ticks = (partial_product + sample_rate - 1) / sample_rate; + if (partial_ticks > std::numeric_limits::max() - ticks) + return false; + ticks += partial_ticks; + return true; +} + +bool account_reference_audio_decode_impl(std::int64_t timestamp, std::uint64_t frame_count, + std::uint32_t sample_rate, + std::uint64_t maximum_padding_frames, + detail::ReferenceAudioDecodeState& state, + ReferenceAudioFrameWindow* retained_window, + const ReferenceMediaDecodePolicy& policy) noexcept { + if (sample_rate == 0 || frame_count == 0) + return false; + + const std::uint64_t public_decoded_frames = + static_cast(policy.maximum_duration_seconds) * sample_rate; + if (maximum_padding_frames > + std::numeric_limits::max() - public_decoded_frames) { + return false; + } + const std::uint64_t maximum_decoded_frames = public_decoded_frames + maximum_padding_frames; + if (state.decoded_frames > maximum_decoded_frames || + frame_count > maximum_decoded_frames - state.decoded_frames) { + return false; + } + + std::uint64_t padding_ticks = 0; + std::uint64_t sample_duration_ticks = 0; + if (!frames_to_ticks_rounded_up(maximum_padding_frames, sample_rate, padding_ticks) || + !frames_to_ticks_rounded_up(frame_count, sample_rate, sample_duration_ticks)) { + return false; + } + std::uint64_t reference_limit = 0; + if (!reference_limit_ticks(policy, reference_limit) || + padding_ticks > std::numeric_limits::max() - reference_limit) { + return false; + } + const std::uint64_t latest_end_tick = reference_limit + padding_ticks; + + std::uint64_t leading_ticks = 0; + if (timestamp < 0) { + leading_ticks = static_cast(-(timestamp + 1)) + 1; + if (leading_ticks > padding_ticks) + return false; + if (sample_duration_ticks > leading_ticks && + sample_duration_ticks - leading_ticks > latest_end_tick) { + return false; + } + } else { + const std::uint64_t start_tick = static_cast(timestamp); + if (start_tick > latest_end_tick || sample_duration_ticks > latest_end_tick - start_tick) { + return false; + } + } + + const std::uint64_t first_frame = + std::min(frame_count, ticks_to_frames(leading_ticks, sample_rate, true)); + std::uint64_t end_frame = first_frame; + if (timestamp < static_cast(reference_limit)) { + const std::uint64_t remaining_ticks = + timestamp < 0 ? reference_limit + leading_ticks + : reference_limit - static_cast(timestamp); + end_frame = std::min(frame_count, ticks_to_frames(remaining_ticks, sample_rate, false)); + end_frame = std::max(end_frame, first_frame); + } + const std::uint64_t padding_frames = frame_count - (end_frame - first_frame); + if (state.decoded_padding_frames > maximum_padding_frames || + padding_frames > maximum_padding_frames - state.decoded_padding_frames) { + return false; + } + + state.decoded_frames += frame_count; + state.decoded_padding_frames += padding_frames; + if (retained_window != nullptr) + *retained_window = {first_frame, end_frame}; + return true; +} + +#if defined(_WIN32) + +using Microsoft::WRL::ComPtr; + +constexpr LONGLONG kMediaTicksPerSecond = 10'000'000; +constexpr std::size_t kMaxConsecutiveEmptyReads = 256; +constexpr std::size_t kMaxTotalEmptyReads = 16'384; + +[[noreturn]] void throw_hresult(const char* operation, HRESULT result) { + std::ostringstream message; + message << operation << " failed with HRESULT 0x" << std::hex << std::setw(8) + << std::setfill('0') << static_cast(result); + throw std::runtime_error(message.str()); +} + +void check_hresult(HRESULT result, const char* operation) { + if (FAILED(result)) + throw_hresult(operation, result); +} + +class MediaFoundationSession { + public: + MediaFoundationSession() { + const HRESULT com_result = CoInitializeEx(nullptr, COINIT_MULTITHREADED); + if (SUCCEEDED(com_result)) + owns_com_ = true; + else if (com_result != RPC_E_CHANGED_MODE) + throw_hresult("CoInitializeEx", com_result); + const HRESULT media_result = MFStartup(MF_VERSION, MFSTARTUP_FULL); + if (FAILED(media_result)) { + if (owns_com_) + CoUninitialize(); + owns_com_ = false; + throw_hresult("MFStartup", media_result); + } + media_started_ = true; + } + + ~MediaFoundationSession() { + if (media_started_) + (void)MFShutdown(); + if (owns_com_) + CoUninitialize(); + } + + MediaFoundationSession(const MediaFoundationSession&) = delete; + MediaFoundationSession& operator=(const MediaFoundationSession&) = delete; + + private: + bool owns_com_{false}; + bool media_started_{false}; +}; + +std::uint32_t checked_u32(std::uint64_t value, const char* label) { + if (value > std::numeric_limits::max()) + throw std::runtime_error(std::string(label) + " exceeds the Media Foundation limit"); + return static_cast(value); +} + +LONGLONG media_time(std::uint64_t numerator, std::uint32_t denominator) { + constexpr std::uint64_t kTicksPerSecond = 10'000'000; + if (denominator == 0 || numerator > std::numeric_limits::max() / kTicksPerSecond) + throw std::runtime_error("media timestamp overflow"); + return static_cast((numerator * kTicksPerSecond) / denominator); +} + +BYTE quantize(float value) { + return static_cast(std::clamp(static_cast(std::lround(value)), 0, 255)); +} + +void rgb_to_nv12(const float* rgb, std::uint32_t width, std::uint32_t height, + std::vector& nv12) { + const std::size_t y_size = static_cast(width) * height; + nv12.resize(y_size + y_size / 2); + auto* y_plane = nv12.data(); + auto* uv_plane = nv12.data() + y_size; + + const auto component = [](const float* pixel, int index) { + return std::clamp(pixel[index], 0.0F, 1.0F); + }; + for (std::uint32_t row = 0; row < height; ++row) { + for (std::uint32_t column = 0; column < width; ++column) { + const float* pixel = rgb + (static_cast(row) * width + column) * 3; + const float red = component(pixel, 0); + const float green = component(pixel, 1); + const float blue = component(pixel, 2); + y_plane[static_cast(row) * width + column] = + quantize(16.0F + 219.0F * (0.2126F * red + 0.7152F * green + 0.0722F * blue)); + } + } + + for (std::uint32_t row = 0; row < height; row += 2) { + for (std::uint32_t column = 0; column < width; column += 2) { + float cb = 0.0F; + float cr = 0.0F; + for (std::uint32_t dy = 0; dy < 2; ++dy) { + for (std::uint32_t dx = 0; dx < 2; ++dx) { + const float* pixel = + rgb + (static_cast(row + dy) * width + column + dx) * 3; + const float red = component(pixel, 0); + const float green = component(pixel, 1); + const float blue = component(pixel, 2); + cb += -0.114572F * red - 0.385428F * green + 0.5F * blue; + cr += 0.5F * red - 0.454153F * green - 0.045847F * blue; + } + } + const auto uv_offset = static_cast(row / 2) * width + column; + uv_plane[uv_offset] = quantize(128.0F + 224.0F * cb / 4.0F); + uv_plane[uv_offset + 1] = quantize(128.0F + 224.0F * cr / 4.0F); + } + } +} + +ComPtr make_sample(const void* bytes, std::size_t byte_count, LONGLONG timestamp, + LONGLONG duration) { + const auto media_bytes = checked_u32(byte_count, "media sample size"); + ComPtr buffer; + check_hresult(MFCreateMemoryBuffer(media_bytes, &buffer), "MFCreateMemoryBuffer"); + BYTE* destination = nullptr; + DWORD maximum_length = 0; + check_hresult(buffer->Lock(&destination, &maximum_length, nullptr), "IMFMediaBuffer::Lock"); + if (maximum_length < media_bytes) { + (void)buffer->Unlock(); + throw std::runtime_error("Media Foundation returned an undersized sample buffer"); + } + std::memcpy(destination, bytes, byte_count); + check_hresult(buffer->Unlock(), "IMFMediaBuffer::Unlock"); + check_hresult(buffer->SetCurrentLength(media_bytes), "IMFMediaBuffer::SetCurrentLength"); + + ComPtr sample; + check_hresult(MFCreateSample(&sample), "MFCreateSample"); + check_hresult(sample->AddBuffer(buffer.Get()), "IMFSample::AddBuffer"); + check_hresult(sample->SetSampleTime(timestamp), "IMFSample::SetSampleTime"); + check_hresult(sample->SetSampleDuration(duration), "IMFSample::SetSampleDuration"); + return sample; +} + +ComPtr make_video_output_type(std::uint32_t width, std::uint32_t height, + std::uint32_t fps) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(video output)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video), "video output major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_H264), "video output subtype"); + const auto bitrate = checked_u32( + std::max(8'000'000, static_cast(width) * height * fps / 3), + "H.264 bitrate"); + check_hresult(type->SetUINT32(MF_MT_AVG_BITRATE, bitrate), "video output bitrate"); + check_hresult(type->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive), + "video output interlace mode"); + check_hresult(type->SetUINT32(MF_MT_MPEG2_PROFILE, eAVEncH264VProfile_Main), + "video output H.264 profile"); + check_hresult(MFSetAttributeSize(type.Get(), MF_MT_FRAME_SIZE, width, height), + "video output frame size"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_FRAME_RATE, fps, 1), + "video output frame rate"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_PIXEL_ASPECT_RATIO, 1, 1), + "video output pixel aspect ratio"); + return type; +} + +ComPtr make_video_input_type(std::uint32_t width, std::uint32_t height, + std::uint32_t fps) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(video input)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video), "video input major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_NV12), "video input subtype"); + check_hresult(type->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive), + "video input interlace mode"); + check_hresult(type->SetUINT32(MF_MT_FIXED_SIZE_SAMPLES, TRUE), "video input fixed samples"); + check_hresult(type->SetUINT32(MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE), + "video input independent samples"); + check_hresult(type->SetUINT32(MF_MT_SAMPLE_SIZE, + checked_u32(static_cast(width) * height * 3 / 2, + "NV12 sample size")), + "video input sample size"); + check_hresult(type->SetUINT32(MF_MT_DEFAULT_STRIDE, width), "video input stride"); + check_hresult(MFSetAttributeSize(type.Get(), MF_MT_FRAME_SIZE, width, height), + "video input frame size"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_FRAME_RATE, fps, 1), + "video input frame rate"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_PIXEL_ASPECT_RATIO, 1, 1), + "video input pixel aspect ratio"); + return type; +} + +ComPtr make_audio_output_type(std::uint32_t sample_rate, std::uint32_t channels) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(audio output)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio), "audio output major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_AAC), "audio output subtype"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, channels), "audio output channels"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate), + "audio output sample rate"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16), "audio output bit depth"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, 24'000), + "audio output bitrate"); + check_hresult(type->SetUINT32(MF_MT_AAC_PAYLOAD_TYPE, 0), "AAC payload type"); + check_hresult(type->SetUINT32(MF_MT_AAC_AUDIO_PROFILE_LEVEL_INDICATION, 0x29), + "AAC profile level"); + return type; +} + +ComPtr make_audio_input_type(std::uint32_t sample_rate, std::uint32_t channels) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(audio input)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio), "audio input major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_PCM), "audio input subtype"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, channels), "audio input channels"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate), + "audio input sample rate"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16), "audio input bit depth"); + const auto block_alignment = + checked_u32(static_cast(channels) * 2, "audio block alignment"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_alignment), + "audio input block alignment"); + check_hresult( + type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, + checked_u32(static_cast(sample_rate) * block_alignment, + "PCM byte rate")), + "audio input byte rate"); + check_hresult(type->SetUINT32(MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE), + "audio input independent samples"); + return type; +} + +void validate_result(const VideoResult& result) { + const auto& frames = result.frames; + if (frames.width <= 0 || frames.height <= 0 || frames.num_frames <= 0 || frames.channels != 3 || + result.fps <= 0) + throw std::runtime_error( + "write_mp4 requires valid THWC RGB video and a positive frame rate"); + if ((frames.width & 1) != 0 || (frames.height & 1) != 0) + throw std::runtime_error("write_mp4 requires even frame dimensions for NV12/H.264"); + const auto pixels_per_frame = static_cast(frames.width) * frames.height * 3; + const auto required_pixels = pixels_per_frame * static_cast(frames.num_frames); + if (required_pixels > frames.pixels.size()) + throw std::runtime_error("write_mp4 frame storage is smaller than its THWC metadata"); + const auto& audio = result.audio; + if (!audio.samples.empty()) { + if (audio.sample_rate <= 0 || (audio.channels != 1 && audio.channels != 2) || + audio.samples.size() % static_cast(audio.channels) != 0) + throw std::runtime_error("write_mp4 requires valid mono or stereo interleaved audio"); + } +} + +std::uint64_t rounded_reference_frame_slot(std::uint64_t frame, std::uint32_t fps_numerator, + std::uint32_t fps_denominator, + const ReferenceMediaDecodePolicy& policy) { + if (fps_numerator == 0 || fps_denominator == 0 || + frame > + std::numeric_limits::max() / policy.target_video_fps / fps_denominator) { + throw std::runtime_error("reference video frame-slot arithmetic overflow"); + } + const std::uint64_t numerator = frame * policy.target_video_fps * fps_denominator; + if (numerator > (std::numeric_limits::max() - fps_numerator) / 2) + throw std::runtime_error("reference video frame-slot rounding overflow"); + return (2 * numerator + fps_numerator) / (2ULL * fps_numerator); +} + +ComPtr source_reader_video_type(std::uint32_t width, std::uint32_t height) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(source video)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video), "source video major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_RGB32), "source video subtype"); + check_hresult(MFSetAttributeSize(type.Get(), MF_MT_FRAME_SIZE, width, height), + "source video decoded frame size"); + return type; +} + +ComPtr source_reader_audio_type(std::uint32_t sample_rate, std::uint32_t channels) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(source audio)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio), "source audio major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_Float), "source audio subtype"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, channels), "source audio channels"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate), + "source audio sample rate"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 32), "source audio bit depth"); + const auto block_alignment = + checked_u32(static_cast(channels) * 4, "source audio block alignment"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_alignment), + "source audio block alignment"); + check_hresult( + type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, + checked_u32(static_cast(sample_rate) * block_alignment, + "source audio byte rate")), + "source audio byte rate"); + return type; +} + +void append_rgb32_frame(IMFSample* sample, std::uint32_t width, std::uint32_t height, LONG stride, + std::vector& pixels, std::size_t maximum_scalars) { + if (sample == nullptr) + throw std::runtime_error("Media Foundation returned a null video sample"); + ComPtr buffer; + check_hresult(sample->ConvertToContiguousBuffer(&buffer), + "IMFSample::ConvertToContiguousBuffer(video)"); + BYTE* data = nullptr; + DWORD maximum_length = 0; + DWORD current_length = 0; + check_hresult(buffer->Lock(&data, &maximum_length, ¤t_length), + "IMFMediaBuffer::Lock(video)"); + const auto unlock = [&]() { (void)buffer->Unlock(); }; + if (current_length == 0) { + unlock(); + detail::require_nonempty_decoded_buffer(current_length, "video"); + } + const auto absolute_stride = + static_cast(stride < 0 ? -static_cast(stride) : stride); + const auto required_bytes = absolute_stride * height; + if (absolute_stride < static_cast(width) * 4 || + required_bytes > current_length) { + unlock(); + throw std::runtime_error("Media Foundation returned an invalid RGB32 stride or buffer"); + } + const auto old_size = pixels.size(); + const auto frame_scalars = static_cast(width) * height * 3; + if (frame_scalars > maximum_scalars || old_size > maximum_scalars - frame_scalars) { + unlock(); + throw std::runtime_error("decoded reference video exceeds its bounded RGB allocation"); + } + if (frame_scalars > std::numeric_limits::max() - old_size) { + unlock(); + throw std::runtime_error("decoded video exceeds host address space"); + } + pixels.resize(old_size + static_cast(frame_scalars)); + float* destination = pixels.data() + old_size; + for (std::uint32_t row = 0; row < height; ++row) { + const auto source_row = stride >= 0 ? row : height - 1 - row; + const BYTE* source = data + static_cast(source_row) * absolute_stride; + for (std::uint32_t column = 0; column < width; ++column) { + const BYTE* bgra = source + static_cast(column) * 4; + float* rgb = destination + (static_cast(row) * width + column) * 3; + rgb[0] = static_cast(bgra[2]) / 255.0F; + rgb[1] = static_cast(bgra[1]) / 255.0F; + rgb[2] = static_cast(bgra[0]) / 255.0F; + } + } + unlock(); +} + +void append_float_audio(IMFSample* sample, std::uint32_t sample_rate, std::uint32_t channels, + LONGLONG timestamp, std::uint64_t maximum_codec_padding_frames, + detail::ReferenceAudioDecodeState& decode_state, + std::vector& samples, std::size_t maximum_scalars, + const ReferenceMediaDecodePolicy& policy) { + if (sample == nullptr) + throw std::runtime_error("Media Foundation returned a null audio sample"); + ComPtr buffer; + check_hresult(sample->ConvertToContiguousBuffer(&buffer), + "IMFSample::ConvertToContiguousBuffer(audio)"); + BYTE* data = nullptr; + DWORD maximum_length = 0; + DWORD current_length = 0; + check_hresult(buffer->Lock(&data, &maximum_length, ¤t_length), + "IMFMediaBuffer::Lock(audio)"); + if (current_length == 0) { + (void)buffer->Unlock(); + detail::require_nonempty_decoded_buffer(current_length, "audio"); + } + if (current_length % (sizeof(float) * channels) != 0) { + (void)buffer->Unlock(); + throw std::runtime_error("Media Foundation returned misaligned float audio"); + } + const auto scalar_count = current_length / sizeof(float); + const std::size_t frame_count = scalar_count / channels; + const bool trim_codec_padding = maximum_codec_padding_frames > 0; + ReferenceAudioFrameWindow retained_window{0, frame_count}; + if (trim_codec_padding) { + if (!account_reference_audio_decode_impl(timestamp, frame_count, sample_rate, + maximum_codec_padding_frames, decode_state, + &retained_window, policy)) { + (void)buffer->Unlock(); + throw std::runtime_error( + "decoded reference audio exceeds the configured duration plus codec padding"); + } + } + if (!trim_codec_padding) { + LONGLONG declared_sample_duration = 0; + if (FAILED(sample->GetSampleDuration(&declared_sample_duration)) || + declared_sample_duration < 0) { + declared_sample_duration = 0; + } + const auto inferred_duration = static_cast( + (static_cast(frame_count) * kMediaTicksPerSecond + sample_rate - 1) / + sample_rate); + const LONGLONG timeline_duration = std::max(declared_sample_duration, inferred_duration); + if (!detail::reference_timeline_within_limit(policy, timestamp, timeline_duration)) { + (void)buffer->Unlock(); + throw std::runtime_error( + "decoded reference audio sample crosses the configured reference timeline"); + } + } + try { + detail::append_reference_audio_frames_on_timeline( + samples, data, frame_count, retained_window.first_frame, retained_window.end_frame, + channels, timestamp, sample_rate, maximum_scalars); + } catch (...) { + (void)buffer->Unlock(); + throw; + } + check_hresult(buffer->Unlock(), "IMFMediaBuffer::Unlock(audio)"); +} + +std::optional presentation_duration(IMFSourceReader& reader) { + PROPVARIANT value; + PropVariantInit(&value); + const HRESULT status = + reader.GetPresentationAttribute(MF_SOURCE_READER_MEDIASOURCE, MF_PD_DURATION, &value); + std::optional result; + if (SUCCEEDED(status)) { + if (value.vt == VT_UI8 && + value.uhVal.QuadPart <= static_cast(std::numeric_limits::max())) { + result = static_cast(value.uhVal.QuadPart); + } else if (value.vt == VT_I8 && value.hVal.QuadPart >= 0) { + result = value.hVal.QuadPart; + } + } + (void)PropVariantClear(&value); + return result; +} + +struct CodecPaddingAllowance { + std::uint64_t frames{0}; + LONGLONG duration{0}; +}; + +CodecPaddingAllowance codec_padding_allowance(IMFMediaType* native_audio) { + if (native_audio == nullptr) + return {}; + GUID subtype{}; + UINT32 sample_rate = 0; + if (FAILED(native_audio->GetGUID(MF_MT_SUBTYPE, &subtype)) || + FAILED(native_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &sample_rate)) || + sample_rate == 0) { + return {}; + } + std::uint64_t padding_frames = 0; + if (subtype == MFAudioFormat_MP3) + padding_frames = 3ULL * 1152; + else if (subtype == MFAudioFormat_AAC) + padding_frames = 3ULL * 1024; + if (padding_frames == 0) + return {}; + return {padding_frames, + static_cast((padding_frames * kMediaTicksPerSecond + sample_rate - 1) / + sample_rate)}; +} + +std::size_t maximum_reference_audio_scalars(std::uint32_t sample_rate, std::uint32_t channels, + const ReferenceMediaDecodePolicy& policy) { + const std::uint64_t public_limit = + static_cast(sample_rate) * channels * policy.maximum_duration_seconds; + return static_cast(std::min( + public_limit, static_cast(std::numeric_limits::max()))); +} + +void reject_source_reader_error(DWORD flags, const char* label) { + if ((flags & MF_SOURCE_READERF_ERROR) != 0) + throw std::runtime_error(std::string(label) + " reported MF_SOURCE_READERF_ERROR"); +} + +void reject_empty_decoded_sample(IMFSample* sample, std::string_view media_kind) { + if (sample == nullptr) + throw std::runtime_error("Media Foundation returned a null decoded sample"); + DWORD total_length = 0; + check_hresult(sample->GetTotalLength(&total_length), "IMFSample::GetTotalLength(decoded)"); + detail::require_nonempty_decoded_buffer(total_length, media_kind); +} + +void reject_reference_timestamp(LONGLONG timestamp, const char* label, + const ReferenceMediaDecodePolicy& policy) { + const auto limit = + static_cast(policy.maximum_duration_seconds) * kMediaTicksPerSecond; + if (timestamp >= static_cast(limit)) + throw std::runtime_error(std::string(label) + + " timestamp exceeds the configured reference duration"); +} + +void reject_video_sample_timeline(IMFSample* sample, LONGLONG timestamp, + std::uint32_t fps_numerator, std::uint32_t fps_denominator, + const ReferenceMediaDecodePolicy& policy) { + LONGLONG duration = 0; + if (sample == nullptr || FAILED(sample->GetSampleDuration(&duration)) || duration <= 0) { + duration = static_cast( + (static_cast(kMediaTicksPerSecond) * fps_denominator + fps_numerator - + 1) / + fps_numerator); + } + if (!detail::reference_timeline_within_limit(policy, timestamp, duration)) + throw std::runtime_error( + "decoded reference video sample crosses the configured reference timeline"); +} + +void note_empty_source_reader_event(DWORD flags, LONGLONG timestamp, LONGLONG& last_tick_timestamp, + std::size_t& consecutive_empty_reads, + std::size_t& total_empty_reads, + std::uint64_t maximum_codec_padding_ticks, const char* label, + const ReferenceMediaDecodePolicy& policy) { + // A stream tick deliberately carries no sample. Treat an advancing tick as + // progress, but bound its timestamp to the public reference duration. All + // other empty reads (including repeated/non-advancing ticks) are bounded so + // a malformed source cannot spin forever. + if (++total_empty_reads > kMaxTotalEmptyReads) + throw std::runtime_error(std::string(label) + " returned too many empty stream events"); + if (maximum_codec_padding_ticks > 0 && !detail::reference_audio_event_timestamp_within_padding( + policy, timestamp, maximum_codec_padding_ticks)) { + throw std::runtime_error(std::string(label) + + " event timestamp exceeds the configured duration plus codec " + "padding"); + } + if ((flags & MF_SOURCE_READERF_STREAMTICK) != 0) { + if (maximum_codec_padding_ticks == 0) + reject_reference_timestamp(timestamp, label, policy); + if (timestamp > last_tick_timestamp) { + last_tick_timestamp = timestamp; + consecutive_empty_reads = 0; + return; + } + } + if (++consecutive_empty_reads > kMaxConsecutiveEmptyReads) + throw std::runtime_error(std::string(label) + " made no progress while decoding"); +} + +#endif + +} // namespace + +void detail::append_reference_audio_frames_on_timeline( + std::vector& destination, const void* source, std::uint64_t source_frame_count, + std::uint64_t first_source_frame, std::uint64_t end_source_frame, std::uint32_t channels, + std::int64_t timestamp, std::uint32_t sample_rate, std::size_t maximum_scalars) { + if (source == nullptr || channels == 0 || sample_rate == 0 || + first_source_frame > end_source_frame || end_source_frame > source_frame_count || + destination.size() % channels != 0) { + throw std::runtime_error("decoded reference audio has an invalid timeline layout"); + } + + const std::uint64_t retained_frames = end_source_frame - first_source_frame; + if (retained_frames == 0) + return; + + // Media Foundation timestamps are expressed in 100 ns units. Map each sample + // independently to the nearest PCM frame so timestamp rounding does not + // accumulate across access units. Negative codec-priming samples have already + // been trimmed to the public timeline and therefore begin at frame zero. + const std::uint64_t target_frame = + timestamp <= 0 + ? 0 + : ticks_to_frames_nearest(static_cast(timestamp), sample_rate); + if (target_frame == std::numeric_limits::max()) + throw std::runtime_error("decoded reference audio timestamp exceeds host range"); + + const std::uint64_t destination_frames = destination.size() / channels; + const std::uint64_t gap_frames = + target_frame > destination_frames ? target_frame - destination_frames : 0; + const std::uint64_t overlap_frames = + target_frame < destination_frames + ? std::min(destination_frames - target_frame, retained_frames) + : 0; + const std::uint64_t append_frames = retained_frames - overlap_frames; + if (gap_frames > std::numeric_limits::max() - append_frames) + throw std::runtime_error("decoded reference audio timeline exceeds host range"); + const std::uint64_t added_frames = gap_frames + append_frames; + if (added_frames > std::numeric_limits::max() / channels) + throw std::runtime_error("decoded audio exceeds host address space"); + const std::size_t added_scalars = static_cast(added_frames) * channels; + const std::size_t old_size = destination.size(); + if (added_scalars > maximum_scalars || old_size > maximum_scalars - added_scalars) + throw std::runtime_error("decoded reference audio exceeds the configured duration limit"); + if (added_scalars > std::numeric_limits::max() - old_size) + throw std::runtime_error("decoded audio exceeds host address space"); + + std::size_t source_byte_offset = 0; + std::size_t append_bytes = 0; + if (append_frames != 0) { + const std::uint64_t source_frame = first_source_frame + overlap_frames; + if (source_frame > std::numeric_limits::max() / channels || + append_frames > std::numeric_limits::max() / channels) { + throw std::runtime_error("decoded audio exceeds host address space"); + } + const std::size_t source_scalar = static_cast(source_frame) * channels; + const std::size_t append_scalars = static_cast(append_frames) * channels; + if (source_scalar > std::numeric_limits::max() / sizeof(float) || + append_scalars > std::numeric_limits::max() / sizeof(float)) { + throw std::runtime_error("decoded audio exceeds host address space"); + } + source_byte_offset = source_scalar * sizeof(float); + append_bytes = append_scalars * sizeof(float); + } + + destination.resize(old_size + added_scalars, 0.0F); + if (append_frames != 0) { + const auto* source_bytes = static_cast(source); + std::memcpy(destination.data() + old_size + static_cast(gap_frames) * channels, + source_bytes + source_byte_offset, append_bytes); + } +} + +void detail::validate_reference_video_soundtrack_format(std::uint32_t sample_rate, + std::uint32_t channels) { + if (sample_rate == 0 || + sample_rate > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("reference video soundtrack has an invalid sample rate"); + } + if (channels != 1 && channels != 2) { + throw std::runtime_error("reference video soundtrack must contain mono or stereo audio"); + } +} + +std::uint64_t detail::reference_video_frame_ceiling(const ReferenceMediaDecodePolicy& policy, + std::uint32_t fps_numerator, + std::uint32_t fps_denominator) { + if (policy.maximum_duration_seconds == 0 || fps_numerator == 0 || fps_denominator == 0) + throw std::invalid_argument("reference video frame rate must be positive"); + const std::uint64_t numerator = + static_cast(policy.maximum_duration_seconds) * fps_numerator; + if (numerator > std::numeric_limits::max() - (fps_denominator - 1ULL)) + throw std::invalid_argument("reference video frame ceiling overflow"); + return (numerator + fps_denominator - 1) / fps_denominator; +} + +std::pair +detail::reference_video_decode_size(const ReferenceMediaDecodePolicy& policy, + std::uint32_t source_width, std::uint32_t source_height) { + validate_reference_media_decode_policy(policy); + const auto size = resolve_reference_video_size(source_width, source_height, policy); + return {size.width, size.height}; +} + +bool detail::reference_timeline_within_limit(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, + std::int64_t duration) noexcept { + std::uint64_t unsigned_limit = 0; + if (!reference_limit_ticks(policy, unsigned_limit) || + unsigned_limit > static_cast(std::numeric_limits::max())) { + return false; + } + const auto limit = static_cast(unsigned_limit); + return timestamp >= 0 && duration >= 0 && timestamp <= limit && duration <= limit - timestamp; +} + +bool detail::reference_audio_event_timestamp_within_padding( + const ReferenceMediaDecodePolicy& policy, std::int64_t timestamp, + std::uint64_t maximum_padding_ticks) noexcept { + std::uint64_t reference_limit = 0; + if (!reference_limit_ticks(policy, reference_limit) || + maximum_padding_ticks > std::numeric_limits::max() - reference_limit) { + return false; + } + if (timestamp < 0) { + const std::uint64_t leading_ticks = static_cast(-(timestamp + 1)) + 1; + return leading_ticks <= maximum_padding_ticks; + } + return static_cast(timestamp) <= reference_limit + maximum_padding_ticks; +} + +bool detail::account_reference_audio_decode(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, std::uint64_t frame_count, + std::uint32_t sample_rate, + std::uint64_t maximum_padding_frames, + ReferenceAudioDecodeState& state) noexcept { + return account_reference_audio_decode_impl(timestamp, frame_count, sample_rate, + maximum_padding_frames, state, nullptr, policy); +} + +void detail::require_nonempty_decoded_buffer(std::uint32_t current_length, + std::string_view media_kind) { + if (current_length == 0) { + throw std::runtime_error("Media Foundation returned an empty decoded " + + std::string(media_kind) + " buffer"); + } +} + +bool is_mp4_path(std::string_view path) { + return lowercase_extension(path) == ".mp4"; +} + +void write_mp4(const VideoResult& result, const std::string& path) { +#if defined(_WIN32) + validate_result(result); + if (!is_mp4_path(path)) + throw std::runtime_error("write_mp4 output path must end in .mp4"); + const auto output_path = std::filesystem::path(path); + if (!output_path.parent_path().empty()) + std::filesystem::create_directories(output_path.parent_path()); + + MediaFoundationSession session; + ComPtr attributes; + check_hresult(MFCreateAttributes(&attributes, 2), "MFCreateAttributes(sink writer)"); + check_hresult(attributes->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE), + "enable Media Foundation hardware transforms"); + check_hresult(attributes->SetUINT32(MF_LOW_LATENCY, FALSE), + "configure Media Foundation latency"); + + ComPtr writer; + check_hresult(MFCreateSinkWriterFromURL(output_path.wstring().c_str(), nullptr, + attributes.Get(), &writer), + "MFCreateSinkWriterFromURL"); + + const auto width = static_cast(result.frames.width); + const auto height = static_cast(result.frames.height); + const auto fps = static_cast(result.fps); + DWORD video_stream = 0; + const auto video_output = make_video_output_type(width, height, fps); + check_hresult(writer->AddStream(video_output.Get(), &video_stream), + "IMFSinkWriter::AddStream(video)"); + const auto video_input = make_video_input_type(width, height, fps); + check_hresult(writer->SetInputMediaType(video_stream, video_input.Get(), nullptr), + "IMFSinkWriter::SetInputMediaType(video)"); + + const bool has_audio = !result.audio.samples.empty(); + DWORD audio_stream = 0; + std::uint32_t audio_rate = 0; + std::uint32_t audio_channels = 0; + if (has_audio) { + audio_rate = static_cast(result.audio.sample_rate); + audio_channels = static_cast(result.audio.channels); + const auto audio_output = make_audio_output_type(audio_rate, audio_channels); + check_hresult(writer->AddStream(audio_output.Get(), &audio_stream), + "IMFSinkWriter::AddStream(audio)"); + const auto audio_input = make_audio_input_type(audio_rate, audio_channels); + check_hresult(writer->SetInputMediaType(audio_stream, audio_input.Get(), nullptr), + "IMFSinkWriter::SetInputMediaType(audio)"); + } + + check_hresult(writer->BeginWriting(), "IMFSinkWriter::BeginWriting"); + + const std::size_t pixels_per_frame = static_cast(width) * height * 3; + const std::uint64_t audio_frames = has_audio ? result.audio.samples.size() / audio_channels : 0; + constexpr std::uint64_t kAudioChunkFrames = 1024; + std::uint64_t video_index = 0; + std::uint64_t audio_offset = 0; + std::vector nv12; + std::vector pcm; + + while (video_index < static_cast(result.frames.num_frames) || + audio_offset < audio_frames) { + const auto next_video_time = + video_index < static_cast(result.frames.num_frames) + ? media_time(video_index, fps) + : std::numeric_limits::max(); + const auto next_audio_time = audio_offset < audio_frames + ? media_time(audio_offset, audio_rate) + : std::numeric_limits::max(); + if (next_video_time <= next_audio_time) { + const float* source = result.frames.pixels.data() + + static_cast(video_index) * pixels_per_frame; + rgb_to_nv12(source, width, height, nv12); + const auto end_time = media_time(video_index + 1, fps); + auto sample = + make_sample(nv12.data(), nv12.size(), next_video_time, end_time - next_video_time); + check_hresult(writer->WriteSample(video_stream, sample.Get()), + "IMFSinkWriter::WriteSample(video)"); + ++video_index; + } else { + const auto chunk_frames = std::min(kAudioChunkFrames, audio_frames - audio_offset); + pcm.resize(static_cast(chunk_frames) * audio_channels); + const auto scalar_offset = static_cast(audio_offset) * audio_channels; + for (std::size_t index = 0; index < pcm.size(); ++index) { + const float value = + std::clamp(result.audio.samples[scalar_offset + index], -1.0F, 1.0F); + pcm[index] = static_cast( + std::lround(value * (value < 0.0F ? 32768.0F : 32767.0F))); + } + const auto end_time = media_time(audio_offset + chunk_frames, audio_rate); + auto sample = make_sample(pcm.data(), pcm.size() * sizeof(std::int16_t), + next_audio_time, end_time - next_audio_time); + check_hresult(writer->WriteSample(audio_stream, sample.Get()), + "IMFSinkWriter::WriteSample(audio)"); + audio_offset += chunk_frames; + } + } + + check_hresult(writer->Finalize(), "IMFSinkWriter::Finalize"); +#else + (void)result; + (void)path; + throw std::runtime_error("native MP4 output is available on Windows through Media Foundation"); +#endif +} + +VideoClipInput read_video_file(const std::string& path, const ReferenceMediaDecodePolicy& policy) { +#if defined(_WIN32) + validate_reference_media_decode_policy(policy); + if (path.empty() || std::filesystem::is_directory(path)) + throw std::runtime_error("read_video_file requires a media file path"); + MediaFoundationSession session; + ComPtr attributes; + check_hresult(MFCreateAttributes(&attributes, 2), "MFCreateAttributes(source reader)"); + check_hresult(attributes->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE), + "enable source-reader hardware transforms"); + check_hresult(attributes->SetUINT32(MF_SOURCE_READER_ENABLE_ADVANCED_VIDEO_PROCESSING, TRUE), + "enable source-reader advanced video processing"); + + ComPtr reader; + check_hresult(MFCreateSourceReaderFromURL(std::filesystem::path(path).wstring().c_str(), + attributes.Get(), &reader), + "MFCreateSourceReaderFromURL"); + const auto declared_duration = presentation_duration(*reader.Get()); + check_hresult(reader->SetStreamSelection(MF_SOURCE_READER_ALL_STREAMS, FALSE), + "disable source-reader streams"); + + DWORD video_stream_index = MAXDWORD; + DWORD audio_stream_index = MAXDWORD; + ComPtr native_video; + ComPtr native_audio; + for (DWORD stream = 0; stream < 64; ++stream) { + ComPtr native_type; + const HRESULT type_result = reader->GetNativeMediaType(stream, 0, &native_type); + if (type_result == MF_E_INVALIDSTREAMNUMBER) + break; + if (FAILED(type_result)) + continue; + GUID major_type{}; + if (FAILED(native_type->GetGUID(MF_MT_MAJOR_TYPE, &major_type))) + continue; + if (major_type == MFMediaType_Video && video_stream_index == MAXDWORD) { + video_stream_index = stream; + native_video = native_type; + } else if (major_type == MFMediaType_Audio && audio_stream_index == MAXDWORD) { + audio_stream_index = stream; + native_audio = native_type; + } + } + if (video_stream_index == MAXDWORD || !native_video) + throw std::runtime_error("reference media contains no video stream"); + const CodecPaddingAllowance padding = codec_padding_allowance(native_audio.Get()); + if (declared_duration && + *declared_duration > + static_cast(policy.maximum_duration_seconds) * kMediaTicksPerSecond + + padding.duration) { + throw std::runtime_error( + "reference media presentation exceeds the configured duration limit"); + } + const bool trim_codec_padding = padding.frames > 0; + UINT32 fps_numerator = 0; + UINT32 fps_denominator = 0; + check_hresult( + MFGetAttributeRatio(native_video.Get(), MF_MT_FRAME_RATE, &fps_numerator, &fps_denominator), + "read source frame rate"); + if (fps_numerator == 0 || fps_denominator == 0 || + static_cast(fps_numerator) > + static_cast(policy.maximum_source_video_fps) * fps_denominator) { + throw std::runtime_error("reference video source frame rate exceeds the configured limit"); + } + UINT32 source_width = 0; + UINT32 source_height = 0; + check_hresult( + MFGetAttributeSize(native_video.Get(), MF_MT_FRAME_SIZE, &source_width, &source_height), + "read source frame size"); + const auto [requested_width, requested_height] = + detail::reference_video_decode_size(policy, source_width, source_height); + const ReferenceVideoSize requested_size{requested_width, requested_height}; + check_hresult(reader->SetStreamSelection(video_stream_index, TRUE), "select source video"); + const auto requested_video = + source_reader_video_type(requested_size.width, requested_size.height); + check_hresult(reader->SetCurrentMediaType(video_stream_index, nullptr, requested_video.Get()), + "SetCurrentMediaType(video RGB32)"); + ComPtr decoded_video; + check_hresult(reader->GetCurrentMediaType(video_stream_index, &decoded_video), + "GetCurrentMediaType(video)"); + UINT32 width = 0; + UINT32 height = 0; + check_hresult(MFGetAttributeSize(decoded_video.Get(), MF_MT_FRAME_SIZE, &width, &height), + "read decoded frame size"); + if (width != requested_size.width || height != requested_size.height) { + throw std::runtime_error( + "Media Foundation did not honor the bounded reference-video decode canvas"); + } + UINT32 raw_stride = 0; + LONG stride = static_cast(width * 4); + if (SUCCEEDED(decoded_video->GetUINT32(MF_MT_DEFAULT_STRIDE, &raw_stride))) + stride = static_cast(raw_stride); + + bool has_audio = false; + std::uint32_t audio_rate = 0; + std::uint32_t audio_channels = 0; + if (audio_stream_index != MAXDWORD && native_audio) { + UINT32 rate = 0; + UINT32 channels = 0; + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &rate), + "read source soundtrack sample rate"); + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &channels), + "read source soundtrack channel count"); + detail::validate_reference_video_soundtrack_format(rate, channels); + audio_rate = rate; + audio_channels = channels; + check_hresult(reader->SetStreamSelection(audio_stream_index, TRUE), "select source audio"); + const auto requested_audio = source_reader_audio_type(audio_rate, audio_channels); + check_hresult( + reader->SetCurrentMediaType(audio_stream_index, nullptr, requested_audio.Get()), + "SetCurrentMediaType(audio float)"); + ComPtr decoded_audio; + check_hresult(reader->GetCurrentMediaType(audio_stream_index, &decoded_audio), + "GetCurrentMediaType(audio)"); + UINT32 decoded_rate = 0; + UINT32 decoded_channels = 0; + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &decoded_rate), + "read decoded soundtrack sample rate"); + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &decoded_channels), + "read decoded soundtrack channel count"); + detail::validate_reference_video_soundtrack_format(decoded_rate, decoded_channels); + if (decoded_rate != audio_rate || decoded_channels != audio_channels) { + throw std::runtime_error( + "Media Foundation did not honor the requested soundtrack PCM format"); + } + has_audio = true; + } + + if (width == 0 || height == 0 || fps_numerator == 0 || fps_denominator == 0 || + width > static_cast(std::numeric_limits::max()) || + height > static_cast(std::numeric_limits::max()) || + fps_numerator > static_cast(std::numeric_limits::max()) || + fps_denominator > static_cast(std::numeric_limits::max())) + throw std::runtime_error("decoded video metadata is outside the public C++ value range"); + + VideoClipInput result; + result.width = static_cast(width); + result.height = static_cast(height); + result.channels = 3; + const bool downsample_video = + static_cast(fps_numerator) > + static_cast(policy.target_video_fps) * fps_denominator; + result.fps_numerator = + static_cast(downsample_video ? policy.target_video_fps : fps_numerator); + result.fps_denominator = static_cast(downsample_video ? 1 : fps_denominator); + result.soundtrack.sample_rate = static_cast(audio_rate); + result.soundtrack.channels = static_cast(audio_channels); + + const std::uint64_t maximum_source_video_frames = std::min( + detail::reference_video_frame_ceiling(policy, fps_numerator, fps_denominator), + static_cast(std::numeric_limits::max())); + const std::uint64_t maximum_video_frames = + downsample_video + ? static_cast(policy.maximum_duration_seconds) * policy.target_video_fps + : maximum_source_video_frames; + const std::uint64_t frame_scalars = static_cast(width) * height * 3; + if (maximum_video_frames > + static_cast(std::numeric_limits::max()) / frame_scalars) + throw std::runtime_error("bounded reference-video RGB allocation exceeds host range"); + const std::size_t maximum_rgb_scalars = + static_cast(maximum_video_frames * frame_scalars); + const std::size_t maximum_soundtrack_scalars = + has_audio ? maximum_reference_audio_scalars(audio_rate, audio_channels, policy) : 0; + + bool video_done = false; + bool audio_done = !has_audio; + LONGLONG last_video_tick = std::numeric_limits::min(); + LONGLONG last_audio_tick = std::numeric_limits::min(); + std::size_t consecutive_empty_reads = 0; + std::size_t total_empty_reads = 0; + std::uint64_t source_video_frames = 0; + detail::ReferenceAudioDecodeState audio_decode_state; + while (!video_done || !audio_done) { + DWORD stream_index = 0; + DWORD flags = 0; + LONGLONG timestamp = 0; + ComPtr sample; + check_hresult(reader->ReadSample(MF_SOURCE_READER_ANY_STREAM, 0, &stream_index, &flags, + ×tamp, &sample), + "IMFSourceReader::ReadSample"); + reject_source_reader_error(flags, "reference media source reader"); + if ((flags & (MF_SOURCE_READERF_NATIVEMEDIATYPECHANGED | + MF_SOURCE_READERF_CURRENTMEDIATYPECHANGED)) != 0) + throw std::runtime_error("reference video changes media type mid-stream"); + if ((flags & MF_SOURCE_READERF_ENDOFSTREAM) != 0) { + if (stream_index == video_stream_index) + video_done = true; + else if (stream_index == audio_stream_index) + audio_done = true; + else + throw std::runtime_error( + "reference media ended an unexpected source-reader stream"); + continue; + } + if (stream_index != video_stream_index && stream_index != audio_stream_index) + throw std::runtime_error("reference media returned an unexpected source-reader stream"); + const bool audio_padding_event = trim_codec_padding && stream_index == audio_stream_index; + const std::uint64_t event_padding_ticks = + audio_padding_event ? static_cast(padding.duration) : 0; + if (!audio_padding_event) + reject_reference_timestamp(timestamp, "reference media source reader", policy); + if ((flags & MF_SOURCE_READERF_STREAMTICK) != 0) { + if (sample) + throw std::runtime_error( + "reference media returned a sample for an empty stream tick"); + auto& last_tick = + stream_index == video_stream_index ? last_video_tick : last_audio_tick; + note_empty_source_reader_event(flags, timestamp, last_tick, consecutive_empty_reads, + total_empty_reads, event_padding_ticks, + "reference media source reader", policy); + continue; + } + if (!sample) { + auto& last_tick = + stream_index == video_stream_index ? last_video_tick : last_audio_tick; + note_empty_source_reader_event(flags, timestamp, last_tick, consecutive_empty_reads, + total_empty_reads, event_padding_ticks, + "reference media source reader", policy); + continue; + } + reject_empty_decoded_sample(sample.Get(), + stream_index == video_stream_index ? "video" : "audio"); + if (stream_index == video_stream_index) { + if (source_video_frames >= maximum_source_video_frames) + throw std::runtime_error( + "decoded reference video exceeds the configured duration limit"); + reject_video_sample_timeline(sample.Get(), timestamp, fps_numerator, fps_denominator, + policy); + bool retain = true; + if (downsample_video) { + const auto begin = rounded_reference_frame_slot(source_video_frames, fps_numerator, + fps_denominator, policy); + const auto end = rounded_reference_frame_slot( + source_video_frames + 1, fps_numerator, fps_denominator, policy); + retain = end > begin; + } + ++source_video_frames; + if (retain) { + if (static_cast(result.num_frames) >= maximum_video_frames) + throw std::runtime_error( + "decoded reference video exceeds the bounded target frame count"); + append_rgb32_frame(sample.Get(), width, height, stride, result.pixels, + maximum_rgb_scalars); + ++result.num_frames; + } + } else if (stream_index == audio_stream_index) { + append_float_audio(sample.Get(), audio_rate, audio_channels, timestamp, padding.frames, + audio_decode_state, result.soundtrack.samples, + maximum_soundtrack_scalars, policy); + } + consecutive_empty_reads = 0; + } + + if (result.num_frames <= 0) + throw std::runtime_error("reference video contains no decoded frames"); + if (!result.soundtrack.samples.empty()) { + if (result.soundtrack.samples.size() > + static_cast(std::numeric_limits::max())) + throw std::runtime_error("decoded soundtrack exceeds the public C++ value range"); + result.soundtrack.num_samples = static_cast(result.soundtrack.samples.size()); + } + return result; +#else + (void)path; + (void)policy; + throw std::runtime_error( + "native media-file input is available on Windows through Media Foundation"); +#endif +} + +AudioResult read_audio_file(const std::string& path, const ReferenceMediaDecodePolicy& policy) { +#if defined(_WIN32) + validate_reference_media_decode_policy(policy); + if (path.empty() || std::filesystem::is_directory(path)) + throw std::runtime_error("read_audio_file requires a media file path"); + MediaFoundationSession session; + ComPtr attributes; + check_hresult(MFCreateAttributes(&attributes, 1), "MFCreateAttributes(audio source reader)"); + check_hresult(attributes->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE), + "enable audio source-reader hardware transforms"); + + ComPtr reader; + check_hresult(MFCreateSourceReaderFromURL(std::filesystem::path(path).wstring().c_str(), + attributes.Get(), &reader), + "MFCreateSourceReaderFromURL(audio)"); + const auto declared_duration = presentation_duration(*reader.Get()); + check_hresult(reader->SetStreamSelection(MF_SOURCE_READER_ALL_STREAMS, FALSE), + "disable audio source-reader streams"); + + DWORD audio_stream_index = MAXDWORD; + ComPtr native_audio; + for (DWORD stream = 0; stream < 64; ++stream) { + ComPtr native_type; + const HRESULT type_result = reader->GetNativeMediaType(stream, 0, &native_type); + if (type_result == MF_E_INVALIDSTREAMNUMBER) + break; + if (FAILED(type_result)) + continue; + GUID major_type{}; + if (SUCCEEDED(native_type->GetGUID(MF_MT_MAJOR_TYPE, &major_type)) && + major_type == MFMediaType_Audio) { + audio_stream_index = stream; + native_audio = native_type; + break; + } + } + if (audio_stream_index == MAXDWORD || !native_audio) + throw std::runtime_error("reference audio file contains no audio stream"); + + UINT32 native_rate = 0; + UINT32 native_channels = 0; + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &native_rate), + "read source audio sample rate"); + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &native_channels), + "read source audio channel count"); + if (native_rate == 0 || + native_rate > static_cast(std::numeric_limits::max()) || + (native_channels != 1 && native_channels != 2)) + throw std::runtime_error( + "reference audio metadata is outside the supported mono/stereo C++ value range"); + const CodecPaddingAllowance padding = codec_padding_allowance(native_audio.Get()); + if (declared_duration && + *declared_duration > + static_cast(policy.maximum_duration_seconds) * kMediaTicksPerSecond + + padding.duration) { + throw std::runtime_error( + "reference audio presentation exceeds the configured duration limit"); + } + const bool trim_codec_padding = padding.frames > 0; + + check_hresult(reader->SetStreamSelection(audio_stream_index, TRUE), "select source audio"); + const auto requested_audio = source_reader_audio_type(native_rate, native_channels); + check_hresult(reader->SetCurrentMediaType(audio_stream_index, nullptr, requested_audio.Get()), + "SetCurrentMediaType(audio float)"); + ComPtr decoded_audio; + check_hresult(reader->GetCurrentMediaType(audio_stream_index, &decoded_audio), + "GetCurrentMediaType(audio)"); + UINT32 decoded_rate = 0; + UINT32 decoded_channels = 0; + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &decoded_rate), + "read decoded audio sample rate"); + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &decoded_channels), + "read decoded audio channel count"); + if (decoded_rate == 0 || + decoded_rate > static_cast(std::numeric_limits::max()) || + (decoded_channels != 1 && decoded_channels != 2)) + throw std::runtime_error( + "decoded audio metadata is outside the supported mono/stereo C++ value range"); + if (decoded_rate != native_rate || decoded_channels != native_channels) { + throw std::runtime_error("Media Foundation did not honor the requested decoded PCM format"); + } + + AudioResult result; + result.sample_rate = static_cast(decoded_rate); + result.channels = static_cast(decoded_channels); + const std::size_t maximum_scalars = + maximum_reference_audio_scalars(decoded_rate, decoded_channels, policy); + LONGLONG last_tick_timestamp = std::numeric_limits::min(); + std::size_t consecutive_empty_reads = 0; + std::size_t total_empty_reads = 0; + detail::ReferenceAudioDecodeState decode_state; + while (true) { + DWORD stream_index = 0; + DWORD flags = 0; + LONGLONG timestamp = 0; + ComPtr sample; + check_hresult( + reader->ReadSample(audio_stream_index, 0, &stream_index, &flags, ×tamp, &sample), + "IMFSourceReader::ReadSample(audio)"); + reject_source_reader_error(flags, "reference audio source reader"); + if ((flags & (MF_SOURCE_READERF_NATIVEMEDIATYPECHANGED | + MF_SOURCE_READERF_CURRENTMEDIATYPECHANGED)) != 0) + throw std::runtime_error("reference audio changes media type mid-stream"); + if ((flags & MF_SOURCE_READERF_ENDOFSTREAM) != 0) + break; + if (!trim_codec_padding) + reject_reference_timestamp(timestamp, "reference audio source reader", policy); + if (stream_index != audio_stream_index) + throw std::runtime_error("reference audio returned an unexpected source-reader stream"); + if ((flags & MF_SOURCE_READERF_STREAMTICK) != 0) { + if (sample) + throw std::runtime_error( + "reference audio returned a sample for an empty stream tick"); + note_empty_source_reader_event(flags, timestamp, last_tick_timestamp, + consecutive_empty_reads, total_empty_reads, + static_cast(padding.duration), + "reference audio source reader", policy); + continue; + } + if (!sample) { + note_empty_source_reader_event(flags, timestamp, last_tick_timestamp, + consecutive_empty_reads, total_empty_reads, + static_cast(padding.duration), + "reference audio source reader", policy); + continue; + } + reject_empty_decoded_sample(sample.Get(), "audio"); + append_float_audio(sample.Get(), decoded_rate, decoded_channels, timestamp, padding.frames, + decode_state, result.samples, maximum_scalars, policy); + consecutive_empty_reads = 0; + } + if (result.samples.empty()) + throw std::runtime_error("reference audio file contains no decoded samples"); + if (result.samples.size() > static_cast(std::numeric_limits::max())) + throw std::runtime_error("decoded audio exceeds the public C++ value range"); + result.num_samples = static_cast(result.samples.size()); + return result; +#else + (void)policy; + if (lowercase_extension(path) == ".wav") + return io::read_wav(path); + throw std::runtime_error( + "compressed native media-file input is available on Windows through Media Foundation"); +#endif +} + +} // namespace trtmc::cli diff --git a/apps/cli/windows_media.h b/apps/cli/windows_media.h new file mode 100644 index 0000000000..23784a3443 --- /dev/null +++ b/apps/cli/windows_media.h @@ -0,0 +1,76 @@ +/* + * 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 +#include +#include + +namespace trtmc::cli { + +namespace detail { +// Overflow-safe decoded-frame allocation ceiling. Presentation timestamps and +// duration remain authoritative for the configured validity decision. +std::uint64_t reference_video_frame_ceiling(const ReferenceMediaDecodePolicy& policy, + std::uint32_t fps_numerator, + std::uint32_t fps_denominator); +std::pair +reference_video_decode_size(const ReferenceMediaDecodePolicy& policy, std::uint32_t source_width, + std::uint32_t source_height); +bool reference_timeline_within_limit(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, std::int64_t duration) noexcept; +bool reference_audio_event_timestamp_within_padding(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, + std::uint64_t maximum_padding_ticks) noexcept; +struct ReferenceAudioDecodeState { + std::uint64_t decoded_frames{0}; + std::uint64_t decoded_padding_frames{0}; +}; +// Account decoded PCM independently of container duration metadata. Compressed +// audio may decode at most maximum_padding_frames beyond the configured sample +// budget, and no more than that many decoded frames may fall outside its timeline. +bool account_reference_audio_decode(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, std::uint64_t frame_count, + std::uint32_t sample_rate, std::uint64_t maximum_padding_frames, + ReferenceAudioDecodeState& state) noexcept; +// Append retained PCM frames at their presentation timestamp. Timeline gaps +// become silence; a later overlapping sample contributes only its new tail. +void append_reference_audio_frames_on_timeline( + std::vector& destination, const void* source, std::uint64_t source_frame_count, + std::uint64_t first_source_frame, std::uint64_t end_source_frame, std::uint32_t channels, + std::int64_t timestamp, std::uint32_t sample_rate, std::size_t maximum_scalars); +// A discovered video soundtrack must fail closed when its native or decoded +// format cannot be represented by the public mono/stereo audio value type. +void validate_reference_video_soundtrack_format(std::uint32_t sample_rate, std::uint32_t channels); +// Reject a decoded Media Foundation sample/buffer that carries no payload. +// media_kind is a trusted internal label such as "audio" or "video". +void require_nonempty_decoded_buffer(std::uint32_t current_length, std::string_view media_kind); +} // namespace detail + +bool is_mp4_path(std::string_view path); + +// Write synchronized H.264 video and optional AAC audio using the Windows +// Media Foundation codecs shipped with the operating system. No FFmpeg or +// other runtime media dependency is involved. +void write_mp4(const VideoResult& result, const std::string& path); + +// Decode a Windows-supported media container into the public THWC RGB video +// and interleaved float-audio value types. Decode fails closed +// as soon as the pipeline's reference-media policy is exceeded. +VideoClipInput read_video_file(const std::string& path, const ReferenceMediaDecodePolicy& policy); + +// Decode the first audio stream in a Windows-supported media file (including +// MP3 and WAV) into interleaved float32 samples. The operating-system Media +// Foundation codecs are the only runtime dependency. Decode fails closed at +// the pipeline's reference-media duration limit. +AudioResult read_audio_file(const std::string& path, const ReferenceMediaDecodePolicy& policy); + +} // namespace trtmc::cli diff --git a/apps/cli/windows_utf8_argv.cpp b/apps/cli/windows_utf8_argv.cpp new file mode 100644 index 0000000000..226cf207bb --- /dev/null +++ b/apps/cli/windows_utf8_argv.cpp @@ -0,0 +1,73 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/windows_utf8_argv.h" + +#if !defined(_WIN32) +#error "windows_utf8_argv.cpp is Windows-only" +#endif + +#define NOMINMAX +#define WIN32_LEAN_AND_MEAN +#include +#include +#include +#include + +namespace trtmc::cli { + +namespace { + +int checked_utf16_size(std::wstring_view value) { + if (value.size() > static_cast(std::numeric_limits::max())) + throw std::length_error("Windows command-line argument is too large to convert to UTF-8"); + return static_cast(value.size()); +} + +[[noreturn]] void throw_conversion_error(DWORD error) { + throw std::runtime_error("WideCharToMultiByte(CP_UTF8) rejected command-line UTF-16 (Windows " + "error " + + std::to_string(error) + ")"); +} + +} // namespace + +std::string utf8_from_utf16(std::wstring_view value) { + if (value.empty()) + return {}; + + const int input_size = checked_utf16_size(value); + const int output_size = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, value.data(), + input_size, nullptr, 0, nullptr, nullptr); + if (output_size <= 0) + throw_conversion_error(GetLastError()); + + std::string result(static_cast(output_size), '\0'); + const int converted = + WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, value.data(), input_size, result.data(), + output_size, nullptr, nullptr); + if (converted != output_size) + throw_conversion_error(GetLastError()); + return result; +} + +Utf8CommandLine::Utf8CommandLine(int argc, wchar_t* const* argv) : argc_(argc) { + if (argc < 0 || (argc > 0 && argv == nullptr)) + throw std::invalid_argument("Windows command line has invalid argc/argv metadata"); + + storage_.reserve(static_cast(argc)); + for (int index = 0; index < argc; ++index) { + if (argv[index] == nullptr) + throw std::invalid_argument("Windows command line contains a null argument"); + storage_.push_back(utf8_from_utf16(argv[index])); + } + + pointers_.reserve(static_cast(argc) + 1U); + for (auto& argument : storage_) + pointers_.push_back(argument.data()); + pointers_.push_back(nullptr); +} + +} // namespace trtmc::cli diff --git a/apps/cli/windows_utf8_argv.h b/apps/cli/windows_utf8_argv.h new file mode 100644 index 0000000000..1c457c73a9 --- /dev/null +++ b/apps/cli/windows_utf8_argv.h @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc::cli { + +// Convert one UTF-16 command-line value to UTF-8. The Windows implementation +// rejects malformed surrogate sequences instead of replacing them. +std::string utf8_from_utf16(std::wstring_view value); + +// Own the UTF-8 storage and the mutable pointer array expected by the existing +// platform-neutral CLI parser. argv()[argc()] is always null. +class Utf8CommandLine { + public: + Utf8CommandLine(int argc, wchar_t* const* argv); + + int argc() const noexcept { return argc_; } + char** argv() noexcept { return pointers_.data(); } + + private: + int argc_{0}; + std::vector storage_; + std::vector pointers_; +}; + +} // namespace trtmc::cli diff --git a/core/builder/tensorrt_model_connect/build.py b/core/builder/tensorrt_model_connect/build.py index d0d5d68638..8884dce612 100644 --- a/core/builder/tensorrt_model_connect/build.py +++ b/core/builder/tensorrt_model_connect/build.py @@ -42,6 +42,7 @@ class BuildRequest: dynamic_kv_cache: bool = False verbose: bool = False graph_transform: GraphTransform | None = None + family_options: tuple[tuple[str, str | int | float | bool | None], ...] = () def __post_init__(self) -> None: if not self.precision: @@ -70,6 +71,19 @@ def __post_init__(self) -> None: raise ValueError("dynamic_kv_cache must be a bool") if self.graph_transform is not None and not callable(self.graph_transform): raise ValueError("graph_transform must be callable when provided") + if not isinstance(self.family_options, tuple): + raise ValueError("family_options must be an immutable tuple") + names: set[str] = set() + for item in self.family_options: + if not isinstance(item, tuple) or len(item) != 2: + raise ValueError("family_options entries must be (name, value) tuples") + name, value = item + _validate_id("family option", name) + if name in names: + raise ValueError(f"duplicate family option: {name}") + if not isinstance(value, (str, int, float, bool, type(None))): + raise ValueError("family option values must be JSON scalar values") + names.add(name) def _validate_id(field: str, value: object) -> str: @@ -136,6 +150,11 @@ def build(request: BuildRequest) -> None: family = _resolve_family(request) _select_backend(request.backend) family_module = _load_family(family) + if request.family_options: + validate_options = getattr(family_module, "validate_build_options", None) + if validate_options is None: + raise ValueError(f"family {family!r} does not accept family_options") + validate_options(dict(request.family_options)) writer = BundleWriter(request.output_path) try: with graph_transform(request.graph_transform): diff --git a/core/builder/tensorrt_model_connect/build_cli.py b/core/builder/tensorrt_model_connect/build_cli.py index 9042a68320..c909bb213d 100644 --- a/core/builder/tensorrt_model_connect/build_cli.py +++ b/core/builder/tensorrt_model_connect/build_cli.py @@ -7,6 +7,7 @@ import argparse import json +import re from pathlib import Path from typing import Sequence @@ -14,6 +15,9 @@ from .model_support import load_model_metadata, resolve_family +_OPTION = re.compile(r"([a-z][a-z0-9_]*)\.([a-z][a-z0-9_]*)=(.*)\Z", re.DOTALL) + + def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="trtmc") commands = parser.add_subparsers(dest="command", required=True) @@ -35,6 +39,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( + "--set", + dest="family_options", + action="append", + default=[], + metavar="FAMILY.KEY=VALUE", + help="Set one family-owned scalar build option", + ) prepare_parser = commands.add_parser( "prepare-structure", help="Prepare one structure request without rebuilding its model bundle", @@ -47,6 +59,33 @@ def _parser() -> argparse.ArgumentParser: return parser +def _parse_family_options( + values: Sequence[str], family: str +) -> tuple[tuple[str, str | int | float | bool | None], ...]: + result: list[tuple[str, str | int | float | bool | None]] = [] + names: set[str] = set() + for value in values: + match = _OPTION.fullmatch(value) + if match is None: + raise ValueError("--set must use FAMILY.KEY=VALUE") + namespace, name, raw = match.groups() + if namespace != family: + raise ValueError( + f"--set namespace {namespace!r} does not match resolved family {family!r}" + ) + if name in names: + raise ValueError(f"duplicate --set option: {family}.{name}") + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + parsed = raw + if not isinstance(parsed, (str, int, float, bool, type(None))): + raise ValueError("--set values must be JSON scalars") + result.append((name, parsed)) + names.add(name) + return tuple(result) + + def main(argv: Sequence[str] | None = None) -> int: args = _parser().parse_args(argv) model_dir = _resolve_model(args.model, args.revision) @@ -93,6 +132,7 @@ def main(argv: Sequence[str] | None = None) -> int: fp32_layers=tuple(args.fp32_layer), dynamic_kv_cache=args.dynamic_kv_cache, verbose=args.verbose, + family_options=_parse_family_options(args.family_options, family), ) ) return 0 diff --git a/core/builder/tensorrt_model_connect/bundle_writer.py b/core/builder/tensorrt_model_connect/bundle_writer.py index fc0944d807..5836009899 100644 --- a/core/builder/tensorrt_model_connect/bundle_writer.py +++ b/core/builder/tensorrt_model_connect/bundle_writer.py @@ -38,6 +38,11 @@ def _validate_nonempty_string(field: str, value: object) -> str: return value +def _file_identity(path: Path) -> tuple[int, int, int, int]: + info = path.stat() + return info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns + + class BundleWriter: """Stage named sections and atomically publish one bundle.""" @@ -50,6 +55,7 @@ def __init__(self, destination: str | Path) -> None: self._header: dict[str, Any] | None = None self._sections: list[tuple[str, Path]] = [] self._section_names: set[str] = set() + self._borrowed_files: dict[str, tuple[int, int, int, int]] = {} self._staging_dir: Path | None = None self._open_sections = 0 self._failed_section = False @@ -89,15 +95,18 @@ def set_header(self, *, family: str, task: str, backend: str) -> None: "backend": _validate_id("backend", backend), } - @contextmanager - def open_section(self, name: str) -> Iterator[BinaryIO]: - """Open one file-backed section for incremental binary writes.""" - + def _validate_section_name(self, name: str) -> str: self._ensure_writable() name = _validate_nonempty_string("section name", name) if name in self._section_names: raise ValueError(f"duplicate bundle section name: {name!r}") + return name + @contextmanager + def open_section(self, name: str) -> Iterator[BinaryIO]: + """Open one file-backed section for incremental binary writes.""" + + name = self._validate_section_name(name) section_path = self._ensure_staging_dir() / f"section-{len(self._sections)}" self._section_names.add(name) self._sections.append((name, section_path)) @@ -111,6 +120,31 @@ def open_section(self, name: str) -> Iterator[BinaryIO]: finally: self._open_sections -= 1 + def add_file(self, name: str, path: str | Path) -> None: + """Borrow an existing file without copying it into section staging. + + The caller must keep the file unchanged until finish or abort. The + writer only reads it; cleanup never removes borrowed source files. + """ + + name = self._validate_section_name(name) + source = Path(path).absolute() + if not source.is_file(): + raise FileNotFoundError(f"bundle section source is not a file: {source}") + resolved = source.resolve() + if resolved == self._destination.resolve() or ( + self._staging_dir is not None and self._staging_dir.resolve() in resolved.parents + ): + raise ValueError("borrowed bundle section must be outside writer-owned paths") + self._borrowed_files[name] = _file_identity(source) + self._section_names.add(name) + self._sections.append((name, source)) + + def _check_borrowed_file(self, name: str, path: Path) -> None: + identity = self._borrowed_files.get(name) + if identity is not None and _file_identity(path) != identity: + raise RuntimeError(f"borrowed bundle section source changed: {name!r}") + def add_bytes(self, name: str, data: bytes) -> None: """Add a complete in-memory binary section.""" @@ -139,6 +173,7 @@ def finish(self) -> None: section_table: dict[str, dict[str, int]] = {} offset = 0 for name, path in self._sections: + self._check_borrowed_file(name, path) length = path.stat().st_size if length > _MAX_UINT64 - offset: raise OverflowError("bundle section table exceeds uint64 range") @@ -162,9 +197,11 @@ def finish(self) -> None: output.write(BUNDLE_MAGIC) output.write(struct.pack(" None: assert request.context_parallel_size == 3 assert request.backend == "trt" assert request.dynamic_kv_cache is False + assert request.family_options == () @pytest.mark.parametrize( @@ -57,6 +58,10 @@ def test_build_request_is_a_plain_frozen_dataclass(tmp_path: Path) -> None: ("dynamic_kv_cache", 1), ("graph_transform", object()), ("backend", "unknown"), + ("family_options", [("option", True)]), + ("family_options", (("bad-name", True),)), + ("family_options", (("option", []),)), + ("family_options", (("option", True), ("option", False))), ], ) def test_build_request_rejects_invalid_direct_inputs( @@ -235,6 +240,22 @@ def transform(network: object, engine_index: int) -> None: assert fake_trt.Builder is FakeTrtBuilder +def test_nonempty_family_options_require_family_validation(monkeypatch, tmp_path: Path) -> None: + request = replace(_request(tmp_path), family_options=(("example_option", True),)) + family = SimpleNamespace(build=lambda _request, _writer: pytest.fail("unexpected build")) + monkeypatch.setattr(build_core, "_load_family", lambda _family: family) + with pytest.raises(ValueError, match="does not accept family_options"): + build_core.build(request) + + def reject_options(options): + assert options == {"example_option": True} + raise ValueError("unsupported example option") + + family.validate_build_options = reject_options + with pytest.raises(ValueError, match="unsupported example option"): + build_core.build(request) + + def test_build_aborts_and_preserves_family_error(monkeypatch, tmp_path: Path) -> None: events: list[str] = [] family_error = RuntimeError("family failed") diff --git a/core/builder/tests/test_build_cli.py b/core/builder/tests/test_build_cli.py index e082ff1ddc..453467c303 100644 --- a/core/builder/tests/test_build_cli.py +++ b/core/builder/tests/test_build_cli.py @@ -56,6 +56,10 @@ def test_build_command_forwards_only_direct_inputs(monkeypatch, tmp_path: Path) "5", "--dynamic-kv-cache", "--verbose", + "--set", + "patchtsmixer.feature=true", + "--set", + "patchtsmixer.path=C:\\models\\weights.bin", ] ) == 0 @@ -78,6 +82,24 @@ def test_build_command_forwards_only_direct_inputs(monkeypatch, tmp_path: Path) assert request.fp32_layers == (2, 5) assert request.dynamic_kv_cache is True assert request.verbose is True + assert request.family_options == ( + ("feature", True), + ("path", "C:\\models\\weights.bin"), + ) + + +@pytest.mark.parametrize( + "values", + [ + ["wrong.option=true"], + ["gpt2.option={}"], + ["gpt2.option=true", "gpt2.option=false"], + ["missing_namespace=true"], + ], +) +def test_family_options_reject_invalid_or_cross_family_values(values) -> None: + with pytest.raises(ValueError): + build_cli._parse_family_options(values, "gpt2") def test_build_command_uses_the_family_owned_default_task(monkeypatch, tmp_path: Path) -> None: diff --git a/core/builder/tests/test_bundle_writer.py b/core/builder/tests/test_bundle_writer.py index b5406b4cec..94f58a354f 100644 --- a/core/builder/tests/test_bundle_writer.py +++ b/core/builder/tests/test_bundle_writer.py @@ -5,6 +5,7 @@ import json import os +import shutil import struct from pathlib import Path @@ -66,9 +67,7 @@ def test_writer_rejects_duplicate_and_empty_section_names(tmp_path: Path) -> Non ("field", "value"), [("family", "../family"), ("task", ""), ("backend", "")], ) -def test_writer_rejects_unsafe_header_ids( - tmp_path: Path, field: str, value: str -) -> None: +def test_writer_rejects_unsafe_header_ids(tmp_path: Path, field: str, value: str) -> None: header = {"family": "family", "task": "text_generation", "backend": "trt"} header[field] = value writer = BundleWriter(tmp_path / "model.bundle") @@ -111,9 +110,7 @@ def test_abort_discards_staging_and_preserves_destination(tmp_path: Path) -> Non writer.finish() -def test_failed_atomic_replace_preserves_destination( - monkeypatch, tmp_path: Path -) -> None: +def test_failed_atomic_replace_preserves_destination(monkeypatch, tmp_path: Path) -> None: destination = tmp_path / "model.bundle" destination.write_bytes(b"previous") writer = BundleWriter(destination) @@ -131,3 +128,124 @@ def fail_replace(_source: Path, _destination: Path) -> None: assert not list(tmp_path.glob(".model.bundle.*.tmp")) writer.abort() + + +def test_borrowed_file_mixes_with_streamed_sections_without_staging_copy(tmp_path: Path) -> None: + source = tmp_path / "existing.plan" + source.write_bytes(b"engine-bytes") + before = source.stat() + destination = tmp_path / "model.bundle" + writer = BundleWriter(destination) + writer.set_header(family="family", task="text_generation", backend="trt") + writer.add_file("engine.plan", source) + assert list(tmp_path.iterdir()) == [source] + writer.add_json("config.json", {"size": 7}) + with writer.open_section("tokenizer.model") as section: + section.write(b"tokens") + writer.finish() + + header, payload = _read_bundle(destination) + assert header["sections"] == { + "engine.plan": {"offset": 0, "length": 12}, + "config.json": {"offset": 12, "length": 10}, + "tokenizer.model": {"offset": 22, "length": 6}, + } + assert payload == b'engine-bytes{"size":7}tokens' + assert source.read_bytes() == b"engine-bytes" + assert source.stat().st_mtime_ns == before.st_mtime_ns + assert set(tmp_path.iterdir()) == {source, destination} + with pytest.raises(RuntimeError, match="already finished"): + writer.add_file("other", source) + + +def test_borrowed_files_reuse_section_guards_and_require_external_files(tmp_path: Path) -> None: + source = tmp_path / "existing.plan" + source.write_bytes(b"plan") + writer = BundleWriter(tmp_path / "model.bundle") + writer.add_file("borrowed", source) + with pytest.raises(ValueError, match="duplicate"): + writer.add_bytes("borrowed", b"other") + writer.add_bytes("staged", b"bytes") + with pytest.raises(ValueError, match="duplicate"): + writer.add_file("staged", source) + with pytest.raises(ValueError, match="section name"): + writer.add_file("", source) + for missing in (tmp_path, tmp_path / "missing.plan"): + with pytest.raises(FileNotFoundError, match="source is not a file"): + writer.add_file("invalid", missing) + with writer.open_section("owned") as section: + section.write(b"owned") + owned_path = Path(section.name) + with pytest.raises(ValueError, match="writer-owned paths"): + writer.add_file("invalid", owned_path) + writer.abort() + assert source.read_bytes() == b"plan" + with pytest.raises(RuntimeError, match="aborted"): + writer.add_file("other", source) + + +def test_borrowing_destination_is_rejected_without_modifying_it(tmp_path: Path) -> None: + destination = tmp_path / "model.bundle" + destination.write_bytes(b"previous") + writer = BundleWriter(destination) + with pytest.raises(ValueError, match="writer-owned paths"): + writer.add_file("previous", destination) + writer.abort() + assert destination.read_bytes() == b"previous" + + +@pytest.mark.parametrize("failure", ("abort", "copy", "replace")) +def test_borrowed_source_and_old_destination_survive_abort_or_finish_failure( + tmp_path: Path, monkeypatch, failure: str +) -> None: + source, destination = tmp_path / "existing.plan", tmp_path / "model.bundle" + source.write_bytes(b"plan") + destination.write_bytes(b"previous") + writer = BundleWriter(destination) + writer.set_header(family="family", task="text_generation", backend="trt") + writer.add_file("engine.plan", source) + writer.add_bytes("metadata", b"metadata") + + def fail(*_args, **_kwargs): + raise OSError("publication failed") + + if failure != "abort": + monkeypatch.setattr( + shutil if failure == "copy" else os, + "copyfileobj" if failure == "copy" else "replace", + fail, + ) + with pytest.raises(OSError, match="publication failed"): + writer.finish() + assert not list(tmp_path.glob(".model.bundle.*.tmp")) + writer.abort() + assert source.read_bytes() == b"plan" + assert destination.read_bytes() == b"previous" + assert set(tmp_path.iterdir()) == {source, destination} + + +@pytest.mark.parametrize("change", ("size", "mtime", "replacement")) +def test_changed_borrowed_source_fails_before_publishing(tmp_path: Path, change: str) -> None: + source, destination = tmp_path / "existing.plan", tmp_path / "model.bundle" + source.write_bytes(b"original") + destination.write_bytes(b"previous") + writer = BundleWriter(destination) + writer.set_header(family="family", task="text_generation", backend="trt") + writer.add_file("engine.plan", source) + before = source.stat() + if change == "size": + source.write_bytes(b"longer source content") + elif change == "mtime": + source.write_bytes(b"modified") + os.utime(source, ns=(before.st_atime_ns, before.st_mtime_ns + 1_000_000_000)) + else: + replacement = tmp_path / "replacement.plan" + replacement.write_bytes(b"modified") + os.utime(replacement, ns=(before.st_atime_ns, before.st_mtime_ns)) + os.replace(replacement, source) + with pytest.raises(RuntimeError, match="borrowed bundle section source changed"): + writer.finish() + writer.abort() + assert source.is_file() + assert destination.read_bytes() == b"previous" + assert set(tmp_path.iterdir()) == {source, destination} diff --git a/core/runtime/bundle/bundle_format.cpp b/core/runtime/bundle/bundle_format.cpp index eb11f20674..f6c48301fe 100644 --- a/core/runtime/bundle/bundle_format.cpp +++ b/core/runtime/bundle/bundle_format.cpp @@ -190,12 +190,18 @@ const BundleSectionInfo* BundleReader::find_section(std::string_view name) const return nullptr; } +std::uint64_t BundleReader::section_file_offset(std::string_view name) const { + const auto* section = find_section(name); + if (section == nullptr) + throw std::runtime_error("Bundle section not found: " + std::string(name)); + return checked_section_file_offset(*section, data_offset_, file_size_, path_); +} + std::vector BundleReader::read_section(std::string_view name) const { const auto* section = find_section(name); if (section == nullptr) throw std::runtime_error("Bundle section not found: " + std::string(name)); - const std::uint64_t file_offset = - checked_section_file_offset(*section, data_offset_, file_size_, path_); + const std::uint64_t file_offset = section_file_offset(name); if (file_offset > static_cast(std::numeric_limits::max())) throw std::runtime_error("Bundle section '" + section->name + "' has an unsupported file offset: " + path_); diff --git a/core/runtime/include/trtmc/bundle.h b/core/runtime/include/trtmc/bundle.h index 154d697ef1..eedcf2c75c 100644 --- a/core/runtime/include/trtmc/bundle.h +++ b/core/runtime/include/trtmc/bundle.h @@ -42,6 +42,10 @@ class BundleReader { const std::string& path() const noexcept { return path_; } const BundleInfo& info() const noexcept { return info_; } const BundleSectionInfo* find_section(std::string_view name) const noexcept; + // Return the validated absolute byte offset of a section in the bundle. + // Large TensorRT plans can be stream-deserialized from this range without + // first duplicating the entire section in host memory. + std::uint64_t section_file_offset(std::string_view name) const; std::vector read_section(std::string_view name) const; private: diff --git a/core/runtime/include/trtmc/runtime/family_factory.h b/core/runtime/include/trtmc/runtime/family_factory.h index c0927b50b0..0a728dd008 100644 --- a/core/runtime/include/trtmc/runtime/family_factory.h +++ b/core/runtime/include/trtmc/runtime/family_factory.h @@ -9,6 +9,7 @@ #include "trtmc/task.h" #include +#include namespace trtmc { @@ -18,6 +19,8 @@ struct FamilyContext { const BundleReader& reader; IBackend& backend; std::uint64_t kv_cache_size_bytes{0}; + std::string runtime_cache_path; + bool cuda_graphs{false}; }; using CreateFamilyFn = ITask* (*)(const FamilyContext& context); diff --git a/core/runtime/include/trtmc/runtime/trt_backend.h b/core/runtime/include/trtmc/runtime/trt_backend.h index ace743553e..d852e3b6c3 100644 --- a/core/runtime/include/trtmc/runtime/trt_backend.h +++ b/core/runtime/include/trtmc/runtime/trt_backend.h @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -23,6 +24,7 @@ struct ModuleCreateOptions { std::shared_ptr distributed_owner; // keeps communicator alive const char* runtime_cache_path{""}; // TensorRT-RTX JIT cache, optional bool cuda_graphs{false}; // TensorRT-RTX whole-graph capture + std::int32_t optimization_profile{0}; // context profile selected by the family }; struct ModuleExternalBinding { @@ -59,6 +61,37 @@ class IBackend { create_dual_profile_modules(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options) = 0; + // Optional TensorRT-RTX capability for multi-GiB bundle sections. The + // backend validates the exact file range and lets TensorRT stream it + // directly instead of materializing a duplicate host buffer. + virtual std::unique_ptr + create_module_from_file(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) { + (void)plan_path; + (void)plan_offset; + (void)plan_size; + (void)options; + (void)external_bindings; + (void)weight_streaming_budget_bytes; + (void)retain_engine; + (void)serial_execution_context; + throw std::runtime_error("backend does not support file-backed TensorRT plans"); + } + + // Optional explicit lifetime for one process-shared runtime cache. A + // family releases its lease after all lazy modules have been destroyed. + virtual std::uint64_t acquire_runtime_cache_lease(const char* path) { + (void)path; + throw std::runtime_error("backend does not support runtime-cache leases"); + } + virtual void release_runtime_cache_lease(std::uint64_t lease) { + (void)lease; + throw std::runtime_error("backend does not support runtime-cache leases"); + } + // Backend identity written into the bundle header (currently "trt"). virtual const char* name() const = 0; }; diff --git a/core/runtime/include/trtmc/task.h b/core/runtime/include/trtmc/task.h index 087c9f02c5..e4e1bd1972 100644 --- a/core/runtime/include/trtmc/task.h +++ b/core/runtime/include/trtmc/task.h @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -107,6 +108,15 @@ struct AudioResult { std::vector samples; std::int32_t num_samples{0}; std::int32_t sample_rate{24000}; + // Interleaved channel count. Existing mono producers keep the default. + std::int32_t channels{1}; +}; + +// Synchronized decoded video and audio produced by an audiovisual model. +struct VideoResult { + ImageResult frames; + AudioResult audio; + std::int32_t fps{0}; }; struct TranscriptionStreamConfig { @@ -379,6 +389,67 @@ struct ImageGenerationConfig { std::string negative_prompt; std::int32_t height{0}; std::int32_t width{0}; + // Requested video frame count. Zero selects the family/bundle default. + std::int32_t video_num_frames{0}; +}; + +// Decoded host-resident inputs for native video conditioning. Image pixels are +// contiguous HWC float32 in [0, 1]; video pixels are contiguous THWC. +struct VideoImageInput { + std::vector pixels; + std::int32_t height{0}; + std::int32_t width{0}; + std::int32_t channels{3}; +}; + +struct VideoClipInput { + std::vector pixels; + std::int32_t num_frames{0}; + std::int32_t height{0}; + std::int32_t width{0}; + std::int32_t channels{3}; + std::int32_t fps_numerator{0}; + std::int32_t fps_denominator{1}; + AudioResult soundtrack; +}; + +struct ReferenceMediaDecodePolicy { + std::uint32_t maximum_duration_seconds{0}; + std::uint32_t target_video_fps{0}; + std::uint32_t maximum_source_video_fps{0}; + std::uint32_t canvas_short_edge{0}; + std::uint64_t canvas_max_pixels{0}; + std::uint32_t canvas_multiple{0}; + double minimum_aspect_ratio{0.0}; + double maximum_aspect_ratio{0.0}; +}; + +enum class VideoReferenceKind { + kImage, + kVideo, + kAudio, +}; + +struct VideoReferenceInput { + VideoReferenceKind kind{VideoReferenceKind::kImage}; + VideoImageInput image; + VideoClipInput video; + AudioResult audio; +}; + +enum class VideoGenerationMode { + kTextToVideoAudio, + kFirstLastFrameToVideoAudio, + kReferenceToVideoAudio, +}; + +struct VideoGenerationRequest { + std::string prompt; + ImageGenerationConfig config; + VideoGenerationMode mode{VideoGenerationMode::kTextToVideoAudio}; + std::optional first_frame; + std::optional last_frame; + std::vector references; }; struct WorldModelRequest { @@ -524,6 +595,19 @@ class IImageGeneration : public virtual ITask { const ImageGenerationConfig& config = {}) = 0; }; +// Optional capability for image-generation families that return synchronized +// video and audio or accept structured video-conditioning inputs. +class IVideoGeneration { + public: + virtual ~IVideoGeneration() = default; + virtual VideoResult generate_video(const std::string& prompt, + const ImageGenerationConfig& config = {}) = 0; + virtual VideoResult generate_video(const VideoGenerationRequest& request) = 0; + virtual std::optional reference_media_decode_policy() const { + return std::nullopt; + } +}; + class IImageEditing : public virtual ITask { public: static constexpr const char* kTask = "image_edit"; diff --git a/core/runtime/loader/family_loader.cpp b/core/runtime/loader/family_loader.cpp index 76c419e633..25daee3c64 100644 --- a/core/runtime/loader/family_loader.cpp +++ b/core/runtime/loader/family_loader.cpp @@ -6,10 +6,10 @@ #include "trtmc/runtime/family_loader.h" #include "runtime/bundle/bundle_format.h" +#include "runtime/platform/dynamic_library.h" #include "trtmc/runtime/family_factory.h" #include "trtmc/runtime/trt_backend.h" -#include #include #include #include @@ -62,12 +62,11 @@ fs::path explicit_runtime_root(const std::string& runtime_root) { class SharedLibrary { public: explicit SharedLibrary(const fs::path& path) : path_(path.string()) { - dlerror(); - handle_ = dlopen(path_.c_str(), RTLD_NOW | RTLD_LOCAL); + std::string error; + handle_ = + internal::open_dynamic_library(path, internal::DynamicLibraryVisibility::local, &error); if (handle_ == nullptr) { - const char* error = dlerror(); - throw std::runtime_error("Unable to load '" + path_ + - "': " + (error != nullptr ? error : "unknown dlopen error")); + throw std::runtime_error("Unable to load '" + path_ + "': " + error); } } @@ -76,29 +75,29 @@ class SharedLibrary { ~SharedLibrary() { if (handle_ != nullptr) - dlclose(handle_); + (void)internal::close_dynamic_library(handle_); } void* require_symbol(const char* name) const { - dlerror(); - void* symbol = dlsym(handle_, name); - const char* error = dlerror(); - if (error != nullptr || symbol == nullptr) { + std::string error; + void* symbol = internal::dynamic_library_symbol(handle_, name, &error); + if (symbol == nullptr) { throw std::runtime_error("Library '" + path_ + "' is missing required symbol '" + name + - "'"); + "': " + error); } return symbol; } private: std::string path_; - void* handle_{nullptr}; + internal::DynamicLibraryHandle handle_{nullptr}; }; class BackendLibrary { public: BackendLibrary(const fs::path& runtime_root, const std::string& backend_id) - : library_(runtime_root / ("libtrtmc_backend_" + backend_id + ".so")) { + : library_(runtime_root / + internal::dynamic_library_filename("trtmc_backend_" + backend_id)) { const auto create = reinterpret_cast(library_.require_symbol("trtmc_create_backend")); destroy_ = @@ -136,7 +135,7 @@ class BackendLibrary { class FamilyLibrary { public: FamilyLibrary(const fs::path& runtime_root, const std::string& family_id) - : library_(runtime_root / ("libtrtmc_model_" + family_id + ".so")), + : library_(runtime_root / internal::dynamic_library_filename("trtmc_model_" + family_id)), create_(reinterpret_cast(library_.require_symbol(kCreateFamilySymbol))) {} FamilyLibrary(const FamilyLibrary&) = delete; @@ -175,6 +174,25 @@ class RuntimeOptionsBackend final : public IBackend { with_runtime_options(options)); } + std::unique_ptr + create_module_from_file(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) override { + return backend_.create_module_from_file( + plan_path, plan_offset, plan_size, with_runtime_options(options), bindings, + weight_streaming_budget_bytes, retain_engine, serial_execution_context); + } + + std::uint64_t acquire_runtime_cache_lease(const char* path) override { + return backend_.acquire_runtime_cache_lease(path); + } + + void release_runtime_cache_lease(std::uint64_t lease) override { + backend_.release_runtime_cache_lease(lease); + } + const char* name() const override { return backend_.name(); } private: @@ -228,7 +246,8 @@ RuntimeLibraryCache& runtime_library_cache() { } IBackend& cached_backend(const fs::path& runtime_root, const std::string& backend_id) { - const std::string path = (runtime_root / ("libtrtmc_backend_" + backend_id + ".so")).string(); + const std::string path = + (runtime_root / internal::dynamic_library_filename("trtmc_backend_" + backend_id)).string(); auto& cache = runtime_library_cache(); std::lock_guard lock(cache.mutex); const auto found = cache.backends.find(path); @@ -258,7 +277,8 @@ IBackend& cached_configured_backend(IBackend& backend, const std::string& runtim } FamilyLibrary& cached_family(const fs::path& runtime_root, const std::string& family_id) { - const std::string path = (runtime_root / ("libtrtmc_model_" + family_id + ".so")).string(); + const std::string path = + (runtime_root / internal::dynamic_library_filename("trtmc_model_" + family_id)).string(); auto& cache = runtime_library_cache(); std::lock_guard lock(cache.mutex); const auto found = cache.families.find(path); @@ -301,7 +321,8 @@ std::unique_ptr load_task(const std::string& bundle_path, const std::stri FamilyLibrary& family = cached_family(root, info.family); IBackend& configured_backend = cached_configured_backend(backend, runtime_cache_path, cuda_graphs); - FamilyContext context{reader, configured_backend, kv_cache_size_bytes}; + FamilyContext context{reader, configured_backend, kv_cache_size_bytes, runtime_cache_path, + cuda_graphs}; std::unique_ptr task(family.create(context)); if (task == nullptr) throw std::runtime_error("trtmc_create_family returned nullptr"); diff --git a/core/runtime/platform/dynamic_library.cpp b/core/runtime/platform/dynamic_library.cpp new file mode 100644 index 0000000000..1b0b7df552 --- /dev/null +++ b/core/runtime/platform/dynamic_library.cpp @@ -0,0 +1,184 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/platform/dynamic_library.h" + +#include +#include + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + +namespace trtmc::internal { +namespace { + +namespace fs = std::filesystem; + +void assign_error(std::string* output, std::string message) noexcept { + if (output == nullptr) + return; + try { + *output = std::move(message); + } catch (...) { + try { + output->clear(); + } catch (...) { + } + } +} + +#if defined(_WIN32) + +std::string utf8_from_wide(const wchar_t* value, int length) { + if (value == nullptr || length <= 0) + return {}; + const int required = + WideCharToMultiByte(CP_UTF8, 0, value, length, nullptr, 0, nullptr, nullptr); + if (required <= 0) + return {}; + std::string result(static_cast(required), '\0'); + if (WideCharToMultiByte(CP_UTF8, 0, value, length, result.data(), required, nullptr, nullptr) <= + 0) { + return {}; + } + return result; +} + +std::string windows_error_message(DWORD code) { + wchar_t* buffer = nullptr; + const DWORD length = FormatMessageW( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, + nullptr, code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + reinterpret_cast(&buffer), 0, nullptr); + std::string message = utf8_from_wide(buffer, static_cast(length)); + if (buffer != nullptr) + LocalFree(buffer); + while (!message.empty() && + (message.back() == '\r' || message.back() == '\n' || message.back() == ' ')) { + message.pop_back(); + } + return message.empty() ? "Windows error " + std::to_string(code) : message; +} + +#else + +std::string loader_error() { + const char* error = dlerror(); + return error == nullptr ? std::string("unknown dynamic-loader error") : std::string(error); +} + +#endif + +} // namespace + +DynamicLibraryHandle open_dynamic_library(const fs::path& path, DynamicLibraryVisibility visibility, + std::string* error) { + if (error != nullptr) + error->clear(); +#if defined(_WIN32) + (void)visibility; + try { + fs::path load_path = path; + DWORD flags = LOAD_LIBRARY_SEARCH_DEFAULT_DIRS; + if (path.is_absolute() || path.has_parent_path()) { + std::error_code ec; + const fs::path absolute = fs::absolute(path, ec); + if (!ec) + load_path = absolute; + flags |= LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR; + } + HMODULE handle = LoadLibraryExW(load_path.c_str(), nullptr, flags); + if (handle == nullptr) { + assign_error(error, windows_error_message(GetLastError())); + return nullptr; + } + return reinterpret_cast(handle); + } catch (const std::exception& exception) { + assign_error(error, exception.what()); + return nullptr; + } +#else + dlerror(); + const int flags = + RTLD_NOW | (visibility == DynamicLibraryVisibility::global ? RTLD_GLOBAL : RTLD_LOCAL); + DynamicLibraryHandle handle = dlopen(path.c_str(), flags); + if (handle == nullptr) + assign_error(error, loader_error()); + return handle; +#endif +} + +void* dynamic_library_symbol(DynamicLibraryHandle handle, const char* name, std::string* error) { + if (error != nullptr) + error->clear(); + if (handle == nullptr || name == nullptr) { + assign_error(error, "invalid dynamic-library handle or symbol name"); + return nullptr; + } +#if defined(_WIN32) + FARPROC symbol = GetProcAddress(reinterpret_cast(handle), name); + if (symbol == nullptr) { + assign_error(error, windows_error_message(GetLastError())); + return nullptr; + } + return reinterpret_cast(symbol); +#else + dlerror(); + void* symbol = dlsym(handle, name); + const char* message = dlerror(); + if (message != nullptr) { + assign_error(error, message); + return nullptr; + } + return symbol; +#endif +} + +bool close_dynamic_library(DynamicLibraryHandle handle, std::string* error) noexcept { + if (error != nullptr) + error->clear(); + if (handle == nullptr) + return true; +#if defined(_WIN32) + if (FreeLibrary(reinterpret_cast(handle)) != 0) + return true; + try { + assign_error(error, windows_error_message(GetLastError())); + } catch (...) { + assign_error(error, "unknown Windows dynamic-loader error"); + } + return false; +#else + if (dlclose(handle) == 0) + return true; + try { + assign_error(error, loader_error()); + } catch (...) { + assign_error(error, "unknown dynamic-loader error"); + } + return false; +#endif +} + +std::string dynamic_library_filename(std::string_view stem) { +#if defined(_WIN32) + return std::string(stem) + ".dll"; +#elif defined(__APPLE__) + return "lib" + std::string(stem) + ".dylib"; +#else + return "lib" + std::string(stem) + ".so"; +#endif +} + +} // namespace trtmc::internal diff --git a/core/runtime/platform/dynamic_library.h b/core/runtime/platform/dynamic_library.h new file mode 100644 index 0000000000..53c8455a00 --- /dev/null +++ b/core/runtime/platform/dynamic_library.h @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc::internal { + +using DynamicLibraryHandle = void*; + +enum class DynamicLibraryVisibility { + local, + global, +}; + +DynamicLibraryHandle +open_dynamic_library(const std::filesystem::path& path, + DynamicLibraryVisibility visibility = DynamicLibraryVisibility::local, + std::string* error = nullptr); +void* dynamic_library_symbol(DynamicLibraryHandle handle, const char* name, + std::string* error = nullptr); +bool close_dynamic_library(DynamicLibraryHandle handle, std::string* error = nullptr) noexcept; + +std::string dynamic_library_filename(std::string_view stem); + +} // namespace trtmc::internal diff --git a/core/runtime/tensorrt/rtx_backend.cpp b/core/runtime/tensorrt/rtx_backend.cpp index 5bb004cd14..0e312fa5c8 100644 --- a/core/runtime/tensorrt/rtx_backend.cpp +++ b/core/runtime/tensorrt/rtx_backend.cpp @@ -4,122 +4,824 @@ */ // RtxBackend: IBackend implementation for TensorRT-RTX. -// Compiled into libtrtmc_backend_trt_rtx.so. Links libtensorrt_rtx.so. +// Compiled into libtrtmc_backend_rtx.so. Links libtensorrt_rtx.so. +// +// Uses the RTX-specific NvInfer.h headers which declare IRuntimeCache, +// CudaGraphStrategy, and DynamicShapesKernelSpecializationStrategy. #include "runtime/primitives/cuda_common.h" +#include "runtime/tensorrt/runtime_cache_persistence.h" #include "runtime/tensorrt/trt_logger.h" #include "trt_module_impl.h" #include "trtmc/runtime/trt_backend.h" #include +#include +#include +#include +#include +#include #include +#include #include +#include #include #include +#include #include #include +#include +#include #include #include +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + namespace trtmc { namespace { +namespace fs = std::filesystem; + struct StreamSetup { cudaStream_t stream{nullptr}; std::shared_ptr owner; }; +class AccumulateWallTime { + public: + explicit AccumulateWallTime(double* milliseconds) + : milliseconds_(milliseconds), + start_(milliseconds ? Clock::now() : Clock::time_point{}) {} + ~AccumulateWallTime() { + if (milliseconds_) + *milliseconds_ += std::chrono::duration(Clock::now() - start_).count(); + } + + private: + using Clock = std::chrono::steady_clock; + double* milliseconds_; + Clock::time_point start_; +}; + +struct FilePlanLoadTiming { + std::chrono::steady_clock::time_point start{std::chrono::steady_clock::now()}; + std::uint64_t read_calls{0}, host_read_bytes{0}, device_read_bytes{0}; + double reader_open_ms{0}, retained_lock_ms{0}, deserialize_ms{0}, read_callback_ms{0}; + double file_read_ms{0}, upload_ms{0}, weight_budget_ms{0}, runtime_config_ms{0}; + double runtime_cache_ms{0}, context_ms{0}, module_init_ms{0}; + + void log(std::uint64_t plan_bytes, std::int32_t profile, bool retained_hit) const { + const double total_ms = + std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); + // Deserialize contains reader callbacks; callbacks contain file reads and uploads. + // Runtime-config contains runtime-cache time. These nested fields are not additive. + std::ostringstream line; + line << "[trtmc.rtx_plan_load] plan_bytes=" << plan_bytes << " profile=" << profile + << " retained_hit=" << retained_hit << " total_ms=" << total_ms + << " reader_open_ms=" << reader_open_ms << " retained_lock_ms=" << retained_lock_ms + << " deserialize_ms=" << deserialize_ms << " read_calls=" << read_calls + << " host_read_bytes=" << host_read_bytes << " device_read_bytes=" << device_read_bytes + << " read_callback_ms=" << read_callback_ms << " file_read_ms=" << file_read_ms + << " upload_ms=" << upload_ms << " weight_budget_ms=" << weight_budget_ms + << " runtime_config_ms=" << runtime_config_ms << " runtime_cache_ms=" << runtime_cache_ms + << " context_ms=" << context_ms << " module_init_ms=" << module_init_ms << '\n'; + std::cerr << line.str(); + } +}; + +#if defined(_WIN32) +// Some Windows CUDA configurations report memory-pool support while rejecting +// the asynchronous allocation path used by TensorRT-RTX. IGpuAsyncAllocator +// permits a synchronizing implementation, so retain the callback surface while +// backing it with cudaMalloc/cudaFree on Windows only. +class SynchronousGpuAllocator final : public nvinfer1::IGpuAsyncAllocator { + public: + void* allocateAsync(std::uint64_t size, std::uint64_t /*alignment*/, + nvinfer1::AllocatorFlags /*flags*/, + cudaStream_t /*stream*/) noexcept override { + if (size == 0 || size > static_cast(std::numeric_limits::max())) + return nullptr; + void* memory = nullptr; + return cudaMalloc(&memory, static_cast(size)) == cudaSuccess ? memory + : nullptr; + } + + bool deallocateAsync(void* memory, cudaStream_t /*stream*/) noexcept override { + return memory == nullptr || cudaFree(memory) == cudaSuccess; + } +}; +#endif + StreamSetup resolve_stream(cudaStream_t requested_stream) { - if (requested_stream) - return {requested_stream, {}}; + if (requested_stream) { + return StreamSetup{requested_stream, {}}; + } auto owned = std::make_shared(); - if (!owned->ok()) + if (!owned->ok()) { throw std::runtime_error("[trtmc] Failed to create CUDA stream"); - return {owned->get(), owned}; + } + + return StreamSetup{owned->get(), owned}; +} + +void append_plan_cache_identity_field(std::string& identity, const std::string& value) { + identity += std::to_string(value.size()); + identity.push_back(':'); + identity += value; +} + +class PlanFileMutationGuard final { + public: + explicit PlanFileMutationGuard(const char* path) { + if (path == nullptr || path[0] == '\0') + throw std::invalid_argument("[trtmc] RTX plan file path must not be empty"); + try { + path_ = fs::canonical(fs::path(path)); + } catch (const fs::filesystem_error& error) { + throw std::runtime_error(std::string("[trtmc] Failed to resolve RTX plan file: ") + + error.what()); + } +#if defined(_WIN32) + // Deny write/delete sharing while the exact file is streamed and deserialized. + handle_ = CreateFileW(path_.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OPEN_REPARSE_POINT, nullptr); + if (handle_ == INVALID_HANDLE_VALUE) { + const std::error_code error(static_cast(GetLastError()), std::system_category()); + throw std::runtime_error("[trtmc] Failed to lock RTX plan file: " + error.message()); + } + BY_HANDLE_FILE_INFORMATION info{}; + if (!GetFileInformationByHandle(handle_, &info)) { + const std::error_code error(static_cast(GetLastError()), std::system_category()); + (void)CloseHandle(handle_); + handle_ = INVALID_HANDLE_VALUE; + throw std::runtime_error("[trtmc] Failed to identify RTX plan file: " + + error.message()); + } + const auto path_utf8 = path_.u8string(); + append_plan_cache_identity_field(file_identity_, path_utf8); + append_plan_cache_identity_field( + file_identity_, std::to_string(static_cast(info.dwVolumeSerialNumber))); + append_plan_cache_identity_field( + file_identity_, + std::to_string((static_cast(info.nFileIndexHigh) << 32U) | + static_cast(info.nFileIndexLow))); + append_plan_cache_identity_field( + file_identity_, std::to_string((static_cast(info.nFileSizeHigh) << 32U) | + static_cast(info.nFileSizeLow))); + append_plan_cache_identity_field( + file_identity_, + std::to_string( + (static_cast(info.ftLastWriteTime.dwHighDateTime) << 32U) | + static_cast(info.ftLastWriteTime.dwLowDateTime))); +#else + struct stat info{}; + if (::stat(path_.c_str(), &info) != 0) { + const std::error_code error(errno, std::generic_category()); + throw std::runtime_error("[trtmc] Failed to identify RTX plan file: " + + error.message()); + } + initial_device_ = static_cast(info.st_dev); + initial_inode_ = static_cast(info.st_ino); + initial_size_ = fs::file_size(path_); + initial_write_time_ = fs::last_write_time(path_); + const auto path_utf8 = path_.u8string(); + append_plan_cache_identity_field(file_identity_, path_utf8); + append_plan_cache_identity_field(file_identity_, std::to_string(initial_device_)); + append_plan_cache_identity_field(file_identity_, std::to_string(initial_inode_)); + append_plan_cache_identity_field(file_identity_, std::to_string(initial_size_)); + append_plan_cache_identity_field( + file_identity_, std::to_string(initial_write_time_.time_since_epoch().count())); +#endif + } + + ~PlanFileMutationGuard() { +#if defined(_WIN32) + if (handle_ != INVALID_HANDLE_VALUE) + (void)CloseHandle(handle_); +#endif + } + + PlanFileMutationGuard(const PlanFileMutationGuard&) = delete; + PlanFileMutationGuard& operator=(const PlanFileMutationGuard&) = delete; + + const fs::path& path() const noexcept { return path_; } + const std::string& cache_identity() const noexcept { return file_identity_; } +#if defined(_WIN32) + HANDLE handle() const noexcept { return handle_; } +#endif + + void verify_unchanged() const { +#if !defined(_WIN32) + struct stat info{}; + if (::stat(path_.c_str(), &info) != 0) { + throw std::runtime_error("[trtmc] RTX plan file disappeared while it was parsed"); + } + if (static_cast(info.st_dev) != initial_device_ || + static_cast(info.st_ino) != initial_inode_ || + fs::file_size(path_) != initial_size_ || + fs::last_write_time(path_) != initial_write_time_) { + throw std::runtime_error("[trtmc] RTX plan file changed while it was parsed"); + } +#endif + } + + private: + fs::path path_; + std::string file_identity_; +#if defined(_WIN32) + HANDLE handle_{INVALID_HANDLE_VALUE}; +#else + std::uintmax_t initial_device_{0}; + std::uintmax_t initial_inode_{0}; + std::uintmax_t initial_size_{0}; + fs::file_time_type initial_write_time_{}; +#endif +}; + +void validate_plan_description(std::uint64_t size) { + if (size == 0 || size > static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("[trtmc] RTX plan file range is invalid"); + } +} + +void validate_plan_file_range(std::uint64_t offset, std::uint64_t size, std::uint64_t file_size) { + const auto max_offset = static_cast(std::numeric_limits::max()); + if (offset > file_size || size > file_size - offset || offset > max_offset || + size > max_offset - offset) { + throw std::invalid_argument("[trtmc] RTX plan section is outside its file"); + } +} + +enum class ReaderDestination { + host, + device, + invalid, +}; + +ReaderDestination classify_reader_destination(void* destination) noexcept { + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, destination); + if (status == cudaErrorInvalidValue) { + (void)cudaGetLastError(); + return ReaderDestination::host; + } + if (status != cudaSuccess) + return ReaderDestination::invalid; + if (attributes.type == cudaMemoryTypeUnregistered || attributes.type == cudaMemoryTypeHost) + return ReaderDestination::host; + if (attributes.type == cudaMemoryTypeDevice || attributes.type == cudaMemoryTypeManaged) + return ReaderDestination::device; + return ReaderDestination::invalid; +} + +#if defined(_WIN32) +using PlanReadFile = HANDLE; +#else +using PlanReadFile = std::ifstream&; +#endif + +std::uint64_t read_plan_bytes(PlanReadFile file, void* destination, std::uint64_t requested, + FilePlanLoadTiming& timing) { + AccumulateWallTime elapsed(&timing.file_read_ms); +#if defined(_WIN32) + // Avoid MSVC filebuf's small fread loop for multi-GiB binary plans. + std::uint64_t completed = 0; + while (completed < requested) { + const auto chunk = static_cast( + std::min(requested - completed, 64U << 20)); + DWORD received = 0; + if (!ReadFile(file, static_cast(destination) + completed, chunk, + &received, nullptr)) { + throw std::runtime_error("[trtmc] Failed to read RTX plan bytes"); + } + completed += received; + if (received != chunk) + break; + } + return completed; +#else + file.read(static_cast(destination), static_cast(requested)); + const auto received = file.gcount(); + return received > 0 ? static_cast(received) : 0; +#endif +} + +void read_host_bytes(PlanReadFile file, void* destination, std::uint64_t requested, + FilePlanLoadTiming& timing) { + const auto received = read_plan_bytes(file, destination, requested, timing); + timing.host_read_bytes += received; + if (received != requested) + throw std::runtime_error("[trtmc] Failed to read RTX plan bytes into host memory"); +} + +void read_device_bytes(PlanReadFile file, std::vector& staging, void* destination, + std::uint64_t requested, FilePlanLoadTiming& timing) { + std::uint64_t copied = 0; + while (copied < requested) { + const auto chunk = std::min(staging.size(), requested - copied); + const auto received = read_plan_bytes(file, staging.data(), chunk, timing); + timing.device_read_bytes += received; + if (received != chunk) + throw std::runtime_error("[trtmc] Failed to stage RTX plan bytes"); + { + AccumulateWallTime elapsed(&timing.upload_ms); + if (cudaMemcpy(static_cast(destination) + copied, staging.data(), + static_cast(chunk), cudaMemcpyHostToDevice) != cudaSuccess) { + throw std::runtime_error("[trtmc] Failed to copy RTX plan bytes to the device"); + } + } + copied += chunk; + } } -std::shared_ptr deserialize_engine(nvinfer1::IRuntime& runtime, - const void* plan_data, size_t plan_size) { - auto* engine = runtime.deserializeCudaEngine(plan_data, plan_size); - if (!engine) - throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); - return {engine, [](nvinfer1::ICudaEngine* value) { delete value; }}; +void read_destination_bytes(PlanReadFile file, std::vector& staging, void* destination, + std::uint64_t requested, FilePlanLoadTiming& timing) { + switch (classify_reader_destination(destination)) { + case ReaderDestination::host: + read_host_bytes(file, destination, requested, timing); + return; + case ReaderDestination::device: + read_device_bytes(file, staging, destination, requested, timing); + return; + case ReaderDestination::invalid: + throw std::runtime_error("[trtmc] Unsupported RTX plan reader destination"); + } } -class RuntimeCacheState { +// TensorRT-RTX may request plan bytes into host or device memory and may seek +// within the stream. Restrict every operation to the validated bundle section. +class BoundedPlanStreamReader final : public nvinfer1::IStreamReaderV2 { public: - RuntimeCacheState(nvinfer1::IRuntimeConfig& config, std::string path) - : cache_(config.createRuntimeCache()), path_(std::move(path)) { - if (cache_ == nullptr) - throw std::runtime_error("[trtmc] Failed to create TensorRT-RTX runtime cache"); + BoundedPlanStreamReader(const char* path, std::uint64_t offset, std::uint64_t size, + FilePlanLoadTiming& timing) + : mutation_guard_(path), offset_(offset), size_(size), staging_(4U << 20), timing_(timing) { + validate_plan_description(size_); +#if defined(_WIN32) + LARGE_INTEGER end{}; + if (!GetFileSizeEx(mutation_guard_.handle(), &end) || end.QuadPart < 0) + throw std::runtime_error("[trtmc] Failed to determine RTX plan file size"); + const auto file_size = static_cast(end.QuadPart); +#else + file_.open(mutation_guard_.path(), std::ios::binary | std::ios::ate); + if (!file_) + throw std::runtime_error("[trtmc] Failed to open RTX plan file"); + const auto end = file_.tellg(); + if (end < 0) + throw std::runtime_error("[trtmc] Failed to determine RTX plan file size"); + const auto file_size = static_cast(static_cast(end)); +#endif + validate_plan_file_range(offset_, size_, file_size); + } + + void verify_unchanged() const { mutation_guard_.verify_unchanged(); } + + std::string cache_identity() const { + std::string identity = mutation_guard_.cache_identity(); + append_plan_cache_identity_field(identity, std::to_string(offset_)); + append_plan_cache_identity_field(identity, std::to_string(size_)); + return identity; + } + + int64_t read(void* destination, int64_t nb_bytes, cudaStream_t stream) noexcept override { + (void)stream; + AccumulateWallTime elapsed(&timing_.read_callback_ms); + ++timing_.read_calls; + if (destination == nullptr || nb_bytes < 0) + return -1; + if (nb_bytes == 0) + return 0; + try { + return read_from_current_position(destination, nb_bytes); + } catch (...) { + return -1; + } + } + + bool seek(int64_t offset, nvinfer1::SeekPosition where) noexcept override { + std::uint64_t base = 0; + switch (where) { + case nvinfer1::SeekPosition::kSET: + break; + case nvinfer1::SeekPosition::kCUR: + base = cursor_; + break; + case nvinfer1::SeekPosition::kEND: + base = size_; + break; + default: + return false; + } + if (offset >= 0) { + const auto delta = static_cast(offset); + if (base > size_ || delta > size_ - base) + return false; + cursor_ = base + delta; + return true; + } + const auto magnitude = static_cast(-(offset + 1)) + 1; + if (magnitude > base) + return false; + cursor_ = base - magnitude; + return true; + } + + private: + int64_t read_from_current_position(void* destination, int64_t nb_bytes) { + const auto requested = + std::min(static_cast(nb_bytes), size_ - cursor_); + if (requested == 0) + return 0; +#if defined(_WIN32) + LARGE_INTEGER offset{}; + offset.QuadPart = static_cast(offset_ + cursor_); + LARGE_INTEGER position{}; + if (!SetFilePointerEx(mutation_guard_.handle(), offset, &position, FILE_BEGIN) || + position.QuadPart != offset.QuadPart) { + throw std::runtime_error("[trtmc] Failed to seek in the RTX plan section"); + } + read_destination_bytes(mutation_guard_.handle(), staging_, destination, requested, timing_); +#else + file_.clear(); + file_.seekg(static_cast(offset_ + cursor_), std::ios::beg); + if (!file_) + throw std::runtime_error("[trtmc] Failed to seek in the RTX plan section"); + read_destination_bytes(file_, staging_, destination, requested, timing_); +#endif + cursor_ += requested; + return static_cast(requested); + } - std::ifstream input(path_, std::ios::binary | std::ios::ate); - if (!input) + PlanFileMutationGuard mutation_guard_; +#if !defined(_WIN32) + std::ifstream file_; +#endif + std::uint64_t offset_{0}; + std::uint64_t size_{0}; + std::uint64_t cursor_{0}; + std::vector staging_; + FilePlanLoadTiming& timing_; +}; + +void apply_weight_streaming_budget(nvinfer1::ICudaEngine& engine, std::int64_t budget, + bool cuda_graphs) { + if (budget < 0) + return; + if (cuda_graphs) + throw std::invalid_argument( + "[trtmc] RTX weight streaming is incompatible with CUDA graph capture"); + const std::int64_t streamable = engine.getStreamableWeightsSize(); + if (streamable <= 0) + throw std::runtime_error("[trtmc] Engine was not built with TensorRT-RTX weight streaming"); + const std::int64_t applied = std::min(budget, streamable); + if (!engine.setWeightStreamingBudgetV2(applied) || + engine.getWeightStreamingBudgetV2() != applied) { + throw std::runtime_error("[trtmc] TensorRT-RTX rejected the weight streaming budget"); + } + std::cerr << "[trtmc.rtx_weight_budget] requested_bytes=" << budget + << " streamable_bytes=" << streamable << " applied_bytes=" << applied + << " streaming_scratch_bytes=" << engine.getWeightStreamingScratchMemorySize() + << '\n'; +} + +void validate_optimization_profile(const nvinfer1::ICudaEngine& engine, int32_t profile) { + if (profile < 0 || profile >= engine.getNbOptimizationProfiles()) + throw std::invalid_argument("[trtmc] Invalid optimization profile index"); +} + +void configure_cuda_graphs(nvinfer1::IRuntimeConfig& config, bool enabled) { + if (enabled && + !config.setCudaGraphStrategy(nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE)) { + throw std::runtime_error("[trtmc] Failed to enable RTX CUDA graph capture"); + } +} + +void validate_serial_execution_context(const ModuleCreateOptions& options, + bool serial_execution_context) { + if (!serial_execution_context) + return; + if (options.stream == nullptr) { + throw std::invalid_argument( + "[trtmc] Shared RTX activation memory requires an explicit CUDA stream"); + } + if (options.cuda_graphs) { + throw std::invalid_argument( + "[trtmc] Shared RTX activation memory is incompatible with CUDA graph capture"); + } +} + +struct ActivationArenaKey { + int device{-1}; + cudaStream_t stream{nullptr}; + + bool operator==(const ActivationArenaKey& other) const noexcept { + return device == other.device && stream == other.stream; + } +}; + +struct ActivationArenaKeyHash { + std::size_t operator()(const ActivationArenaKey& key) const noexcept { + const auto stream = reinterpret_cast(key.stream); + return std::hash{}(stream) ^ + (std::hash{}(key.device) + 0x9e3779b9U + (stream << 6U) + (stream >> 2U)); + } +}; + +// One arena is shared only by explicitly marked contexts on one CUDA stream. +// Context creation records each profile upper bound. Allocation is deferred to +// the first enqueue, after that context has received its request shapes. Each +// serial context caches its shape-sized requirement until its input shapes +// change, so repeated forwards avoid repartitioning activation memory. +class RtxActivationArena final : public ITrtActivationArena { + public: + RtxActivationArena(int device, cudaStream_t stream) : device_(device), stream_(stream) {} + + ~RtxActivationArena() override { + if (memory_ == nullptr) return; - const auto end = input.tellg(); - if (end <= 0) + int previous_device = -1; + const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess && + previous_device >= 0 && previous_device != device_; + if (restore_device && cudaSetDevice(device_) != cudaSuccess) return; - std::vector bytes(static_cast(end)); - input.seekg(0); - input.read(bytes.data(), static_cast(bytes.size())); - if (!input) - throw std::runtime_error("[trtmc] Failed to read TensorRT-RTX runtime cache"); - if (!cache_->deserialize(bytes.data(), bytes.size())) - throw std::runtime_error("[trtmc] TensorRT-RTX rejected the runtime cache"); - } - - RuntimeCacheState(const RuntimeCacheState&) = delete; - RuntimeCacheState& operator=(const RuntimeCacheState&) = delete; - - ~RuntimeCacheState() { - auto* serialized = cache_ != nullptr ? cache_->serialize() : nullptr; - if (serialized != nullptr && serialized->size() > 0) { - std::ofstream output(path_, std::ios::binary | std::ios::trunc); - if (output) { - output.write(static_cast(serialized->data()), - static_cast(serialized->size())); - } else { - std::cerr << "[trtmc] Failed to write TensorRT-RTX runtime cache: " << path_ - << '\n'; + (void)cudaStreamSynchronize(stream_); + if (memory_) + (void)cudaFree(memory_); + if (restore_device) + (void)cudaSetDevice(previous_device); + } + + void attach(nvinfer1::IExecutionContext* context, cudaStream_t stream, + std::int64_t required_bytes) override { + if (context == nullptr) + throw std::invalid_argument("[trtmc] Cannot attach a null RTX execution context"); + if (stream == nullptr || stream != stream_) + throw std::invalid_argument("[trtmc] RTX activation arena stream mismatch"); + if (required_bytes < 0 || + static_cast(required_bytes) > + static_cast(std::numeric_limits::max())) { + throw std::overflow_error("[trtmc] RTX activation memory requirement is invalid"); + } + require_current_device(); + + std::lock_guard lock(mutex_); + if (contexts_.count(context) != 0) + throw std::logic_error("[trtmc] RTX execution context is already attached to an arena"); + contexts_.emplace(context, required_bytes); + required_bytes_ = std::max(required_bytes_, required_bytes); + } + + void begin_enqueue(nvinfer1::IExecutionContext* context) override { + mutex_.lock(); + try { + require_current_device(); + const auto found = contexts_.find(context); + if (found == contexts_.end()) { + throw std::logic_error( + "[trtmc] RTX execution context is not attached to its activation arena"); } + auto shaped = shaped_requirements_.find(context); + if (shaped == shaped_requirements_.end()) { + const std::int64_t requested = shaped_requirement(*context, found->second); + shaped = shaped_requirements_.emplace(context, requested).first; + } + ensure_capacity_locked(shaped->second); + // setDeviceMemoryV2 has no status return. Rebind on every enqueue + // because a later serial context may have raised the arena's + // high-water mark and replaced its allocation. + context->setDeviceMemoryV2(memory_, capacity_bytes_); + } catch (...) { + mutex_.unlock(); + throw; } - delete serialized; - delete cache_; } - nvinfer1::IRuntimeCache& cache() const { return *cache_; } - const std::string& path() const { return path_; } + void end_enqueue() noexcept override { mutex_.unlock(); } + + void invalidate_shapes(nvinfer1::IExecutionContext* context) override { + std::lock_guard lock(mutex_); + shaped_requirements_.erase(context); + } + + void detach(nvinfer1::IExecutionContext* context) noexcept override { + if (context == nullptr) + return; + std::lock_guard lock(mutex_); + const auto found = contexts_.find(context); + if (found == contexts_.end()) + return; + // TensorRT permits the activation buffer and execution context to be + // released only after the enqueued work completes. + const cudaError_t status = cudaStreamSynchronize(stream_); + if (status != cudaSuccess) { + std::cerr << "[trtmc] Failed to synchronize shared RTX activation memory during " + "context teardown: " + << cudaGetErrorString(status) << '\n'; + } + contexts_.erase(found); + shaped_requirements_.erase(context); + if (memory_ == nullptr) { + required_bytes_ = 0; + for (const auto& [attached_context, attached_bytes] : contexts_) { + (void)attached_context; + required_bytes_ = std::max(required_bytes_, attached_bytes); + } + } + } private: - nvinfer1::IRuntimeCache* cache_{nullptr}; - std::string path_; + int device_{-1}; + cudaStream_t stream_{nullptr}; + std::mutex mutex_; + std::unordered_map contexts_; + std::unordered_map shaped_requirements_; + void* memory_{nullptr}; + std::int64_t capacity_bytes_{0}; + std::int64_t required_bytes_{0}; + + void require_current_device() const { + int current_device = -1; + const cudaError_t status = cudaGetDevice(¤t_device); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("[trtmc] Failed to identify the CUDA device: ") + + cudaGetErrorString(status)); + } + if (current_device != device_) + throw std::runtime_error("[trtmc] RTX activation arena CUDA device mismatch"); + } + + static std::int64_t shaped_requirement(nvinfer1::IExecutionContext& context, + std::int64_t profile_upper_bound) { + const std::size_t shaped_bytes = context.updateDeviceMemorySizeForShapes(); + // Zero reports an error (including incomplete shapes). Preserve the + // validated profile bound as a correctness-first fallback. + if (shaped_bytes == 0) + return profile_upper_bound; + if (shaped_bytes > static_cast(profile_upper_bound)) { + throw std::runtime_error( + "[trtmc] RTX shape-sized activation memory exceeds the profile bound"); + } + return static_cast(shaped_bytes); + } + + void ensure_capacity_locked(std::int64_t requested_bytes) { + if (requested_bytes <= capacity_bytes_) + return; + if (memory_ != nullptr) { + // Serial work using the old high-water allocation is ordered on + // this stream. Synchronize once before replacing it for a larger + // context encountered later in the same request. + const cudaError_t synchronize_status = cudaStreamSynchronize(stream_); + if (synchronize_status != cudaSuccess) { + throw std::runtime_error( + std::string("[trtmc] Failed to synchronize before growing the RTX activation " + "arena: ") + + cudaGetErrorString(synchronize_status)); + } + const cudaError_t free_status = cudaFree(memory_); + if (free_status != cudaSuccess) { + throw std::runtime_error( + std::string("[trtmc] Failed to release the previous RTX activation arena: ") + + cudaGetErrorString(free_status)); + } + memory_ = nullptr; + capacity_bytes_ = 0; + } + + const cudaError_t allocation_status = + cudaMalloc(&memory_, static_cast(requested_bytes)); + if (allocation_status != cudaSuccess) { + memory_ = nullptr; + throw std::runtime_error( + std::string("[trtmc] Failed to allocate shared RTX activation memory (") + + std::to_string(requested_bytes) + + " bytes): " + cudaGetErrorString(allocation_status)); + } + capacity_bytes_ = requested_bytes; + std::cerr << "[trtmc.rtx_activation_arena] device=" << device_ + << " contexts=" << contexts_.size() << " capacity_bytes=" << capacity_bytes_ + << " profile_capacity_bytes=" << required_bytes_ << '\n'; + } }; -void keep_backend_resources(ITrtModule& module, - const std::shared_ptr& engine, - const StreamSetup& stream_setup, - const std::shared_ptr& distributed_owner) { - module.keep_alive(engine); - if (stream_setup.owner) - module.keep_alive(stream_setup.owner); - if (distributed_owner) - module.keep_alive(distributed_owner); +std::ifstream open_runtime_cache_file(const char* path) { + std::error_code exists_error; + const bool cache_exists = fs::exists(path, exists_error); + if (exists_error) { + throw std::system_error(exists_error, + "failed to inspect RTX runtime cache " + std::string(path)); + } + + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (cache_exists && !file) { + throw std::runtime_error("[trtmc] Failed to open existing RTX runtime cache: " + + std::string(path)); + } + return file; +} + +std::streamoff runtime_cache_file_size(std::ifstream& file, const char* path) { + const std::streamoff size = file.tellg(); + if (size < 0) { + throw std::runtime_error("[trtmc] Failed to size RTX runtime cache: " + std::string(path)); + } + if (static_cast(size) > + static_cast(std::numeric_limits::max()) || + size > std::numeric_limits::max()) { + throw std::overflow_error("[trtmc] RTX runtime cache is too large to load: " + + std::string(path)); + } + return size; +} + +std::vector read_runtime_cache_bytes(std::ifstream& file, std::streamoff size, + const char* path) { + std::vector bytes(static_cast(size)); + file.seekg(0, std::ios::beg); + if (!file) { + throw std::runtime_error("[trtmc] Failed to seek RTX runtime cache: " + std::string(path)); + } + file.read(bytes.data(), static_cast(size)); + if (!file || file.gcount() != static_cast(size)) { + throw std::runtime_error("[trtmc] Failed to read RTX runtime cache: " + std::string(path)); + } + return bytes; +} + +void load_runtime_cache(nvinfer1::IRuntimeCache& runtime_cache, const char* path) { + std::ifstream file = open_runtime_cache_file(path); + if (!file) + return; + + const std::streamoff size = runtime_cache_file_size(file, path); + if (size == 0) + return; + + const std::vector bytes = read_runtime_cache_bytes(file, size, path); + if (!runtime_cache.deserialize(bytes.data(), bytes.size())) { + throw std::runtime_error("[trtmc] TensorRT-RTX rejected the serialized runtime cache: " + + std::string(path)); + } + std::cerr << "[trtmc] RTX runtime cache loaded: " << path << " (" << size << " bytes)\n"; +} + +std::unique_ptr create_runtime_cache(nvinfer1::IRuntimeConfig& config, + const char* path) { + std::unique_ptr runtime_cache(config.createRuntimeCache()); + if (runtime_cache == nullptr) + throw std::runtime_error("[trtmc] Failed to create RTX runtime cache"); + load_runtime_cache(*runtime_cache, path); + return runtime_cache; } } // namespace class RtxBackend final : public IBackend { public: - RtxBackend() : runtime_(create_trt_runtime()) { + RtxBackend() + : runtime_(create_trt_runtime()) +#if defined(_WIN32) + , + staged_runtime_(create_trt_runtime()) +#endif + { if (!runtime_) throw std::runtime_error("[trtmc] Failed to create TRT-RTX runtime"); +#if defined(_WIN32) + if (!staged_runtime_) + throw std::runtime_error("[trtmc] Failed to create staged TRT-RTX runtime"); + // This dedicated runtime serves only file-backed staged plans, so the + // synchronizing allocator cannot change ordinary in-memory RTX paths. + staged_runtime_->setGpuAllocator(&gpu_allocator_); +#endif + } + + ~RtxBackend() override { + if (!runtime_cache_leases_.empty()) { + std::cerr << "[trtmc] RTX runtime cache not persisted: " << runtime_cache_leases_.size() + << " pipeline lease(s) are still active\n"; + } + delete runtime_cache_; } std::unique_ptr create_module(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options) override { - return create_module_impl(plan_data, plan_size, options, {}); + auto* engine = runtime_->deserializeCudaEngine(plan_data, plan_size); + if (!engine) + throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); + return create_single_module(engine, options, {}, -1); } std::unique_ptr @@ -128,91 +830,333 @@ class RtxBackend final : public IBackend { const std::vector& external_bindings) override { if (external_bindings.empty()) throw std::invalid_argument("[trtmc] External I/O prebindings must not be empty"); - return create_module_impl(plan_data, plan_size, options, external_bindings); + auto* engine = runtime_->deserializeCudaEngine(plan_data, plan_size); + if (!engine) + throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); + return create_single_module(engine, options, external_bindings, -1); + } + + std::unique_ptr + create_module_from_file(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) override { + return create_module_from_file_impl(plan_path, plan_offset, plan_size, options, + external_bindings, weight_streaming_budget_bytes, + retain_engine, serial_execution_context); } BackendDualProfileModules create_dual_profile_modules(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options) override { - auto engine = deserialize_engine(*runtime_, plan_data, plan_size); - const auto stream_setup = resolve_stream(options.stream); - if (engine->getNbOptimizationProfiles() < 2) - throw std::runtime_error("[trtmc] Dual-profile engine requires two profiles"); + auto* engine_raw = runtime_->deserializeCudaEngine(plan_data, plan_size); + if (!engine_raw) + throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); + std::shared_ptr engine(engine_raw, + [](nvinfer1::ICudaEngine* p) { delete p; }); + + StreamSetup stream_setup = resolve_stream(options.stream); - BackendDualProfileModules modules; - modules.prefill = create_context_module(engine, stream_setup, options, 0, {}); - modules.decode = create_context_module(engine, stream_setup, options, 1, {}); - return modules; + const int32_t nprofiles = engine->getNbOptimizationProfiles(); + if (nprofiles < 2) + throw std::runtime_error("[trtmc] Dual-profile modules require at least two profiles"); + auto make_ctx_module = [&](int32_t profile_idx) -> std::unique_ptr { + auto profile_options = options; + profile_options.optimization_profile = profile_idx; + auto config = create_runtime_config(*engine, profile_options, false); + return create_execution_module(engine, config, stream_setup, profile_options, {}, false, + false); + }; + + BackendDualProfileModules out; + out.prefill = make_ctx_module(0); + out.decode = make_ctx_module(1); + return out; } const char* name() const override { return "trt_rtx"; } + std::uint64_t acquire_runtime_cache_lease(const char* path) override { + std::lock_guard lock(runtime_cache_mutex_); + return runtime_cache_leases_.acquire(path); + } + + void release_runtime_cache_lease(std::uint64_t lease) override { + std::lock_guard lock(runtime_cache_mutex_); + runtime_cache_leases_.release(lease, runtime_cache_ != nullptr, + [this] { persist_runtime_cache_locked(); }); + } + private: +#if defined(_WIN32) + // Declared before runtime_ so it outlives the runtime and every engine + // deserialized through it (members are destroyed in reverse order). + SynchronousGpuAllocator gpu_allocator_; +#endif + TrtUniquePtr runtime_; +#if defined(_WIN32) + TrtUniquePtr staged_runtime_; +#endif + nvinfer1::IRuntimeCache* runtime_cache_{nullptr}; + std::mutex runtime_cache_mutex_; + internal::RuntimeCacheLeaseState runtime_cache_leases_; + std::mutex retained_engines_mutex_; + std::unordered_map> retained_engines_; + std::mutex activation_arenas_mutex_; + std::unordered_map, + ActivationArenaKeyHash> + activation_arenas_; + std::unique_ptr - create_module_impl(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options, - const std::vector& external_bindings) { - auto engine = deserialize_engine(*runtime_, plan_data, plan_size); - return create_context_module(engine, resolve_stream(options.stream), options, 0, - external_bindings); + create_module_from_file_impl(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) { + FilePlanLoadTiming timing; + validate_serial_execution_context(options, serial_execution_context); + if (weight_streaming_budget_bytes >= 0 && options.cuda_graphs) { + throw std::invalid_argument( + "[trtmc] RTX weight streaming is incompatible with CUDA graph capture"); + } + BoundedPlanStreamReader reader = [&] { + AccumulateWallTime elapsed(&timing.reader_open_ms); + return BoundedPlanStreamReader(plan_path, plan_offset, plan_size, timing); + }(); + // Stable file and section identity prevents retained-engine reuse across + // bundle files or after an in-place plan replacement. + const std::string cache_key = + retain_engine ? retained_engine_key(reader, weight_streaming_budget_bytes) + : std::string{}; + std::unique_lock retained_lock; + if (retain_engine) { + // Serialize retained-engine creation. Two concurrent deserializations + // of the same multi-GiB plan can exceed the device-memory envelope + // before the losing insertion is released. + { + AccumulateWallTime elapsed(&timing.retained_lock_ms); + retained_lock = std::unique_lock(retained_engines_mutex_); + } + const auto hit = retained_engines_.find(cache_key); + if (hit != retained_engines_.end()) { + reader.verify_unchanged(); + std::cerr << "[trtmc.rtx_engine_cache] hit=1\n"; + auto module = create_single_module_from_engine( + hit->second, options, external_bindings, true, serial_execution_context, &timing); + timing.log(plan_size, options.optimization_profile, true); + return module; + } + } + TrtUniquePtr engine; + { + AccumulateWallTime elapsed(&timing.deserialize_ms); + engine.reset(file_backed_runtime().deserializeCudaEngine(reader)); + } + reader.verify_unchanged(); + if (!engine) + throw std::runtime_error("[trtmc] Failed to stream-deserialize engine (RTX)"); + if (retain_engine) { + std::shared_ptr retained( + engine.release(), [](nvinfer1::ICudaEngine* value) { delete value; }); + { + AccumulateWallTime elapsed(&timing.weight_budget_ms); + apply_weight_streaming_budget(*retained, weight_streaming_budget_bytes, + options.cuda_graphs); + } + auto module = create_single_module_from_engine(retained, options, external_bindings, + true, serial_execution_context, &timing); + retained_engines_.emplace(cache_key, retained); + std::cerr << "[trtmc.rtx_engine_cache] hit=0 retained=1\n"; + timing.log(plan_size, options.optimization_profile, false); + return module; + } + auto module = create_single_module(engine.release(), options, external_bindings, + weight_streaming_budget_bytes, true, + serial_execution_context, &timing); + timing.log(plan_size, options.optimization_profile, false); + return module; + } + + nvinfer1::IRuntime& file_backed_runtime() { +#if defined(_WIN32) + return *staged_runtime_; +#else + return *runtime_; +#endif + } + + static std::string retained_engine_key(const BoundedPlanStreamReader& reader, + std::int64_t weight_streaming_budget_bytes) { + int device = -1; + const cudaError_t status = cudaGetDevice(&device); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("[trtmc] Failed to identify the CUDA device: ") + + cudaGetErrorString(status)); + } + return std::to_string(device) + ":" + reader.cache_identity() + ":" + + std::to_string(weight_streaming_budget_bytes); + } + + std::unique_ptr create_single_module_from_engine( + const std::shared_ptr& engine, const ModuleCreateOptions& options, + const std::vector& external_bindings, + bool use_synchronous_allocator = false, bool serial_execution_context = false, + FilePlanLoadTiming* timing = nullptr) { + validate_optimization_profile(*engine, options.optimization_profile); + const StreamSetup stream_setup = resolve_stream(options.stream); + auto config = create_runtime_config(*engine, options, serial_execution_context, timing); + return create_execution_module(engine, config, stream_setup, options, external_bindings, + use_synchronous_allocator, serial_execution_context, timing); } std::unique_ptr - create_context_module(const std::shared_ptr& engine, - const StreamSetup& stream_setup, const ModuleCreateOptions& options, - int32_t profile_idx, - const std::vector& external_bindings) { - TrtUniquePtr runtime_config(engine->createRuntimeConfig()); - if (!runtime_config) - throw std::runtime_error("[trtmc] Failed to create RTX runtime config"); + create_single_module(nvinfer1::ICudaEngine* engine_raw, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, + bool use_synchronous_allocator = false, + bool serial_execution_context = false, + FilePlanLoadTiming* timing = nullptr) { + std::shared_ptr engine( + engine_raw, [](nvinfer1::ICudaEngine* value) { delete value; }); + { + AccumulateWallTime elapsed(timing ? &timing->weight_budget_ms : nullptr); + apply_weight_streaming_budget(*engine, weight_streaming_budget_bytes, options.cuda_graphs); + } + return create_single_module_from_engine(engine, options, external_bindings, + use_synchronous_allocator, + serial_execution_context, timing); + } - auto runtime_cache = resolve_runtime_cache(*runtime_config, options.runtime_cache_path); - if (runtime_cache && !runtime_config->setRuntimeCache(runtime_cache->cache())) - throw std::runtime_error("[trtmc] TensorRT-RTX rejected the runtime cache"); - if (options.cuda_graphs) { - if (!runtime_config->setCudaGraphStrategy( - nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE)) { - throw std::runtime_error("[trtmc] TensorRT-RTX rejected whole-graph capture"); + std::shared_ptr + create_runtime_config(nvinfer1::ICudaEngine& engine, const ModuleCreateOptions& options, + bool serial_execution_context, FilePlanLoadTiming* timing = nullptr) { + AccumulateWallTime elapsed(timing ? &timing->runtime_config_ms : nullptr); + std::shared_ptr config( + engine.createRuntimeConfig(), [](nvinfer1::IRuntimeConfig* value) { delete value; }); + if (!config) + throw std::runtime_error("[trtmc] Failed to create RTX runtime config"); + if (options.runtime_cache_path && options.runtime_cache_path[0] != '\0') { + AccumulateWallTime cache_elapsed(timing ? &timing->runtime_cache_ms : nullptr); + ensure_runtime_cache(config.get(), options.runtime_cache_path); + } + if (serial_execution_context) { + config->setExecutionContextAllocationStrategy( + nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED); + if (config->getExecutionContextAllocationStrategy() != + nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED) { + throw std::runtime_error( + "[trtmc] TensorRT-RTX rejected user-managed activation memory"); } } + configure_cuda_graphs(*config, options.cuda_graphs); + return config; + } + + std::shared_ptr resolve_activation_arena(bool serial_execution_context, + cudaStream_t stream) { + if (!serial_execution_context) + return {}; + int device = -1; + const cudaError_t status = cudaGetDevice(&device); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("[trtmc] Failed to identify the CUDA device: ") + + cudaGetErrorString(status)); + } - TrtUniquePtr context( - engine->createExecutionContext(runtime_config.get())); + const ActivationArenaKey key{device, stream}; + std::lock_guard lock(activation_arenas_mutex_); + const auto found = activation_arenas_.find(key); + if (found != activation_arenas_.end()) { + if (auto arena = found->second.lock()) + return arena; + activation_arenas_.erase(found); + } + auto arena = std::make_shared(device, stream); + activation_arenas_.emplace(key, arena); + return arena; + } + + std::unique_ptr + create_execution_module(const std::shared_ptr& engine, + const std::shared_ptr& config, + const StreamSetup& stream_setup, const ModuleCreateOptions& options, + const std::vector& external_bindings, + bool use_synchronous_allocator, bool serial_execution_context, + FilePlanLoadTiming* timing = nullptr) { + const auto activation_arena = + resolve_activation_arena(serial_execution_context, stream_setup.stream); + const std::int64_t activation_memory_bytes = + activation_arena ? engine->getDeviceMemorySizeForProfileV2(options.optimization_profile) + : 0; + if (activation_memory_bytes < 0) + throw std::runtime_error("[trtmc] Invalid RTX activation memory requirement"); + std::unique_ptr context; + { + AccumulateWallTime elapsed(timing ? &timing->context_ms : nullptr); + context.reset(engine->createExecutionContext(config.get())); + } if (!context) throw std::runtime_error("[trtmc] Failed to create RTX execution context"); - - auto module = std::make_unique( - engine.get(), context.release(), stream_setup.stream, profile_idx, - options.distributed_communicator, external_bindings, true); +#if defined(_WIN32) + if (use_synchronous_allocator && !context->setTemporaryStorageAllocator(&gpu_allocator_)) + throw std::runtime_error( + "[trtmc] Failed to set synchronous RTX temporary-storage allocator"); +#else + (void)use_synchronous_allocator; +#endif + auto module = [&] { + // Includes profile selection, activation attachment, and I/O-buffer initialization. + AccumulateWallTime elapsed(timing ? &timing->module_init_ms : nullptr); + return std::make_unique( + engine.get(), context.get(), stream_setup.stream, options.optimization_profile, + nullptr, external_bindings, options.cuda_graphs, activation_arena, + activation_memory_bytes); + }(); + // A completed TrtModuleImpl owns (or has already rejected and deleted) + // the context. Before that point the local owner handles exceptions. + (void)context.release(); if (!module->ok()) throw std::runtime_error("[trtmc] TrtModuleImpl creation failed (RTX)"); - - keep_backend_resources(*module, engine, stream_setup, options.distributed_owner); - if (runtime_cache) - module->keep_alive(runtime_cache); + module->keep_alive(engine); + module->keep_alive(config); + if (stream_setup.owner) + module->keep_alive(stream_setup.owner); return module; } - std::shared_ptr resolve_runtime_cache(nvinfer1::IRuntimeConfig& config, - const char* raw_path) { - if (raw_path == nullptr || *raw_path == '\0') - return {}; - const std::string path(raw_path); + void ensure_runtime_cache(nvinfer1::IRuntimeConfig* cfg, const char* path) { std::lock_guard lock(runtime_cache_mutex_); - if (auto existing = runtime_cache_.lock()) { - if (existing->path() != path) { - throw std::invalid_argument( - "one TensorRT-RTX backend cannot use two runtime cache paths concurrently"); - } - return existing; + if (runtime_cache_leases_.empty()) { + throw std::runtime_error( + "[trtmc] RTX runtime cache use requires an active pipeline lease"); + } + if (runtime_cache_leases_.path() != path) { + throw std::invalid_argument( + "[trtmc] RTX backend cannot share one runtime cache across different paths"); } - auto created = std::make_shared(config, path); - runtime_cache_ = created; - return created; + if (!runtime_cache_) + runtime_cache_ = create_runtime_cache(*cfg, path).release(); + if (!cfg->setRuntimeCache(*runtime_cache_)) + throw std::runtime_error("[trtmc] Failed to attach the TensorRT-RTX runtime cache"); } - TrtUniquePtr runtime_; - std::mutex runtime_cache_mutex_; - std::weak_ptr runtime_cache_; + void persist_runtime_cache_locked() { + auto* serialized = runtime_cache_->serialize(); + if (serialized == nullptr) + throw std::runtime_error( + "[trtmc] TensorRT-RTX returned a null serialized runtime cache"); + std::unique_ptr memory(serialized); + if (memory->size() > + static_cast(std::numeric_limits::max())) { + throw std::overflow_error("[trtmc] RTX runtime cache is too large to persist"); + } + + internal::persist_runtime_cache_file(runtime_cache_leases_.path(), memory->data(), + memory->size()); + std::cerr << "[trtmc] RTX runtime cache saved: " << runtime_cache_leases_.path() << " (" + << memory->size() << " bytes)\n"; + } }; } // namespace trtmc @@ -220,12 +1164,26 @@ class RtxBackend final : public IBackend { extern "C" trtmc::IBackend* trtmc_create_backend() { try { return new trtmc::RtxBackend(); - } catch (const std::exception& error) { - std::cerr << "[trtmc] RTX backend init failed: " << error.what() << std::endl; + } catch (const std::exception& e) { + std::cerr << "[trtmc] RTX backend init failed: " << e.what() << std::endl; return nullptr; } } -extern "C" void trtmc_destroy_backend(trtmc::IBackend* backend) { - delete backend; +extern "C" void trtmc_destroy_backend(trtmc::IBackend* b) { + delete b; +} + +extern "C" const char* trtmc_backend_abi() { + static const std::string abi = + std::to_string(getInferLibMajorVersion()) + "." + std::to_string(getInferLibMinorVersion()); + return abi.c_str(); +} + +extern "C" const char* trtmc_backend_runtime_version() { + static const std::string version = std::to_string(getInferLibMajorVersion()) + "." + + std::to_string(getInferLibMinorVersion()) + "." + + std::to_string(getInferLibPatchVersion()) + "." + + std::to_string(getInferLibBuildVersion()); + return version.c_str(); } diff --git a/core/runtime/tensorrt/runtime_cache_persistence.cpp b/core/runtime/tensorrt/runtime_cache_persistence.cpp new file mode 100644 index 0000000000..c168859871 --- /dev/null +++ b/core/runtime/tensorrt/runtime_cache_persistence.cpp @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/tensorrt/runtime_cache_persistence.h" + +#include +#include +#include +#include + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + +namespace trtmc::internal { + +namespace { + +namespace fs = std::filesystem; + +std::uint64_t current_process_id() { +#if defined(_WIN32) + return static_cast(GetCurrentProcessId()); +#else + return static_cast(getpid()); +#endif +} + +fs::path runtime_cache_temporary_path(const fs::path& target) { + fs::path temporary = target; + temporary += ".tmp." + std::to_string(current_process_id()); + return temporary; +} + +void durable_flush(const fs::path& path) { +#if defined(_WIN32) + HANDLE handle = CreateFileW(path.c_str(), GENERIC_WRITE, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (handle == INVALID_HANDLE_VALUE) { + const auto error = static_cast(GetLastError()); + throw std::system_error(error, std::system_category(), + "failed to open RTX runtime cache for durable flush " + + path.string()); + } + DWORD error = ERROR_SUCCESS; + if (!FlushFileBuffers(handle)) + error = GetLastError(); + if (!CloseHandle(handle) && error == ERROR_SUCCESS) + error = GetLastError(); + if (error != ERROR_SUCCESS) { + throw std::system_error(static_cast(error), std::system_category(), + "failed to durably flush RTX runtime cache " + path.string()); + } +#else + // POSIX rename retains the existing best-effort durability contract. + (void)path; +#endif +} + +void atomic_replace(const fs::path& temporary, const fs::path& target) { +#if defined(_WIN32) + if (!MoveFileExW(temporary.c_str(), target.c_str(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH)) { + const auto error = static_cast(GetLastError()); + throw std::system_error(error, std::system_category(), + "failed to atomically replace RTX runtime cache " + + target.string()); + } +#else + std::error_code error; + fs::rename(temporary, target, error); + if (error) { + throw std::system_error(error, "failed to atomically replace RTX runtime cache " + + target.string()); + } +#endif +} + +} // namespace + +void persist_runtime_cache_file(const fs::path& target, const void* data, std::size_t size) { + if (target.empty()) + throw std::invalid_argument("[trtmc] RTX runtime cache path must not be empty"); + if (size != 0 && data == nullptr) + throw std::invalid_argument("[trtmc] Serialized RTX runtime cache data is null"); + if (size > static_cast(std::numeric_limits::max())) + throw std::overflow_error("[trtmc] RTX runtime cache is too large to persist"); + + const fs::path temporary = runtime_cache_temporary_path(target); + try { + std::ofstream output(temporary, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("[trtmc] Failed to open RTX runtime cache temporary file: " + + temporary.string()); + } + if (size != 0) + output.write(static_cast(data), static_cast(size)); + output.flush(); + if (!output) { + throw std::runtime_error("[trtmc] Failed to write or flush RTX runtime cache: " + + temporary.string()); + } + output.close(); + if (!output) { + throw std::runtime_error("[trtmc] Failed to close RTX runtime cache: " + + temporary.string()); + } + durable_flush(temporary); + atomic_replace(temporary, target); + } catch (...) { + std::error_code ignored; + fs::remove(temporary, ignored); + throw; + } +} + +std::uint64_t RuntimeCacheLeaseState::acquire(const char* path) { + if (path == nullptr || path[0] == '\0') + throw std::invalid_argument("[trtmc] RTX runtime cache lease requires a path"); + if (path_.empty()) { + path_ = path; + } else if (path_ != path) { + throw std::invalid_argument( + "[trtmc] RTX backend cannot share one runtime cache across different paths"); + } + const std::uint64_t lease = next_lease_++; + if (lease == 0 || !active_.insert(lease).second) + throw std::runtime_error("[trtmc] RTX runtime cache lease space exhausted"); + return lease; +} + +void RuntimeCacheLeaseState::release(std::uint64_t lease, bool cache_materialized, + const std::function& persist_final_cache) { + const auto active = active_.find(lease); + if (lease == 0 || active == active_.end()) + throw std::invalid_argument("[trtmc] RTX runtime cache lease is invalid or inactive"); + if (active_.size() == 1 && cache_materialized) + persist_final_cache(); + active_.erase(active); +} + +} // namespace trtmc::internal diff --git a/core/runtime/tensorrt/runtime_cache_persistence.h b/core/runtime/tensorrt/runtime_cache_persistence.h new file mode 100644 index 0000000000..fce4196a8c --- /dev/null +++ b/core/runtime/tensorrt/runtime_cache_persistence.h @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace trtmc::internal { + +// Write a complete serialized cache to a same-directory temporary file, make +// its contents durable, then atomically replace the target. The temporary is +// removed on every reported failure. +void persist_runtime_cache_file(const std::filesystem::path& target, const void* data, + std::size_t size); + +// Lease bookkeeping is deliberately independent of TensorRT-RTX so the +// failure contract can be regression-tested without constructing an engine. +// Callers must serialize access to this object. +class RuntimeCacheLeaseState final { + public: + std::uint64_t acquire(const char* path); + + // The final active lease is erased only after persistence succeeds. If the + // callback throws, the exact lease remains active and can be retried. + void release(std::uint64_t lease, bool cache_materialized, + const std::function& persist_final_cache); + + bool empty() const noexcept { return active_.empty(); } + std::size_t size() const noexcept { return active_.size(); } + const std::string& path() const noexcept { return path_; } + + private: + std::string path_; + std::uint64_t next_lease_{1}; + std::unordered_set active_; +}; + +} // namespace trtmc::internal diff --git a/core/runtime/tensorrt/trt_backend.cpp b/core/runtime/tensorrt/trt_backend.cpp index caf20f03a0..44b29f761e 100644 --- a/core/runtime/tensorrt/trt_backend.cpp +++ b/core/runtime/tensorrt/trt_backend.cpp @@ -124,8 +124,9 @@ class TrtBackend final : public IBackend { stream_owner = owned; } - auto module = std::make_unique( - engine, ctx, stream, 0, options.distributed_communicator, external_bindings); + auto module = + std::make_unique(engine, ctx, stream, options.optimization_profile, + options.distributed_communicator, external_bindings); if (!module->ok()) { delete engine; throw std::runtime_error("[trtmc] TrtModuleImpl creation failed"); diff --git a/core/runtime/tensorrt/trt_logger.cpp b/core/runtime/tensorrt/trt_logger.cpp index acd902a8b7..976a4fd363 100644 --- a/core/runtime/tensorrt/trt_logger.cpp +++ b/core/runtime/tensorrt/trt_logger.cpp @@ -54,8 +54,11 @@ void TrtLogger::clear_error() { } TrtUniquePtr create_trt_runtime() { - static TrtLogger logger; - return TrtUniquePtr(nvinfer1::createInferRuntime(logger)); + // Backends may be released from an atexit handler after function-static + // destruction has begun. Keep the small logger process-lifetime so it + // outlives every TensorRT runtime and engine on every platform. + static TrtLogger* logger = new TrtLogger(); + return TrtUniquePtr(nvinfer1::createInferRuntime(*logger)); } } // namespace trtmc diff --git a/core/runtime/tensorrt/trt_module_impl.cpp b/core/runtime/tensorrt/trt_module_impl.cpp index 6aaaef4f23..f51c4ad2f7 100644 --- a/core/runtime/tensorrt/trt_module_impl.cpp +++ b/core/runtime/tensorrt/trt_module_impl.cpp @@ -62,37 +62,58 @@ TrtModuleImpl::TrtModuleImpl(nvinfer1::ICudaEngine* engine, nvinfer1::IExecution cudaStream_t stream, int32_t profile_idx, void* distributed_communicator, const std::vector& external_bindings, - bool backend_managed_cuda_graph) + bool backend_managed_cuda_graph, + std::shared_ptr activation_arena, + std::int64_t activation_memory_bytes) : engine_(engine), ctx_(ctx), stream_(stream), profile_idx_(profile_idx), distributed_communicator_(distributed_communicator), - backend_managed_cuda_graph_(backend_managed_cuda_graph), - cuda_graph_(std::make_unique()) { + activation_arena_(std::move(activation_arena)), + backend_managed_cuda_graph_(backend_managed_cuda_graph) { if (!ctx_) return; try { - discover_tensor_aliases(engine); - validate_initial_external_bindings(engine, external_bindings); + initialize_module(engine, external_bindings, activation_memory_bytes); } catch (const std::exception& error) { - std::cerr << "[trt_module] Invalid external binding: " << error.what() << '\n'; - delete ctx_; - ctx_ = nullptr; - return; + std::cerr << "[trt_module] Module initialization failed: " << error.what() << '\n'; + cleanup_failed_initialization(); + } catch (...) { + std::cerr << "[trt_module] Module initialization failed\n"; + cleanup_failed_initialization(); } - if (!attach_distributed_communicator()) { - delete ctx_; - ctx_ = nullptr; +} + +void TrtModuleImpl::initialize_module(nvinfer1::ICudaEngine* engine, + const std::vector& external_bindings, + std::int64_t activation_memory_bytes) { + cuda_graph_ = std::make_unique(); + discover_tensor_aliases(engine); + validate_initial_external_bindings(engine, external_bindings); + if (!attach_distributed_communicator()) + throw std::runtime_error("failed to attach the distributed communicator"); + select_optimization_profile(); + if (activation_arena_) + activation_arena_->attach(ctx_, stream_, activation_memory_bytes); + allocate_buffers(engine); +} + +void TrtModuleImpl::select_optimization_profile() { + if (profile_idx_ <= 0) return; + if (!ctx_->setOptimizationProfileAsync(profile_idx_, stream_)) + throw std::runtime_error("failed to select the optimization profile"); + const cudaError_t status = cudaStreamSynchronize(stream_); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("failed to synchronize the optimization ") + + "profile selection: " + cudaGetErrorString(status)); } - if (profile_idx_ > 0) { - if (!ctx_->setOptimizationProfileAsync(profile_idx_, stream_)) { - std::cerr << "[trt_module] Failed to set optimization profile " << profile_idx_ << "\n"; - delete ctx_; - ctx_ = nullptr; - return; - } - cudaStreamSynchronize(stream_); +} + +void TrtModuleImpl::cleanup_failed_initialization() noexcept { + try { + free_buffers(); + } catch (...) { } - allocate_buffers(engine); + destroy_execution_context(); } void TrtModuleImpl::discover_tensor_aliases(nvinfer1::ICudaEngine* engine) { @@ -122,39 +143,96 @@ void TrtModuleImpl::discover_tensor_aliases(nvinfer1::ICudaEngine* engine) { void TrtModuleImpl::validate_initial_external_bindings( nvinfer1::ICudaEngine* engine, const std::vector& external_bindings) { for (const auto& binding : external_bindings) { - if (binding.tensor_name.empty()) - throw std::invalid_argument("tensor name must not be empty"); - if (binding.device_ptr == nullptr) - throw std::invalid_argument("buffer for '" + binding.tensor_name + "' is null"); - if (initial_external_bindings_.count(binding.tensor_name) != 0) - throw std::invalid_argument("duplicate tensor '" + binding.tensor_name + "'"); - - if (!engine_has_io_tensor(engine, binding.tensor_name)) - throw std::invalid_argument("unknown tensor '" + binding.tensor_name + "'"); - if (alias_input_by_output_.count(binding.tensor_name) != 0 || - alias_outputs_by_input_.count(binding.tensor_name) != 0) { - throw std::invalid_argument("TensorRT alias tensor '" + binding.tensor_name + - "' must be bound after module creation"); - } - - const auto dims = engine->getTensorShape(binding.tensor_name.c_str()); - if (dims_are_dynamic(dims)) { - throw std::invalid_argument("tensor '" + binding.tensor_name + - "' is dynamic; prebinding requires a static shape"); - } - std::vector shape; - const auto required_bytes = compute_alloc_bytes( - dims, from_trt_dtype(engine->getTensorDataType(binding.tensor_name.c_str())), shape); - if (binding.capacity_bytes < required_bytes) { - throw std::invalid_argument("buffer for '" + binding.tensor_name + "' has " + - std::to_string(binding.capacity_bytes) + - " bytes; expected at least " + - std::to_string(required_bytes)); - } + validate_initial_external_binding_identity(engine, binding); + const auto required_bytes = initial_external_binding_required_bytes(engine, binding); + validate_external_binding_capacity(binding.tensor_name, binding.capacity_bytes, + required_bytes); initial_external_bindings_.emplace(binding.tensor_name, binding.device_ptr); + initial_external_binding_capacities_.emplace(binding.tensor_name, binding.capacity_bytes); + } +} + +void TrtModuleImpl::validate_initial_external_binding_identity( + nvinfer1::ICudaEngine* engine, const ModuleExternalBinding& binding) const { + if (binding.tensor_name.empty()) + throw std::invalid_argument("tensor name must not be empty"); + if (binding.device_ptr == nullptr) + throw std::invalid_argument("buffer for '" + binding.tensor_name + "' is null"); + if (initial_external_bindings_.count(binding.tensor_name) != 0) + throw std::invalid_argument("duplicate tensor '" + binding.tensor_name + "'"); + if (!engine_has_io_tensor(engine, binding.tensor_name)) + throw std::invalid_argument("unknown tensor '" + binding.tensor_name + "'"); + if (alias_input_by_output_.count(binding.tensor_name) != 0 || + alias_outputs_by_input_.count(binding.tensor_name) != 0) { + throw std::invalid_argument("TensorRT alias tensor '" + binding.tensor_name + + "' must be bound after module creation"); + } +} + +std::size_t +TrtModuleImpl::initial_external_binding_required_bytes(nvinfer1::ICudaEngine* engine, + const ModuleExternalBinding& binding) const { + auto dims = engine->getTensorShape(binding.tensor_name.c_str()); + const bool is_input = + engine->getTensorIOMode(binding.tensor_name.c_str()) == nvinfer1::TensorIOMode::kINPUT; + if (dims_are_dynamic(dims) && is_input) { + if (profile_idx_ < 0 || engine->getNbOptimizationProfiles() <= profile_idx_) + throw std::invalid_argument("dynamic tensor '" + binding.tensor_name + + "' has no selected optimization profile"); + dims = engine->getProfileShape(binding.tensor_name.c_str(), profile_idx_, + nvinfer1::OptProfileSelector::kMAX); + } + if (dims_are_dynamic(dims)) + return 0; + std::vector shape; + return compute_alloc_bytes( + dims, from_trt_dtype(engine->getTensorDataType(binding.tensor_name.c_str())), shape); +} + +void TrtModuleImpl::validate_external_binding_capacity(const std::string& name, + std::size_t available, + std::size_t required) { + if (required > 0 && available < required) { + throw std::invalid_argument("buffer for '" + name + "' has " + std::to_string(available) + + " bytes; expected at least " + std::to_string(required)); + } +} + +bool TrtModuleImpl::attach_initial_external_binding(const std::string& name, BufferEntry& entry) { + const auto external = initial_external_bindings_.find(name); + if (external == initial_external_bindings_.end()) + return false; + const auto capacity = initial_external_binding_capacities_.find(name); + const auto available = + capacity == initial_external_binding_capacities_.end() ? std::size_t{0} : capacity->second; + validate_external_binding_capacity(name, available, entry.nbytes); + entry.d_ptr = external->second; + entry.is_external = true; + return true; +} + +void TrtModuleImpl::assign_initial_input_storage(const std::string& name, BufferEntry& entry) { + if (attach_initial_external_binding(name, entry)) + return; + if (entry.is_dynamic || !should_allocate_input(name, entry.nbytes)) + return; + void* allocation = nullptr; + if (cudaMalloc(&allocation, entry.nbytes) == cudaSuccess) { + entry.d_ptr = allocation; + cudaMemsetAsync(entry.d_ptr, 0, entry.nbytes, stream_); } } +void TrtModuleImpl::assign_initial_output_storage(const std::string& name, BufferEntry& entry) { + if (attach_initial_external_binding(name, entry) || entry.nbytes == 0 || entry.lazy_output) + return; + if (cudaMalloc(&entry.d_ptr, entry.nbytes) != cudaSuccess) { + entry.d_ptr = nullptr; + return; + } + cudaMemsetAsync(entry.d_ptr, 0, entry.nbytes, stream_); +} + void TrtModuleImpl::bind_external(const std::string& name, void* ptr, const std::vector& shape) { bind_external(name, ptr); @@ -213,9 +291,22 @@ TrtModuleImpl::~TrtModuleImpl() { // CUDA Graphs may contain TensorRT collective launches that retain the // distributed communicator. Destroy the captured graph before member // teardown releases distributed_owner from keep_alive_. - cuda_graph_->reset(); + if (cuda_graph_) + cuda_graph_->reset(); + auto activation_arena = activation_arena_; + if (activation_arena && ctx_) + activation_arena->detach(ctx_); free_buffers(); delete ctx_; + ctx_ = nullptr; + activation_arena_.reset(); +} + +void TrtModuleImpl::destroy_execution_context() noexcept { + if (activation_arena_ && ctx_) + activation_arena_->detach(ctx_); + delete ctx_; + ctx_ = nullptr; } void TrtModuleImpl::keep_alive(std::shared_ptr resource) { @@ -257,8 +348,11 @@ void TrtModuleImpl::update_dynamic_shape(const std::string& name, BufferEntry& e dims.nbDims = static_cast(new_shape.size()); for (int32_t d = 0; d < dims.nbDims; ++d) dims.d[d] = new_shape[d]; - ctx_->setInputShape(name.c_str(), dims); + if (!ctx_->setInputShape(name.c_str(), dims)) + throw std::runtime_error("TensorRT rejected dynamic input shape for '" + name + "'"); entry.shape = new_shape; + if (activation_arena_) + activation_arena_->invalidate_shapes(ctx_); } std::size_t TrtModuleImpl::compute_alloc_bytes(const nvinfer1::Dims& dims, DType dtype, @@ -298,7 +392,9 @@ void TrtModuleImpl::set_dynamic_input_shapes(nvinfer1::ICudaEngine* engine, int3 continue; if (dims_are_dynamic(engine->getTensorShape(name.c_str()))) { auto dims = engine->getProfileShape(name.c_str(), profile_idx_, selector); - ctx_->setInputShape(name.c_str(), dims); + if (!ctx_->setInputShape(name.c_str(), dims)) + throw std::runtime_error("TensorRT rejected profile input shape for '" + name + + "'"); } } } @@ -328,41 +424,95 @@ void TrtModuleImpl::allocate_single_input(nvinfer1::ICudaEngine* engine, const s entry.nbytes = nbytes; entry.is_input = true; entry.is_dynamic = is_dynamic; + entry.lazy_input = activation_arena_ && !backend_managed_cuda_graph_ && is_dynamic && + should_allocate_input(name, nbytes); entry.shape = is_dynamic ? dims_to_shape(init_dims) : shape; - const auto external = initial_external_bindings_.find(name); - if (external != initial_external_bindings_.end()) { - entry.d_ptr = external->second; - entry.is_external = true; - } else if (!is_dynamic && should_allocate_input(name, nbytes)) { - void* allocation = nullptr; - if (cudaMalloc(&allocation, nbytes) == cudaSuccess) { - entry.d_ptr = allocation; - cudaMemsetAsync(entry.d_ptr, 0, nbytes, stream_); - } - } + assign_initial_input_storage(name, entry); if (entry.d_ptr) bind_tensor_address(name, entry); - if (is_dynamic) - ctx_->setInputShape(name.c_str(), init_dims); - buffers_[name] = std::move(entry); + + if (is_dynamic && !ctx_->setInputShape(name.c_str(), init_dims)) + throw std::runtime_error("TensorRT rejected initial dynamic input shape for '" + name + + "'"); +} + +void TrtModuleImpl::prepare_input_buffer(const std::string& name, BufferEntry& entry, + const std::vector& shape, DType dtype) { + if (!ctx_) + throw std::runtime_error("TensorRT execution context is unavailable"); + if (entry.lazy_input && !entry.is_external) { + const auto minimum = engine_->getProfileShape(name.c_str(), profile_idx_, + nvinfer1::OptProfileSelector::kMIN); + const auto maximum = engine_->getProfileShape(name.c_str(), profile_idx_, + nvinfer1::OptProfileSelector::kMAX); + if (dtype != entry.dtype || minimum.nbDims < 0 || minimum.nbDims != maximum.nbDims || + shape.size() != static_cast(maximum.nbDims)) { + throw std::invalid_argument("TensorRT input dtype or rank mismatch for '" + name + + "'"); + } + for (int32_t axis = 0; axis < maximum.nbDims; ++axis) { + if (shape[axis] < minimum.d[axis] || shape[axis] > maximum.d[axis]) { + throw std::invalid_argument("TensorRT input shape exceeds its profile for '" + + name + "'"); + } + } + } + const auto previous_shape = entry.shape; + update_dynamic_shape(name, entry, shape); + try { + ensure_input_buffer(name, entry); + } catch (...) { + // A failed growth must not leave a larger live shape bound to the + // smaller, still-owned allocation from the preceding forward. + if (entry.lazy_input && !entry.is_external) { + try { + update_dynamic_shape(name, entry, previous_shape); + } catch (...) { + destroy_execution_context(); + } + } + throw; + } } void TrtModuleImpl::ensure_input_buffer(const std::string& name, BufferEntry& entry) { - if (entry.d_ptr != nullptr || entry.is_external || !should_allocate_input(name, entry.nbytes)) + if (entry.is_external || !should_allocate_input(name, entry.nbytes)) + return; + std::size_t required_bytes = entry.nbytes; + if (entry.lazy_input) { + required_bytes = dtype_size(entry.dtype); + for (const auto dim : entry.shape) + required_bytes *= static_cast(std::max(dim, int64_t{1})); + if (required_bytes > entry.nbytes) + throw std::runtime_error("TensorRT input shape exceeds its profile capacity for '" + + name + "'"); + } + if (entry.d_ptr && (!entry.lazy_input || required_bytes <= entry.input_capacity_bytes)) return; - const cudaError_t error = cudaMalloc(&entry.d_ptr, entry.nbytes); - if (error != cudaSuccess) + + // The previous forward may still reference this input. Preserve its + // allocation and binding until a replacement has been initialized. + if (entry.d_ptr && cudaStreamSynchronize(stream_) != cudaSuccess) + throw std::runtime_error("Unable to synchronize TensorRT input growth for '" + name + + "'"); + BufferEntry candidate = entry; + candidate.d_ptr = nullptr; + if (cudaMalloc(&candidate.d_ptr, required_bytes) != cudaSuccess) throw std::runtime_error("Unable to allocate TensorRT input buffer for '" + name + "'"); - cudaMemsetAsync(entry.d_ptr, 0, entry.nbytes, stream_); - if (!bind_tensor_address(name, entry)) { - cudaFree(entry.d_ptr); - entry.d_ptr = nullptr; - throw std::runtime_error("TensorRT rejected input buffer for '" + name + "'"); + if (cudaMemsetAsync(candidate.d_ptr, 0, required_bytes, stream_) != cudaSuccess || + !bind_tensor_address(name, candidate)) { + cudaFree(candidate.d_ptr); + throw std::runtime_error("Unable to initialize TensorRT input buffer for '" + name + "'"); } + void* const previous_ptr = entry.d_ptr; + candidate.input_capacity_bytes = required_bytes; + entry = std::move(candidate); + if (previous_ptr) + cudaFree(previous_ptr); } void TrtModuleImpl::allocate_input_buffers(nvinfer1::ICudaEngine* engine, int32_t num_io, @@ -386,50 +536,91 @@ void TrtModuleImpl::allocate_output_buffers(nvinfer1::ICudaEngine* engine, int32 const std::string name(raw_name); if (engine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT) continue; + allocate_single_output(engine, name); + } +} - auto dtype = from_trt_dtype(engine->getTensorDataType(name.c_str())); +void TrtModuleImpl::allocate_single_output(nvinfer1::ICudaEngine* engine, const std::string& name) { + auto dtype = from_trt_dtype(engine->getTensorDataType(name.c_str())); - // For dynamic engines, query the context for inferred output shape - // (based on the max input shapes set by the caller). - // For static engines, use the engine shape directly. - nvinfer1::Dims out_dims = has_dynamic_shapes_ ? ctx_->getTensorShape(name.c_str()) - : engine->getTensorShape(name.c_str()); + // For dynamic engines, query the context for inferred output shape + // (based on the max input shapes set by the caller). + // For static engines, use the engine shape directly. + nvinfer1::Dims out_dims = has_dynamic_shapes_ ? ctx_->getTensorShape(name.c_str()) + : engine->getTensorShape(name.c_str()); - std::vector shape; - std::size_t nbytes = compute_alloc_bytes(out_dims, dtype, shape); + std::vector shape; + std::size_t nbytes = compute_alloc_bytes(out_dims, dtype, shape); - BufferEntry entry; - entry.shape = shape; - entry.dtype = dtype; - entry.nbytes = nbytes; - entry.is_input = false; + BufferEntry entry; + entry.shape = shape; + entry.dtype = dtype; + entry.nbytes = nbytes; + entry.is_input = false; - if (alias_input_by_output_.count(name) != 0) { - buffers_[name] = std::move(entry); - continue; - } + if (alias_input_by_output_.count(name) != 0) { + buffers_[name] = std::move(entry); + return; + } - const auto external = initial_external_bindings_.find(name); - if (external != initial_external_bindings_.end()) { - entry.d_ptr = external->second; - entry.is_external = true; - } else if (nbytes > 0) { - auto err = cudaMalloc(&entry.d_ptr, nbytes); - if (err != cudaSuccess) - entry.d_ptr = nullptr; - else - cudaMemsetAsync(entry.d_ptr, 0, nbytes, stream_); - } + // Only serial, non-graph contexts opt into live-shape output storage. + // Keep MAX metadata for profile validation and external-buffer contracts. + // Data-dependent outputs retain their existing allocation policy. + entry.lazy_output = activation_arena_ && !backend_managed_cuda_graph_ && + dims_are_dynamic(engine->getTensorShape(name.c_str())) && + out_dims.nbDims >= 0 && !dims_are_dynamic(out_dims); + assign_initial_output_storage(name, entry); - if (entry.d_ptr) - bind_tensor_address(name, entry); + if (entry.d_ptr) + bind_tensor_address(name, entry); + if (!entry.lazy_output) allocate_host_output_staging(host_output_staging_, name, nbytes, entry.is_external); - buffers_[name] = std::move(entry); + buffers_[name] = std::move(entry); +} + +void TrtModuleImpl::ensure_output_buffers() { + for (auto& [name, entry] : buffers_) { + if (entry.lazy_output && !entry.is_external) + ensure_output_buffer(name, entry); } } +void TrtModuleImpl::ensure_output_buffer(const std::string& name, BufferEntry& entry) { + const auto dims = ctx_->getTensorShape(name.c_str()); + if (dims.nbDims < 0 || dims_are_dynamic(dims)) + throw std::runtime_error("TensorRT output shape is unresolved for '" + name + "'"); + std::vector runtime_shape; + const auto required_bytes = compute_alloc_bytes(dims, entry.dtype, runtime_shape); + if (required_bytes > entry.nbytes) + throw std::runtime_error("TensorRT output shape exceeds its profile capacity for '" + + name + "'"); + if (entry.d_ptr && required_bytes <= entry.output_capacity_bytes) + return; + + // An earlier async forward may still use the old address. Keep it valid + // until completion, and retain the old allocation if growth fails. + if (entry.d_ptr && cudaStreamSynchronize(stream_) != cudaSuccess) + throw std::runtime_error("Unable to synchronize TensorRT output growth for '" + name + + "'"); + BufferEntry candidate = entry; + candidate.d_ptr = nullptr; + if (cudaMalloc(&candidate.d_ptr, required_bytes) != cudaSuccess) + throw std::runtime_error("Unable to allocate TensorRT output buffer for '" + name + "'"); + if (cudaMemsetAsync(candidate.d_ptr, 0, required_bytes, stream_) != cudaSuccess || + !bind_tensor_address(name, candidate)) { + cudaFree(candidate.d_ptr); + throw std::runtime_error("Unable to initialize TensorRT output buffer for '" + name + + "'"); + } + void* const previous_ptr = entry.d_ptr; + candidate.output_capacity_bytes = required_bytes; + entry = std::move(candidate); + if (previous_ptr) + cudaFree(previous_ptr); +} + // --- Buffer allocation --- void TrtModuleImpl::allocate_buffers(nvinfer1::ICudaEngine* engine) { @@ -437,7 +628,8 @@ void TrtModuleImpl::allocate_buffers(nvinfer1::ICudaEngine* engine) { const int32_t num_profiles = engine->getNbOptimizationProfiles(); detect_dynamic_shapes(engine, num_io); - // Pass 1: allocate input buffers (use profile-0 max shape for dynamic inputs). + // Pass 1: record selected-profile input capacity. Dynamic input storage is + // deferred until forward; serial contexts allocate only the live shape. allocate_input_buffers(engine, num_io, num_profiles); // Pass 2: allocate output buffers. For dynamic shapes, temporarily set @@ -451,6 +643,7 @@ void TrtModuleImpl::allocate_buffers(nvinfer1::ICudaEngine* engine) { set_dynamic_input_shapes(engine, num_io, nvinfer1::OptProfileSelector::kOPT); initial_external_bindings_.clear(); + initial_external_binding_capacities_.clear(); cudaStreamSynchronize(stream_); } @@ -479,9 +672,9 @@ void TrtModuleImpl::bind_alias_outputs_or_invalidate(const std::vectorreset(); - delete ctx_; - ctx_ = nullptr; + if (cuda_graph_) + cuda_graph_->reset(); + destroy_execution_context(); throw std::runtime_error("TensorRT rejected external alias output '" + output_name + "'"); } @@ -583,6 +776,8 @@ TensorMap TrtModuleImpl::forward(const TensorMap& inputs) { } auto& staging = host_output_staging_[name]; + if (entry.lazy_output) + staging.resize(runtime_nbytes); cudaMemcpy(staging.data(), entry.d_ptr, runtime_nbytes, cudaMemcpyDeviceToHost); Tensor t; @@ -599,6 +794,10 @@ TensorMap TrtModuleImpl::forward(const TensorMap& inputs) { void TrtModuleImpl::enable_cuda_graph() { if (backend_managed_cuda_graph_) return; + if (activation_arena_) { + throw std::logic_error( + "CUDA graph capture is incompatible with a shared TensorRT activation arena"); + } use_cuda_graph_ = true; cuda_graph_->reset(); } @@ -613,13 +812,13 @@ void TrtModuleImpl::forward_async(const TensorMap& inputs) { if (!entry.is_input) continue; - ensure_input_buffer(name, entry); + prepare_input_buffer(name, entry, tensor.shape, tensor.dtype); if (!entry.d_ptr) continue; - update_dynamic_shape(name, entry, tensor.shape); - - auto copy_bytes = std::min(tensor.nbytes(), entry.nbytes); + const auto capacity = entry.lazy_input && !entry.is_external ? entry.input_capacity_bytes + : entry.nbytes; + auto copy_bytes = std::min(tensor.nbytes(), capacity); if (copy_bytes > 0 && tensor.data) { cudaMemcpyAsync(entry.d_ptr, tensor.data, copy_bytes, cudaMemcpyHostToDevice, stream_); } @@ -635,7 +834,20 @@ void TrtModuleImpl::execute_enqueue() { validate_alias_groups_bound(); alias_groups_ready_ = true; } - record_timed_enqueue(); + if (activation_arena_) + ensure_output_buffers(); + if (!activation_arena_) { + record_timed_enqueue(); + return; + } + activation_arena_->begin_enqueue(ctx_); + try { + record_timed_enqueue(); + } catch (...) { + activation_arena_->end_enqueue(); + throw; + } + activation_arena_->end_enqueue(); } bool TrtModuleImpl::cuda_graph_captured() const { @@ -756,14 +968,15 @@ void TrtModuleImpl::forward_device_async(const DeviceTensorMap& inputs) { if (!entry.is_input) continue; - ensure_input_buffer(name, entry); + prepare_input_buffer(name, entry, dt_ptr->shape(), dt_ptr->dtype()); if (!entry.d_ptr) continue; - update_dynamic_shape(name, entry, dt_ptr->shape()); - if (dt_ptr->data() != entry.d_ptr) { - auto copy_bytes = std::min(dt_ptr->nbytes(), entry.nbytes); + const auto capacity = entry.lazy_input && !entry.is_external + ? entry.input_capacity_bytes + : entry.nbytes; + auto copy_bytes = std::min(dt_ptr->nbytes(), capacity); if (copy_bytes > 0) { cudaMemcpyAsync(entry.d_ptr, dt_ptr->data(), copy_bytes, cudaMemcpyDeviceToDevice, stream_); diff --git a/core/runtime/tensorrt/trt_module_impl.h b/core/runtime/tensorrt/trt_module_impl.h index 58f11a33f2..b17be1edb5 100644 --- a/core/runtime/tensorrt/trt_module_impl.h +++ b/core/runtime/tensorrt/trt_module_impl.h @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -24,6 +25,19 @@ namespace trtmc { class CudaGraphExec; class TrtModuleImplTestPeer; +// Optional TensorRT-RTX coordinator for execution contexts which run serially +// on one stream and share a user-managed activation allocation. +class ITrtActivationArena { + public: + virtual ~ITrtActivationArena() = default; + virtual void attach(nvinfer1::IExecutionContext* context, cudaStream_t stream, + std::int64_t required_bytes) = 0; + virtual void begin_enqueue(nvinfer1::IExecutionContext* context) = 0; + virtual void end_enqueue() noexcept = 0; + virtual void invalidate_shapes(nvinfer1::IExecutionContext* context) = 0; + virtual void detach(nvinfer1::IExecutionContext* context) noexcept = 0; +}; + class TrtModuleImpl final : public ITrtModule { public: // Backend creates engine + context, passes them in. @@ -32,7 +46,9 @@ class TrtModuleImpl final : public ITrtModule { cudaStream_t stream, int32_t profile_idx = 0, void* distributed_communicator = nullptr, const std::vector& external_bindings = {}, - bool backend_managed_cuda_graph = false); + bool backend_managed_cuda_graph = false, + std::shared_ptr activation_arena = {}, + std::int64_t activation_memory_bytes = 0); ~TrtModuleImpl() override; TrtModuleImpl(const TrtModuleImpl&) = delete; @@ -80,6 +96,10 @@ class TrtModuleImpl final : public ITrtModule { bool is_input{true}; bool is_external{false}; bool is_dynamic{false}; + bool lazy_input{false}; + std::size_t input_capacity_bytes{0}; + bool lazy_output{false}; + std::size_t output_capacity_bytes{0}; }; struct TimingEvent { cudaEvent_t start{nullptr}; @@ -91,6 +111,7 @@ class TrtModuleImpl final : public ITrtModule { cudaStream_t stream_{nullptr}; int32_t profile_idx_{0}; void* distributed_communicator_{nullptr}; + std::shared_ptr activation_arena_; bool has_dynamic_shapes_{false}; bool use_cuda_graph_{false}; bool backend_managed_cuda_graph_{false}; @@ -98,6 +119,7 @@ class TrtModuleImpl final : public ITrtModule { std::unique_ptr cuda_graph_; std::vector> keep_alive_; std::unordered_map initial_external_bindings_; + std::unordered_map initial_external_binding_capacities_; std::unordered_map buffers_; std::unordered_map alias_input_by_output_; std::unordered_map> alias_outputs_by_input_; @@ -106,11 +128,25 @@ class TrtModuleImpl final : public ITrtModule { std::string timing_label_{"engine"}; std::vector timing_events_; + void initialize_module(nvinfer1::ICudaEngine* engine, + const std::vector& external_bindings, + std::int64_t activation_memory_bytes); + void select_optimization_profile(); + void cleanup_failed_initialization() noexcept; void allocate_buffers(nvinfer1::ICudaEngine* engine); void discover_tensor_aliases(nvinfer1::ICudaEngine* engine); void validate_initial_external_bindings(nvinfer1::ICudaEngine* engine, const std::vector& external_bindings); + void validate_initial_external_binding_identity(nvinfer1::ICudaEngine* engine, + const ModuleExternalBinding& binding) const; + std::size_t initial_external_binding_required_bytes(nvinfer1::ICudaEngine* engine, + const ModuleExternalBinding& binding) const; + static void validate_external_binding_capacity(const std::string& name, std::size_t available, + std::size_t required); + bool attach_initial_external_binding(const std::string& name, BufferEntry& entry); + void assign_initial_input_storage(const std::string& name, BufferEntry& entry); + void assign_initial_output_storage(const std::string& name, BufferEntry& entry); void bind_alias_group(const std::string& input_name, void* ptr); bool should_allocate_input(const std::string& name, std::size_t nbytes) const; void validate_alias_outputs_exist(const std::vector& output_names) const; @@ -123,13 +159,19 @@ class TrtModuleImpl final : public ITrtModule { int32_t num_profiles); void allocate_single_input(nvinfer1::ICudaEngine* engine, const std::string& name, int32_t num_profiles); + void prepare_input_buffer(const std::string& name, BufferEntry& entry, + const std::vector& shape, DType dtype); void ensure_input_buffer(const std::string& name, BufferEntry& entry); void allocate_output_buffers(nvinfer1::ICudaEngine* engine, int32_t num_io); + void allocate_single_output(nvinfer1::ICudaEngine* engine, const std::string& name); + void ensure_output_buffers(); + void ensure_output_buffer(const std::string& name, BufferEntry& entry); void set_dynamic_input_shapes(nvinfer1::ICudaEngine* engine, int32_t num_io, nvinfer1::OptProfileSelector selector); void update_dynamic_shape(const std::string& name, BufferEntry& entry, const std::vector& new_shape); void execute_enqueue(); + void destroy_execution_context() noexcept; void flush_timing_events(); bool begin_timing_event(TimingEvent& event); void finish_timing_event(TimingEvent event); diff --git a/core/runtime/tests/test_bundle_format_v1.cpp b/core/runtime/tests/test_bundle_format_v1.cpp index 0da636b49b..f61cca4d0a 100644 --- a/core/runtime/tests/test_bundle_format_v1.cpp +++ b/core/runtime/tests/test_bundle_format_v1.cpp @@ -5,13 +5,13 @@ #include "runtime/bundle/bundle_format.h" +#include #include #include #include #include #include #include -#include namespace { @@ -25,10 +25,11 @@ void check(bool condition, const char* name) { } std::filesystem::path temp_dir() { - char pattern[] = "/tmp/trtmc_bundle_v1_XXXXXX"; - char* path = mkdtemp(pattern); - if (path == nullptr) - throw std::runtime_error("mkdtemp failed"); + const auto nonce = std::chrono::steady_clock::now().time_since_epoch().count(); + const auto path = + std::filesystem::temp_directory_path() / ("trtmc-bundle-v1-" + std::to_string(nonce)); + if (!std::filesystem::create_directory(path)) + throw std::runtime_error("failed to create temporary test directory"); return path; } diff --git a/core/runtime/tests/test_dynamic_library.cpp b/core/runtime/tests/test_dynamic_library.cpp new file mode 100644 index 0000000000..6cc0043970 --- /dev/null +++ b/core/runtime/tests/test_dynamic_library.cpp @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/platform/dynamic_library.h" + +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +} // namespace + +int main() { +#if defined(_WIN32) + check(trtmc::internal::dynamic_library_filename("trtmc_example") == "trtmc_example.dll", + "Windows library filename"); +#elif defined(__APPLE__) + check(trtmc::internal::dynamic_library_filename("trtmc_example") == "libtrtmc_example.dylib", + "macOS library filename"); +#else + check(trtmc::internal::dynamic_library_filename("trtmc_example") == "libtrtmc_example.so", + "Linux library filename"); +#endif + + std::string error; + auto missing = trtmc::internal::open_dynamic_library( + std::filesystem::temp_directory_path() / "trtmc-missing-dynamic-library", + trtmc::internal::DynamicLibraryVisibility::local, &error); + check(missing == nullptr, "missing dynamic library is rejected"); + check(!error.empty(), "missing dynamic library reports an error"); + + auto handle = trtmc::internal::open_dynamic_library( + TRTMC_TEST_LIBRARY, trtmc::internal::DynamicLibraryVisibility::local, &error); + check(handle != nullptr, "test dynamic library opens"); + if (handle != nullptr) { + auto symbol = + trtmc::internal::dynamic_library_symbol(handle, "trtmc_create_backend", &error); + check(symbol != nullptr, "exported C ABI symbol resolves"); + check(trtmc::internal::close_dynamic_library(handle, &error), + "test dynamic library closes"); + } + + std::cerr << (failures == 0 ? "ALL PASSED" : "SOME FAILED") << '\n'; + return failures; +} diff --git a/core/runtime/tests/test_family_loader.cpp b/core/runtime/tests/test_family_loader.cpp index 1f77d55db2..e7ae6a62fe 100644 --- a/core/runtime/tests/test_family_loader.cpp +++ b/core/runtime/tests/test_family_loader.cpp @@ -4,10 +4,10 @@ */ #include "runtime/bundle/bundle_format.h" +#include "runtime/platform/dynamic_library.h" #include "trtmc/runtime/family_loader.h" #include -#include #include #include #include @@ -62,18 +62,19 @@ bool rtx_options_throw(const std::filesystem::path& bundle, const std::string& r void check_rtx_options(const std::filesystem::path& runtime_root, const std::string& expected_cache_path, bool expected_cuda_graphs) { - const auto library_path = runtime_root / "libtrtmc_backend_trt_rtx.so"; - void* handle = dlopen(library_path.c_str(), RTLD_NOW | RTLD_LOCAL); + const auto library_path = + runtime_root / trtmc::internal::dynamic_library_filename("trtmc_backend_trt_rtx"); + void* handle = trtmc::internal::open_dynamic_library(library_path); check(handle != nullptr, "fake RTX backend remains loaded"); if (handle == nullptr) return; using CachePathFn = const char* (*)(); using CudaGraphsFn = bool (*)(); - const auto cache_path = - reinterpret_cast(dlsym(handle, "trtmc_test_backend_last_runtime_cache_path")); - const auto cuda_graphs = - reinterpret_cast(dlsym(handle, "trtmc_test_backend_last_cuda_graphs")); + const auto cache_path = reinterpret_cast(trtmc::internal::dynamic_library_symbol( + handle, "trtmc_test_backend_last_runtime_cache_path")); + const auto cuda_graphs = reinterpret_cast( + trtmc::internal::dynamic_library_symbol(handle, "trtmc_test_backend_last_cuda_graphs")); check(cache_path != nullptr, "fake RTX cache-path probe is exported"); check(cuda_graphs != nullptr, "fake RTX CUDA-graphs probe is exported"); if (cache_path != nullptr) @@ -81,7 +82,7 @@ void check_rtx_options(const std::filesystem::path& runtime_root, if (cuda_graphs != nullptr) check(cuda_graphs() == expected_cuda_graphs, "CUDA-graphs option reaches delayed module creation"); - dlclose(handle); + trtmc::internal::close_dynamic_library(handle); } } // namespace diff --git a/core/runtime/tests/test_runtime_cache_persistence.cpp b/core/runtime/tests/test_runtime_cache_persistence.cpp new file mode 100644 index 0000000000..9fcc442166 --- /dev/null +++ b/core/runtime/tests/test_runtime_cache_persistence.cpp @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/tensorrt/runtime_cache_persistence.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +namespace fs = std::filesystem; + +class TempDirectory { + public: + TempDirectory() { + const auto nonce = std::chrono::steady_clock::now().time_since_epoch().count(); + path_ = fs::temp_directory_path() / ("trtmc-runtime-cache-test-" + std::to_string(nonce)); + fs::create_directory(path_); + } + + ~TempDirectory() { + std::error_code error; + fs::remove_all(path_, error); + } + + const fs::path& path() const { return path_; } + + private: + fs::path path_; +}; + +void check(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(message); +} + +std::string read_file(const fs::path& path) { + std::ifstream input(path, std::ios::binary); + if (!input) + throw std::runtime_error("failed to read test file"); + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +void test_only_final_lease_persists_with_durable_atomic_replace() { + TempDirectory directory; + const fs::path target = directory.path() / "cache.bin"; + { + std::ofstream output(target, std::ios::binary); + output << "old-cache"; + } + + trtmc::internal::RuntimeCacheLeaseState leases; + const std::string target_string = target.string(); + const auto first = leases.acquire(target_string.c_str()); + const auto final = leases.acquire(target_string.c_str()); + int serialization_attempts = 0; + const std::string new_cache = "new-serialized-cache"; + const auto persist = [&] { + ++serialization_attempts; + trtmc::internal::persist_runtime_cache_file(target, new_cache.data(), new_cache.size()); + }; + + leases.release(first, true, persist); + check(serialization_attempts == 0, "non-final lease unexpectedly persisted the cache"); + check(read_file(target) == "old-cache", "non-final lease changed the cache file"); + check(leases.size() == 1, "final cache lease was lost before persistence"); + + leases.release(final, true, persist); + check(serialization_attempts == 1, "final lease did not serialize exactly once"); + check(leases.empty(), "successful final persistence retained its lease"); + check(read_file(target) == new_cache, "atomic cache replacement wrote incorrect bytes"); +} + +void test_write_failure_preserves_final_lease_for_successful_retry() { + TempDirectory directory; + const fs::path parent = directory.path() / "not-created"; + const fs::path target = parent / "cache.bin"; + const std::string target_string = target.string(); + const std::string serialized = "cache-after-write-retry"; + + trtmc::internal::RuntimeCacheLeaseState leases; + const auto lease = leases.acquire(target_string.c_str()); + bool failure_observed = false; + try { + leases.release(lease, true, [&] { + trtmc::internal::persist_runtime_cache_file(target, serialized.data(), + serialized.size()); + }); + } catch (const std::runtime_error& error) { + failure_observed = std::string(error.what()).find("temporary file") != std::string::npos; + } + + check(failure_observed, "runtime-cache write failure was not reported"); + check(leases.size() == 1, "runtime-cache write failure consumed the final lease"); + fs::create_directory(parent); + leases.release(lease, true, [&] { + trtmc::internal::persist_runtime_cache_file(target, serialized.data(), serialized.size()); + }); + check(leases.empty(), "successful write retry retained the final lease"); + check(read_file(target) == serialized, "write retry persisted incorrect cache bytes"); +} + +} // namespace + +int main() { + try { + test_only_final_lease_persists_with_durable_atomic_replace(); + test_write_failure_preserves_final_lease_for_successful_retry(); + std::cout << "Runtime-cache persistence tests passed\n"; + return 0; + } catch (const std::exception& error) { + std::cerr << "Runtime-cache persistence test failed: " << error.what() << '\n'; + return 1; + } +} diff --git a/core/runtime/tests/test_task_api.cpp b/core/runtime/tests/test_task_api.cpp index f302c9fa0e..f7544c8211 100644 --- a/core/runtime/tests/test_task_api.cpp +++ b/core/runtime/tests/test_task_api.cpp @@ -14,6 +14,7 @@ static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); +static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); @@ -70,6 +71,39 @@ void test_robot_observation() { throw std::runtime_error("RobotObservation did not preserve its input spans"); } +void test_structured_video_generation_contract() { + trtmc::AudioResult stereo; + stereo.samples = {0.25F, -0.25F, 0.5F, -0.5F}; + stereo.num_samples = static_cast(stereo.samples.size()); + stereo.sample_rate = 32000; + stereo.channels = 2; + + trtmc::VideoGenerationRequest request; + request.prompt = "ordered references"; + request.mode = trtmc::VideoGenerationMode::kReferenceToVideoAudio; + request.config.video_num_frames = 124; + + trtmc::VideoReferenceInput image; + image.kind = trtmc::VideoReferenceKind::kImage; + trtmc::VideoReferenceInput audio; + audio.kind = trtmc::VideoReferenceKind::kAudio; + audio.audio = stereo; + trtmc::VideoReferenceInput video; + video.kind = trtmc::VideoReferenceKind::kVideo; + video.video.soundtrack = stereo; + request.references = {image, audio, video}; + + if (stereo.channels != 2 || stereo.num_samples != 4 || + request.mode != trtmc::VideoGenerationMode::kReferenceToVideoAudio || + request.config.video_num_frames != 124 || request.references.size() != 3 || + request.references[0].kind != trtmc::VideoReferenceKind::kImage || + request.references[1].kind != trtmc::VideoReferenceKind::kAudio || + request.references[2].kind != trtmc::VideoReferenceKind::kVideo || + request.references[2].video.soundtrack.channels != 2) { + throw std::runtime_error("structured video request did not preserve mode or media order"); + } +} + } // namespace int main() { @@ -87,6 +121,7 @@ int main() { return 1; test_robot_observation(); + test_structured_video_generation_contract(); std::unique_ptr task = std::make_unique(); if (std::string(task->task()) != trtmc::ITextGeneration::kTask) diff --git a/core/runtime/tests/test_trt_module_dynamic_input.cpp b/core/runtime/tests/test_trt_module_dynamic_input.cpp index af3db11ad4..cfe664a42a 100644 --- a/core/runtime/tests/test_trt_module_dynamic_input.cpp +++ b/core/runtime/tests/test_trt_module_dynamic_input.cpp @@ -8,12 +8,48 @@ #include "runtime/tensorrt/trt_module_impl.h" #include +#include +#include #include #include #include +#include +#include #include #include +namespace trtmc { + +class TrtModuleImplTestPeer { + public: + static std::size_t input_capacity(const TrtModuleImpl& module, const std::string& name) { + return module.buffers_.at(name).input_capacity_bytes; + } + + static std::size_t profile_capacity(const TrtModuleImpl& module, const std::string& name) { + return module.buffers_.at(name).nbytes; + } + + static std::size_t output_capacity(const TrtModuleImpl& module, const std::string& name) { + return module.buffers_.at(name).output_capacity_bytes; + } + + static std::size_t host_staging_size(const TrtModuleImpl& module, const std::string& name) { + const auto found = module.host_output_staging_.find(name); + return found == module.host_output_staging_.end() ? 0 : found->second.size(); + } + + static bool shares_activation_arena(const TrtModuleImpl& first, const TrtModuleImpl& second) { + return first.activation_arena_ && first.activation_arena_ == second.activation_arena_; + } + + static std::int64_t activation_profile_bytes(const TrtModuleImpl& module) { + return module.engine_->getDeviceMemorySizeForProfileV2(module.profile_idx_); + } +}; + +} // namespace trtmc + namespace { int failures = 0; @@ -25,6 +61,19 @@ void check(bool condition, const char* name) { } } +// The small test context owns its activation allocation. This coordinator +// exercises the serial opt-in and records shape invalidation independently. +class RecordingActivationArena final : public trtmc::ITrtActivationArena { + public: + void attach(nvinfer1::IExecutionContext*, cudaStream_t, std::int64_t) override {} + void begin_enqueue(nvinfer1::IExecutionContext*) override {} + void end_enqueue() noexcept override {} + void invalidate_shapes(nvinfer1::IExecutionContext*) override { ++invalidations; } + void detach(nvinfer1::IExecutionContext*) noexcept override {} + + int invalidations{0}; +}; + trtmc::TrtUniquePtr build_dynamic_identity() { static trtmc::TrtLogger logger; auto builder = trtmc::TrtUniquePtr(nvinfer1::createInferBuilder(logger)); @@ -57,6 +106,8 @@ void test_first_forward_allocates_unbound_dynamic_input(nvinfer1::ICudaEngine& e trtmc::TrtModuleImpl module(&engine, engine.createExecutionContext(), stream); check(module.device_ptr("input") == nullptr, "dynamic input is not allocated during module construction"); + check(module.device_ptr("output") != nullptr, + "ordinary contexts retain eager output allocation"); float values[] = {1.0F, 2.0F}; const auto outputs = module.forward({{"input", trtmc::Tensor{values, {2}, trtmc::DType::kFloat32}}}); @@ -100,6 +151,351 @@ void test_backend_managed_graph_does_not_enable_nested_capture(nvinfer1::ICudaEn check(!module.cuda_graph_active(), "backend-managed graph leaves stream capture disabled"); } +void check_identity_output(const trtmc::TensorMap& outputs, const float* expected, + int64_t count) { + const auto found = outputs.find("output"); + check(found != outputs.end() && found->second.shape == std::vector{count}, + "lazy output returns the actual runtime shape"); + if (found == outputs.end() || found->second.numel() != count) + return; + const auto* output = static_cast(found->second.data); + for (int64_t index = 0; index < count; ++index) + check(std::fabs(output[index] - expected[index]) < 1.0e-6F, + "lazy output preserves every output value"); +} + +void test_serial_dynamic_output_grows_and_reuses(nvinfer1::ICudaEngine& engine, + cudaStream_t stream) { + auto arena = std::make_shared(); + trtmc::TrtModuleImpl module(&engine, engine.createExecutionContext(), stream, 0, nullptr, {}, + false, arena); + check(module.device_ptr("output") == nullptr, + "serial dynamic output does not allocate at profile MAX during construction"); + check(trtmc::TrtModuleImplTestPeer::host_staging_size(module, "output") == 0, + "serial dynamic output does not allocate host staging during construction"); + check(module.tensor_shape("output") == std::vector{4}, + "lazy allocation preserves profile MAX output metadata"); + bool graph_rejected = false; + try { + module.enable_cuda_graph(); + } catch (const std::logic_error&) { + graph_rejected = true; + } + check(graph_rejected, "serial live-shape outputs cannot enter stream graph capture"); + + float values[] = {2.0F, -3.0F, 4.0F, 5.0F}; + check_identity_output( + module.forward({{"input", trtmc::Tensor{values, {1}, trtmc::DType::kFloat32}}}), + values, 1); + check(trtmc::TrtModuleImplTestPeer::output_capacity(module, "output") == sizeof(float), + "first enqueue allocates only its actual output shape"); + check(trtmc::TrtModuleImplTestPeer::input_capacity(module, "input") == sizeof(float), + "first host forward allocates only its actual input shape"); + check(trtmc::TrtModuleImplTestPeer::profile_capacity(module, "input") == 4 * sizeof(float), + "actual input allocation preserves the profile MAX capacity contract"); + check(trtmc::TrtModuleImplTestPeer::host_staging_size(module, "output") == sizeof(float), + "first download allocates only its actual host output shape"); + + check_identity_output( + module.forward({{"input", trtmc::Tensor{values, {4}, trtmc::DType::kFloat32}}}), + values, 4); + void* const grown_output = module.device_ptr("output"); + void* const grown_input = module.device_ptr("input"); + check(trtmc::TrtModuleImplTestPeer::output_capacity(module, "output") == 4 * sizeof(float), + "larger enqueue grows output storage through profile MAX"); + check(trtmc::TrtModuleImplTestPeer::input_capacity(module, "input") == 4 * sizeof(float), + "larger host forward grows input storage through profile MAX"); + + values[0] = -7.0F; + check_identity_output( + module.forward({{"input", trtmc::Tensor{values, {1}, trtmc::DType::kFloat32}}}), + values, 1); + check(module.device_ptr("output") == grown_output, + "smaller enqueue reuses the high-water output allocation"); + check(module.device_ptr("input") == grown_input && + trtmc::TrtModuleImplTestPeer::input_capacity(module, "input") == 4 * sizeof(float), + "smaller host forward reuses the high-water input allocation"); + check(trtmc::TrtModuleImplTestPeer::host_staging_size(module, "output") == sizeof(float), + "smaller download reports only its actual host output shape"); + check(arena->invalidations == 3, "A to B to A invalidates activation shape requirements"); + (void)module.forward({{"input", trtmc::Tensor{values, {1}, trtmc::DType::kFloat32}}}); + check(arena->invalidations == 3, "unchanged input shape preserves activation requirement cache"); +} + +void test_serial_device_forward_defers_host_staging(nvinfer1::ICudaEngine& engine, + cudaStream_t stream) { + auto arena = std::make_shared(); + trtmc::TrtModuleImpl module(&engine, engine.createExecutionContext(), stream, 0, nullptr, {}, + false, arena); + float values[] = {8.0F, 9.0F, -10.0F, 11.0F}; + std::size_t high_water = 0; + void* grown_input = nullptr; + for (const int64_t count : {1, 4, 1}) { + trtmc::DeviceTensor input({count}, trtmc::DType::kFloat32, stream); + check(input.ok(), "device-forward input allocation succeeds"); + if (!input.ok()) + return; + const auto bytes = static_cast(count) * sizeof(float); + cudaMemcpyAsync(input.data(), values, bytes, cudaMemcpyHostToDevice, stream); + module.forward_device_async({{"input", &input}}); + module.sync(); + high_water = std::max(high_water, bytes); + check(trtmc::TrtModuleImplTestPeer::input_capacity(module, "input") == high_water, + "device forward allocates actual input capacity and reuses its high-water mark"); + if (count == 4) + grown_input = module.device_ptr("input"); + else if (grown_input) + check(module.device_ptr("input") == grown_input, + "smaller device forward preserves the grown input allocation"); + check(module.device_ptr("output") != nullptr, + "device-forward path allocates its live output before enqueue"); + check(trtmc::TrtModuleImplTestPeer::host_staging_size(module, "output") == 0, + "device-only forward never allocates host output staging"); + float output[4]{}; + cudaMemcpy(output, module.device_ptr("output"), bytes, cudaMemcpyDeviceToHost); + for (int64_t index = 0; index < count; ++index) + check(output[index] == values[index], "device-forward lazy output preserves values"); + values[0] += 1.0F; + } + check(arena->invalidations == 3, + "device A to B to A invalidates activation shape requirements"); +} + +void test_serial_input_rejects_invalid_requests_before_allocation(nvinfer1::ICudaEngine& engine, + cudaStream_t stream) { + auto arena = std::make_shared(); + trtmc::TrtModuleImpl module(&engine, engine.createExecutionContext(), stream, 0, nullptr, {}, + false, arena); + float values[] = {1.0F, 2.0F, 3.0F, 4.0F, 5.0F}; + for (const auto& shape : {std::vector{0}, std::vector{5}, + std::vector{1, 1}}) { + bool rejected = false; + try { + module.forward_async({{"input", trtmc::Tensor{values, shape, trtmc::DType::kFloat32}}}); + } catch (const std::invalid_argument&) { + rejected = true; + } + check(rejected && module.device_ptr("input") == nullptr, + "invalid dynamic input shape is rejected before allocation"); + } + bool dtype_rejected = false; + try { + module.forward_async({{"input", trtmc::Tensor{values, {1}, trtmc::DType::kInt32}}}); + } catch (const std::invalid_argument&) { + dtype_rejected = true; + } + check(dtype_rejected && module.device_ptr("input") == nullptr, + "invalid dynamic input dtype is rejected before allocation"); + check_identity_output( + module.forward({{"input", trtmc::Tensor{values, {1}, trtmc::DType::kFloat32}}}), values, 1); +} + +void test_serial_external_input_stays_external(nvinfer1::ICudaEngine& engine, + cudaStream_t stream) { + void* external = nullptr; + if (cudaMalloc(&external, 4 * sizeof(float)) != cudaSuccess) { + ++failures; + return; + } + for (const bool initial_binding : {false, true}) { + { + auto arena = std::make_shared(); + const std::vector bindings = + initial_binding ? std::vector{ + {"input", external, 4 * sizeof(float)}} + : std::vector{}; + trtmc::TrtModuleImpl module(&engine, engine.createExecutionContext(), stream, 0, + nullptr, bindings, false, arena); + if (!initial_binding) + module.bind_external("input", external, {1}); + float values[] = {-1.0F, 2.0F, 3.0F, 4.0F}; + for (const int64_t count : {1, 4, 1}) { + check_identity_output( + module.forward( + {{"input", trtmc::Tensor{values, {count}, trtmc::DType::kFloat32}}}), + values, count); + check(module.device_ptr("input") == external && + trtmc::TrtModuleImplTestPeer::input_capacity(module, "input") == 0, + "serial dynamic forward never allocates or replaces an external input"); + } + } + check(cudaMemset(external, 0, 4 * sizeof(float)) == cudaSuccess, + "destroying a serial module does not free externally owned input memory"); + } + cudaFree(external); +} + +void test_serial_external_output_stays_external(nvinfer1::ICudaEngine& engine, + cudaStream_t stream) { + void* external = nullptr; + if (cudaMalloc(&external, 4 * sizeof(float)) != cudaSuccess) { + ++failures; + return; + } + { + auto arena = std::make_shared(); + const std::vector bindings{ + {"output", external, 4 * sizeof(float)}}; + trtmc::TrtModuleImpl module(&engine, engine.createExecutionContext(), stream, 0, nullptr, + bindings, false, arena); + float values[] = {-1.0F, 2.0F, 3.0F, 4.0F}; + const auto outputs = + module.forward({{"input", trtmc::Tensor{values, {4}, trtmc::DType::kFloat32}}}); + check(module.device_ptr("output") == external, + "serial enqueue preserves the initial external output binding"); + check(outputs.count("output") == 0, + "external output is not silently downloaded or returned as host data"); + check(trtmc::TrtModuleImplTestPeer::host_staging_size(module, "output") == 0, + "external output never receives host staging"); + float output[4]{}; + cudaMemcpy(output, external, sizeof(output), cudaMemcpyDeviceToHost); + for (int index = 0; index < 4; ++index) + check(output[index] == values[index], "external output preserves all result values"); + } + check(cudaMemset(external, 0, 4 * sizeof(float)) == cudaSuccess, + "destroying a serial module does not free externally owned output memory"); + cudaFree(external); +} + +// Own only this uniquely created directory and its one generated plan. Never +// overwrite an existing path or remove anything recursively. +struct TemporaryArenaPlan { + std::filesystem::path directory; + std::filesystem::path path; + + TemporaryArenaPlan() { + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + for (int attempt = 0; attempt < 16; ++attempt) { + const auto candidate = std::filesystem::temp_directory_path() / + ("trtmc-rtx-arena-" + std::to_string(stamp) + "-" + + std::to_string(attempt)); + if (std::filesystem::create_directory(candidate)) { + directory = candidate; + path = directory / "dynamic.plan"; + return; + } + } + throw std::runtime_error("unable to create a unique test plan directory"); + } + + ~TemporaryArenaPlan() { + std::error_code ignored; + std::filesystem::remove(path, ignored); + std::filesystem::remove(directory, ignored); + } +}; + +trtmc::TrtUniquePtr build_dynamic_softmax_plan() { + static trtmc::TrtLogger logger; + auto builder = trtmc::TrtUniquePtr(nvinfer1::createInferBuilder(logger)); + if (!builder) + throw std::runtime_error("cannot create the RTX test builder"); + auto network = trtmc::TrtUniquePtr(builder->createNetworkV2(0)); + auto config = trtmc::TrtUniquePtr(builder->createBuilderConfig()); + if (!network || !config) + throw std::runtime_error("cannot create the RTX test network/config"); + auto* input = network->addInput("input", nvinfer1::DataType::kFLOAT, + nvinfer1::Dims{2, {-1, 16}}); + if (!input) + throw std::runtime_error("cannot create the RTX test input"); + auto* scores = network->addMatrixMultiply(*input, nvinfer1::MatrixOperation::kNONE, *input, + nvinfer1::MatrixOperation::kTRANSPOSE); + if (!scores) + throw std::runtime_error("cannot create the RTX test matrix product"); + auto* softmax = network->addSoftMax(*scores->getOutput(0)); + if (!softmax) + throw std::runtime_error("cannot create the RTX test softmax"); + softmax->setAxes(1U << 1U); + softmax->getOutput(0)->setName("output"); + network->markOutput(*softmax->getOutput(0)); + auto* profile = builder->createOptimizationProfile(); + if (!profile || + !profile->setDimensions("input", nvinfer1::OptProfileSelector::kMIN, + nvinfer1::Dims{2, {8, 16}}) || + !profile->setDimensions("input", nvinfer1::OptProfileSelector::kOPT, + nvinfer1::Dims{2, {32, 16}}) || + !profile->setDimensions("input", nvinfer1::OptProfileSelector::kMAX, + nvinfer1::Dims{2, {512, 16}}) || + config->addOptimizationProfile(profile) < 0) + throw std::runtime_error("cannot create the RTX test dynamic profile"); + return trtmc::TrtUniquePtr( + builder->buildSerializedNetwork(*network, *config)); +} + +void test_real_rtx_activation_arena(cudaStream_t stream) { + try { + std::unique_ptr backend( + trtmc_create_backend(), &trtmc_destroy_backend); + if (!backend) + throw std::runtime_error("cannot create the selected backend"); + if (std::string(backend->name()) != "trt_rtx") { + std::cout << "SKIP: real activation arena requires the RTX backend\n"; + return; + } + auto plan = build_dynamic_softmax_plan(); + if (!plan) + throw std::runtime_error("cannot build the real-arena test plan"); + TemporaryArenaPlan temporary; + { + std::ofstream file(temporary.path, std::ios::binary); + file.write(static_cast(plan->data()), + static_cast(plan->size())); + file.close(); + if (!file) + throw std::runtime_error("cannot write the real-arena test plan"); + } + trtmc::ModuleCreateOptions options; + options.stream = stream; + options.cuda_graphs = false; + const auto path = temporary.path.u8string(); + auto first = backend->create_module_from_file(path.c_str(), 0, plan->size(), options, {}, + -1, false, true); + auto second = backend->create_module_from_file(path.c_str(), 0, plan->size(), options, {}, + -1, false, true); + if (!first || !second || !first->ok() || !second->ok()) + throw std::runtime_error("cannot create the real-arena serial modules"); + const auto* first_impl = dynamic_cast(first.get()); + const auto* second_impl = dynamic_cast(second.get()); + if (!first_impl || !second_impl) + throw std::runtime_error("unexpected RTX module implementation"); + check(trtmc::TrtModuleImplTestPeer::shares_activation_arena(*first_impl, *second_impl), + "public RTX backend modules share a real arena on the explicit stream"); + check(trtmc::TrtModuleImplTestPeer::activation_profile_bytes(*first_impl) > 0, + "real-arena test plan requires nonzero activation memory"); + + auto run = [&](trtmc::ITrtModule& module, int64_t rows) { + std::vector values(static_cast(rows) * 16, 1.0F); + const auto outputs = module.forward( + {{"input", trtmc::Tensor{values.data(), {rows, 16}, trtmc::DType::kFloat32}}}); + const auto found = outputs.find("output"); + const bool correct_shape = + found != outputs.end() && found->second.shape == std::vector{rows, rows}; + check(correct_shape, "real arena preserves dynamic matrix/softmax output shape"); + if (!correct_shape) + return; + const auto* output = static_cast(found->second.data); + const float expected = 1.0F / static_cast(rows); + bool correct_values = true; + for (int64_t index = 0; index < rows * rows; ++index) + correct_values &= std::fabs(output[index] - expected) <= 1.0e-6F; + check(correct_values, "real arena preserves all expected softmax probabilities"); + }; + // Exercise another context raising the shared arena high-water mark + // before the first context itself has ever used the larger shape. + run(*first, 8); + run(*second, 512); + run(*first, 8); + // The first context also transitions A -> B -> A without recreation. + run(*first, 512); + run(*first, 8); + check(cudaStreamSynchronize(stream) == cudaSuccess, + "real-arena shape/context transitions finish without a CUDA error"); + } catch (const std::exception& error) { + std::cerr << "FAIL: real RTX activation arena: " << error.what() << '\n'; + ++failures; + } +} + } // namespace int main() { @@ -115,6 +511,12 @@ int main() { test_first_forward_allocates_unbound_dynamic_input(*engine, stream); test_external_dynamic_input_stays_external(*engine, stream); test_backend_managed_graph_does_not_enable_nested_capture(*engine, stream); + test_serial_dynamic_output_grows_and_reuses(*engine, stream); + test_serial_device_forward_defers_host_staging(*engine, stream); + test_serial_input_rejects_invalid_requests_before_allocation(*engine, stream); + test_serial_external_input_stays_external(*engine, stream); + test_serial_external_output_stays_external(*engine, stream); + test_real_rtx_activation_arena(stream); cudaStreamDestroy(stream); return failures; } diff --git a/families/minimax_h3/adaln_builder.py b/families/minimax_h3/adaln_builder.py index fc543c3cdd..05c9e823b8 100644 --- a/families/minimax_h3/adaln_builder.py +++ b/families/minimax_h3/adaln_builder.py @@ -12,8 +12,10 @@ import gc import sys +from pathlib import Path import tensorrt as trt +from . import trt_compat from . import graph_ops as op from .config import ( @@ -41,6 +43,7 @@ def checkpoint_keys( return tuple(names) +@op.cleanup_failed_build def build_adaln_precompute_engine( weights: dict, profile: MiniMaxH3Config, @@ -48,19 +51,24 @@ def build_adaln_precompute_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: profile.validate() logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) config = builder.create_builder_config() config.builder_optimization_level = 1 - op.configure_builder(config) - op.configure_workspace( - config, - workspace_bytes, - default_bytes=ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES, - ) + op.configure_builder(config, weight_streaming=weight_streaming) + # Let TensorRT-RTX use its native maximum pools by default. An explicit + # override remains useful for constrained build hosts and unit tests. + if workspace_bytes is not None: + op.configure_workspace( + config, + workspace_bytes, + default_bytes=ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES, + ) # Match PyTorch's default FP32 matmul policy used by the reference. config.clear_flag(trt.BuilderFlag.TF32) @@ -126,14 +134,21 @@ def build_adaln_precompute_engine( f"timesteps={profile.max_timestep_count}", file=sys.stderr, ) + plan = None + record = None try: - plan = builder.build_serialized_network(network, config) + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) finally: op.release_weight_buffers(network) if consume_weights: weights.clear() - if plan is None: + if output_path is None and plan is None: raise RuntimeError("TensorRT failed to build MiniMax-H3 AdaLN precompute engine") del network, config, builder gc.collect() - return bytes(plan) + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/audio_vae_builder.py b/families/minimax_h3/audio_vae_builder.py new file mode 100644 index 0000000000..2e25705786 --- /dev/null +++ b/families/minimax_h3/audio_vae_builder.py @@ -0,0 +1,771 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native FP32 TensorRT decoder for the MiniMax-H3 audio VAE. + +The released checkpoint uses a BigVGAN decoder. Stereo is represented as two +batch items, so this plan maps denormalized ``[2, 32, frames]`` latents to +``[2, 1, frames * 800]`` mono waveforms. Interleaving the two batch items and +applying the checkpoint's per-channel latent normalization belong to the +pipeline boundary, not to this decoder. + +TensorRT-RTX convolution layers are two-dimensional. The graph therefore +keeps a singleton height dimension and expresses every Conv1d/ConvTranspose1d +as a native 2-D convolution with a ``(1, kernel)`` kernel. No plugin or +framework runtime is required by the serialized engine. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import gc +import math +import sys +from pathlib import Path +from typing import Mapping + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .config import ( + AUDIO_LATENT_FRAMES_MAX, + AUDIO_LATENT_FRAMES_MIN, + AUDIO_LATENT_FRAMES_OPT, + AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES, +) + + +trt = trt_compat.get_trt() + + +@dataclass(frozen=True) +class AudioVAEDecoderConfig: + """Static architecture and input shape of one audio decoder plan.""" + + batch_size: int = 2 + latent_channels: int = 32 + latent_frames: int = AUDIO_LATENT_FRAMES_OPT + min_latent_frames: int = AUDIO_LATENT_FRAMES_MIN + max_latent_frames: int = AUDIO_LATENT_FRAMES_MAX + latent_dim: int = 2048 + decoder_dim: int = 1024 + decoder_rates: tuple[int, ...] = (5, 5, 2, 2, 2, 2, 2) + decoder_kernel_sizes: tuple[int, ...] = (9, 9, 4, 4, 4, 4, 4) + resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11) + resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ( + (1, 3, 5), + (1, 3, 5), + (1, 3, 5), + ) + sampling_rate: int = 32000 + + @property + def num_upsamples(self) -> int: + return len(self.decoder_rates) + + @property + def num_kernels(self) -> int: + return len(self.resblock_kernel_sizes) + + @property + def hop_length(self) -> int: + return math.prod(self.decoder_rates) + + @property + def output_samples(self) -> int: + return self.latent_frames * self.hop_length + + @property + def min_output_samples(self) -> int: + return self.min_latent_frames * self.hop_length + + @property + def max_output_samples(self) -> int: + return self.max_latent_frames * self.hop_length + + def validate(self) -> None: + integer_fields = { + "batch_size": self.batch_size, + "latent_channels": self.latent_channels, + "latent_frames": self.latent_frames, + "min_latent_frames": self.min_latent_frames, + "max_latent_frames": self.max_latent_frames, + "latent_dim": self.latent_dim, + "decoder_dim": self.decoder_dim, + "sampling_rate": self.sampling_rate, + } + if any( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + for value in integer_fields.values() + ): + raise ValueError("MiniMax-H3 AudioVAE dimensions must be positive integers") + if self.batch_size != 2: + raise ValueError("MiniMax-H3 AudioVAE stereo requires batch_size=2") + if not self.min_latent_frames <= self.latent_frames <= self.max_latent_frames: + raise ValueError( + "MiniMax-H3 AudioVAE latent frames must satisfy min <= opt <= max" + ) + if len(self.decoder_rates) != len(self.decoder_kernel_sizes): + raise ValueError("MiniMax-H3 AudioVAE decoder rates and kernels must align") + if not self.decoder_rates or not self.resblock_kernel_sizes: + raise ValueError("MiniMax-H3 AudioVAE decoder stages and residual kernels are required") + if len(self.resblock_kernel_sizes) != len(self.resblock_dilation_sizes): + raise ValueError("MiniMax-H3 AudioVAE residual kernels and dilations must align") + if any(rate <= 0 for rate in self.decoder_rates): + raise ValueError("MiniMax-H3 AudioVAE decoder rates must be positive") + if any(kernel <= 0 for kernel in self.decoder_kernel_sizes): + raise ValueError("MiniMax-H3 AudioVAE decoder kernels must be positive") + if any(kernel <= 0 or kernel % 2 == 0 for kernel in self.resblock_kernel_sizes): + raise ValueError("MiniMax-H3 AudioVAE residual kernels must be positive and odd") + if any(not dilations or any(value <= 0 for value in dilations) for dilations in self.resblock_dilation_sizes): + raise ValueError("MiniMax-H3 AudioVAE residual dilations must be positive") + divisor = 1 << self.num_upsamples + if self.decoder_dim % divisor: + raise ValueError("MiniMax-H3 AudioVAE decoder_dim must halve at every stage") + if self.decoder_dim // divisor <= 0: + raise ValueError("MiniMax-H3 AudioVAE final decoder width must be positive") + + +DEFAULT_AUDIO_VAE_DECODER_CONFIG = AudioVAEDecoderConfig() + + +def decoder_config_from_checkpoint( + raw: Mapping[str, object], + *, + latent_frames: int, + min_latent_frames: int | None = None, + max_latent_frames: int | None = None, +) -> AudioVAEDecoderConfig: + """Validate and convert the public Diffusers audio-VAE config.""" + + required = ( + "latent_channels", + "latent_dim", + "decoder_dim", + "decoder_rates", + "decoder_kernel_sizes", + "resblock_kernel_sizes", + "resblock_dilation_sizes", + "sampling_rate", + ) + missing = [name for name in required if name not in raw] + if missing: + raise ValueError(f"MiniMax-H3 AudioVAE config is missing fields: {missing}") + + try: + profile = AudioVAEDecoderConfig( + latent_channels=int(raw["latent_channels"]), + latent_frames=latent_frames, + min_latent_frames=( + latent_frames if min_latent_frames is None else min_latent_frames + ), + max_latent_frames=( + latent_frames if max_latent_frames is None else max_latent_frames + ), + latent_dim=int(raw["latent_dim"]), + decoder_dim=int(raw["decoder_dim"]), + decoder_rates=tuple(int(value) for value in raw["decoder_rates"]), + decoder_kernel_sizes=tuple(int(value) for value in raw["decoder_kernel_sizes"]), + resblock_kernel_sizes=tuple(int(value) for value in raw["resblock_kernel_sizes"]), + resblock_dilation_sizes=tuple( + tuple(int(value) for value in values) + for values in raw["resblock_dilation_sizes"] + ), + sampling_rate=int(raw["sampling_rate"]), + ) + except (TypeError, ValueError) as error: + raise ValueError("MiniMax-H3 AudioVAE config has invalid decoder fields") from error + profile.validate() + + encoder_rates = raw.get("encoder_rates") + if not isinstance(encoder_rates, (list, tuple)): + raise ValueError("MiniMax-H3 AudioVAE config is missing encoder_rates") + try: + encoder_hop = math.prod(int(value) for value in encoder_rates) + except (TypeError, ValueError) as error: + raise ValueError("MiniMax-H3 AudioVAE encoder_rates are invalid") from error + if encoder_hop != profile.hop_length: + raise ValueError( + "MiniMax-H3 AudioVAE encoder/decoder hop mismatch: " + f"encoder={encoder_hop}, decoder={profile.hop_length}" + ) + return profile + + +def checkpoint_keys(profile: AudioVAEDecoderConfig = DEFAULT_AUDIO_VAE_DECODER_CONFIG) -> tuple[str, ...]: + """Return exactly the checkpoint tensors consumed by the decoder plan.""" + + profile.validate() + names = [ + "dec_in_proj.weight", + "dec_in_proj.bias", + "decoder.conv_pre.weight_g", + "decoder.conv_pre.weight_v", + "decoder.conv_pre.bias", + ] + for stage in range(profile.num_upsamples): + prefix = f"decoder.ups.{stage}.0" + names.extend((f"{prefix}.weight_g", f"{prefix}.weight_v", f"{prefix}.bias")) + + for block in range(profile.num_upsamples * profile.num_kernels): + dilation_count = len( + profile.resblock_dilation_sizes[block % profile.num_kernels] + ) + prefix = f"decoder.resblocks.{block}" + for activation in range(2 * dilation_count): + activation_prefix = f"{prefix}.activations.{activation}" + names.extend( + ( + f"{activation_prefix}.act.alpha", + f"{activation_prefix}.act.beta", + f"{activation_prefix}.upsample.filter", + f"{activation_prefix}.downsample.lowpass.filter", + ) + ) + for group in ("convs1", "convs2"): + for index in range(dilation_count): + convolution_prefix = f"{prefix}.{group}.{index}" + names.extend( + ( + f"{convolution_prefix}.weight_g", + f"{convolution_prefix}.weight_v", + f"{convolution_prefix}.bias", + ) + ) + + names.extend( + ( + "decoder.activation_post.act.alpha", + "decoder.activation_post.act.beta", + "decoder.activation_post.upsample.filter", + "decoder.activation_post.downsample.lowpass.filter", + "decoder.conv_post.weight_g", + "decoder.conv_post.weight_v", + ) + ) + return tuple(names) + + +def _require_array(weights: Mapping[str, object], name: str, shape: tuple[int, ...]) -> np.ndarray: + try: + value = np.asarray(weights[name]) + except KeyError as error: + raise ValueError(f"MiniMax-H3 AudioVAE checkpoint is missing tensor: {name}") from error + if tuple(value.shape) != shape: + raise ValueError( + f"MiniMax-H3 AudioVAE tensor {name} has shape {tuple(value.shape)}, expected {shape}" + ) + if value.dtype != np.float32: + raise ValueError( + f"MiniMax-H3 AudioVAE tensor {name} must remain FP32, got {value.dtype}" + ) + return np.ascontiguousarray(value) + + +def fold_weight_norm(weight_g: object, weight_v: object) -> np.ndarray: + """Fold PyTorch ``weight_norm(..., dim=0)`` without changing FP32 dtype.""" + + gain = np.asarray(weight_g) + direction = np.asarray(weight_v) + if gain.dtype != np.float32 or direction.dtype != np.float32: + raise ValueError("MiniMax-H3 AudioVAE weight-normalized tensors must be FP32") + if gain.ndim != direction.ndim or gain.shape[0] != direction.shape[0]: + raise ValueError("MiniMax-H3 AudioVAE weight-normalized tensor shapes do not align") + if gain.shape[1:] != (1,) * (direction.ndim - 1): + raise ValueError("MiniMax-H3 AudioVAE weight_g must normalize dimension zero") + axes = tuple(range(1, direction.ndim)) + norm = np.sqrt( + np.sum(direction * direction, axis=axes, keepdims=True, dtype=np.float32) + ) + if np.any(norm == 0.0): + raise ValueError("MiniMax-H3 AudioVAE weight_v contains a zero-norm filter") + return np.ascontiguousarray(direction * (gain / norm), dtype=np.float32) + + +def _fold_checkpoint_weight( + weights: Mapping[str, object], prefix: str, shape: tuple[int, ...] +) -> np.ndarray: + gain_shape = (shape[0],) + (1,) * (len(shape) - 1) + gain = _require_array(weights, f"{prefix}.weight_g", gain_shape) + direction = _require_array(weights, f"{prefix}.weight_v", shape) + return fold_weight_norm(gain, direction) + + +def _conv1d( + network, + hidden, + weight: np.ndarray, + bias: np.ndarray | None, + *, + stride: int = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + name: str, + owned_weights: list[np.ndarray], +): + out_channels, in_channels_per_group, kernel_size = weight.shape + kernel = np.ascontiguousarray(weight[:, :, None, :], dtype=np.float32) + owned_weights.append(kernel) + if bias is not None: + bias = np.ascontiguousarray(bias, dtype=np.float32) + owned_weights.append(bias) + layer = network.add_convolution_nd( + hidden, + out_channels, + (1, kernel_size), + kernel, + bias, + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE convolution {name}") + layer.name = name + layer.stride_nd = (1, stride) + layer.padding_nd = (0, padding) + layer.dilation_nd = (1, dilation) + layer.num_groups = groups + expected_inputs = int(tuple(hidden.shape)[1]) // groups + if in_channels_per_group != expected_inputs: + raise ValueError( + f"MiniMax-H3 AudioVAE convolution {name} expects {in_channels_per_group} " + f"channels per group, got {expected_inputs}" + ) + return layer.get_output(0) + + +def _deconv1d( + network, + hidden, + weight: np.ndarray, + bias: np.ndarray | None, + *, + out_channels: int, + stride: int = 1, + padding: int = 0, + groups: int = 1, + name: str, + owned_weights: list[np.ndarray], +): + in_channels, out_channels_per_group, kernel_size = weight.shape + kernel = np.ascontiguousarray(weight[:, :, None, :], dtype=np.float32) + owned_weights.append(kernel) + if bias is not None: + bias = np.ascontiguousarray(bias, dtype=np.float32) + owned_weights.append(bias) + if in_channels != int(tuple(hidden.shape)[1]): + raise ValueError( + f"MiniMax-H3 AudioVAE deconvolution {name} expects {in_channels} input channels, " + f"got {tuple(hidden.shape)[1]}" + ) + if out_channels_per_group * groups != out_channels: + raise ValueError( + f"MiniMax-H3 AudioVAE deconvolution {name} has an invalid grouped output width" + ) + layer = network.add_deconvolution_nd( + hidden, + out_channels, + (1, kernel_size), + kernel, + bias, + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE deconvolution {name}") + layer.name = name + layer.stride_nd = (1, stride) + layer.padding_nd = (0, padding) + layer.num_groups = groups + return layer.get_output(0) + + +def _replicate_pad(network, hidden, left: int, right: int, *, name: str): + batch, channels, height, _width = tuple(int(value) for value in hidden.shape) + layer = network.add_slice( + hidden, + (0, 0, 0, -left), + (batch, channels, height, 1), + (1, 1, 1, 1), + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE replicate pad {name}") + layer.name = name + layer.mode = trt.SampleMode.CLAMP + shape = network.add_shape(hidden).get_output(0) + delta = op.constant( + network, + np.asarray([0, 0, 0, left + right], dtype=np.int64), + dtype=np.int64, + ) + size = network.add_elementwise(shape, delta, trt.ElementWiseOperation.SUM).get_output(0) + layer.set_input(2, size) + return layer.get_output(0) + + +def _crop_time(network, hidden, left: int, right: int, *, name: str): + batch, channels, height, _width = tuple(int(value) for value in hidden.shape) + layer = network.add_slice( + hidden, + (0, 0, 0, left), + (batch, channels, height, 1), + (1, 1, 1, 1), + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE crop {name}") + layer.name = name + shape = network.add_shape(hidden).get_output(0) + delta = op.constant( + network, + np.asarray([0, 0, 0, -(left + right)], dtype=np.int64), + dtype=np.int64, + ) + size = network.add_elementwise(shape, delta, trt.ElementWiseOperation.SUM).get_output(0) + layer.set_input(2, size) + return layer.get_output(0) + + +def _snake_beta(network, hidden, weights: Mapping[str, object], prefix: str): + channels = int(tuple(hidden.shape)[1]) + alpha = _require_array(weights, f"{prefix}.alpha", (channels,)).reshape(1, channels, 1, 1) + beta = _require_array(weights, f"{prefix}.beta", (channels,)).reshape(1, channels, 1, 1) + alpha = network.add_unary(op.weight_constant(network, alpha), trt.UnaryOperation.EXP).get_output(0) + beta = network.add_unary(op.weight_constant(network, beta), trt.UnaryOperation.EXP).get_output(0) + phase = network.add_elementwise(hidden, alpha, trt.ElementWiseOperation.PROD).get_output(0) + sine = network.add_unary(phase, trt.UnaryOperation.SIN).get_output(0) + sine_square = network.add_elementwise(sine, sine, trt.ElementWiseOperation.PROD).get_output(0) + epsilon = op.constant(network, np.full((1, 1, 1, 1), 1.0e-9, np.float32)) + denominator = network.add_elementwise(beta, epsilon, trt.ElementWiseOperation.SUM).get_output(0) + periodic = network.add_elementwise( + sine_square, denominator, trt.ElementWiseOperation.DIV + ).get_output(0) + return network.add_elementwise(hidden, periodic, trt.ElementWiseOperation.SUM).get_output(0) + + +def _alias_free_activation( + network, + hidden, + weights: Mapping[str, object], + prefix: str, + *, + owned_weights: list[np.ndarray], +): + ratio = 2 + kernel_size = 12 + channels = int(tuple(hidden.shape)[1]) + pad = kernel_size // ratio - 1 + pad_left = pad * ratio + (kernel_size - ratio) // 2 + pad_right = pad * ratio + (kernel_size - ratio + 1) // 2 + + padded = _replicate_pad(network, hidden, pad, pad, name=f"{prefix}.upsample.pad") + upsample_filter = _require_array(weights, f"{prefix}.upsample.filter", (1, 1, kernel_size)) + upsample_filter = np.broadcast_to( + upsample_filter, (channels, 1, kernel_size) + ).copy() + hidden = _deconv1d( + network, + padded, + upsample_filter, + None, + out_channels=channels, + stride=ratio, + groups=channels, + name=f"{prefix}.upsample", + owned_weights=owned_weights, + ) + scale = op.constant(network, np.full((1, 1, 1, 1), float(ratio), np.float32)) + hidden = network.add_elementwise(hidden, scale, trt.ElementWiseOperation.PROD).get_output(0) + hidden = _crop_time( + network, + hidden, + pad_left, + pad_right, + name=f"{prefix}.upsample.crop", + ) + hidden = _snake_beta(network, hidden, weights, f"{prefix}.act") + + even = kernel_size % 2 == 0 + down_left = kernel_size // 2 - int(even) + down_right = kernel_size // 2 + hidden = _replicate_pad( + network, + hidden, + down_left, + down_right, + name=f"{prefix}.downsample.pad", + ) + downsample_filter = _require_array( + weights, f"{prefix}.downsample.lowpass.filter", (1, 1, kernel_size) + ) + downsample_filter = np.broadcast_to( + downsample_filter, (channels, 1, kernel_size) + ).copy() + return _conv1d( + network, + hidden, + downsample_filter, + None, + stride=ratio, + groups=channels, + name=f"{prefix}.downsample", + owned_weights=owned_weights, + ) + + +def _amp_block( + network, + hidden, + weights: Mapping[str, object], + prefix: str, + *, + kernel_size: int, + dilations: tuple[int, ...], + owned_weights: list[np.ndarray], +): + channels = int(tuple(hidden.shape)[1]) + for index, dilation in enumerate(dilations): + residual = _alias_free_activation( + network, + hidden, + weights, + f"{prefix}.activations.{2 * index}", + owned_weights=owned_weights, + ) + conv1_prefix = f"{prefix}.convs1.{index}" + conv1_weight = _fold_checkpoint_weight( + weights, conv1_prefix, (channels, channels, kernel_size) + ) + conv1_bias = _require_array(weights, f"{conv1_prefix}.bias", (channels,)) + residual = _conv1d( + network, + residual, + conv1_weight, + conv1_bias, + padding=(kernel_size * dilation - dilation) // 2, + dilation=dilation, + name=conv1_prefix, + owned_weights=owned_weights, + ) + residual = _alias_free_activation( + network, + residual, + weights, + f"{prefix}.activations.{2 * index + 1}", + owned_weights=owned_weights, + ) + conv2_prefix = f"{prefix}.convs2.{index}" + conv2_weight = _fold_checkpoint_weight( + weights, conv2_prefix, (channels, channels, kernel_size) + ) + conv2_bias = _require_array(weights, f"{conv2_prefix}.bias", (channels,)) + residual = _conv1d( + network, + residual, + conv2_weight, + conv2_bias, + padding=(kernel_size - 1) // 2, + name=conv2_prefix, + owned_weights=owned_weights, + ) + hidden = network.add_elementwise( + hidden, residual, trt.ElementWiseOperation.SUM + ).get_output(0) + return hidden + + +@op.cleanup_failed_build +def build_audio_vae_decoder_engine( + weights: dict, + profile: AudioVAEDecoderConfig = DEFAULT_AUDIO_VAE_DECODER_CONFIG, + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the released MiniMax-H3 BigVGAN decoder as a native FP32 plan.""" + + profile.validate() + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES, + ) + missing = sorted(set(checkpoint_keys(profile)) - set(weights)) + if missing: + raise ValueError(f"MiniMax-H3 AudioVAE checkpoint is missing tensors: {missing}") + owned_weights: list[np.ndarray] = [] + + latents = network.add_input( + "audio_latents", + trt.float32, + (profile.batch_size, profile.latent_channels, -1), + ) + shape_profile = builder.create_optimization_profile() + shape_profile.set_shape( + "audio_latents", + (profile.batch_size, profile.latent_channels, profile.min_latent_frames), + (profile.batch_size, profile.latent_channels, profile.latent_frames), + (profile.batch_size, profile.latent_channels, profile.max_latent_frames), + ) + profile_index = config.add_optimization_profile(shape_profile) + if profile_index != 0: + raise RuntimeError("TensorRT rejected MiniMax-H3 AudioVAE dynamic shape profile") + expanded = network.add_shuffle(latents) + expanded.reshape_dims = ( + profile.batch_size, + profile.latent_channels, + 1, + -1, + ) + hidden = expanded.get_output(0) + + dec_in_weight = _require_array( + weights, + "dec_in_proj.weight", + (profile.latent_dim, profile.latent_channels, 1), + ) + dec_in_bias = _require_array(weights, "dec_in_proj.bias", (profile.latent_dim,)) + hidden = _conv1d( + network, + hidden, + dec_in_weight, + dec_in_bias, + name="dec_in_proj", + owned_weights=owned_weights, + ) + + conv_pre_weight = _fold_checkpoint_weight( + weights, + "decoder.conv_pre", + (profile.decoder_dim, profile.latent_dim, 7), + ) + conv_pre_bias = _require_array(weights, "decoder.conv_pre.bias", (profile.decoder_dim,)) + hidden = _conv1d( + network, + hidden, + conv_pre_weight, + conv_pre_bias, + padding=3, + name="decoder.conv_pre", + owned_weights=owned_weights, + ) + + for stage, (rate, kernel_size) in enumerate( + zip(profile.decoder_rates, profile.decoder_kernel_sizes) + ): + input_channels = profile.decoder_dim // (1 << stage) + output_channels = profile.decoder_dim // (1 << (stage + 1)) + upsample_prefix = f"decoder.ups.{stage}.0" + upsample_weight = _fold_checkpoint_weight( + weights, + upsample_prefix, + (input_channels, output_channels, kernel_size), + ) + upsample_bias = _require_array(weights, f"{upsample_prefix}.bias", (output_channels,)) + hidden = _deconv1d( + network, + hidden, + upsample_weight, + upsample_bias, + out_channels=output_channels, + stride=rate, + padding=(kernel_size - rate) // 2, + name=upsample_prefix, + owned_weights=owned_weights, + ) + + block_outputs = [] + for kernel_index, (resblock_kernel, dilations) in enumerate( + zip(profile.resblock_kernel_sizes, profile.resblock_dilation_sizes) + ): + block_index = stage * profile.num_kernels + kernel_index + block_outputs.append( + _amp_block( + network, + hidden, + weights, + f"decoder.resblocks.{block_index}", + kernel_size=resblock_kernel, + dilations=dilations, + owned_weights=owned_weights, + ) + ) + combined = block_outputs[0] + for block_output in block_outputs[1:]: + combined = network.add_elementwise( + combined, block_output, trt.ElementWiseOperation.SUM + ).get_output(0) + divisor = op.constant( + network, np.full((1, 1, 1, 1), float(profile.num_kernels), np.float32) + ) + hidden = network.add_elementwise( + combined, divisor, trt.ElementWiseOperation.DIV + ).get_output(0) + + hidden = _alias_free_activation( + network, + hidden, + weights, + "decoder.activation_post", + owned_weights=owned_weights, + ) + final_channels = profile.decoder_dim // (1 << profile.num_upsamples) + conv_post_weight = _fold_checkpoint_weight( + weights, + "decoder.conv_post", + (1, final_channels, 7), + ) + hidden = _conv1d( + network, + hidden, + conv_post_weight, + None, + padding=3, + name="decoder.conv_post", + owned_weights=owned_weights, + ) + clamp = network.add_activation(hidden, trt.ActivationType.CLIP) + if clamp is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 AudioVAE output clamp") + clamp.alpha = -1.0 + clamp.beta = 1.0 + output = network.add_shuffle(clamp.get_output(0)) + output.reshape_dims = (profile.batch_size, 1, -1) + decoded = output.get_output(0) + decoded.name = "decoded_audio" + network.mark_output(decoded) + op.validate_native_network(network, expected_attentions=0, label="AudioVAE decoder") + + print( + "[minimax-h3] building native FP32 AudioVAE decoder: " + f"latents={profile.min_latent_frames}/{profile.latent_frames}/" + f"{profile.max_latent_frames}, samples={profile.min_output_samples}/" + f"{profile.output_samples}/{profile.max_output_samples}, " + f"upsamples={profile.num_upsamples}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 AudioVAE decoder") + del owned_weights, network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/checkpoint.py b/families/minimax_h3/checkpoint.py index 8221812c45..c9f489c548 100644 --- a/families/minimax_h3/checkpoint.py +++ b/families/minimax_h3/checkpoint.py @@ -24,56 +24,45 @@ def __array_finalize__(self, source) -> None: self._tensor_owner = getattr(source, "_tensor_owner", None) -def load_component_state_dict(component_dir: str | Path) -> dict[str, Any]: - """Load a safetensors component without materializing duplicate tensors.""" - - from safetensors.torch import load_file - - root = Path(component_dir) - indexes = sorted(root.glob("*.safetensors.index.json")) - if indexes: - if len(indexes) != 1: - raise ValueError(f"Expected one safetensors index in {root}, found {len(indexes)}") - weight_map = json.loads(indexes[0].read_text())["weight_map"] - paths = [root / name for name in sorted(set(weight_map.values()))] - else: - paths = sorted(root.glob("*.safetensors")) - if not paths: - raise FileNotFoundError(f"No safetensors checkpoint found in {root}") - - state: dict[str, Any] = {} - for path in paths: - for name, tensor in load_file(path, device="cpu").items(): - if name in state: - raise ValueError(f"Duplicate MiniMax-H3 tensor {name!r}") - state[name] = tensor - return state - - def load_selected_component_state_dict( component_dir: str | Path, names: Iterable[str] ) -> dict[str, Any]: - """Load only selected indexed tensors, avoiding unused H3 language/vision weights.""" + """Load selected tensors from one indexed or unsharded component.""" from safetensors import safe_open root = Path(component_dir) indexes = sorted(root.glob("*.safetensors.index.json")) - if len(indexes) != 1: - raise ValueError(f"Selective loading requires one safetensors index in {root}") - weight_map = json.loads(indexes[0].read_text())["weight_map"] requested = tuple(names) - missing = sorted(set(requested) - set(weight_map)) - if missing: - raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") - by_file: dict[str, list[str]] = {} - for name in requested: - by_file.setdefault(weight_map[name], []).append(name) state: dict[str, Any] = {} - for filename, tensor_names in sorted(by_file.items()): - with safe_open(root / filename, framework="pt", device="cpu") as reader: - for name in tensor_names: - state[name] = reader.get_tensor(name) + if indexes: + if len(indexes) != 1: + raise ValueError(f"Selective loading requires one safetensors index in {root}") + weight_map = json.loads(indexes[0].read_text())["weight_map"] + missing = sorted(set(requested) - set(weight_map)) + if missing: + raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") + by_file: dict[str, list[str]] = {} + for name in requested: + by_file.setdefault(weight_map[name], []).append(name) + for filename, tensor_names in sorted(by_file.items()): + with safe_open(root / filename, framework="pt", device="cpu") as reader: + for name in tensor_names: + state[name] = reader.get_tensor(name) + return state + + paths = sorted(root.glob("*.safetensors")) + if len(paths) != 1: + raise ValueError( + f"Selective loading requires one safetensors index or one unsharded file in {root}" + ) + with safe_open(paths[0], framework="pt", device="cpu") as reader: + available = set(reader.keys()) + missing = sorted(set(requested) - available) + if missing: + raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") + for name in requested: + state[name] = reader.get_tensor(name) return state @@ -104,12 +93,6 @@ def validate_component_key_partition( ) -def require_keys(state: dict[str, Any], names: Iterable[str]) -> None: - missing = sorted(set(names) - set(state)) - if missing: - raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") - - def numpy_state(state: dict[str, Any]) -> dict[str, Any]: """Expose CPU tensors to NumPy without expanding checkpoint-native BF16. diff --git a/families/minimax_h3/config.py b/families/minimax_h3/config.py index d038603de5..efdd8c067c 100644 --- a/families/minimax_h3/config.py +++ b/families/minimax_h3/config.py @@ -3,11 +3,12 @@ """Validated native TensorRT profiles for MiniMax-H3. -The default profile is the 124-frame, 1344x768 shape used by the public -Sol-Engine H3 benchmark. Structural row counts are explicit because prompt -packing is part of the engine ABI and must match the Hugging Face reference. -Text rows are dynamic; ``text_rows`` remains the maximum for compatible -bundle metadata. +Production uses one dynamic profile for the complete released 5--15 second +T2VA and FL2VA media envelope. It supports dynamic prompt lengths, the public +continuous 1:4--4:1 canvas resolver, and the documented explicit performance +canvases, including the optional landscape 480p super-resolution source geometry. +Structural row counts are explicit because prompt/media packing is part of +the engine ABI and must match the Hugging Face reference. """ from __future__ import annotations @@ -16,15 +17,49 @@ TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +VISION_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES = 64 << 30 DENOISER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 VAE_TILE_DECODER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +AUDIO_LATENT_FRAMES_MIN = 207 +AUDIO_LATENT_FRAMES_OPT = 207 +AUDIO_LATENT_FRAMES_MAX = 575 +VIDEO_NUM_FRAMES_MIN = 124 +VIDEO_NUM_FRAMES_OPT = 124 +VIDEO_NUM_FRAMES_MAX = 345 +VIDEO_ROWS_MIN = 14_985 +VIDEO_ROWS_OPT = 37_296 +VIDEO_ROWS_MAX = 108_576 +CANVAS_MULTIPLE = 32 +CANVAS_SHORT_EDGE = 768 +CANVAS_MAX_PIXELS = 768 * 1344 +CANVAS_MIN_ASPECT_RATIO = 0.25 +CANVAS_MAX_ASPECT_RATIO = 4.0 +# Extra explicit performance canvases, stored as (height, width). The TensorRT +# runtime remains a finite allowlist: these orientations are in addition to, +# not a replacement for, the 95 resolver-produced canvases. The landscape +# 480p canvas is the source geometry for the optional native SR stage. +NATIVE_EXPLICIT_CANVAS_SIZES = ( + (480, 864), + (544, 960), + (960, 544), +) +# The RTX path builds each plan in a fresh process, so one conservative +# workspace and runtime budget cover every stage without coupling the public +# artifact to a particular workstation identity. +RTX_STAGED_WORKSPACE_BYTES = 16 << 30 +RTX_WEIGHT_STREAMING_BUDGET_BYTES = 32 << 30 +TRT_DEFAULT_WORKSPACE_POLICY = "trt_default_max" DEFAULT_WORKSPACE_LIMIT_BYTES = { "text_encoder.plan": TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES, + "vision_encoder.plan": VISION_ENCODER_DEFAULT_WORKSPACE_BYTES, "adaln_precompute.plan": ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES, - "denoiser.plan": DENOISER_DEFAULT_WORKSPACE_BYTES, + "fl2va_keyframe_vae_encoder.plan": KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES, "vae_tile_decoder.plan": VAE_TILE_DECODER_DEFAULT_WORKSPACE_BYTES, + "audio_vae_decoder.plan": AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES, } FIRST_BLOCK_CACHE_DENOISER_PLAN_FILENAMES = ( @@ -34,32 +69,30 @@ ) -def native_plan_filenames(*, first_block_cache: bool) -> tuple[str, ...]: - """Return the exact plan set selected by the native denoiser profile.""" +def native_plan_filenames() -> tuple[str, ...]: + """Return the singular original-weight dense FirstBlockCache plan set.""" - if not isinstance(first_block_cache, bool): - raise ValueError("MiniMax-H3 first_block_cache must be a boolean") - denoiser_plans = ( - FIRST_BLOCK_CACHE_DENOISER_PLAN_FILENAMES if first_block_cache else ("denoiser.plan",) - ) return ( "text_encoder.plan", + "vision_encoder.plan", "adaln_precompute.plan", - *denoiser_plans, + *FIRST_BLOCK_CACHE_DENOISER_PLAN_FILENAMES, + "fl2va_keyframe_vae_encoder.plan", "vae_tile_decoder.plan", + "audio_vae_decoder.plan", ) -def default_workspace_limit_bytes(*, first_block_cache: bool) -> dict[str, int]: - """Return per-plan tactic workspace limits for one denoiser layout.""" +def default_workspace_limit_bytes() -> dict[str, int | str]: + """Return per-plan tactic workspace limits for the dense FBC layout.""" return { filename: ( - DENOISER_DEFAULT_WORKSPACE_BYTES - if filename.startswith("denoiser_") or filename == "denoiser.plan" + TRT_DEFAULT_WORKSPACE_POLICY + if filename.startswith("denoiser_") or filename.startswith("adaln_") else DEFAULT_WORKSPACE_LIMIT_BYTES[filename] ) - for filename in native_plan_filenames(first_block_cache=first_block_cache) + for filename in native_plan_filenames() } @@ -89,7 +122,11 @@ class MiniMaxH3Config: timestep_embed_dim: int = 2688 rope_freq_dim: int = 16 norm_eps: float = 1.0e-5 + min_video_rows: int = 37296 + opt_video_rows: int = 37296 video_rows: int = 37296 + min_audio_rows: int = 414 + opt_audio_rows: int = 414 audio_rows: int = 414 min_text_rows: int = 1 opt_text_rows: int = 128 @@ -97,7 +134,7 @@ class MiniMaxH3Config: padded_sequence_length: int = 38247 max_timestep_count: int = 4 context_parallel_size: int = 1 - first_block_cache: bool = False + first_block_cache: bool = True @property def sequence_length(self) -> int: @@ -105,11 +142,27 @@ def sequence_length(self) -> int: @property def min_sequence_length(self) -> int: - return self.video_rows + self.audio_rows + self.min_text_rows + return self.min_video_rows + self.min_audio_rows + self.min_text_rows @property def opt_sequence_length(self) -> int: - return self.video_rows + self.audio_rows + self.opt_text_rows + return self.opt_video_rows + self.opt_audio_rows + self.opt_text_rows + + @property + def video_row_profile(self) -> tuple[int, int, int]: + return self.min_video_rows, self.opt_video_rows, self.video_rows + + @property + def audio_row_profile(self) -> tuple[int, int, int]: + return self.min_audio_rows, self.opt_audio_rows, self.audio_rows + + @property + def packed_row_profile(self) -> tuple[int, int, int]: + return self.min_sequence_length, self.opt_sequence_length, self.sequence_length + + @property + def text_row_profile(self) -> tuple[int, int, int]: + return self.min_text_rows, self.opt_text_rows, self.text_rows @property def padding_rows(self) -> int: @@ -137,6 +190,12 @@ def validate(self) -> None: raise ValueError("MiniMax-H3 attention width must exceed its residual width") if not 1 <= self.min_text_rows <= self.opt_text_rows <= self.text_rows: raise ValueError("MiniMax-H3 text rows must satisfy 1 <= min <= opt <= max") + if not 0 < self.min_video_rows <= self.opt_video_rows <= self.video_rows: + raise ValueError("MiniMax-H3 video rows must satisfy 0 < min <= opt <= max") + if not 0 < self.min_audio_rows <= self.opt_audio_rows <= self.audio_rows: + raise ValueError("MiniMax-H3 audio rows must satisfy 0 < min <= opt <= max") + if any(rows % 2 for rows in self.audio_row_profile): + raise ValueError("MiniMax-H3 audio_rows must contain two equal stereo channels") if self.sequence_length != self.padded_sequence_length: raise ValueError( "MiniMax-H3 requires no packed-sequence padding: " @@ -149,3 +208,18 @@ def validate(self) -> None: SOL_ENGINE_1344X768_124F = MiniMaxH3Config() + +# Each production FirstBlockCache engine has one profile covering the complete +# public 5--15 second envelope, including dynamic prompts and keyframes. +# The released local pipeline aligns requested frame counts to ``17 * n + 5``. +# At 24 fps its supported 5--15 second endpoints are therefore 124 and 345 +# frames. Video tokens use 1,008 rows per latent frame at 1344x768; audio is +# packed as two stereo row groups of 207 through 575 latent frames. +SOL_ENGINE_1344X768_124_TO_345F = MiniMaxH3Config( + min_video_rows=VIDEO_ROWS_MIN, + opt_video_rows=VIDEO_ROWS_OPT, + video_rows=VIDEO_ROWS_MAX, + audio_rows=1150, + text_rows=2641, + padded_sequence_length=112367, +) diff --git a/families/minimax_h3/delivery.py b/families/minimax_h3/delivery.py new file mode 100644 index 0000000000..d3d04603b2 --- /dev/null +++ b/families/minimax_h3/delivery.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time source resolution for the single quantized H3 delivery.""" + +from __future__ import annotations + +from pathlib import Path + + +COMFY_REPOSITORY = "Comfy-Org/MiniMax-H3" +COMFY_REVISION = "4cc1d817b6184899b41293954329f576cb5ae86b" +SR_SOURCE_SHAPE = (480, 864) +QUANTIZED_SOURCES = { + "quantized_transformer": "diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors", + "quantized_ref_transformer": "diffusion_models/minimax_h3_ref2va_int8_convrot.safetensors", + "quantized_text_encoder": "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", +} + + +def resolve_quantized_sources(model_dir: Path, options: dict) -> dict[str, Path]: + """Prefer explicit/offline files, otherwise use the pinned public HF cache.""" + + from huggingface_hub import hf_hub_download + + for option in QUANTIZED_SOURCES: + if options.get(option) and not Path(options[option]).is_file(): + raise FileNotFoundError(f"MiniMax-H3 {option} checkpoint is missing: {options[option]}") + sources = {} + for option, filename in QUANTIZED_SOURCES.items(): + explicit = options.get(option) + local = Path(model_dir) / filename + if explicit: + source = Path(explicit).absolute() + elif local.is_file(): + source = local.absolute() + else: + source = Path(hf_hub_download(COMFY_REPOSITORY, filename, revision=COMFY_REVISION)) + if not source.is_file(): + raise FileNotFoundError(f"MiniMax-H3 {option} checkpoint is missing: {source}") + # Preserve the public filename when HF cache entries are symlinks. + sources[option] = source.absolute() + return sources + + +def resolve_super_resolution_sources(options: dict) -> tuple[Path | None, Path | None]: + """Only the explicit build flag may enable the fixed-base SR delivery.""" + + enabled = options.get("super_resolution", False) + keys = ("super_resolution_model", "super_resolution_weak_model") + if not enabled: + if any(options.get(key) for key in keys): + raise ValueError("MiniMax-H3 SR checkpoint overrides require super_resolution=true") + return None, None + from torch.hub import download_url_to_file, get_dir + + result = [] + for key, filename in zip(keys, ("realesr-general-x4v3.pth", "realesr-general-wdn-x4v3.pth")): + explicit = options.get(key) + path = Path(explicit).absolute() if explicit else Path(get_dir()) / "checkpoints" / filename + if explicit and not path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 SR checkpoint is missing: {path}") + if not path.is_file(): + path.parent.mkdir(parents=True, exist_ok=True) + download_url_to_file( + "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/" + filename, + str(path), + ) + result.append(path.absolute()) + return result[0], result[1] diff --git a/families/minimax_h3/dit_builder.py b/families/minimax_h3/dit_builder.py index 0113411375..94e781df8e 100644 --- a/families/minimax_h3/dit_builder.py +++ b/families/minimax_h3/dit_builder.py @@ -8,10 +8,12 @@ import gc import math import sys +from pathlib import Path import numpy as np import tensorrt as trt +from . import trt_compat from . import graph_ops as op from .config import ( @@ -85,9 +87,7 @@ def tail_checkpoint_keys( return _block_checkpoint_keys(range(1, profile.num_layers)) -def finish_checkpoint_keys( - profile: MiniMaxH3Config = SOL_ENGINE_1344X768_124F, -) -> tuple[str, ...]: +def finish_checkpoint_keys() -> tuple[str, ...]: """Weights used by the final norm and modality-specific projections.""" return ( @@ -107,11 +107,11 @@ def checkpoint_keys( return ( *head_checkpoint_keys(profile), *tail_checkpoint_keys(profile), - *finish_checkpoint_keys(profile), + *finish_checkpoint_keys(), ) -def _slice_modulation(network, selected, index: int, rows: int, width: int): +def _slice_modulation(network, selected, index: int, width: int): value = op.dynamic_slice(network, selected, (0, index, 0), (None, 1, width)) reshape = network.add_shuffle(value) reshape.reshape_dims = (-1, width) @@ -129,7 +129,7 @@ def _per_head_norm(network, tensor, weight, profile: MiniMaxH3Config, rows: int) return flatten.get_output(0) -def _rope_tables(network, position_ids, profile: MiniMaxH3Config, rows: int): +def _rope_tables(network, position_ids, profile: MiniMaxH3Config): positions = op.cast(network, position_ids, trt.float32) position_shape = network.add_shuffle(positions) position_shape.reshape_dims = (-1, 3, 1) @@ -152,9 +152,9 @@ def _rope_tables(network, position_ids, profile: MiniMaxH3Config, rows: int): def _native_attention(network, q, k, v, *, rows: int, profile: MiniMaxH3Config, name: str): - q4 = op.rows_to_heads(network, q, rows, profile.num_heads, profile.head_dim) - k4 = op.rows_to_heads(network, k, rows, profile.num_heads, profile.head_dim) - v4 = op.rows_to_heads(network, v, rows, profile.num_heads, profile.head_dim) + q4 = op.rows_to_heads(network, q, profile.num_heads, profile.head_dim) + k4 = op.rows_to_heads(network, k, profile.num_heads, profile.head_dim) + v4 = op.rows_to_heads(network, v, profile.num_heads, profile.head_dim) scale = op.constant( network, np.full((1, 1, 1, 1), 1.0 / math.sqrt(profile.head_dim), dtype=np.float32), @@ -168,7 +168,7 @@ def _native_attention(network, q, k, v, *, rows: int, profile: MiniMaxH3Config, attention.metadata = f"trtmc.native_op=IAttention;source={name}" attention.get_output(0).name = f"{name}.output" attention.decomposable = False - return op.heads_to_rows(network, attention.get_output(0), rows, profile.attention_size) + return op.heads_to_rows(network, attention.get_output(0), profile.attention_size) def _attention_block( @@ -181,8 +181,15 @@ def _attention_block( *, cos=None, sin=None, + consume_weights: bool = False, ): - q, k, v = op.fused_qkv(network, hidden, weights, f"{prefix}.attn") + q, k, v = op.fused_qkv( + network, + hidden, + weights, + f"{prefix}.attn", + consume_weights=consume_weights, + ) q = _per_head_norm(network, q, weights[f"{prefix}.attn.norm_q.weight"], profile, rows) k = _per_head_norm(network, k, weights[f"{prefix}.attn.norm_k.weight"], profile, rows) if cos is not None: @@ -192,7 +199,6 @@ def _attention_block( q, cos, sin, - rows=rows, heads=profile.num_heads, head_dim=profile.head_dim, rotary_dim=rotary_dim, @@ -202,7 +208,6 @@ def _attention_block( k, cos, sin, - rows=rows, heads=profile.num_heads, head_dim=profile.head_dim, rotary_dim=rotary_dim, @@ -212,7 +217,6 @@ def _attention_block( q, k, v, - rows=rows, heads=profile.num_heads, head_dim=profile.head_dim, name=f"{prefix}.attn.native_attention", @@ -230,7 +234,14 @@ def _attention_block( return op.linear(network, attended, weights[f"{prefix}.attn.to_out.0.weight"]) -def _refine_text(network, text, weights, profile: MiniMaxH3Config): +def _refine_text( + network, + text, + weights, + profile: MiniMaxH3Config, + *, + consume_weights: bool = False, +): hidden = op.linear( network, text, weights["context_embedder.weight"], weights["context_embedder.bias"] ) @@ -244,7 +255,15 @@ def _refine_text(network, text, weights, profile: MiniMaxH3Config): profile.hidden_size, profile.norm_eps, ) - update = _attention_block(network, normalized, weights, prefix, profile, rows) + update = _attention_block( + network, + normalized, + weights, + prefix, + profile, + rows, + consume_weights=consume_weights, + ) hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) normalized = op.rms_norm( network, @@ -270,10 +289,25 @@ def _refine_text(network, text, weights, profile: MiniMaxH3Config): ) -def _packed_hidden(network, video, audio, text, weights, profile: MiniMaxH3Config): +def _packed_hidden( + network, + video, + audio, + text, + weights, + profile: MiniMaxH3Config, + *, + consume_weights: bool = False, +): """Project and pack text | audio | video exactly like the Diffusers model.""" - text_hidden = _refine_text(network, text, weights, profile) + text_hidden = _refine_text( + network, + text, + weights, + profile, + consume_weights=consume_weights, + ) audio_hidden = op.linear( network, audio, weights["audio_proj_in.weight"], weights["audio_proj_in.bias"], bf16=False ) @@ -297,6 +331,8 @@ def _transformer_block( weights, profile: MiniMaxH3Config, index: int, + *, + consume_weights: bool = False, ): """Add one native H3 transformer block and return its residual stream.""" @@ -304,7 +340,7 @@ def _transformer_block( prefix = f"transformer_blocks.{index}" selected = op.gather_rows(network, block_modulation, adaln_indices) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - _slice_modulation(network, selected, part, rows, profile.hidden_size) for part in range(6) + _slice_modulation(network, selected, part, profile.hidden_size) for part in range(6) ) normalized = op.rms_norm( network, @@ -323,6 +359,7 @@ def _transformer_block( rows, cos=cos, sin=sin, + consume_weights=consume_weights, ) hidden = op.gated_residual(network, hidden, update, gate_msa) @@ -345,10 +382,9 @@ def _transformer_block( def _final_hidden(network, hidden, timestep_indices, final_modulation, weights, profile): - rows = -1 selected = op.gather_rows(network, final_modulation, timestep_indices) - final_shift = _slice_modulation(network, selected, 0, rows, profile.hidden_size) - final_scale = _slice_modulation(network, selected, 1, rows, profile.hidden_size) + final_shift = _slice_modulation(network, selected, 0, profile.hidden_size) + final_scale = _slice_modulation(network, selected, 1, profile.hidden_size) hidden = op.rms_norm( network, hidden, weights["norm_out.norm.weight"], profile.hidden_size, profile.norm_eps ) @@ -356,37 +392,13 @@ def _final_hidden(network, hidden, timestep_indices, final_modulation, weights, return op.cast(network, hidden, trt.float32) -def _mark_full_velocity_outputs(network, hidden, weights): - """Preserve the original monolithic plan's full packed-sequence outputs.""" - - video_all = op.linear( - network, hidden, weights["proj_out.weight"], weights["proj_out.bias"], bf16=False - ) - audio_all = op.linear( - network, - hidden, - weights["audio_proj_out.weight"], - weights["audio_proj_out.bias"], - bf16=False, - ) - video_all.name = "video_velocity" - audio_all.name = "audio_velocity" - network.mark_output(video_all) - network.mark_output(audio_all) - - -def _mark_sliced_velocity_outputs(network, hidden, weights, profile: MiniMaxH3Config): +def _mark_sliced_velocity_outputs(network, hidden, weights, video_reference, audio_reference): """Project only rows consumed by the audio and video scheduler updates.""" - audio_hidden = op.slice_rows_from_end( - network, - hidden, - offset=profile.audio_rows + profile.video_rows, - rows=profile.audio_rows, - ) - video_hidden = op.slice_rows_from_end( - network, hidden, offset=profile.video_rows, rows=profile.video_rows + audio_hidden = op.slice_rows_like_from_end( + network, hidden, audio_reference, trailing_reference=video_reference ) + video_hidden = op.slice_rows_like_from_end(network, hidden, video_reference) video = op.linear( network, video_hidden, @@ -407,18 +419,27 @@ def _mark_sliced_velocity_outputs(network, hidden, weights, profile: MiniMaxH3Co network.mark_output(audio) -def _native_builder(verbose: bool, workspace_bytes: int | None): +def _native_builder( + verbose: bool, + workspace_bytes: int | None, + *, + weight_streaming: bool = False, +): logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) config = builder.create_builder_config() config.builder_optimization_level = 1 - op.configure_builder(config) - op.configure_workspace( - config, - workspace_bytes, - default_bytes=DENOISER_DEFAULT_WORKSPACE_BYTES, - ) + op.configure_builder(config, weight_streaming=weight_streaming) + # Leave TensorRT-RTX at its device-derived maximum unless the caller + # explicitly asks for a smaller tactic workspace. The largest dense plans + # use the TensorRT default workspace limit. + if workspace_bytes is not None: + op.configure_workspace( + config, + workspace_bytes, + default_bytes=DENOISER_DEFAULT_WORKSPACE_BYTES, + ) # Hugging Face keeps TF32 disabled for the FP32 input/output projections. config.clear_flag(trt.BuilderFlag.TF32) return logger, builder, network, config @@ -429,6 +450,8 @@ def _add_dynamic_profile( config, profile: MiniMaxH3Config, *, + video_inputs: tuple[str, ...] = (), + audio_inputs: tuple[str, ...] = (), text_inputs: tuple[str, ...] = (), packed_inputs: tuple[str, ...] = (), ) -> None: @@ -443,6 +466,20 @@ def _add_dynamic_profile( profile.opt_sequence_length, profile.sequence_length, ) + for name in video_inputs: + optimization.set_shape( + name, + min=(profile.min_video_rows, profile.video_patch_dim), + opt=(profile.opt_video_rows, profile.video_patch_dim), + max=(profile.video_rows, profile.video_patch_dim), + ) + for name in audio_inputs: + optimization.set_shape( + name, + min=(profile.min_audio_rows, profile.audio_in_channels), + opt=(profile.opt_audio_rows, profile.audio_in_channels), + max=(profile.audio_rows, profile.audio_in_channels), + ) for name in text_inputs: optimization.set_shape( name, @@ -460,6 +497,17 @@ def _add_dynamic_profile( config.add_optimization_profile(optimization) +def _add_first_block_cache_profiles( + builder, + config, + profile: MiniMaxH3Config, + **binding_groups, +) -> None: + """Use one dynamic profile for the complete requested media envelope.""" + + _add_dynamic_profile(builder, config, profile, **binding_groups) + + def _serialize( *, logger, @@ -469,109 +517,26 @@ def _serialize( weights: dict, consume_weights: bool, label: str, -) -> bytes: + output_path: str | Path | None, +) -> bytes | dict[str, int | str]: + plan = None + record = None try: - plan = builder.build_serialized_network(network, config) + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) finally: op.release_weight_buffers(network) if consume_weights: weights.clear() - if plan is None: + if output_path is None and plan is None: raise RuntimeError(f"TensorRT failed to build MiniMax-H3 {label} engine") del network, config, builder, logger gc.collect() - return bytes(plan) - - -def build_dit_engine( - weights: dict, - profile: MiniMaxH3Config, - *, - verbose: bool = False, - consume_weights: bool = False, - workspace_bytes: int | None = None, -) -> bytes: - """Build the full-sequence single-device H3 TensorRT plan.""" - - profile.validate() - if profile.first_block_cache: - raise ValueError("MiniMax-H3 first_block_cache profile requires the split DiT builders") - rows = -1 - logger, builder, network, config = _native_builder(verbose, workspace_bytes) - - video = network.add_input( - "video_hidden_states", trt.float32, (profile.video_rows, profile.video_patch_dim) - ) - audio = network.add_input( - "audio_hidden_states", trt.float32, (profile.audio_rows, profile.audio_in_channels) - ) - text = network.add_input("encoder_hidden_states", trt.float32, (-1, profile.text_dim)) - positions = network.add_input("position_ids", trt.float32, (-1, 3)) - adaln_indices = network.add_input("adaln_indices", trt.int32, (-1,)) - timestep_indices = network.add_input("timestep_indices", trt.int32, (-1,)) - _add_dynamic_profile( - builder, - config, - profile, - text_inputs=("encoder_hidden_states",), - packed_inputs=("position_ids", "adaln_indices", "timestep_indices"), - ) - block_modulations = [ - network.add_input( - f"block_modulation_{index}", - trt.bfloat16, - (profile.adaln_table_rows, 6, profile.hidden_size), - ) - for index in range(profile.num_layers) - ] - final_modulation = network.add_input( - "final_modulation", trt.bfloat16, (profile.max_timestep_count, 2, profile.hidden_size) - ) - # The public single-device FL2VA profile is packed as text | audio | video. - # Projection, text refinement, packing, and full-sequence attention all - # remain native TensorRT operations on one device. - hidden = _packed_hidden(network, video, audio, text, weights, profile) - - cos, sin = _rope_tables(network, positions, profile, rows) - # The dynamic packed sequence contains live rows only, like Diffusers, so - # its attention mask remains None for every supported prompt length. - for index in range(profile.num_layers): - hidden = _transformer_block( - network, - hidden, - block_modulations[index], - adaln_indices, - cos, - sin, - weights, - profile, - index, - ) - - hidden = _final_hidden(network, hidden, timestep_indices, final_modulation, weights, profile) - _mark_sliced_velocity_outputs(network, hidden, weights, profile) - - op.validate_native_network( - network, - expected_attentions=profile.num_refiner_layers + profile.num_layers, - label="DiT", - ) - - print( - f"[minimax-h3] building native DiT: layers={profile.num_layers}, " - f"packed={profile.min_sequence_length}..{profile.sequence_length} " - f"(opt={profile.opt_sequence_length}), devices=1", - file=sys.stderr, - ) - return _serialize( - logger=logger, - builder=builder, - network=network, - config=config, - weights=weights, - consume_weights=consume_weights, - label="DiT", - ) + return record if record is not None else bytes(plan) def _require_first_block_cache_profile(profile: MiniMaxH3Config) -> None: @@ -580,6 +545,7 @@ def _require_first_block_cache_profile(profile: MiniMaxH3Config) -> None: raise ValueError("MiniMax-H3 split DiT plans require profile.first_block_cache=True") +@op.cleanup_failed_build def build_dit_head_engine( weights: dict, profile: MiniMaxH3Config, @@ -587,17 +553,20 @@ def build_dit_head_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: """Build packing, text refinement, block zero, and the native cache metric.""" _require_first_block_cache_profile(profile) - rows = -1 - logger, builder, network, config = _native_builder(verbose, workspace_bytes) + logger, builder, network, config = _native_builder( + verbose, workspace_bytes, weight_streaming=weight_streaming + ) video = network.add_input( - "video_hidden_states", trt.float32, (profile.video_rows, profile.video_patch_dim) + "video_hidden_states", trt.float32, (-1, profile.video_patch_dim) ) audio = network.add_input( - "audio_hidden_states", trt.float32, (profile.audio_rows, profile.audio_in_channels) + "audio_hidden_states", trt.float32, (-1, profile.audio_in_channels) ) text = network.add_input("encoder_hidden_states", trt.float32, (-1, profile.text_dim)) positions = network.add_input("position_ids", trt.float32, (-1, 3)) @@ -610,16 +579,26 @@ def build_dit_head_engine( previous_head_residual = network.add_input( "previous_head_residual", trt.bfloat16, (-1, profile.hidden_size) ) - _add_dynamic_profile( + _add_first_block_cache_profiles( builder, config, profile, + video_inputs=("video_hidden_states",), + audio_inputs=("audio_hidden_states",), text_inputs=("encoder_hidden_states",), packed_inputs=("position_ids", "adaln_indices", "previous_head_residual"), ) - pre_block_hidden = _packed_hidden(network, video, audio, text, weights, profile) - cos, sin = _rope_tables(network, positions, profile, rows) + pre_block_hidden = _packed_hidden( + network, + video, + audio, + text, + weights, + profile, + consume_weights=consume_weights, + ) + cos, sin = _rope_tables(network, positions, profile) head_hidden = _transformer_block( network, pre_block_hidden, @@ -630,6 +609,7 @@ def build_dit_head_engine( weights, profile, 0, + consume_weights=consume_weights, ) head_residual = network.add_elementwise( head_hidden, pre_block_hidden, trt.ElementWiseOperation.SUB @@ -687,9 +667,11 @@ def build_dit_head_engine( weights=weights, consume_weights=consume_weights, label="DiT FirstBlockCache head", + output_path=output_path, ) +@op.cleanup_failed_build def build_dit_tail_engine( weights: dict, profile: MiniMaxH3Config, @@ -697,16 +679,19 @@ def build_dit_tail_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: """Build blocks one through 49 and expose their reusable total residual.""" _require_first_block_cache_profile(profile) - rows = -1 - logger, builder, network, config = _native_builder(verbose, workspace_bytes) + logger, builder, network, config = _native_builder( + verbose, workspace_bytes, weight_streaming=weight_streaming + ) head_hidden = network.add_input("head_hidden", trt.bfloat16, (-1, profile.hidden_size)) positions = network.add_input("position_ids", trt.float32, (-1, 3)) adaln_indices = network.add_input("adaln_indices", trt.int32, (-1,)) - _add_dynamic_profile( + _add_first_block_cache_profiles( builder, config, profile, @@ -720,7 +705,7 @@ def build_dit_tail_engine( ) for index in range(1, profile.num_layers) } - cos, sin = _rope_tables(network, positions, profile, rows) + cos, sin = _rope_tables(network, positions, profile) hidden = head_hidden for index in range(1, profile.num_layers): hidden = _transformer_block( @@ -733,6 +718,7 @@ def build_dit_tail_engine( weights, profile, index, + consume_weights=consume_weights, ) tail_residual = network.add_elementwise( hidden, head_hidden, trt.ElementWiseOperation.SUB @@ -757,9 +743,11 @@ def build_dit_tail_engine( weights=weights, consume_weights=consume_weights, label="DiT FirstBlockCache tail", + output_path=output_path, ) +@op.cleanup_failed_build def build_dit_finish_engine( weights: dict, profile: MiniMaxH3Config, @@ -767,18 +755,30 @@ def build_dit_finish_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: """Apply a selected tail residual, final norm, and consumed-row projections.""" _require_first_block_cache_profile(profile) - logger, builder, network, config = _native_builder(verbose, workspace_bytes) + logger, builder, network, config = _native_builder( + verbose, workspace_bytes, weight_streaming=weight_streaming + ) head_hidden = network.add_input("head_hidden", trt.bfloat16, (-1, profile.hidden_size)) tail_residual = network.add_input("tail_residual", trt.bfloat16, (-1, profile.hidden_size)) timestep_indices = network.add_input("timestep_indices", trt.int32, (-1,)) - _add_dynamic_profile( + video = network.add_input( + "video_hidden_states", trt.float32, (-1, profile.video_patch_dim) + ) + audio = network.add_input( + "audio_hidden_states", trt.float32, (-1, profile.audio_in_channels) + ) + _add_first_block_cache_profiles( builder, config, profile, + video_inputs=("video_hidden_states",), + audio_inputs=("audio_hidden_states",), packed_inputs=("head_hidden", "tail_residual", "timestep_indices"), ) final_modulation = network.add_input( @@ -788,7 +788,7 @@ def build_dit_finish_engine( head_hidden, tail_residual, trt.ElementWiseOperation.SUM ).get_output(0) hidden = _final_hidden(network, hidden, timestep_indices, final_modulation, weights, profile) - _mark_sliced_velocity_outputs(network, hidden, weights, profile) + _mark_sliced_velocity_outputs(network, hidden, weights, video, audio) op.validate_native_network( network, expected_attentions=0, @@ -807,4 +807,5 @@ def build_dit_finish_engine( weights=weights, consume_weights=consume_weights, label="DiT FirstBlockCache finish", + output_path=output_path, ) diff --git a/families/minimax_h3/fl2va_contract.py b/families/minimax_h3/fl2va_contract.py new file mode 100644 index 0000000000..9e4d78d4c3 --- /dev/null +++ b/families/minimax_h3/fl2va_contract.py @@ -0,0 +1,665 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time contracts for MiniMax-H3 first/last-frame conditioning. + +The runtime consumes the plans described here without importing this module. +Keeping the public FL2VA geometry, Qwen presentation, and row accounting in a +small module makes it possible to validate a bundle before any multi-gigabyte +engine is built. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + + +CANVAS_MULTIPLE = 32 +CANVAS_SHORT_EDGE = 768 +CANVAS_MAX_PIXELS = 768 * 1344 +VAE_SPATIAL_COMPRESSION = 16 +VAE_TILE_SIZE = 256 +VAE_TILE_MIN_OVERLAP = 64 +VAE_TILE_OPT_BATCH = 28 +VAE_TILE_MAX_BATCH = 33 +VIDEO_PATCH_SIZE = (1, 2, 2) +QWEN_VISION_PATCH_SIZE = 16 +QWEN_VISION_TEMPORAL_PATCH_SIZE = 2 +QWEN_VISION_MERGE_SIZE = 2 +QWEN_VISION_PATCH_WIDTH = 3 * QWEN_VISION_TEMPORAL_PATCH_SIZE * QWEN_VISION_PATCH_SIZE**2 +QWEN_VISION_HIDDEN_SIZE = 1152 +QWEN_TEXT_HIDDEN_SIZE = 5120 +QWEN_LABEL_TOKENS = 6 +QWEN_VISION_BOUNDARY_TOKENS = 2 +QWEN_VISION_START_TOKEN_ID = 151652 +QWEN_VISION_END_TOKEN_ID = 151653 +QWEN_IMAGE_PAD_TOKEN_ID = 151655 +MIN_FRAMES = 124 +MAX_FRAMES = 345 +FRAMES_PER_CHUNK = 17 +LATENTS_PER_CHUNK = 5 +AUDIO_LATENTS_PER_SECOND = 40 +FPS = 24 +LATENTS_MEAN = ( + 0.858090341091156, + -0.9606591463088989, + 1.0661640167236328, + -0.5090325474739075, + -0.2727581858634949, + -1.3675414323806763, + -0.2553254961967468, + -0.26907554268836975, + -0.5376840829849243, + -0.0464097298681736, + 0.6657370328903198, + 0.19690127670764923, + -0.5460608005523682, + -0.4035342037677765, + -0.23683024942874908, + 0.25928452610969543, + -0.30133944749832153, + 0.211341992020607, + -1.1206848621368408, + 0.3581933379173279, + -0.04225143790245056, + 0.2604829967021942, + 0.22864092886447906, + 0.7056031823158264, +) +LATENTS_STD = ( + 1.2223774194717407, + 1.2767263650894165, + 1.6831774711608887, + 1.7549455165863037, + 1.5636216402053833, + 2.194143533706665, + 0.9653137922286987, + 1.0569885969161987, + 0.841948926448822, + 0.7729952931404114, + 1.8955937623977661, + 0.946841835975647, + 0.7996809482574463, + 0.44988900423049927, + 0.7197399735450745, + 0.6936293244361877, + 2.961095094680786, + 2.7694199085235596, + 3.0496184825897217, + 2.1088054180145264, + 3.276226282119751, + 3.1627357006073, + 2.2816812992095947, + 2.6127843856811523, +) + + +@dataclass(frozen=True) +class TensorAbi: + """One TensorRT binding, including every optimization-profile shape.""" + + name: str + dtype: str + min_shape: tuple[int, ...] + opt_shape: tuple[int, ...] + max_shape: tuple[int, ...] + + +@dataclass(frozen=True) +class PlanAbi: + """The complete externally visible binding contract of one native plan.""" + + filename: str + inputs: tuple[TensorAbi, ...] + outputs: tuple[TensorAbi, ...] + + +@dataclass(frozen=True) +class VisionEncoderProfile: + """Patch-row profile for one Qwen3-VL image or temporal video block.""" + + min_patches: int = 1620 + opt_patches: int = 4032 + max_patches: int = 4176 + + def validate(self) -> None: + values = (self.min_patches, self.opt_patches, self.max_patches) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + raise ValueError("MiniMax-H3 vision profile dimensions must be integers") + if not 4 <= self.min_patches <= self.opt_patches <= self.max_patches: + raise ValueError("MiniMax-H3 vision profile must satisfy 4 <= min <= opt <= max") + if any(value % (QWEN_VISION_MERGE_SIZE**2) for value in values): + raise ValueError("MiniMax-H3 vision profile dimensions must be divisible by four") + + +@dataclass(frozen=True) +class MultimodalTextProfile: + """Sequence profile shared by T2VA, FL2VA, and Ref2VA conditioning.""" + + min_sequence_length: int = 1 + opt_sequence_length: int = 1144 + max_sequence_length: int = 2641 + # TensorRT profiles cannot bind a zero-length tensor. Text-only T2VA uses + # one ignored dummy visual row with ``vision_count == 0``. + min_vision_rows: int = 1 + opt_vision_rows: int = 1008 + max_vision_rows: int = 2088 + + def validate(self) -> None: + values = ( + self.min_sequence_length, + self.opt_sequence_length, + self.max_sequence_length, + self.min_vision_rows, + self.opt_vision_rows, + self.max_vision_rows, + ) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + raise ValueError("MiniMax-H3 text profile dimensions must be integers") + if ( + not 1 + <= self.min_sequence_length + <= self.opt_sequence_length + <= self.max_sequence_length + ): + raise ValueError("MiniMax-H3 text profile must satisfy 1 <= min <= opt <= max") + if not 1 <= self.min_vision_rows <= self.opt_vision_rows <= self.max_vision_rows: + raise ValueError("MiniMax-H3 visual-row profile must satisfy 1 <= min <= opt <= max") + if self.min_vision_rows > self.min_sequence_length: + raise ValueError("MiniMax-H3 minimum visual rows cannot exceed minimum sequence rows") + if self.opt_vision_rows > self.opt_sequence_length: + raise ValueError("MiniMax-H3 optimum visual rows cannot exceed optimum sequence rows") + if self.max_vision_rows > self.max_sequence_length: + raise ValueError("MiniMax-H3 maximum visual rows cannot exceed maximum sequence rows") + + +@dataclass(frozen=True) +class TileAxis: + starts: tuple[int, ...] + lengths: tuple[int, ...] + overlaps: tuple[int, ...] + + +@dataclass(frozen=True) +class KeyframeResize: + """Integer geometry for one exact LANCZOS resize/crop operation.""" + + mode: str + resized_height: int + resized_width: int + crop_top: int = 0 + crop_left: int = 0 + + +@dataclass(frozen=True) +class PackedRows: + text: int + condition_video: int + target_audio: int + target_video: int + + @property + def total(self) -> int: + return self.text + self.condition_video + self.target_audio + self.target_video + + +def resolve_canvas_size( + aspect_width: float, + aspect_height: float, + *, + multiple: int = CANVAS_MULTIPLE, + short_edge: int = CANVAS_SHORT_EDGE, + max_pixels: int = CANVAS_MAX_PIXELS, +) -> tuple[int, int]: + """Reproduce the public canvas resolver and return ``(height, width)``.""" + + values = (aspect_width, aspect_height, multiple, short_edge, max_pixels) + if any(isinstance(value, bool) for value in values): + raise ValueError("MiniMax-H3 canvas geometry cannot use booleans") + if not all(isinstance(value, (int, float)) for value in (aspect_width, aspect_height)): + raise ValueError("MiniMax-H3 aspect dimensions must be numeric") + if not all(isinstance(value, int) for value in (multiple, short_edge, max_pixels)): + raise ValueError("MiniMax-H3 canvas limits must be integers") + if not math.isfinite(aspect_width) or not math.isfinite(aspect_height): + raise ValueError("MiniMax-H3 aspect dimensions must be finite") + if ( + aspect_width <= 0 + or aspect_height <= 0 + or multiple <= 0 + or short_edge <= 0 + or max_pixels <= 0 + ): + raise ValueError("MiniMax-H3 canvas geometry must be positive") + + ratio = aspect_width / aspect_height + if not 0.25 <= ratio <= 4.0: + raise ValueError( + f"MiniMax-H3 supports aspect ratios from 1:4 to 4:1, got {aspect_width}:{aspect_height}" + ) + if ratio >= 1.0: + width, height = short_edge * ratio, float(short_edge) + else: + width, height = float(short_edge), short_edge / ratio + area = width * height + if area > max_pixels: + scale = math.sqrt(max_pixels / area) + width *= scale + height *= scale + return ( + max(multiple, round(height / multiple) * multiple), + max(multiple, round(width / multiple) * multiple), + ) + + +def validate_canvas(height: int, width: int) -> None: + if isinstance(height, bool) or isinstance(width, bool): + raise ValueError("MiniMax-H3 canvas dimensions must be integers") + if not isinstance(height, int) or not isinstance(width, int) or height <= 0 or width <= 0: + raise ValueError(f"MiniMax-H3 canvas must be positive, got {height}x{width}") + if height % CANVAS_MULTIPLE or width % CANVAS_MULTIPLE: + raise ValueError( + f"MiniMax-H3 canvas must be a multiple of {CANVAS_MULTIPLE}, got {height}x{width}" + ) + if width > 4 * height or height > 4 * width: + raise ValueError(f"MiniMax-H3 canvas must be within 1:4 and 4:1, got {height}x{width}") + + +def resolve_keyframe_anchors(*, has_first: bool, has_last: bool) -> tuple[str, ...]: + """Return the reference's packed keyframe order and endpoint anchors.""" + + if not isinstance(has_first, bool) or not isinstance(has_last, bool): + raise ValueError("MiniMax-H3 keyframe presence flags must be booleans") + anchors = tuple( + anchor for anchor, present in (("first", has_first), ("last", has_last)) if present + ) + if not anchors: + raise ValueError("MiniMax-H3 FL2VA needs a first frame, a last frame, or both") + return anchors + + +def keyframe_resize_geometry( + source_height: int, + source_width: int, + target_height: int, + target_width: int, + *, + packed_index: int, +) -> KeyframeResize: + """Resolve the exact stretch/cover-crop geometry for one packed keyframe. + + Packed index zero is always the geometry anchor, including a last-only + request. It is stretched to the canvas. Only the second keyframe is a + cover-cropped follower. + """ + + values = (source_height, source_width, target_height, target_width, packed_index) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + raise ValueError("MiniMax-H3 resize geometry must use integer values") + if min(source_height, source_width, target_height, target_width) <= 0: + raise ValueError("MiniMax-H3 resize dimensions must be positive") + if packed_index not in (0, 1): + raise ValueError("MiniMax-H3 FL2VA has at most two packed keyframes") + validate_canvas(target_height, target_width) + if (source_height, source_width) == (target_height, target_width): + return KeyframeResize("identity", target_height, target_width) + if packed_index == 0: + return KeyframeResize("stretch", target_height, target_width) + + scale = max(target_width / source_width, target_height / source_height) + resized_width = max(target_width, round(source_width * scale)) + resized_height = max(target_height, round(source_height * scale)) + return KeyframeResize( + "cover_crop", + resized_height, + resized_width, + crop_top=max(0, (resized_height - target_height) // 2), + crop_left=max(0, (resized_width - target_width) // 2), + ) + + +def split_tile_axis( + length: int, + tile_size: int = VAE_TILE_SIZE, + min_overlap: int = VAE_TILE_MIN_OVERLAP, + alignment: int = VAE_SPATIAL_COMPRESSION, +) -> TileAxis: + """Reproduce ``AutoencoderKLMiniMaxH3._split_tiles`` exactly.""" + + if any(isinstance(value, bool) for value in (length, tile_size, min_overlap, alignment)): + raise ValueError("MiniMax-H3 tile geometry must use integer values") + if not all(isinstance(value, int) for value in (length, tile_size, min_overlap, alignment)): + raise ValueError("MiniMax-H3 tile geometry must use integer values") + if length <= 0 or tile_size <= 0 or alignment <= 0 or not 0 <= min_overlap < tile_size: + raise ValueError("invalid MiniMax-H3 tile geometry") + if length % alignment: + raise ValueError(f"MiniMax-H3 tile axis {length} is not latent-aligned to {alignment}") + if tile_size >= length: + return TileAxis((0,), (length,), ()) + + num_tiles = math.ceil(length / tile_size) + while tile_size * num_tiles - min_overlap * (num_tiles - 1) - length < 0: + num_tiles += 1 + + overlaps = [min_overlap] * (num_tiles - 1) + remaining = tile_size * num_tiles - sum(overlaps) - length + if remaining % alignment: + raise ValueError("MiniMax-H3 tile slack is not latent-aligned") + for index in range(remaining // alignment): + overlaps[index % (num_tiles - 1)] += alignment + + starts = [0] + for overlap in overlaps: + starts.append(starts[-1] + tile_size - overlap) + result = TileAxis(tuple(starts), (tile_size,) * num_tiles, tuple(overlaps)) + if result.starts[-1] + result.lengths[-1] != length: + raise ValueError("MiniMax-H3 tile split does not cover the canvas exactly") + return result + + +def latent_tile_axis(length: int) -> TileAxis: + """Map an exact pixel-space tile split onto the VAE latent grid. + + Stitching blends each overlap of ``n`` latent cells at positions + ``i = 0..n-1`` as ``previous * (1 - i/n) + current * (i/n)``. Vertical + blending is applied before horizontal blending, then every non-final tile + drops its trailing overlap exactly as the reference ``_stitch_tiles`` does. + """ + + pixels = split_tile_axis(length) + return TileAxis( + tuple(start // VAE_SPATIAL_COMPRESSION for start in pixels.starts), + tuple(size // VAE_SPATIAL_COMPRESSION for size in pixels.lengths), + tuple(overlap // VAE_SPATIAL_COMPRESSION for overlap in pixels.overlaps), + ) + + +def keyframe_tile_count(height: int, width: int) -> int: + validate_canvas(height, width) + if height < VAE_TILE_SIZE or width < VAE_TILE_SIZE: + raise ValueError( + f"MiniMax-H3 keyframe plan requires both canvas axes to be at least " + f"{VAE_TILE_SIZE}, got {height}x{width}" + ) + count = len(split_tile_axis(height).starts) * len(split_tile_axis(width).starts) + if count > VAE_TILE_MAX_BATCH: + raise ValueError( + f"MiniMax-H3 keyframe needs {count} VAE tiles, plan capacity is {VAE_TILE_MAX_BATCH}" + ) + return count + + +def qwen_vision_patch_rows(height: int, width: int) -> int: + validate_canvas(height, width) + if height % QWEN_VISION_PATCH_SIZE or width % QWEN_VISION_PATCH_SIZE: + raise ValueError("MiniMax-H3 Qwen vision canvas is not patch-aligned") + return (height // QWEN_VISION_PATCH_SIZE) * (width // QWEN_VISION_PATCH_SIZE) + + +def qwen_vision_token_rows(height: int, width: int) -> int: + patches = qwen_vision_patch_rows(height, width) + merge_unit = QWEN_VISION_MERGE_SIZE**2 + if patches % merge_unit: + raise ValueError("MiniMax-H3 Qwen vision patches are not merge-aligned") + return patches // merge_unit + + +def fl2va_text_rows(prompt_tokens: int, keyframes: int, *, height: int, width: int) -> int: + if isinstance(prompt_tokens, bool) or not isinstance(prompt_tokens, int) or prompt_tokens < 0: + raise ValueError("MiniMax-H3 prompt token count must be a non-negative integer") + if keyframes not in (1, 2): + raise ValueError(f"MiniMax-H3 FL2VA needs one or two keyframes, got {keyframes}") + per_keyframe = ( + QWEN_LABEL_TOKENS + QWEN_VISION_BOUNDARY_TOKENS + qwen_vision_token_rows(height, width) + ) + return prompt_tokens + keyframes * per_keyframe + + +def fl2va_mrope_position_ids( + prompt_tokens: int, + keyframes: int, + *, + height: int, + width: int, +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """Build Qwen3-VL's exact temporal/height/width position IDs. + + Qwen marks only ``image_pad`` tokens as image modality. Each label and + ``vision_start`` therefore forms a text run before the image grid, while a + preceding ``vision_end`` joins the next label when two images are present. + This mirrors ``Qwen3VLModel.get_rope_index`` without requiring Transformers + at bundle-build time. + """ + + expected_rows = fl2va_text_rows(prompt_tokens, keyframes, height=height, width=width) + llm_height = height // (QWEN_VISION_PATCH_SIZE * QWEN_VISION_MERGE_SIZE) + llm_width = width // (QWEN_VISION_PATCH_SIZE * QWEN_VISION_MERGE_SIZE) + axes: tuple[list[int], list[int], list[int]] = ([], [], []) + current_position = 0 + + def append_text(length: int) -> None: + nonlocal current_position + positions = range(current_position, current_position + length) + for axis in axes: + axis.extend(positions) + current_position += length + + for index in range(keyframes): + # The second text run begins with the previous image's vision_end. + append_text(QWEN_LABEL_TOKENS + 1 + int(index > 0)) + for row in range(llm_height): + for column in range(llm_width): + axes[0].append(current_position) + axes[1].append(current_position + row) + axes[2].append(current_position + column) + current_position += max(llm_height, llm_width) + append_text(1 + prompt_tokens) # final vision_end followed by the prompt + + result = tuple(tuple(axis) for axis in axes) + if any(len(axis) != expected_rows for axis in result): + raise RuntimeError("MiniMax-H3 FL2VA MRoPE row accounting mismatch") + return result + + +def align_num_frames(num_frames: int) -> int: + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError("MiniMax-H3 frame count must be a positive integer") + aligned = num_frames + (LATENTS_PER_CHUNK - num_frames % FRAMES_PER_CHUNK) % FRAMES_PER_CHUNK + if not MIN_FRAMES <= aligned <= MAX_FRAMES: + raise ValueError( + f"MiniMax-H3 aligned frame count must be {MIN_FRAMES}..{MAX_FRAMES}, got {aligned}" + ) + return aligned + + +def video_latent_frames(num_frames: int) -> int: + aligned = align_num_frames(num_frames) + return ((aligned - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * LATENTS_PER_CHUNK + 2 + + +def audio_latent_frames(num_frames: int) -> int: + aligned = align_num_frames(num_frames) + return int(round(aligned / FPS * AUDIO_LATENTS_PER_SECOND)) + + +def packed_rows( + *, + prompt_tokens: int, + keyframes: int, + height: int, + width: int, + num_frames: int, +) -> PackedRows: + text = fl2va_text_rows(prompt_tokens, keyframes, height=height, width=width) + rows_per_frame = (height // VAE_SPATIAL_COMPRESSION // VIDEO_PATCH_SIZE[1]) * ( + width // VAE_SPATIAL_COMPRESSION // VIDEO_PATCH_SIZE[2] + ) + return PackedRows( + text=text, + condition_video=keyframes * rows_per_frame, + target_audio=2 * audio_latent_frames(num_frames), + target_video=video_latent_frames(num_frames) * rows_per_frame, + ) + + +def last_keyframe_rotary_time(text_rows: int, num_frames: int) -> float: + """The reference's NumPy-summed last-frame rotary anchor.""" + + latent_frames = video_latent_frames(num_frames) + spans = [5.0 / 3.0 * (1, 4, 4, 4, 4)[index % 5] for index in range(latent_frames)] + # Python's sum is deliberately not used: NumPy pairwise summation is the + # reference. The native runtime implements the same reduction order. + import numpy as np + + return float(text_rows) + float(np.asarray(spans, dtype=np.float64).sum()) - 5.0 / 3.0 + + +def keyframe_vae_encoder_abi() -> PlanAbi: + """Tile-plan ABI for the exact released keyframe posterior recipe. + + The caller supplies row-major tiles after ImageNet normalization. It must + stitch all 48 posterior-parameter channels before splitting mean/logvar, + clamp logvar to ``[-30, 20]``, sample with a fresh CPU generator seeded 42, + round the sample through FP16, and apply ``LATENTS_MEAN/LATENTS_STD``. Each + keyframe repeats that fresh-generator recipe independently. + """ + + return PlanAbi( + filename="fl2va_keyframe_vae_encoder.plan", + inputs=( + TensorAbi( + "pixel_tiles", + "float32", + (1, 3, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + (VAE_TILE_OPT_BATCH, 3, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + (VAE_TILE_MAX_BATCH, 3, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + ), + ), + outputs=( + TensorAbi( + "posterior_parameter_tiles", + "float32", + ( + 1, + 48, + 1, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + ), + ( + VAE_TILE_OPT_BATCH, + 48, + 1, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + ), + ( + VAE_TILE_MAX_BATCH, + 48, + 1, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + ), + ), + ), + ) + + +def vision_encoder_abi(profile: VisionEncoderProfile = VisionEncoderProfile()) -> PlanAbi: + """One-image/frame dynamic-aspect Qwen3-VL vision-tower plan. + + A reference video is evaluated one temporal patch block at a time. This + preserves Qwen3-VL's ``cu_seqlens`` boundary (vision attention never crosses + temporal blocks) while allowing the same weights and engine to serve FL2VA + images, Ref2VA images, and Ref2VA video frames. + """ + + profile.validate() + min_patches = profile.min_patches + max_patches = profile.max_patches + opt_patches = profile.opt_patches + bindings = ( + TensorAbi( + "pixel_values", + "float32", + (min_patches, QWEN_VISION_PATCH_WIDTH), + (opt_patches, QWEN_VISION_PATCH_WIDTH), + (max_patches, QWEN_VISION_PATCH_WIDTH), + ), + TensorAbi("interp_indices", "int32", (min_patches, 4), (opt_patches, 4), (max_patches, 4)), + TensorAbi( + "interp_weights", "float32", (min_patches, 4), (opt_patches, 4), (max_patches, 4) + ), + TensorAbi( + "vision_position_ids", "int32", (min_patches, 2), (opt_patches, 2), (max_patches, 2) + ), + ) + + def output(name: str) -> TensorAbi: + return TensorAbi( + name, + "float32", + (min_patches // 4, QWEN_TEXT_HIDDEN_SIZE), + (opt_patches // 4, QWEN_TEXT_HIDDEN_SIZE), + (max_patches // 4, QWEN_TEXT_HIDDEN_SIZE), + ) + + return PlanAbi( + filename="vision_encoder.plan", + inputs=bindings, + outputs=tuple( + output(name) for name in ("vision_embeds", "deepstack_0", "deepstack_1", "deepstack_2") + ), + ) + + +def text_encoder_abi( + profile: MultimodalTextProfile = MultimodalTextProfile(), +) -> PlanAbi: + """Unified Qwen3-VL language-stack ABI for all three public workflows. + + Visual features remain compact at the ABI and are scattered through + ``vision_row_indices``. Text-only requests bind one dummy row, set + ``vision_count`` to zero, and clear ``vision_mask``. Thus the plan owns one + copy of language weights instead of shipping separate T2VA and FL2VA + language plans. + """ + + profile.validate() + min_rows = profile.min_sequence_length + opt_rows = profile.opt_sequence_length + max_rows = profile.max_sequence_length + min_vision = profile.min_vision_rows + opt_vision = profile.opt_vision_rows + max_vision = profile.max_vision_rows + + def rows(name: str, dtype: str, width: int | None = None) -> TensorAbi: + suffix = () if width is None else (width,) + return TensorAbi(name, dtype, (min_rows, *suffix), (opt_rows, *suffix), (max_rows, *suffix)) + + def visual_rows(name: str, dtype: str, width: int | None = None) -> TensorAbi: + suffix = () if width is None else (width,) + return TensorAbi( + name, + dtype, + (min_vision, *suffix), + (opt_vision, *suffix), + (max_vision, *suffix), + ) + + inputs = ( + rows("input_ids", "int32"), + TensorAbi("mrope_position_ids", "int32", (3, min_rows), (3, opt_rows), (3, max_rows)), + rows("vision_mask", "float32", 1), + TensorAbi("vision_count", "int32", (1,), (1,), (1,)), + visual_rows("vision_row_indices", "int32"), + visual_rows("vision_embeds", "float32", QWEN_TEXT_HIDDEN_SIZE), + visual_rows("deepstack_0", "float32", QWEN_TEXT_HIDDEN_SIZE), + visual_rows("deepstack_1", "float32", QWEN_TEXT_HIDDEN_SIZE), + visual_rows("deepstack_2", "float32", QWEN_TEXT_HIDDEN_SIZE), + ) + return PlanAbi( + filename="text_encoder.plan", + inputs=inputs, + outputs=(rows("encoder_hidden_states", "float32", QWEN_TEXT_HIDDEN_SIZE),), + ) diff --git a/families/minimax_h3/fl2va_vae_encoder_builder.py b/families/minimax_h3/fl2va_vae_encoder_builder.py new file mode 100644 index 0000000000..0bb885d399 --- /dev/null +++ b/families/minimax_h3/fl2va_vae_encoder_builder.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native TensorRT encoder for MiniMax-H3 FL2VA keyframe tiles. + +The released VAE always encodes keyframes as one temporal frame. Causal 3-D +convolution therefore has one non-zero temporal tap: the last tap of every +3x3x3 kernel (or the only tap of a 1x1x1 kernel). Expressing that exact case +as 2-D convolutions avoids manufacturing zero temporal frames while retaining +the released spatial reflection padding, isolated group norm, and FP32 math. + +The plan deliberately emits posterior parameters. Posterior sampling, +FP16-rounding, per-channel latent normalization, tile stitching, and the +request-generator draw order belong to the native runtime rather than the +engine and are stated in :mod:`fl2va_contract`. +""" + +from __future__ import annotations + +import gc +import sys +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .fl2va_contract import ( + VAE_SPATIAL_COMPRESSION, + VAE_TILE_MAX_BATCH, + VAE_TILE_SIZE, + keyframe_vae_encoder_abi, +) + + +trt = trt_compat.get_trt() + +IN_CHANNELS = 3 +MOMENT_CHANNELS = 48 +BLOCK_OUT_CHANNELS = (128, 256, 256, 512, 512, 1024) +SPATIAL_DOWNSAMPLE_FACTORS = (2, 2, 2, 2, 1, 1) +TEMPORAL_DOWNSAMPLE_FACTORS = (1, 2, 2, 1, 1, 1) +LAYERS_PER_BLOCK = 2 +NORM_GROUPS = 32 +NORM_EPS = 1.0e-6 +KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 + + +def checkpoint_keys() -> tuple[str, ...]: + """Return the exhaustive VAE checkpoint partition consumed by this plan.""" + + names = ["encoder.conv_in.weight", "encoder.conv_in.bias"] + block_in_channels = (BLOCK_OUT_CHANNELS[0], *BLOCK_OUT_CHANNELS[:-1]) + for block_index, (in_channels, out_channels) in enumerate( + zip(block_in_channels, BLOCK_OUT_CHANNELS) + ): + for layer_index in range(LAYERS_PER_BLOCK): + prefix = f"encoder.down_blocks.{block_index}.resnets.{layer_index}" + names.extend( + [ + f"{prefix}.norm1.weight", + f"{prefix}.norm1.bias", + f"{prefix}.conv1.weight", + f"{prefix}.conv1.bias", + f"{prefix}.norm2.weight", + f"{prefix}.norm2.bias", + f"{prefix}.conv2.weight", + f"{prefix}.conv2.bias", + ] + ) + if layer_index == 0 and in_channels != out_channels: + names.extend([f"{prefix}.conv_shortcut.weight", f"{prefix}.conv_shortcut.bias"]) + if SPATIAL_DOWNSAMPLE_FACTORS[block_index] * TEMPORAL_DOWNSAMPLE_FACTORS[block_index] > 1: + prefix = f"encoder.down_blocks.{block_index}.downsamplers.0.conv" + names.extend([f"{prefix}.weight", f"{prefix}.bias"]) + names.extend( + [ + "encoder.norm_out.weight", + "encoder.norm_out.bias", + "encoder.conv_out.weight", + "encoder.conv_out.bias", + "quant_conv.weight", + "quant_conv.bias", + ] + ) + return tuple(names) + + +def _shape_dim(network, tensor, axis: int): + shape = network.add_shape(tensor).get_output(0) + return network.add_slice(shape, (axis,), (1,), (1,)).get_output(0) + + +def _shape_vector(network, values): + tensors = [] + for value in values: + if isinstance(value, int): + tensors.append( + op.constant(network, np.asarray([value], dtype=np.int64), dtype=np.int64) + ) + else: + tensors.append(value) + if len(tensors) == 1: + return tensors[0] + layer = network.add_concatenation(tensors) + layer.axis = 0 + return layer.get_output(0) + + +def _reflect_pad_2d( + network, + hidden, + *, + channels: int, + height: int, + width: int, + top: int = 0, + bottom: int = 0, + left: int = 0, + right: int = 0, +): + """PyTorch-compatible reflection padding for static CHW and dynamic batch.""" + + if min(top, bottom, left, right) < 0: + raise ValueError("MiniMax-H3 reflect padding cannot be negative") + if top >= height or bottom >= height or left >= width or right >= width: + raise ValueError("MiniMax-H3 reflect padding must be smaller than its input axis") + if top == bottom == left == right == 0: + return hidden + output_height = height + top + bottom + output_width = width + left + right + layer = network.add_slice( + hidden, + (0, 0, -top, -left), + (1, channels, output_height, output_width), + (1, 1, 1, 1), + ) + if layer is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 VAE reflection padding") + layer.set_input( + 2, + _shape_vector( + network, + (_shape_dim(network, hidden, 0), channels, output_height, output_width), + ), + ) + layer.mode = trt.SampleMode.REFLECT + return layer.get_output(0) + + +def _spatial_kernel(weight: np.ndarray) -> np.ndarray: + """Select the only non-zero causal temporal tap for a one-frame clip.""" + + value = np.asarray(weight) + if value.ndim != 5 or value.shape[2] not in (1, 3): + raise ValueError(f"MiniMax-H3 VAE convolution weight has invalid shape {value.shape}") + return np.ascontiguousarray(value[:, :, -1]) + + +def _conv( + network, + hidden, + weights: dict[str, np.ndarray], + prefix: str, + *, + in_channels: int, + out_channels: int, + height: int, + width: int, + stride: int = 1, + symmetric_reflect: int = 0, + bottom_right_reflect: int = 0, +): + if symmetric_reflect and bottom_right_reflect: + raise ValueError("MiniMax-H3 VAE convolution has ambiguous padding") + if symmetric_reflect: + hidden = _reflect_pad_2d( + network, + hidden, + channels=in_channels, + height=height, + width=width, + top=symmetric_reflect, + bottom=symmetric_reflect, + left=symmetric_reflect, + right=symmetric_reflect, + ) + height += 2 * symmetric_reflect + width += 2 * symmetric_reflect + elif bottom_right_reflect: + hidden = _reflect_pad_2d( + network, + hidden, + channels=in_channels, + height=height, + width=width, + bottom=bottom_right_reflect, + right=bottom_right_reflect, + ) + height += bottom_right_reflect + width += bottom_right_reflect + + kernel = _spatial_kernel(weights[f"{prefix}.weight"]) + layer = network.add_convolution_nd( + hidden, + out_channels, + tuple(kernel.shape[-2:]), + kernel, + np.ascontiguousarray(weights[f"{prefix}.bias"]), + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 VAE convolution {prefix}") + layer.name = prefix + layer.stride_nd = (stride, stride) + output_height = (height - kernel.shape[-2]) // stride + 1 + output_width = (width - kernel.shape[-1]) // stride + 1 + return layer.get_output(0), output_height, output_width + + +def _group_norm(network, hidden, weights, prefix: str, channels: int): + rank = len(tuple(hidden.shape)) + gamma = op.weight_constant( + network, np.asarray(weights[f"{prefix}.weight"]).reshape(1, channels, 1, 1) + ) + beta = op.weight_constant( + network, np.asarray(weights[f"{prefix}.bias"]).reshape(1, channels, 1, 1) + ) + gamma = op.cast(network, gamma, hidden.dtype) + beta = op.cast(network, beta, hidden.dtype) + # TensorRT group normalization treats the second NCHW dimension as the + # grouped channel axis. When num_groups != 1, the reduction mask must + # therefore contain only dimensions after C (NvInfer.h, INormalizationLayer + # contract), rather than C itself. Including bit 1 is rejected by + # TensorRT-RTX before the graph can be serialized. + axes = sum(1 << axis for axis in range(2, rank)) + layer = network.add_normalization_v2(hidden, gamma, beta, axes) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 VAE group norm {prefix}") + layer.name = prefix + layer.epsilon = NORM_EPS + layer.num_groups = NORM_GROUPS + return layer.get_output(0) + + +def _resnet( + network, + hidden, + weights, + prefix: str, + *, + in_channels: int, + out_channels: int, + height: int, + width: int, +): + residual = hidden + update = _group_norm(network, hidden, weights, f"{prefix}.norm1", in_channels) + update = op.silu(network, update) + update, update_height, update_width = _conv( + network, + update, + weights, + f"{prefix}.conv1", + in_channels=in_channels, + out_channels=out_channels, + height=height, + width=width, + symmetric_reflect=1, + ) + if (update_height, update_width) != (height, width): + raise RuntimeError("MiniMax-H3 VAE resnet conv1 changed spatial geometry") + update = _group_norm(network, update, weights, f"{prefix}.norm2", out_channels) + update = op.silu(network, update) + update, update_height, update_width = _conv( + network, + update, + weights, + f"{prefix}.conv2", + in_channels=out_channels, + out_channels=out_channels, + height=height, + width=width, + symmetric_reflect=1, + ) + if in_channels != out_channels: + residual, shortcut_height, shortcut_width = _conv( + network, + residual, + weights, + f"{prefix}.conv_shortcut", + in_channels=in_channels, + out_channels=out_channels, + height=height, + width=width, + ) + if (shortcut_height, shortcut_width) != (height, width): + raise RuntimeError("MiniMax-H3 VAE shortcut changed spatial geometry") + result = network.add_elementwise(residual, update, trt.ElementWiseOperation.SUM) + if result is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 VAE residual {prefix}") + return result.get_output(0), update_height, update_width + + +@op.cleanup_failed_build +def build_keyframe_vae_encoder_engine( + weights: dict[str, np.ndarray], + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the dynamic-tile FL2VA VAE encoder plan.""" + + missing = sorted(set(checkpoint_keys()) - set(weights)) + unexpected = sorted(set(weights) - set(checkpoint_keys())) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 FL2VA VAE encoder checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + + abi = keyframe_vae_encoder_abi() + binding = abi.inputs[0] + pixels = network.add_input( + binding.name, + trt.float32, + (-1, IN_CHANNELS, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + ) + profile = builder.create_optimization_profile() + profile.set_shape(binding.name, binding.min_shape, binding.opt_shape, binding.max_shape) + config.add_optimization_profile(profile) + + # T=1 is removed only after the ABI boundary. Every causal 3-D + # convolution below selects its last temporal kernel tap. + squeeze = network.add_shuffle(pixels) + squeeze.reshape_dims = (-1, IN_CHANNELS, VAE_TILE_SIZE, VAE_TILE_SIZE) + hidden = squeeze.get_output(0) + height = width = VAE_TILE_SIZE + hidden, height, width = _conv( + network, + hidden, + weights, + "encoder.conv_in", + in_channels=IN_CHANNELS, + out_channels=BLOCK_OUT_CHANNELS[0], + height=height, + width=width, + symmetric_reflect=1, + ) + + current_channels = BLOCK_OUT_CHANNELS[0] + for block_index, out_channels in enumerate(BLOCK_OUT_CHANNELS): + for layer_index in range(LAYERS_PER_BLOCK): + prefix = f"encoder.down_blocks.{block_index}.resnets.{layer_index}" + hidden, height, width = _resnet( + network, + hidden, + weights, + prefix, + in_channels=current_channels, + out_channels=out_channels, + height=height, + width=width, + ) + current_channels = out_channels + if SPATIAL_DOWNSAMPLE_FACTORS[block_index] * TEMPORAL_DOWNSAMPLE_FACTORS[block_index] > 1: + prefix = f"encoder.down_blocks.{block_index}.downsamplers.0.conv" + spatial_stride = SPATIAL_DOWNSAMPLE_FACTORS[block_index] + hidden, height, width = _conv( + network, + hidden, + weights, + prefix, + in_channels=current_channels, + out_channels=current_channels, + height=height, + width=width, + stride=spatial_stride, + bottom_right_reflect=int(spatial_stride == 2), + ) + + expected_side = VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION + if (height, width) != (expected_side, expected_side): + raise RuntimeError( + f"MiniMax-H3 VAE encoder produced {height}x{width}, expected {expected_side}x{expected_side}" + ) + hidden = _group_norm(network, hidden, weights, "encoder.norm_out", current_channels) + hidden = op.silu(network, hidden) + hidden, height, width = _conv( + network, + hidden, + weights, + "encoder.conv_out", + in_channels=current_channels, + out_channels=MOMENT_CHANNELS, + height=height, + width=width, + symmetric_reflect=1, + ) + hidden, height, width = _conv( + network, + hidden, + weights, + "quant_conv", + in_channels=MOMENT_CHANNELS, + out_channels=MOMENT_CHANNELS, + height=height, + width=width, + ) + expand = network.add_shuffle(hidden) + expand.reshape_dims = (-1, MOMENT_CHANNELS, 1, expected_side, expected_side) + output = expand.get_output(0) + output.name = abi.outputs[0].name + network.mark_output(output) + + op.validate_native_network(network, expected_attentions=0, label="FL2VA keyframe VAE encoder") + print( + f"[minimax-h3] building native FL2VA VAE encoder: tiles=1..{VAE_TILE_MAX_BATCH}, " + f"tile={VAE_TILE_SIZE}x{VAE_TILE_SIZE}, moments={MOMENT_CHANNELS}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 FL2VA VAE encoder") + del network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/graph_ops.py b/families/minimax_h3/graph_ops.py index d81c8ab133..6948560141 100644 --- a/families/minimax_h3/graph_ops.py +++ b/families/minimax_h3/graph_ops.py @@ -10,6 +10,8 @@ from __future__ import annotations from collections import Counter +from contextvars import ContextVar +from functools import lru_cache, wraps import math import ml_dtypes @@ -18,18 +20,55 @@ import tensorrt as trt from .config import resolve_workspace_bytes +from .quantized_checkpoint import ConvRotInt8Weight # TensorRT's explicit BF16 ``Weights`` constructor stores a pointer rather # than owning its input. Retain every backing array until the associated # network has finished building, including temporary packed QKV buffers. _WEIGHT_BUFFER_KEEPALIVE: dict[int, list[np.ndarray]] = {} +_ACTIVE_BUILD_NETWORKS: ContextVar[set[int] | None] = ContextVar( + "minimax_h3_active_build_networks", default=None +) -def configure_builder(config) -> None: - """Retain enough engine metadata to audit native TensorRT lowering.""" +def cleanup_failed_build(function): + """Release retained arrays and consumed weights after graph-build failures.""" + + @wraps(function) + def wrapped(weights, *args, **kwargs): + network_ids: set[int] = set() + token = _ACTIVE_BUILD_NETWORKS.set(network_ids) + try: + return function(weights, *args, **kwargs) + except BaseException: + for network_id in network_ids: + _WEIGHT_BUFFER_KEEPALIVE.pop(network_id, None) + if kwargs.get("consume_weights", False): + weights.clear() + raise + finally: + _ACTIVE_BUILD_NETWORKS.reset(token) + + return wrapped + + +def configure_builder(config, *, weight_streaming: bool = False) -> None: + """Retain graph metadata and enable RTX weight streaming when requested.""" config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED + if not isinstance(weight_streaming, bool): + raise ValueError("MiniMax-H3 weight_streaming must be a boolean") + if not weight_streaming: + return + flag = getattr(trt.BuilderFlag, "WEIGHT_STREAMING", None) + if flag is None: + raise RuntimeError( + "Selected TensorRT-RTX bindings do not expose BuilderFlag.WEIGHT_STREAMING" + ) + config.set_flag(flag) + if not config.get_flag(flag): + raise RuntimeError("TensorRT-RTX did not enable MiniMax-H3 weight streaming") def configure_workspace(config, workspace_bytes: int | None, *, default_bytes: int) -> int: @@ -47,8 +86,13 @@ def configure_workspace(config, workspace_bytes: int | None, *, default_bytes: i return applied -def validate_native_network(network, *, expected_attentions: int, label: str) -> dict[str, int]: - """Fail closed if a native H3 graph gains plugins or collectives.""" +def validate_native_network( + network, + *, + expected_attentions: int, + label: str, +) -> dict[str, int]: + """Fail closed on any layer outside the selected native H3 contract.""" counts = Counter(network.get_layer(index).type for index in range(network.num_layers)) expected = { @@ -79,7 +123,11 @@ def validate_native_network(network, *, expected_attentions: int, label: str) -> def _add_constant(network, array: np.ndarray): array = np.ascontiguousarray(array) - _WEIGHT_BUFFER_KEEPALIVE.setdefault(id(network), []).append(array) + network_id = id(network) + _WEIGHT_BUFFER_KEEPALIVE.setdefault(network_id, []).append(array) + active_networks = _ACTIVE_BUILD_NETWORKS.get() + if active_networks is not None: + active_networks.add(network_id) if array.dtype == np.dtype(ml_dtypes.bfloat16): weights = trt.Weights(trt.bfloat16, array.ctypes.data, array.size) else: @@ -113,6 +161,247 @@ def cast(network, tensor, dtype): return network.add_cast(tensor, dtype).get_output(0) +def _name_convrot_layer(network, layer, role: str) -> None: + """Assign a stable, network-unique diagnostic name to a ConvRot layer.""" + + layer.name = f"minimax_h3.convrot.{role}.{network.num_layers - 1}" + + +@lru_cache(maxsize=None) +def _regular_hadamard(group_size: int) -> np.ndarray: + """Return Comfy's normalized regular Hadamard matrix in checkpoint dtype. + + Comfy ConvRot starts from its symmetric 4x4 basis and takes Kronecker + powers. H3 currently uses groups of 64 and 256. Keeping the coefficients + in BF16 exactly matches the activation dtype used by the published + checkpoint and avoids introducing a second activation quantizer. + """ + + if not isinstance(group_size, int) or isinstance(group_size, bool) or group_size < 4: + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + basis = np.asarray( + ( + (1.0, 1.0, 1.0, -1.0), + (1.0, 1.0, -1.0, 1.0), + (1.0, -1.0, 1.0, 1.0), + (-1.0, 1.0, 1.0, 1.0), + ), + dtype=np.float32, + ) + matrix = basis + while matrix.shape[0] < group_size: + matrix = np.kron(matrix, basis) + if matrix.shape != (group_size, group_size): + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + return np.ascontiguousarray( + matrix / math.sqrt(float(group_size)), dtype=ml_dtypes.bfloat16 + ) + + +def _convrot_activation(network, tensor, *, in_features: int, group_size: int): + """Apply the checkpoint's grouped right-Hadamard rotation with native TRT.""" + + if in_features % group_size: + raise ValueError( + "MiniMax-H3 ConvRot input width must be divisible by group_size: " + f"width={in_features}, group_size={group_size}" + ) + if len(tuple(tensor.shape)) != 2: + raise ValueError("MiniMax-H3 ConvRot linears require rank-2 activations") + static_width = int(tensor.shape[1]) + if static_width >= 0 and static_width != in_features: + raise ValueError( + "MiniMax-H3 ConvRot activation/weight width mismatch: " + f"activation={static_width}, weight={in_features}" + ) + + value = cast(network, tensor, trt.bfloat16) + grouped = network.add_shuffle(value) + # Collapse rows and groups into one GEMM M dimension. This is materially + # faster than issuing one tiny batched GEMM per sequence row while keeping + # the exact same contiguous group boundaries. + grouped.reshape_dims = (-1, group_size) + hadamard = weight_constant(network, _regular_hadamard(group_size)) + rotation = network.add_matrix_multiply( + grouped.get_output(0), + trt.MatrixOperation.NONE, + hadamard, + trt.MatrixOperation.NONE, + ) + if rotation is None: + raise RuntimeError("TensorRT rejected the MiniMax-H3 ConvRot activation rotation") + _name_convrot_layer(network, rotation, f"group_{group_size}") + rotation.metadata = ( + "trtmc.quantization=int8_tensorwise_convrot;" + f"activation=bf16;group_size={group_size}" + ) + restored = network.add_shuffle(rotation.get_output(0)) + restored.reshape_dims = (-1, in_features) + return restored.get_output(0) + + +def _convrot_int8_linear(network, tensor, weight: ConvRotInt8Weight, bias=None): + """Build Comfy-compatible dynamic-rowwise ConvRot W8A8 with native TRT. + + TensorRT's ordinary Quantize layer requires a build-time-constant scale. + To retain the published checkpoint's exact dynamic semantics without a + plugin, express the BF16 rowwise quantizer as native math, then place unit + Q/DQ markers around the resulting INT8 tensors. TensorRT recognizes those + markers and lowers the MatrixMultiply to an INT8 GEMM; the dynamic row and + per-output weight scales are applied explicitly in FP32 before BF16 output. + """ + + qweight = np.asarray(weight.qweight) + scale = np.asarray(weight.scale) + if qweight.dtype != np.int8 or qweight.ndim != 2: + raise ValueError("MiniMax-H3 ConvRot qweight must be rank-2 INT8") + if scale.dtype != np.float32 or scale.shape not in { + (qweight.shape[0],), + (qweight.shape[0], 1), + }: + raise ValueError( + "MiniMax-H3 ConvRot scale must be FP32 [out] or [out,1]: " + f"weight={qweight.shape}, scale={scale.shape}" + ) + + rotated = _convrot_activation( + network, + tensor, + in_features=int(qweight.shape[1]), + group_size=int(weight.group_size), + ) + + # Match comfy-kitchen >= 0.2.15 exactly: absmax is computed at the source + # BF16 dtype, scale storage is FP32, scale math and division return to BF16, + # and ROUND is round-to-nearest-even before the signed INT8 clamp. + absolute = network.add_unary(rotated, trt.UnaryOperation.ABS) + if absolute is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation abs") + _name_convrot_layer(network, absolute, "activation_abs_bf16") + row_max = network.add_reduce( + absolute.get_output(0), + trt.ReduceOperation.MAX, + 1 << 1, + True, + ) + if row_max is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot rowwise maximum") + _name_convrot_layer(network, row_max, "activation_row_max_bf16") + row_max_f32 = cast(network, row_max.get_output(0), trt.float32) + divisor = constant(network, np.full((1, 1), 127.0, dtype=np.float32)) + row_scale = network.add_elementwise( + row_max_f32, + divisor, + trt.ElementWiseOperation.DIV, + ) + if row_scale is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot row scale") + _name_convrot_layer(network, row_scale, "activation_scale_f32") + minimum_scale = constant(network, np.full((1, 1), 1.0e-30, dtype=np.float32)) + clamped_scale = network.add_elementwise( + row_scale.get_output(0), + minimum_scale, + trt.ElementWiseOperation.MAX, + ) + if clamped_scale is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot row scale clamp") + _name_convrot_layer(network, clamped_scale, "activation_scale_clamped_f32") + row_scale_f32 = clamped_scale.get_output(0) + row_scale_bf16 = cast(network, row_scale_f32, trt.bfloat16) + divided = network.add_elementwise( + rotated, + row_scale_bf16, + trt.ElementWiseOperation.DIV, + ) + if divided is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot BF16 quantization division") + _name_convrot_layer(network, divided, "activation_divide_bf16") + rounded = network.add_unary(divided.get_output(0), trt.UnaryOperation.ROUND) + if rounded is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation rounding") + _name_convrot_layer(network, rounded, "activation_round_bf16") + minimum_code = weight_constant( + network, + np.full((1, 1), -128.0, dtype=ml_dtypes.bfloat16), + ) + maximum_code = weight_constant( + network, + np.full((1, 1), 127.0, dtype=ml_dtypes.bfloat16), + ) + clipped_minimum = network.add_elementwise( + rounded.get_output(0), + minimum_code, + trt.ElementWiseOperation.MAX, + ) + if clipped_minimum is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation lower clamp") + clipped = network.add_elementwise( + clipped_minimum.get_output(0), + maximum_code, + trt.ElementWiseOperation.MIN, + ) + if clipped is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation upper clamp") + activation_int8 = cast(network, clipped.get_output(0), trt.int8) + + quantized = weight_constant(network, qweight) + unit_scale = constant(network, np.asarray(1.0, dtype=np.float32)) + activation_dequantize = network.add_dequantize( + activation_int8, + unit_scale, + trt.float32, + ) + weight_dequantize = network.add_dequantize( + quantized, + unit_scale, + trt.float32, + ) + if activation_dequantize is None or weight_dequantize is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot unit INT8 Q/DQ markers") + _name_convrot_layer(network, activation_dequantize, "activation_unit_dequantize") + _name_convrot_layer(network, weight_dequantize, "weight_unit_dequantize") + matmul = network.add_matrix_multiply( + activation_dequantize.get_output(0), + trt.MatrixOperation.NONE, + weight_dequantize.get_output(0), + trt.MatrixOperation.TRANSPOSE, + ) + if matmul is None: + raise RuntimeError("TensorRT rejected the MiniMax-H3 ConvRot INT8 linear") + _name_convrot_layer(network, matmul, "dynamic_rowwise_int8_gemm") + matmul.metadata = ( + "trtmc.quantization=int8_tensorwise_convrot;" + "activation=dynamic_int8_rowwise;weight=int8;accumulator=fp32" + ) + accumulator = cast(network, matmul.get_output(0), trt.float32) + weight_scale = weight_constant(network, scale.reshape(1, -1)) + output_scale = network.add_elementwise( + row_scale_f32, + weight_scale, + trt.ElementWiseOperation.PROD, + ) + if output_scale is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot combined output scale") + _name_convrot_layer(network, output_scale, "combined_scale_f32") + scaled = network.add_elementwise( + accumulator, + output_scale.get_output(0), + trt.ElementWiseOperation.PROD, + ) + if scaled is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot output scaling") + _name_convrot_layer(network, scaled, "output_scale_f32") + output = scaled.get_output(0) + if bias is not None: + shape = (1,) * (len(tuple(output.shape)) - 1) + (-1,) + bias_tensor = weight_constant(network, np.asarray(bias).reshape(shape)) + bias_tensor = cast(network, bias_tensor, trt.float32) + output = network.add_elementwise( + output, bias_tensor, trt.ElementWiseOperation.SUM + ).get_output(0) + return cast(network, output, trt.bfloat16) + + def linear( network, tensor, @@ -124,6 +413,11 @@ def linear( ): """PyTorch ``[out, in]`` linear expressed as native TensorRT GEMM.""" + if isinstance(weight, ConvRotInt8Weight): + if compute_dtype not in (None, trt.bfloat16) or not bf16: + raise ValueError("MiniMax-H3 ConvRot linears require BF16 activation compute") + return _convrot_int8_linear(network, tensor, weight, bias) + tensor_rank = len(tuple(tensor.shape)) rhs_value = np.asarray(weight) if tensor_rank > 2: @@ -218,26 +512,50 @@ def dynamic_slice(network, tensor, starts: tuple[int, ...], sizes: tuple[int | N if len(starts) != len(sizes): raise ValueError("MiniMax-H3 dynamic slice rank mismatch") - runtime_sizes = [ - _shape_dim(network, tensor, axis) if size is None else size - for axis, size in enumerate(sizes) - ] + runtime_sizes = [] + for axis, size in enumerate(sizes): + if size is not None: + runtime_sizes.append(size) + continue + static_size = int(tensor.shape[axis]) + runtime_sizes.append( + static_size if static_size >= 0 else _shape_dim(network, tensor, axis) + ) + if all(isinstance(size, (int, np.integer)) for size in runtime_sizes): + return network.add_slice( + tensor, + starts, + tuple(int(size) for size in runtime_sizes), + (1,) * len(sizes), + ).get_output(0) initial_sizes = tuple(1 if size is None else size for size in sizes) layer = network.add_slice(tensor, starts, initial_sizes, (1,) * len(sizes)) layer.set_input(2, _shape_vector(network, runtime_sizes)) return layer.get_output(0) -def slice_rows_from_end(network, tensor, *, offset: int, rows: int): - """Take fixed rows from a dynamic 2-D tensor, measured from its end.""" +def slice_rows_like_from_end(network, tensor, reference, *, trailing_reference=None): + """Take runtime-sized rows from ``tensor`` using modality input shapes. + + ``reference`` supplies the number of rows to return. When a trailing + modality is supplied, its runtime row count is included in the offset + measured from the packed sequence end. Only shape tensors participate in + this operation, so the reference contents are never copied or consumed. + """ total_rows = _shape_dim(network, tensor, 0) - offset_tensor = constant(network, np.asarray([offset], dtype=np.int64), dtype=np.int64) + rows = _shape_dim(network, reference, 0) + offset = rows + if trailing_reference is not None: + trailing_rows = _shape_dim(network, trailing_reference, 0) + offset = network.add_elementwise( + offset, trailing_rows, trt.ElementWiseOperation.SUM + ).get_output(0) start_row = network.add_elementwise( - total_rows, offset_tensor, trt.ElementWiseOperation.SUB + total_rows, offset, trt.ElementWiseOperation.SUB ).get_output(0) width = int(tensor.shape[1]) - layer = network.add_slice(tensor, (0, 0), (rows, width), (1, 1)) + layer = network.add_slice(tensor, (0, 0), (1, width), (1, 1)) layer.set_input(1, _shape_vector(network, (start_row, 0))) layer.set_input(2, _shape_vector(network, (rows, width))) return layer.get_output(0) @@ -283,20 +601,50 @@ def swiglu(network, tensor, weight_in, weight_out, ffn_dim: int): return linear(network, hidden, weight_out) -def fused_qkv(network, tensor, weights: dict, prefix: str): +def fused_qkv( + network, + tensor, + weights: dict, + prefix: str, + *, + consume_weights: bool = False, +): """Pack Q/K/V into one TensorRT GEMM, matching Sol-Engine's lossless path.""" - packed_weight = np.concatenate( - [weights[f"{prefix}.to_{name}.weight"] for name in ("q", "k", "v")], axis=0 - ) + keys = tuple(f"{prefix}.to_{name}.weight" for name in ("q", "k", "v")) + sources = tuple(weights[key] for key in keys) + if all(isinstance(source, ConvRotInt8Weight) for source in sources): + parents = tuple(source.packed_parent for source in sources) + if parents[0] is None or not all(parent is parents[0] for parent in parents): + raise ValueError( + "MiniMax-H3 quantized Q/K/V must share one packed parent" + ) + packed_weight = parents[0] + if not packed_weight.is_full_fused_qkv: + raise ValueError("MiniMax-H3 quantized QKV parent is not marked as full fused QKV") + expected_width = int(packed_weight.qweight.shape[0]) // 3 + expected_slices = tuple( + (index * expected_width, (index + 1) * expected_width) for index in range(3) + ) + if tuple(source.row_slice for source in sources) != expected_slices: + raise ValueError("MiniMax-H3 quantized Q/K/V row slices are not q-k-v ordered") + elif any(isinstance(source, ConvRotInt8Weight) for source in sources): + raise ValueError("MiniMax-H3 fused QKV cannot mix quantized and BF16 weights") + else: + packed_weight = np.concatenate(sources, axis=0) packed = linear(network, tensor, packed_weight) + if consume_weights: + # The packed array is retained by the network. Its three sources are + # no longer referenced and can be released before serialization. + for key in keys: + weights.pop(key) width = int(packed.shape[1]) // 3 return tuple( dynamic_slice(network, packed, (0, index * width), (None, width)) for index in range(3) ) -def rows_to_heads(network, tensor, rows: int, heads: int, head_dim: int): +def rows_to_heads(network, tensor, heads: int, head_dim: int): reshape = network.add_shuffle(tensor) reshape.reshape_dims = (-1, heads, head_dim) reshape.second_transpose = trt.Permutation([1, 0, 2]) @@ -305,7 +653,7 @@ def rows_to_heads(network, tensor, rows: int, heads: int, head_dim: int): return batch.get_output(0) -def heads_to_rows(network, tensor, rows: int, width: int): +def heads_to_rows(network, tensor, width: int): reshape = network.add_shuffle(tensor) reshape.first_transpose = trt.Permutation([0, 2, 1, 3]) reshape.reshape_dims = (-1, width) @@ -318,7 +666,6 @@ def partial_rope( cos_half, sin_half, *, - rows: int, heads: int, head_dim: int, rotary_dim: int, @@ -326,7 +673,7 @@ def partial_rope( ): """Apply H3's 96-channel rotate-half MM-RoPE with native layers.""" - value = rows_to_heads(network, tensor, rows, heads, head_dim) + value = rows_to_heads(network, tensor, heads, head_dim) if interleaved: raise ValueError("MiniMax-H3 uses rotate-half, non-interleaved RoPE") rotary = dynamic_slice(network, value, (0, 0, 0, 0), (1, heads, None, rotary_dim)) @@ -357,21 +704,17 @@ def duplicate_table(table): rotated = network.add_elementwise(left, right, trt.ElementWiseOperation.SUM).get_output(0) result = network.add_concatenation((rotated, passthrough)) result.axis = 3 - return heads_to_rows(network, result.get_output(0), rows, heads * head_dim) + return heads_to_rows(network, result.get_output(0), heads * head_dim) -def native_attention(network, q, k, v, *, rows: int, heads: int, head_dim: int, name: str): +def native_attention(network, q, k, v, *, heads: int, head_dim: int, name: str): """Full-sequence single-device fused TensorRT attention.""" - q = rows_to_heads(network, q, rows, heads, head_dim) - k = rows_to_heads(network, k, rows, heads, head_dim) - v = rows_to_heads(network, v, rows, heads, head_dim) - # TensorRT's fused FP16 attention has three additional mantissa bits over - # BF16. Keep the surrounding block in checkpoint-native BF16 while - # reducing recurrent numerical drift across 49 denoising evaluations. - q = cast(network, q, trt.float16) - k = cast(network, k, trt.float16) - v = cast(network, v, trt.float16) + q = rows_to_heads(network, q, heads, head_dim) + k = rows_to_heads(network, k, heads, head_dim) + v = rows_to_heads(network, v, heads, head_dim) + # Preserve BF16's exponent range. H3 residuals can exceed FP16's finite + # limit before a later block normalizes them. scale = constant( network, np.full((1, 1, 1, 1), 1.0 / math.sqrt(head_dim), dtype=np.float32), @@ -385,5 +728,4 @@ def native_attention(network, q, k, v, *, rows: int, heads: int, head_dim: int, layer.metadata = f"trtmc.native_op=IAttention;source={name}" layer.get_output(0).name = f"{name}.output" layer.decomposable = False - context = cast(network, layer.get_output(0), trt.bfloat16) - return heads_to_rows(network, context, rows, heads * head_dim) + return heads_to_rows(network, layer.get_output(0), heads * head_dim) diff --git a/families/minimax_h3/model.py b/families/minimax_h3/model.py index e33e746f5a..c6f6cbb697 100644 --- a/families/minimax_h3/model.py +++ b/families/minimax_h3/model.py @@ -5,37 +5,61 @@ from __future__ import annotations -import gc +from fractions import Fraction import json +import math from dataclasses import replace from pathlib import Path from types import SimpleNamespace from typing import TYPE_CHECKING -from .checkpoint import ( - load_selected_component_state_dict, - numpy_state, - validate_component_key_partition, -) from .config import ( - SOL_ENGINE_1344X768_124F, - default_workspace_limit_bytes, + CANVAS_MAX_ASPECT_RATIO, + CANVAS_MAX_PIXELS, + CANVAS_MIN_ASPECT_RATIO, + CANVAS_MULTIPLE, + CANVAS_SHORT_EDGE, + NATIVE_EXPLICIT_CANVAS_SIZES, + SOL_ENGINE_1344X768_124_TO_345F, + VIDEO_NUM_FRAMES_MAX, + VIDEO_NUM_FRAMES_MIN, + VIDEO_NUM_FRAMES_OPT, ) +from .runtime_config_schema import normalize_build_options as validate_build_options +from .delivery import SR_SOURCE_SHAPE if TYPE_CHECKING: from tensorrt_model_connect.build import BuildRequest from tensorrt_model_connect.bundle_writer import BundleWriter -def _fixed_profile(raw: dict): +def _effective_build_config(raw: dict) -> dict: + family_options = raw.get("_family_build_options", {}) + options = family_options.get("minimax_h3", {}) if isinstance(family_options, dict) else {} + if not isinstance(options, dict): + raise ValueError("minimax_h3 build options must be an object") + return {**raw, **validate_build_options(options)} + + +def _public_dynamic_profile(raw: dict): + profile = SOL_ENGINE_1344X768_124_TO_345F expected = { - "text_rows": SOL_ENGINE_1344X768_124F.text_rows, - "text_rows_min": SOL_ENGINE_1344X768_124F.min_text_rows, - "text_rows_opt": SOL_ENGINE_1344X768_124F.opt_text_rows, - "text_rows_max": SOL_ENGINE_1344X768_124F.text_rows, - "audio_rows": SOL_ENGINE_1344X768_124F.audio_rows, - "video_rows": SOL_ENGINE_1344X768_124F.video_rows, - "padded_sequence_length": SOL_ENGINE_1344X768_124F.padded_sequence_length, + "text_rows": profile.text_rows, + "text_rows_min": profile.min_text_rows, + "text_rows_opt": profile.opt_text_rows, + "text_rows_max": profile.text_rows, + "audio_rows": profile.opt_audio_rows, + "audio_rows_min": profile.min_audio_rows, + "audio_rows_opt": profile.opt_audio_rows, + "audio_rows_max": profile.audio_rows, + "video_rows": profile.opt_video_rows, + "video_rows_min": profile.min_video_rows, + "video_rows_opt": profile.opt_video_rows, + "video_rows_max": profile.video_rows, + "packed_sequence_length_min": profile.min_sequence_length, + "packed_sequence_length_opt": profile.opt_sequence_length, + "packed_sequence_length_max": profile.sequence_length, + "padded_sequence_length": profile.padded_sequence_length, } mismatches = { name: (raw[name], value) @@ -44,28 +68,141 @@ def _fixed_profile(raw: dict): } if mismatches: raise ValueError(f"Unsupported MiniMax-H3 packed-row profile: {mismatches}") - explicit_flag = raw.get("first_block_cache") - mode = raw.get( - "denoiser_cache_mode", - "first_block" if explicit_flag is True else "monolithic", - ) - if mode not in ("monolithic", "first_block"): - raise ValueError(f"Unsupported MiniMax-H3 denoiser_cache_mode: {mode!r}") - if explicit_flag is not None and not isinstance(explicit_flag, bool): - raise ValueError("MiniMax-H3 first_block_cache must be a boolean") - mode_flag = mode == "first_block" - if explicit_flag is not None and explicit_flag != mode_flag: - raise ValueError("MiniMax-H3 cache mode and first_block_cache flag disagree") - if not mode_flag: - return SOL_ENGINE_1344X768_124F - return replace(SOL_ENGINE_1344X768_124F, first_block_cache=True) - - -class _MiniMaxH3Model: - def load_weights(self, model_dir: str, config) -> dict: + explicit_flag = raw.get("first_block_cache", True) + if explicit_flag is not True: + raise ValueError("MiniMax-H3 only supports the dense FirstBlockCache build") + mode = raw.get("denoiser_cache_mode", "first_block") + if mode != "first_block": + raise ValueError("MiniMax-H3 only supports denoiser_cache_mode='first_block'") + return replace(profile, first_block_cache=True) + + +def _default_num_frames(raw: dict) -> int: + value = raw.get("video_num_frames", raw.get("num_frames", VIDEO_NUM_FRAMES_OPT)) + if isinstance(value, bool): + raise ValueError("MiniMax-H3 video_num_frames must be a valid 5--15 second geometry") + try: + frames = int(value) + except (TypeError, ValueError) as error: + raise ValueError( + "MiniMax-H3 video_num_frames must be a valid 5--15 second geometry" + ) from error + if not VIDEO_NUM_FRAMES_MIN <= frames <= VIDEO_NUM_FRAMES_MAX or frames % 17 != 5: + raise ValueError("MiniMax-H3 video_num_frames must be a valid 5--15 second geometry") + return frames + + +def _resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: + """Mirror the public H3 resolver with Python ties-to-even 32 rounding.""" + + if not math.isfinite(aspect_width) or not math.isfinite(aspect_height): + raise ValueError("MiniMax-H3 canvas aspect must be finite and positive") + if aspect_width <= 0.0 or aspect_height <= 0.0: + raise ValueError("MiniMax-H3 canvas aspect must be finite and positive") + ratio = aspect_width / aspect_height + if not CANVAS_MIN_ASPECT_RATIO <= ratio <= CANVAS_MAX_ASPECT_RATIO: + raise ValueError("MiniMax-H3 canvas aspect must be within 1:4 through 4:1") + if ratio >= 1.0: + height = float(CANVAS_SHORT_EDGE) + width = height * ratio + else: + width = float(CANVAS_SHORT_EDGE) + height = width / ratio + pixels = width * height + if pixels > CANVAS_MAX_PIXELS: + scale = math.sqrt(CANVAS_MAX_PIXELS / pixels) + width *= scale + height *= scale + resolved_width = int(round(width / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE + resolved_height = int(round(height / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE + return resolved_height, resolved_width + + +def _reachable_canvas_sizes() -> tuple[tuple[int, int], ...]: + """Enumerate the exact finite image of the continuous 32-rounded resolver.""" + + landscape = { + (CANVAS_SHORT_EDGE, width) + for width in range( + CANVAS_SHORT_EDGE, + int(CANVAS_SHORT_EDGE * 1.75) + CANVAS_MULTIPLE, + CANVAS_MULTIPLE, + ) + } + half = CANVAS_MULTIPLE // 2 + maximum_dimension = CANVAS_SHORT_EDGE * int(CANVAS_MAX_ASPECT_RATIO) + for height in range(CANVAS_MULTIPLE, CANVAS_SHORT_EDGE + 1, CANVAS_MULTIPLE): + for width in range(CANVAS_SHORT_EDGE, maximum_dimension + 1, CANVAS_MULTIPLE): + # In the area-limited landscape branch raw dimensions satisfy + # h*w=max_pixels and r=w/h. Intersect the two nearest-multiple + # rounding cells with the resolver's exact r interval. + lower = max( + Fraction(7, 4), + Fraction(CANVAS_MAX_PIXELS, (height + half) ** 2), + Fraction((width - half) ** 2, CANVAS_MAX_PIXELS), + ) + upper = min( + Fraction(4, 1), + Fraction(CANVAS_MAX_PIXELS, (height - half) ** 2), + Fraction((width + half) ** 2, CANVAS_MAX_PIXELS), + ) + if lower >= upper: + continue + sample = float((lower + upper) / 2) + if _resolve_canvas_size(sample, 1.0) == (height, width): + landscape.add((height, width)) + reachable = landscape | {(width, height) for height, width in landscape} + return tuple(sorted(reachable)) + + +def _default_canvas_size(raw: dict) -> tuple[int, int]: + sr = raw.get("super_resolution", False) + default_height, default_width = SR_SOURCE_SHAPE if sr else (768, 1344) + height_value = raw.get("video_height", raw.get("height", default_height)) + width_value = raw.get("video_width", raw.get("width", default_width)) + if isinstance(height_value, bool) or isinstance(width_value, bool): + raise ValueError("MiniMax-H3 video dimensions must match the public canvas resolver") + try: + height = int(height_value) + width = int(width_value) + except (TypeError, ValueError) as error: + raise ValueError( + "MiniMax-H3 video dimensions must match the public canvas resolver" + ) from error + if sr and (height, width) != SR_SOURCE_SHAPE: + raise ValueError("MiniMax-H3 super_resolution=true requires height=480 and width=864") + if ( + height <= 0 + or width <= 0 + or ( + (height, width) not in NATIVE_EXPLICIT_CANVAS_SIZES + and (height, width) != _resolve_canvas_size(width, height) + ) + ): + raise ValueError("MiniMax-H3 video dimensions must match the public canvas resolver") + return height, width + + +def _first_block_cache_threshold(raw: dict, *, default: float = 0.08) -> float: + value = raw.get("first_block_cache_threshold", default) + if isinstance(value, bool): + raise ValueError("MiniMax-H3 first_block_cache_threshold must be finite and positive") + try: + threshold = float(value) + except (TypeError, ValueError) as error: + raise ValueError( + "MiniMax-H3 first_block_cache_threshold must be finite and positive" + ) from error + if not math.isfinite(threshold) or threshold <= 0.0: + raise ValueError("MiniMax-H3 first_block_cache_threshold must be finite and positive") + return threshold + + +class MiniMaxH3Plugin: + def load_weights(self, model_dir: str, config, **_kwargs) -> dict: del config root = Path(model_dir) - required_dirs = ("transformer", "text_encoder", "vae", "audio_vae", "tokenizer") + required_dirs = ("transformer", "vae", "audio_vae", "tokenizer") missing = [str(root / name) for name in required_dirs if not (root / name).is_dir()] if missing: raise FileNotFoundError( @@ -86,241 +223,131 @@ def load_weights(self, model_dir: str, config) -> dict: } if mismatches: raise ValueError(f"Unsupported MiniMax-H3 transformer architecture: {mismatches}") - return { - "_model_dir": str(root), - "_transformer_dir": str(root / "transformer"), - "_text_encoder_dir": str(root / "text_encoder"), - "_vae_dir": str(root / "vae"), - "_audio_vae_dir": str(root / "audio_vae"), - "_tokenizer_dir": str(root / "tokenizer"), - } + return {"_model_dir": str(root)} - def build_components( + def build_staged_bundle( self, model_dir: str, + writer: "BundleWriter", config, weights: dict, *, - precision: str = "bf16", + plans_dir: str | Path, + precision: str, verbose: bool = False, - ) -> dict: - del model_dir - if precision.lower() != "bf16": - raise ValueError("MiniMax-H3 native builds require BF16 checkpoint weights") - raw = getattr(config, "raw", {}) - profile = _fixed_profile(raw) - profile.validate() - workspace_limits = default_workspace_limit_bytes( - first_block_cache=profile.first_block_cache - ) - from .adaln_builder import build_adaln_precompute_engine - from .adaln_builder import checkpoint_keys as adaln_checkpoint_keys - from .dit_builder import ( - build_dit_engine, - build_dit_finish_engine, - build_dit_head_engine, - build_dit_tail_engine, - checkpoint_keys as dit_checkpoint_keys, - finish_checkpoint_keys, - head_checkpoint_keys, - tail_checkpoint_keys, - ) - from .text_encoder_builder import ( - build_text_encoder_engine, - checkpoint_keys as text_encoder_checkpoint_keys, - ) - - if profile.first_block_cache: - denoiser_specs = ( - ( - "denoiser_head", - "denoiser_head.plan", - build_dit_head_engine, - head_checkpoint_keys(profile), - ), - ( - "denoiser_tail", - "denoiser_tail.plan", - build_dit_tail_engine, - tail_checkpoint_keys(profile), - ), - ( - "denoiser_finish", - "denoiser_finish.plan", - build_dit_finish_engine, - finish_checkpoint_keys(profile), - ), - ) - checkpoint_groups = ( - adaln_checkpoint_keys(profile), - *(spec[3] for spec in denoiser_specs), - ) - else: - denoiser_specs = ( - ( - "denoiser", - "denoiser.plan", - build_dit_engine, - dit_checkpoint_keys(profile), - ), - ) - checkpoint_groups = ( - adaln_checkpoint_keys(profile), - dit_checkpoint_keys(profile), - ) - validate_component_key_partition(weights["_transformer_dir"], checkpoint_groups) + parallel_config=None, + max_batch_size: int = 1, + ) -> None: + """Build a staged RTX bundle without retaining serialized plans in RAM.""" - text_state = load_selected_component_state_dict( - weights["_text_encoder_dir"], text_encoder_checkpoint_keys() - ) - text_weights = numpy_state(text_state) - del text_state - text_encoder_plan = build_text_encoder_engine( - text_weights, - sequence_length=profile.text_rows, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits["text_encoder.plan"], - ) - del text_weights - gc.collect() + if precision.lower() != "bf16": + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds require BF16") + if max_batch_size != 1: + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds require max_batch_size=1") + mode = str(getattr(parallel_config, "mode", "single")) + if mode != "single": + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds require one GPU") - adaln_state = load_selected_component_state_dict( - weights["_transformer_dir"], adaln_checkpoint_keys(profile) - ) - adaln_weights = numpy_state(adaln_state) - del adaln_state - adaln_plan = build_adaln_precompute_engine( - adaln_weights, - profile, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits["adaln_precompute.plan"], - ) - del adaln_weights - gc.collect() + raw = _effective_build_config(getattr(config, "raw", {})) + if raw.get("_fp32_layers"): + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds do not support FP32 layers") + from .delivery import resolve_quantized_sources, resolve_super_resolution_sources - denoiser_components = {} - for component_name, filename, denoiser_builder, selected_keys in denoiser_specs: - dit_state = load_selected_component_state_dict( - weights["_transformer_dir"], selected_keys - ) - dit_weights = numpy_state(dit_state) - del dit_state - denoiser_plan = denoiser_builder( - dit_weights, - profile, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits[filename], - ) - del dit_weights - gc.collect() - denoiser_components[component_name] = denoiser_plan + root = Path(weights.get("_model_dir", model_dir)) + staged_raw = dict(raw) + staged_raw.setdefault("first_block_cache", True) + staged_raw.setdefault("denoiser_cache_mode", "first_block") + _public_dynamic_profile(staged_raw) + expected_request = { + "num_inference_steps": 50, + "seed": 0, + } + mismatches = { + name: (raw[name], value) + for name, value in expected_request.items() + if name in raw and int(raw[name]) != value + } + if mismatches: + raise ValueError(f"Unsupported MiniMax-H3 staged profile: {mismatches}") + _default_canvas_size(raw) + _default_num_frames(raw) - from .vae_builder import ( - build_vae_tile_decoder_engine, - checkpoint_keys as vae_checkpoint_keys, - ) + from .staged_build import build_staged_bundle - vae_state = load_selected_component_state_dict(weights["_vae_dir"], vae_checkpoint_keys()) - vae_weights = numpy_state(vae_state) - del vae_state - vae_decoder_plan = build_vae_tile_decoder_engine( - vae_weights, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits["vae_tile_decoder.plan"], - ) - tokenizer_json = (Path(weights["_tokenizer_dir"]) / "tokenizer.json").read_bytes() - - return { - "text_encoder": text_encoder_plan, - "adaln_precompute": adaln_plan, - **denoiser_components, - "vae_decoder": vae_decoder_plan, - "profile": profile, - # Text/VAE paths remain explicit so follow-on native component - # builders cannot silently substitute a different checkpoint. - "vae_dir": weights["_vae_dir"], - "audio_vae_dir": weights["_audio_vae_dir"], - "tokenizer_dir": weights["_tokenizer_dir"], - "tokenizer_json": tokenizer_json, + # All public builds use the same quantized checkpoints for all three + # workflows. Only the explicit SR flag changes the delivery mode. + staged_options = {"verbose": verbose, **resolve_quantized_sources(root, raw)} + super_resolution_model, super_resolution_weak_model = resolve_super_resolution_sources(raw) + staged_options["runtime_defaults"] = { + "height": _default_canvas_size(raw)[0], + "width": _default_canvas_size(raw)[1], + "num_frames": _default_num_frames(raw), + "first_block_cache_threshold": _first_block_cache_threshold(raw), } + staged_options["runtime_defaults"].update( + ref2va_first_block_cache=raw.get("ref2va_first_block_cache", True), + ref2va_first_block_cache_threshold=raw.get("ref2va_first_block_cache_threshold", 0.08), + ) + if super_resolution_model is not None: + staged_options["super_resolution_model"] = super_resolution_model + if super_resolution_weak_model is not None: + staged_options["super_resolution_weak_model"] = super_resolution_weak_model + build_staged_bundle(root, writer, plans_dir=plans_dir, **staged_options) -def build(request: "BuildRequest", writer: "BundleWriter") -> None: - """Build one MiniMax-H3 image-generation bundle.""" - if request.dynamic_kv_cache: - raise NotImplementedError("minimax_h3 does not support dynamic_kv_cache") +plugin = MiniMaxH3Plugin() - if request.max_sequence_length is not None: - raise NotImplementedError("minimax_h3 does not support max_sequence_length") - if request.context_parallel_size != 1: - raise ValueError("this family does not support context parallelism") +def build(request: "BuildRequest", writer: "BundleWriter") -> None: + """Build the unified T2VA/FL2VA/Ref2VA bundle through TensorRT-RTX.""" if request.task != "image_generation": raise ValueError("minimax_h3 supports only task=image_generation") + if request.backend != "trt_rtx": + raise ValueError("MiniMax-H3 requires backend=trt_rtx") if request.precision != "bf16": raise ValueError("MiniMax-H3 requires precision=bf16") - if request.tensor_parallel_size != 1: - raise NotImplementedError("MiniMax-H3 requires tensor_parallel_size=1") + if request.dynamic_kv_cache: + raise NotImplementedError("minimax_h3 does not support dynamic_kv_cache") + if request.max_sequence_length is not None: + raise NotImplementedError("minimax_h3 does not support max_sequence_length") + if request.tensor_parallel_size != 1 or request.context_parallel_size != 1: + raise ValueError("MiniMax-H3 requires single-device parallel settings") if request.max_batch_size != 1: - raise NotImplementedError("MiniMax-H3 requires max_batch_size=1") - if request.quantization not in {None, "none"}: - raise NotImplementedError("MiniMax-H3 does not support quantization") + raise ValueError("MiniMax-H3 requires max_batch_size=1") if request.fp32_layers: raise NotImplementedError("MiniMax-H3 does not support fp32_layers") - if int(request.image_height or 768) != 768 or int(request.image_width or 1344) != 1344: - raise ValueError("MiniMax-H3 requires image_height=768 and image_width=1344") - if int(request.video_num_frames or 124) != 124: - raise ValueError("MiniMax-H3 requires video_num_frames=124") - - config = SimpleNamespace(raw={}) - model_dir = Path(request.model_dir) - model = _MiniMaxH3Model() - weights = model.load_weights(str(model_dir), config) - components = model.build_components( - str(model_dir), + + family_options = validate_build_options(dict(getattr(request, "family_options", ()))) + if request.quantization not in {None, "int8_tensorwise_convrot"}: + raise ValueError("MiniMax-H3 delivery requires Comfy INT8 denoisers and the NVFP4 text checkpoint") + if (request.image_height is None) != (request.image_width is None): + raise ValueError("MiniMax-H3 requires both image height and width, or neither") + sr = family_options.get("super_resolution", False) + default_height, default_width = SR_SOURCE_SHAPE if sr else (768, 1344) + + raw = { + "_family_build_options": {"minimax_h3": family_options}, + "height": int(request.image_height or default_height), + "width": int(request.image_width or default_width), + "video_num_frames": int(request.video_num_frames or VIDEO_NUM_FRAMES_OPT), + "num_inference_steps": 50, + "seed": 0, + } + config = SimpleNamespace(raw=raw) + _default_canvas_size(_effective_build_config(raw)) + _default_num_frames(raw) + weights = plugin.load_weights(str(request.model_dir), config) + + writer.set_header(family="minimax_h3", task=request.task, backend=request.backend) + plans_dir = request.output_path.with_name(f"{request.output_path.name}.plans") + plugin.build_staged_bundle( + str(request.model_dir), + writer, config, weights, + plans_dir=plans_dir, precision=request.precision, verbose=request.verbose, - ) - profile = components["profile"] - if profile.first_block_cache: - raise RuntimeError("MiniMax-H3 minimal build uses the monolithic denoiser profile") - - writer.set_header(family="minimax_h3", task=request.task, backend=request.backend) - writer.add_bytes("text_encoder.plan", components["text_encoder"]) - writer.add_bytes("adaln.plan", components["adaln_precompute"]) - writer.add_bytes("denoiser.plan", components["denoiser"]) - writer.add_bytes("vae.plan", components["vae_decoder"]) - writer.add_bytes("tokenizer.json", components["tokenizer_json"]) - writer.add_json( - "runtime.json", - { - "height": 768, - "width": 1344, - "num_frames": 124, - "fps": 24, - "num_inference_steps": 50, - "seed": 0, - "first_block_cache": False, - "denoiser_cache_mode": "monolithic", - "first_block_cache_threshold": 0.025, - "text_rows": profile.text_rows, - "text_rows_min": profile.min_text_rows, - "text_rows_opt": profile.opt_text_rows, - "text_rows_max": profile.text_rows, - "audio_rows": profile.audio_rows, - "video_rows": profile.video_rows, - "padded_sequence_length": profile.padded_sequence_length, - "max_timestep_count": profile.max_timestep_count, - "context_parallel_size": profile.context_parallel_size, - "vae_tile_batch": 28, - "vae_tile_size": 256, - "vae_tile_overlap": 64, - }, + max_batch_size=request.max_batch_size, ) diff --git a/families/minimax_h3/multimodal_text_encoder_builder.py b/families/minimax_h3/multimodal_text_encoder_builder.py new file mode 100644 index 0000000000..dd9ae604b4 --- /dev/null +++ b/families/minimax_h3/multimodal_text_encoder_builder.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared native Qwen3-VL language stack for MiniMax-H3 public workflows. + +The plan contains one copy of the Qwen embedding and decoder layers 0..49. +T2VA binds one disabled dummy visual row; FL2VA and Ref2VA bind compact vision +features plus their presentation-row indices. The engine scatters the main +features before layer zero, injects DeepStack features after layers 0..2, and +emits ``hidden_states[50]`` before the full model's final normalization. +""" + +from __future__ import annotations + +import gc +import math +import sys +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .config import TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES +from .fl2va_contract import MultimodalTextProfile, QWEN_TEXT_HIDDEN_SIZE, text_encoder_abi + + +trt = trt_compat.get_trt() + +HIDDEN_SIZE = QWEN_TEXT_HIDDEN_SIZE +NUM_LAYERS = 50 +NUM_HEADS = 64 +NUM_KV_HEADS = 8 +HEAD_DIM = 128 +NORM_EPS = 1.0e-6 +# Bit-exact output of Qwen3VLTextRotaryEmbedding's Torch FP32 construction. +# Recomputing ``theta**exponent`` with NumPy differs by one ULP at index 19. +ROPE_INV_FREQ_BITS = ( + 0x3F800000, + 0x3F492C28, + 0x3F1E165D, + 0x3EF875A7, + 0x3EC33F3A, + 0x3E996E52, + 0x3E712429, + 0x3E3D7EFC, + 0x3E14E963, + 0x3DEA09DB, + 0x3DB7EA1A, + 0x3D908687, + 0x3D63251B, + 0x3D327F50, + 0x3D0C44BF, + 0x3CDC7456, + 0x3CAD3D5E, + 0x3C88230F, + 0x3C55F605, + 0x3C282311, + 0x3C042088, + 0x3BCFA8A9, + 0x3BA32F3E, + 0x3B803C3D, + 0x3B498AD3, + 0x3B1E60C3, + 0x3AF8EA93, + 0x3AC39B1C, + 0x3A99B686, + 0x3A7195A5, + 0x3A3DD829, + 0x3A152F76, + 0x39EA77FF, + 0x39B840A8, + 0x3990CA8B, + 0x39639000, + 0x3932D34F, + 0x390C86C1, + 0x38DCDC14, + 0x38AD8EE4, + 0x38886320, + 0x38565AB5, + 0x38287230, + 0x38045EB6, + 0x37D00A62, + 0x37A37C09, + 0x37807896, + 0x3749E9AB, + 0x371EAB4B, + 0x36F95FB6, + 0x36C3F72A, + 0x3699FEDC, + 0x36720755, + 0x363E3180, + 0x361575AB, + 0x35EAE655, + 0x35B8975D, + 0x35910EAE, + 0x3563FB18, + 0x35332777, + 0x350CC8E3, + 0x34DD4405, + 0x34ADE091, + 0x3488A34F, +) + + +def checkpoint_keys() -> tuple[str, ...]: + """Return the exhaustive shared language-weight partition.""" + + names = ["model.language_model.embed_tokens.weight"] + for index in range(NUM_LAYERS): + prefix = f"model.language_model.layers.{index}" + names.extend( + [ + f"{prefix}.input_layernorm.weight", + f"{prefix}.post_attention_layernorm.weight", + f"{prefix}.self_attn.q_proj.weight", + f"{prefix}.self_attn.k_proj.weight", + f"{prefix}.self_attn.v_proj.weight", + f"{prefix}.self_attn.o_proj.weight", + f"{prefix}.self_attn.q_norm.weight", + f"{prefix}.self_attn.k_norm.weight", + f"{prefix}.mlp.gate_proj.weight", + f"{prefix}.mlp.up_proj.weight", + f"{prefix}.mlp.down_proj.weight", + ] + ) + return tuple(names) + + +def _network_shape(binding) -> tuple[int, ...]: + return tuple( + minimum if minimum == maximum else -1 + for minimum, maximum in zip(binding.min_shape, binding.max_shape) + ) + + +def _per_head_norm(network, tensor, weight, heads: int): + reshape = network.add_shuffle(tensor) + reshape.reshape_dims = (-1, heads, HEAD_DIM) + normalized = op.rms_norm(network, reshape.get_output(0), weight, HEAD_DIM, NORM_EPS) + flatten = network.add_shuffle(normalized) + flatten.reshape_dims = (-1, heads * HEAD_DIM) + return flatten.get_output(0) + + +def _repeat_kv(network, tensor): + repeated = [] + repeat = NUM_HEADS // NUM_KV_HEADS + for index in range(NUM_KV_HEADS): + head = op.dynamic_slice(network, tensor, (0, index, 0, 0), (1, 1, None, HEAD_DIM)) + repeated.extend([head] * repeat) + concatenation = network.add_concatenation(repeated) + concatenation.axis = 1 + return concatenation.get_output(0) + + +def _mrope_cache(network, position_ids): + """Build Qwen3-VL interleaved ``[T,H,W]`` MRoPE in FP32.""" + + positions = op.cast(network, position_ids, trt.float32) + inverse = np.asarray(ROPE_INV_FREQ_BITS, dtype=np.uint32).view(np.float32) + inverse = op.constant(network, inverse.reshape(1, HEAD_DIM // 2)) + frequencies = [] + for axis in range(3): + coordinate = op.dynamic_slice(network, positions, (axis, 0), (1, None)) + reshape = network.add_shuffle(coordinate) + reshape.reshape_dims = (-1, 1) + frequencies.append( + network.add_elementwise( + reshape.get_output(0), inverse, trt.ElementWiseOperation.PROD + ).get_output(0) + ) + + height_mask = np.zeros((1, HEAD_DIM // 2), dtype=np.float32) + width_mask = np.zeros_like(height_mask) + height_mask[:, 1:60:3] = 1.0 + width_mask[:, 2:60:3] = 1.0 + temporal_mask = 1.0 - height_mask - width_mask + selected = [] + for frequency, mask in zip(frequencies, (temporal_mask, height_mask, width_mask)): + selected.append( + network.add_elementwise( + frequency, + op.constant(network, mask), + trt.ElementWiseOperation.PROD, + ).get_output(0) + ) + frequency = network.add_elementwise( + selected[0], selected[1], trt.ElementWiseOperation.SUM + ).get_output(0) + frequency = network.add_elementwise( + frequency, selected[2], trt.ElementWiseOperation.SUM + ).get_output(0) + cos = network.add_unary(frequency, trt.UnaryOperation.COS).get_output(0) + sin = network.add_unary(frequency, trt.UnaryOperation.SIN).get_output(0) + return op.cast(network, cos, trt.bfloat16), op.cast(network, sin, trt.bfloat16) + + +def _visual_count_active(network, vision_count): + zero = op.constant(network, np.zeros((1,), dtype=np.int32), dtype=np.int32) + active = network.add_elementwise( + vision_count, zero, trt.ElementWiseOperation.GREATER + ).get_output(0) + reshape = network.add_shuffle(active) + reshape.reshape_dims = (1, 1) + return reshape.get_output(0) + + +def _scatter_visual_rows(network, base_like, row_indices, compact, active): + """Scatter compact valid features into a dynamic sequence-sized zero base.""" + + zero = op.constant(network, np.zeros((1, 1), dtype=np.float32)) + zero = op.cast(network, zero, base_like.dtype) + base = network.add_elementwise(base_like, zero, trt.ElementWiseOperation.PROD).get_output(0) + compact = op.cast(network, compact, base_like.dtype) + compact = network.add_select(active, compact, zero).get_output(0) + indices = network.add_shuffle(row_indices) + indices.reshape_dims = (-1, 1) + scatter = network.add_scatter( + base, + indices.get_output(0), + compact, + trt.ScatterMode.ND, + ) + if scatter is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 compact visual-row scatter") + return scatter.get_output(0) + + +def _linear(network, hidden, weights, name: str): + prescale = weights.get(f"{name}.pre_quant_scale") + if prescale is not None: + prescale = op.weight_constant(network, np.asarray(prescale).reshape(1, -1)) + prescale = op.cast(network, prescale, hidden.dtype) + hidden = network.add_elementwise( + hidden, prescale, trt.ElementWiseOperation.PROD + ).get_output(0) + return op.linear(network, hidden, weights[f"{name}.weight"]) + + +@op.cleanup_failed_build +def build_multimodal_text_encoder_engine( + weights: dict[str, np.ndarray], + profile: MultimodalTextProfile = MultimodalTextProfile(), + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the unified T2VA/FL2VA/Ref2VA Qwen language plan.""" + + profile.validate() + expected_keys = set(checkpoint_keys()) + optional_keys = { + f"model.language_model.layers.{index}.{name}.pre_quant_scale" + for index in range(NUM_LAYERS) + for name in ("self_attn.o_proj", "mlp.down_proj") + } + missing = sorted(expected_keys - set(weights)) + unexpected = sorted(set(weights) - expected_keys - optional_keys) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 multimodal text checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + for name in optional_keys & set(weights): + width = np.asarray(weights[f"{name.removesuffix('.pre_quant_scale')}.weight"]).shape[1] + value = np.asarray(weights[name]) + if value.shape != (width,) or not np.isfinite(value).all(): + raise ValueError(f"Invalid MiniMax-H3 AWQ input scale: {name}") + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + + abi = text_encoder_abi(profile) + dtype_map = {"float32": trt.float32, "int32": trt.int32} + inputs = { + binding.name: network.add_input( + binding.name, dtype_map[binding.dtype], _network_shape(binding) + ) + for binding in abi.inputs + } + optimization = builder.create_optimization_profile() + for binding in abi.inputs: + if binding.min_shape != binding.max_shape: + optimization.set_shape( + binding.name, binding.min_shape, binding.opt_shape, binding.max_shape + ) + config.add_optimization_profile(optimization) + + table = op.weight_constant(network, weights["model.language_model.embed_tokens.weight"]) + table = op.cast(network, table, trt.bfloat16) + hidden = network.add_gather(table, inputs["input_ids"], 0).get_output(0) + + active = _visual_count_active(network, inputs["vision_count"]) + visual = _scatter_visual_rows( + network, + hidden, + inputs["vision_row_indices"], + inputs["vision_embeds"], + active, + ) + zero_mask = op.constant(network, np.zeros((1, 1), dtype=np.float32)) + visual_mask = network.add_elementwise( + inputs["vision_mask"], zero_mask, trt.ElementWiseOperation.GREATER + ).get_output(0) + hidden = network.add_select(visual_mask, visual, hidden).get_output(0) + cos, sin = _mrope_cache(network, inputs["mrope_position_ids"]) + + for index in range(NUM_LAYERS): + prefix = f"model.language_model.layers.{index}" + normalized = op.rms_norm( + network, + hidden, + weights[f"{prefix}.input_layernorm.weight"], + HIDDEN_SIZE, + NORM_EPS, + ) + query = _linear(network, normalized, weights, f"{prefix}.self_attn.q_proj") + key = _linear(network, normalized, weights, f"{prefix}.self_attn.k_proj") + value = _linear(network, normalized, weights, f"{prefix}.self_attn.v_proj") + query = _per_head_norm( + network, query, weights[f"{prefix}.self_attn.q_norm.weight"], NUM_HEADS + ) + key = _per_head_norm( + network, key, weights[f"{prefix}.self_attn.k_norm.weight"], NUM_KV_HEADS + ) + query = op.partial_rope( + network, + query, + cos, + sin, + heads=NUM_HEADS, + head_dim=HEAD_DIM, + rotary_dim=HEAD_DIM, + ) + key = op.partial_rope( + network, + key, + cos, + sin, + heads=NUM_KV_HEADS, + head_dim=HEAD_DIM, + rotary_dim=HEAD_DIM, + ) + query_heads = op.rows_to_heads(network, query, NUM_HEADS, HEAD_DIM) + key_heads = op.rows_to_heads(network, key, NUM_KV_HEADS, HEAD_DIM) + value_heads = op.rows_to_heads(network, value, NUM_KV_HEADS, HEAD_DIM) + key_heads = _repeat_kv(network, key_heads) + value_heads = _repeat_kv(network, value_heads) + scale = op.constant( + network, + np.full((1, 1, 1, 1), 1.0 / math.sqrt(HEAD_DIM), dtype=np.float32), + ) + scale = op.cast(network, scale, query_heads.dtype) + query_heads = network.add_elementwise( + query_heads, scale, trt.ElementWiseOperation.PROD + ).get_output(0) + attention = network.add_attention( + query_heads, + key_heads, + value_heads, + trt.AttentionNormalizationOp.SOFTMAX, + True, + ) + if attention is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 Qwen attention layer {index}") + attention.name = f"{prefix}.self_attn.native_attention" + attention.metadata = f"trtmc.native_op=IAttention;source={attention.name}" + attention.get_output(0).name = f"{attention.name}.output" + attention.decomposable = False + update = op.heads_to_rows( + network, + attention.get_output(0), + NUM_HEADS * HEAD_DIM, + ) + update = _linear(network, update, weights, f"{prefix}.self_attn.o_proj") + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + normalized = op.rms_norm( + network, + hidden, + weights[f"{prefix}.post_attention_layernorm.weight"], + HIDDEN_SIZE, + NORM_EPS, + ) + gate = _linear(network, normalized, weights, f"{prefix}.mlp.gate_proj") + up = _linear(network, normalized, weights, f"{prefix}.mlp.up_proj") + gate = op.silu(network, gate) + gated = network.add_elementwise(gate, up, trt.ElementWiseOperation.PROD).get_output(0) + update = _linear(network, gated, weights, f"{prefix}.mlp.down_proj") + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + if index < 3: + deepstack = _scatter_visual_rows( + network, + hidden, + inputs["vision_row_indices"], + inputs[f"deepstack_{index}"], + active, + ) + hidden = network.add_elementwise( + hidden, deepstack, trt.ElementWiseOperation.SUM + ).get_output(0) + + output = op.cast(network, hidden, trt.float32) + output.name = abi.outputs[0].name + network.mark_output(output) + op.validate_native_network( + network, expected_attentions=NUM_LAYERS, label="multimodal text encoder" + ) + print( + "[minimax-h3] building shared Qwen3-VL language stack: " + f"layers={NUM_LAYERS}, sequence={profile.min_sequence_length}.." + f"{profile.max_sequence_length}, compact_vision={profile.min_vision_rows}.." + f"{profile.max_vision_rows}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 multimodal text encoder") + del network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/multimodal_vision_builder.py b/families/minimax_h3/multimodal_vision_builder.py new file mode 100644 index 0000000000..bcd80e7291 --- /dev/null +++ b/families/minimax_h3/multimodal_vision_builder.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native Qwen3-VL vision tower shared by MiniMax-H3 FL2VA and Ref2VA. + +The engine evaluates one image or one temporal video-patch block. Ref2VA video +conditioning invokes it once per temporal block and concatenates outputs, which +is exactly Qwen3-VL's packed ``cu_seqlens`` attention boundary. Interpolation +indices/weights and 2-D rotary positions are runtime inputs because reference +images have dynamic aspect and Ref2VA images use a larger 2048-short-edge grid. +The learned position table and every vision weight remain inside this one plan. +""" + +from __future__ import annotations + +import gc +import math +import sys +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .fl2va_contract import ( + QWEN_VISION_HIDDEN_SIZE, + QWEN_VISION_MERGE_SIZE, + QWEN_VISION_PATCH_WIDTH, + VisionEncoderProfile, + vision_encoder_abi, +) + + +trt = trt_compat.get_trt() + +DEPTH = 27 +NUM_HEADS = 16 +HEAD_DIM = QWEN_VISION_HIDDEN_SIZE // NUM_HEADS +INTERMEDIATE_SIZE = 4304 +POSITION_TABLE_SIDE = 48 +MERGE_UNIT = QWEN_VISION_MERGE_SIZE**2 +MERGED_HIDDEN_SIZE = QWEN_VISION_HIDDEN_SIZE * MERGE_UNIT +DEEPSTACK_VISUAL_INDEXES = (8, 16, 24) +NORM_EPS = 1.0e-6 +VISION_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 +ROPE_INV_FREQ_BITS = ( + 0x3F800000, + 0x3F1977CC, + 0x3EB800D6, + 0x3E5C9D35, + 0x3E044133, + 0x3D9E91B6, + 0x3D3E1E95, + 0x3CE3F280, + 0x3C88A69B, + 0x3C23D70A, + 0x3BC47060, + 0x3B6B8631, + 0x3B0D3169, + 0x3AA94938, + 0x3A4AF7F3, + 0x39F35A5C, + 0x3991E2E1, + 0x392EE9BF, +) + + +def checkpoint_keys() -> tuple[str, ...]: + names = ["model.visual.patch_embed.proj.weight", "model.visual.patch_embed.proj.bias"] + names.append("model.visual.pos_embed.weight") + for index in range(DEPTH): + prefix = f"model.visual.blocks.{index}" + names.extend( + [ + f"{prefix}.norm1.weight", + f"{prefix}.norm1.bias", + f"{prefix}.attn.qkv.weight", + f"{prefix}.attn.qkv.bias", + f"{prefix}.attn.proj.weight", + f"{prefix}.attn.proj.bias", + f"{prefix}.norm2.weight", + f"{prefix}.norm2.bias", + f"{prefix}.mlp.linear_fc1.weight", + f"{prefix}.mlp.linear_fc1.bias", + f"{prefix}.mlp.linear_fc2.weight", + f"{prefix}.mlp.linear_fc2.bias", + ] + ) + for prefix in [ + "model.visual.merger", + *(f"model.visual.deepstack_merger_list.{index}" for index in range(3)), + ]: + names.extend( + [ + f"{prefix}.norm.weight", + f"{prefix}.norm.bias", + f"{prefix}.linear_fc1.weight", + f"{prefix}.linear_fc1.bias", + f"{prefix}.linear_fc2.weight", + f"{prefix}.linear_fc2.bias", + ] + ) + return tuple(names) + + +def _layer_norm(network, hidden, weights, prefix: str, width: int): + rank = len(tuple(hidden.shape)) + shape = (1,) * (rank - 1) + (width,) + gamma = op.weight_constant(network, np.asarray(weights[f"{prefix}.weight"]).reshape(shape)) + beta = op.weight_constant(network, np.asarray(weights[f"{prefix}.bias"]).reshape(shape)) + gamma = op.cast(network, gamma, hidden.dtype) + beta = op.cast(network, beta, hidden.dtype) + layer = network.add_normalization_v2(hidden, gamma, beta, 1 << (rank - 1)) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 Qwen vision LayerNorm {prefix}") + layer.name = prefix + layer.epsilon = NORM_EPS + return layer.get_output(0) + + +def _gelu(network, hidden, activation_type): + layer = network.add_activation(hidden, activation_type) + if layer is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen vision GELU") + return layer.get_output(0) + + +def _interpolated_position_embeddings(network, weights, indices, interp_weights): + table = op.weight_constant(network, weights["model.visual.pos_embed.weight"]) + table = op.cast(network, table, trt.bfloat16) + gathered = network.add_gather(table, indices, 0) + if gathered is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen position-table gather") + # The learned table is BF16, but Transformers multiplies it by FP32 + # bilinear weights and performs the four-way sum in FP32 before rounding + # the completed positional embedding back to the patch-embedding dtype. + gathered_value = op.cast(network, gathered.get_output(0), trt.float32) + coefficients = network.add_shuffle(interp_weights) + coefficients.reshape_dims = (-1, 4, 1) + coefficients_value = op.cast(network, coefficients.get_output(0), trt.float32) + weighted = network.add_elementwise( + gathered_value, coefficients_value, trt.ElementWiseOperation.PROD + ) + if weighted is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen position interpolation") + reduced = network.add_reduce(weighted.get_output(0), trt.ReduceOperation.SUM, 1 << 1, False) + if reduced is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen position interpolation reduce") + return reduced.get_output(0) + + +def _vision_rope_cache(network, position_ids): + positions = op.cast(network, position_ids, trt.float32) + # Preserve Qwen3VLVisionRotaryEmbedding's Torch-created FP32 buffer + # without making Torch a builder dependency. + inverse = np.asarray(ROPE_INV_FREQ_BITS, dtype=np.uint32).view(np.float32) + inverse = op.constant(network, inverse.reshape(1, -1)) + axes = [] + for axis in range(2): + coordinate = op.dynamic_slice(network, positions, (0, axis), (None, 1)) + axes.append( + network.add_elementwise(coordinate, inverse, trt.ElementWiseOperation.PROD).get_output( + 0 + ) + ) + frequency = network.add_concatenation(axes) + frequency.axis = 1 + cos = network.add_unary(frequency.get_output(0), trt.UnaryOperation.COS).get_output(0) + sin = network.add_unary(frequency.get_output(0), trt.UnaryOperation.SIN).get_output(0) + # Qwen3-VL explicitly performs vision rotary embedding in FP32 and rounds + # q/k back to their source dtype only after both products are added. + return cos, sin + + +def _vision_partial_rope(network, tensor, cos_half, sin_half): + """Qwen3-VL's FP32 vision rotate-half operation.""" + + value = op.rows_to_heads(network, tensor, NUM_HEADS, HEAD_DIM) + source_dtype = value.dtype + value = op.cast(network, value, trt.float32) + first = op.dynamic_slice(network, value, (0, 0, 0, 0), (1, NUM_HEADS, None, HEAD_DIM // 2)) + second = op.dynamic_slice( + network, + value, + (0, 0, 0, HEAD_DIM // 2), + (1, NUM_HEADS, None, HEAD_DIM // 2), + ) + negative_second = network.add_unary(second, trt.UnaryOperation.NEG).get_output(0) + rotated_half = network.add_concatenation((negative_second, first)) + rotated_half.axis = 3 + + def duplicate(table): + reshape = network.add_shuffle(op.cast(network, table, trt.float32)) + reshape.reshape_dims = (1, 1, -1, HEAD_DIM // 2) + result = network.add_concatenation((reshape.get_output(0), reshape.get_output(0))) + result.axis = 3 + return result.get_output(0) + + left = network.add_elementwise( + value, duplicate(cos_half), trt.ElementWiseOperation.PROD + ).get_output(0) + right = network.add_elementwise( + rotated_half.get_output(0), duplicate(sin_half), trt.ElementWiseOperation.PROD + ).get_output(0) + result = network.add_elementwise(left, right, trt.ElementWiseOperation.SUM).get_output(0) + result = op.cast(network, result, source_dtype) + return op.heads_to_rows(network, result, QWEN_VISION_HIDDEN_SIZE) + + +def _attention(network, hidden, weights, prefix: str, cos, sin): + projected = op.linear( + network, + hidden, + weights[f"{prefix}.qkv.weight"], + weights[f"{prefix}.qkv.bias"], + ) + query, key, value = tuple( + op.dynamic_slice( + network, + projected, + (0, index * QWEN_VISION_HIDDEN_SIZE), + (None, QWEN_VISION_HIDDEN_SIZE), + ) + for index in range(3) + ) + query = _vision_partial_rope(network, query, cos, sin) + key = _vision_partial_rope(network, key, cos, sin) + query_heads = op.rows_to_heads(network, query, NUM_HEADS, HEAD_DIM) + key_heads = op.rows_to_heads(network, key, NUM_HEADS, HEAD_DIM) + value_heads = op.rows_to_heads(network, value, NUM_HEADS, HEAD_DIM) + scale = op.cast( + network, + op.constant(network, np.full((1, 1, 1, 1), 1.0 / math.sqrt(HEAD_DIM), np.float32)), + query_heads.dtype, + ) + query_heads = network.add_elementwise( + query_heads, scale, trt.ElementWiseOperation.PROD + ).get_output(0) + layer = network.add_attention( + query_heads, + key_heads, + value_heads, + trt.AttentionNormalizationOp.SOFTMAX, + False, + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 Qwen vision attention {prefix}") + layer.name = f"{prefix}.native_attention" + layer.metadata = f"trtmc.native_op=IAttention;source={layer.name}" + layer.get_output(0).name = f"{layer.name}.output" + layer.decomposable = False + rows = op.heads_to_rows(network, layer.get_output(0), QWEN_VISION_HIDDEN_SIZE) + return op.linear( + network, + rows, + weights[f"{prefix}.proj.weight"], + weights[f"{prefix}.proj.bias"], + ) + + +def _mlp(network, hidden, weights, prefix: str): + hidden = op.linear( + network, + hidden, + weights[f"{prefix}.linear_fc1.weight"], + weights[f"{prefix}.linear_fc1.bias"], + ) + hidden = _gelu(network, hidden, trt.ActivationType.GELU_TANH) + return op.linear( + network, + hidden, + weights[f"{prefix}.linear_fc2.weight"], + weights[f"{prefix}.linear_fc2.bias"], + ) + + +def _merge(network, hidden, weights, prefix: str, *, postshuffle_norm: bool): + if postshuffle_norm: + reshape = network.add_shuffle(hidden) + reshape.reshape_dims = (-1, MERGED_HIDDEN_SIZE) + merged = _layer_norm( + network, reshape.get_output(0), weights, f"{prefix}.norm", MERGED_HIDDEN_SIZE + ) + else: + normalized = _layer_norm( + network, hidden, weights, f"{prefix}.norm", QWEN_VISION_HIDDEN_SIZE + ) + reshape = network.add_shuffle(normalized) + reshape.reshape_dims = (-1, MERGED_HIDDEN_SIZE) + merged = reshape.get_output(0) + merged = op.linear( + network, + merged, + weights[f"{prefix}.linear_fc1.weight"], + weights[f"{prefix}.linear_fc1.bias"], + ) + # Qwen3-VL patch mergers use nn.GELU(approximate="none"), unlike the + # tanh-approximate GELU in each vision MLP. + merged = _gelu(network, merged, trt.ActivationType.GELU_ERF) + return op.linear( + network, + merged, + weights[f"{prefix}.linear_fc2.weight"], + weights[f"{prefix}.linear_fc2.bias"], + ) + + +@op.cleanup_failed_build +def build_multimodal_vision_encoder_engine( + weights: dict[str, np.ndarray], + profile: VisionEncoderProfile = VisionEncoderProfile(), + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build one shared dynamic-aspect Qwen3-VL vision tower.""" + + profile.validate() + expected_keys = set(checkpoint_keys()) + missing = sorted(expected_keys - set(weights)) + unexpected = sorted(set(weights) - expected_keys) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 multimodal vision checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=VISION_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + + abi = vision_encoder_abi(profile) + inputs = { + binding.name: network.add_input( + binding.name, + { + "float32": trt.float32, + "int32": trt.int32, + }[binding.dtype], + (-1, binding.max_shape[1]), + ) + for binding in abi.inputs + } + optimization = builder.create_optimization_profile() + for binding in abi.inputs: + optimization.set_shape( + binding.name, binding.min_shape, binding.opt_shape, binding.max_shape + ) + config.add_optimization_profile(optimization) + + patch_weight = np.asarray(weights["model.visual.patch_embed.proj.weight"]).reshape( + QWEN_VISION_HIDDEN_SIZE, QWEN_VISION_PATCH_WIDTH + ) + hidden = op.linear( + network, + inputs["pixel_values"], + patch_weight, + weights["model.visual.patch_embed.proj.bias"], + ) + position = _interpolated_position_embeddings( + network, weights, inputs["interp_indices"], inputs["interp_weights"] + ) + position = op.cast(network, position, hidden.dtype) + hidden = network.add_elementwise(hidden, position, trt.ElementWiseOperation.SUM).get_output(0) + cos, sin = _vision_rope_cache(network, inputs["vision_position_ids"]) + + deepstack = [] + for index in range(DEPTH): + prefix = f"model.visual.blocks.{index}" + normalized = _layer_norm( + network, hidden, weights, f"{prefix}.norm1", QWEN_VISION_HIDDEN_SIZE + ) + update = _attention(network, normalized, weights, f"{prefix}.attn", cos, sin) + update = op.cast(network, update, hidden.dtype) + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + normalized = _layer_norm( + network, hidden, weights, f"{prefix}.norm2", QWEN_VISION_HIDDEN_SIZE + ) + update = _mlp(network, normalized, weights, f"{prefix}.mlp") + update = op.cast(network, update, hidden.dtype) + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + if index in DEEPSTACK_VISUAL_INDEXES: + merger_index = DEEPSTACK_VISUAL_INDEXES.index(index) + deepstack.append( + _merge( + network, + hidden, + weights, + f"model.visual.deepstack_merger_list.{merger_index}", + postshuffle_norm=True, + ) + ) + + main = _merge( + network, + hidden, + weights, + "model.visual.merger", + postshuffle_norm=False, + ) + for tensor, binding in zip((main, *deepstack), abi.outputs): + output = op.cast(network, tensor, trt.float32) + output.name = binding.name + network.mark_output(output) + + op.validate_native_network( + network, expected_attentions=DEPTH, label="multimodal vision encoder" + ) + print( + "[minimax-h3] building shared Qwen3-VL vision tower: " + f"layers={DEPTH}, patches={profile.min_patches}..{profile.max_patches}, " + f"merged={profile.min_patches // MERGE_UNIT}..{profile.max_patches // MERGE_UNIT}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 multimodal vision encoder") + del network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/nvfp4_text_checkpoint.py b/families/minimax_h3/nvfp4_text_checkpoint.py new file mode 100644 index 0000000000..20767bb438 --- /dev/null +++ b/families/minimax_h3/nvfp4_text_checkpoint.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time import of Comfy's public Qwen3-VL NVFP4 AWQ checkpoint. + +The released checkpoint explicitly requests full-precision matrix products. +Decode its packed weights to BF16 for the native TensorRT language engine; +this is checkpoint compatibility, not FP4 activation or matrix computation. +The same file contains the unquantized BF16 vision tower. AWQ input scales +remain separate so the language graph preserves their BF16 rounding point. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +import math +import os +from pathlib import Path +from typing import Iterable + +import ml_dtypes +import numpy as np + +from .quantized_checkpoint import ( + CHECKPOINT_REVISION, + MODEL_ID, + QuantizedSourceFileIdentity, + _read_header, + _read_marker, + _source_file_identity, +) + +CHECKPOINT_FILENAME = "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors" +CHECKPOINT_BYTES = 15_687_142_551 +_NVFP4_MARKER = {"format": "nvfp4", "full_precision_matrix_mult": True} +_EMBEDDING = "model.embed_tokens" +_DTYPES = { + "BF16": np.dtype(ml_dtypes.bfloat16), + "F32": np.dtype(" dict[str, object]: + """Public source and compute identity, without local paths.""" + return { + "schema_version": 1, + "model_id": self.model_id, + "revision": self.revision, + "filename": self.filename, + "size_bytes": self.size_bytes, + "tensor_count": self.tensor_count, + "quantized_weight_count": self.quantized_weight_count, + "quantization": "nvfp4_awq", + "embedding_quantization": "int8_tensorwise", + "full_precision_matrix_mult": True, + "engine_weights_dtype": "bfloat16", + "matmul_dtype": "bfloat16", + "runtime_framework": None, + } + + +QUANTIZED_TEXT_CHECKPOINT_IDENTITY = QuantizedTextCheckpointIdentity() + + +def _physical_specs() -> dict[str, tuple[str, tuple[int, ...]]]: + specs = {} + + def tensor(name, shape, dtype="BF16"): + specs[name] = (dtype, tuple(shape)) + + tensor(f"{_EMBEDDING}.weight", (151936, 5120), "I8") + tensor(f"{_EMBEDDING}.weight_scale", (151936, 1), "F32") + tensor(f"{_EMBEDDING}.comfy_quant", (29,), "U8") + for index in range(50): + prefix = f"model.layers.{index}" + for name, width in (("input_layernorm", 5120), ("post_attention_layernorm", 5120), + ("self_attn.q_norm", 128), ("self_attn.k_norm", 128)): + tensor(f"{prefix}.{name}.weight", (width,)) + for name, rows, cols in ( + ("self_attn.q_proj", 8192, 5120), ("self_attn.k_proj", 1024, 5120), + ("self_attn.v_proj", 1024, 5120), ("self_attn.o_proj", 5120, 8192), + ("mlp.gate_proj", 25600, 5120), ("mlp.up_proj", 25600, 5120), + ("mlp.down_proj", 5120, 25600), + ): + base = f"{prefix}.{name}" + tensor(f"{base}.weight", (rows, cols // 2), "U8") + tensor(f"{base}.weight_scale", (rows, cols // 16), "F8_E4M3") + tensor(f"{base}.weight_scale_2", (), "F32") + tensor(f"{base}.comfy_quant", (55,), "U8") + if name in ("self_attn.o_proj", "mlp.down_proj"): + tensor(f"{base}.pre_quant_scale", (cols,)) + + def linear(base, rows, cols): + tensor(f"{base}.weight", (rows, cols)) + tensor(f"{base}.bias", (rows,)) + + def norm(base, width): + tensor(f"{base}.weight", (width,)) + tensor(f"{base}.bias", (width,)) + + tensor("visual.patch_embed.proj.weight", (1152, 3, 2, 16, 16)) + tensor("visual.patch_embed.proj.bias", (1152,)) + tensor("visual.pos_embed.weight", (2304, 1152)) + for index in range(27): + prefix = f"visual.blocks.{index}" + norm(f"{prefix}.norm1", 1152) + norm(f"{prefix}.norm2", 1152) + linear(f"{prefix}.attn.qkv", 3456, 1152) + linear(f"{prefix}.attn.proj", 1152, 1152) + linear(f"{prefix}.mlp.linear_fc1", 4304, 1152) + linear(f"{prefix}.mlp.linear_fc2", 1152, 4304) + for prefix in ("visual.merger", *(f"visual.deepstack_merger_list.{i}" for i in range(3))): + norm(f"{prefix}.norm", 1152 if prefix == "visual.merger" else 4608) + linear(f"{prefix}.linear_fc1", 4608, 4608) + linear(f"{prefix}.linear_fc2", 5120, 4608) + return specs + + +_PHYSICAL_SPECS = _physical_specs() + + +def _authenticate_source(path: Path) -> None: + identity = QUANTIZED_TEXT_CHECKPOINT_IDENTITY + relative = Path(identity.filename) + root = path.parents[len(relative.parts) - 1] + if os.path.normcase(os.path.abspath(path)) != os.path.normcase(os.path.abspath(root / relative)): + raise ValueError("MiniMax-H3 text checkpoint must retain its released text_encoders path") + if root.parent.name == "snapshots": + if root.name != identity.revision or root.parent.parent.name != "models--Comfy-Org--MiniMax-H3": + raise ValueError("MiniMax-H3 text checkpoint is not the pinned Hugging Face snapshot") + return + metadata_path = root / ".cache/huggingface/download" / f"{relative.as_posix()}.metadata" + try: + lines = metadata_path.read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError) as error: + raise ValueError("Download the MiniMax-H3 text checkpoint with hf download and pinned revision metadata") from error + if not lines or lines[0] != identity.revision: + raise ValueError("MiniMax-H3 text checkpoint does not match the pinned Hugging Face revision") + + +def _validate_inventory(header, data_offset: int, file_size: int): + tensors = {key: value for key, value in header.items() if key != "__metadata__"} + if set(tensors) != set(_PHYSICAL_SPECS): + raise ValueError("MiniMax-H3 NVFP4 text checkpoint tensor inventory mismatch") + regions = [] + for name, (dtype, shape) in _PHYSICAL_SPECS.items(): + entry = tensors[name] + if not isinstance(entry, dict) or set(entry) != {"dtype", "shape", "data_offsets"}: + raise ValueError(f"Invalid MiniMax-H3 text tensor header: {name}") + if entry["dtype"] != dtype or entry["shape"] != list(shape): + raise ValueError(f"MiniMax-H3 text tensor contract mismatch: {name}") + offsets = entry["data_offsets"] + if (not isinstance(offsets, list) or len(offsets) != 2 + or any(type(value) is not int for value in offsets)): + raise ValueError(f"Invalid MiniMax-H3 text tensor offsets: {name}") + start, end = offsets + if start < 0 or end - start != math.prod(shape) * _DTYPES[dtype].itemsize: + raise ValueError(f"Invalid MiniMax-H3 text tensor byte range: {name}") + regions.append((start, end)) + cursor = 0 + for start, end in sorted(regions): + if start != cursor: + raise ValueError("MiniMax-H3 text tensor storage has gaps or overlaps") + cursor = end + if data_offset + cursor != file_size: + raise ValueError("MiniMax-H3 text checkpoint payload size mismatch") + return tensors + + +def validate_quantized_text_checkpoint(checkpoint_file: str | Path) -> QuantizedTextCheckpointIdentity: + """Check the pinned file, header and all tiny markers; never hash its payload.""" + path = Path(checkpoint_file) + identity = QUANTIZED_TEXT_CHECKPOINT_IDENTITY + before = _source_file_identity(path) + if path.name != Path(identity.filename).name or before.size_bytes != identity.size_bytes: + raise ValueError("MiniMax-H3 NVFP4 text checkpoint filename or size mismatch") + _authenticate_source(path) + header, offset = _read_header(path) + tensors = _validate_inventory(header, offset, before.size_bytes) + for name, entry in tensors.items(): + if name.endswith(".comfy_quant"): + expected = {"format": "int8_tensorwise"} if name == f"{_EMBEDDING}.comfy_quant" else _NVFP4_MARKER + if _read_marker(path, offset, entry) != expected: + raise ValueError(f"Unsupported MiniMax-H3 text comfy_quant marker: {name}") + if _source_file_identity(path) != before: + raise ValueError("MiniMax-H3 text checkpoint changed during validation") + return replace(identity, source_file_identity=before) + + +def _unblock_scales(scales: np.ndarray, rows: int, cols: int) -> np.ndarray: + """Invert Comfy's cuBLAS 128x4 block-scale storage on the build CPU. + + Format reference: Comfy-Org/comfy-kitchen, float_utils.py::from_blocked. + """ + row_blocks, col_blocks = (rows + 127) // 128, (cols + 3) // 4 + if scales.shape != (row_blocks * 128, col_blocks * 4): + raise ValueError("MiniMax-H3 NVFP4 block-scale shape mismatch") + tiles = scales.reshape(row_blocks, col_blocks, 32, 4, 4) + plain = tiles.transpose(0, 1, 3, 2, 4).reshape(row_blocks, col_blocks, 128, 4) + plain = plain.transpose(0, 2, 1, 3).reshape(row_blocks * 128, col_blocks * 4) + return plain[:rows, :cols] + + +def _dequantize_nvfp4(packed, scales, global_scale) -> np.ndarray: + """Match Comfy CUDA decode: FP32 scale/product, then one BF16 rounding. + + Comfy stores the even element in the HIGH nibble. Its CUDA reference is + backends/cuda/ops/quantize_nvfp4.cu::dequantize_nvfp4_kernel; the eager CPU + backend has extra intermediate BF16 rounding and is not used as an oracle. + """ + if packed.dtype != np.uint8 or packed.ndim != 2 or packed.shape[1] % 8: + raise ValueError("MiniMax-H3 NVFP4 weights require packed uint8 [out,in/2], in divisible by 16") + rows, cols = packed.shape[0], packed.shape[1] * 2 + if np.asarray(global_scale).shape != () or not np.isfinite(global_scale) or global_scale <= 0: + raise ValueError("MiniMax-H3 NVFP4 global scale must be a positive finite scalar") + plain_scales = _unblock_scales(scales, rows, cols // 16) + lut = np.asarray((0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6), np.float32) + output = np.empty((rows, cols), dtype=ml_dtypes.bfloat16) + for start in range(0, rows, 128): + end = min(rows, start + 128) + scale = plain_scales[start:end].astype(np.float32) * np.float32(global_scale) + if not np.isfinite(scale).all() or (scale < 0).any(): + raise ValueError("MiniMax-H3 NVFP4 block scales must be finite and non-negative") + values = np.empty((end - start, cols), np.float32) + values[:, 0::2] = lut[packed[start:end] >> 4] + values[:, 1::2] = lut[packed[start:end] & 15] + values.reshape(end - start, cols // 16, 16)[:] *= scale[..., None] + output[start:end] = values + return output + + +def _physical_name(logical_name: str) -> str: + for logical, physical in (("model.language_model.", "model."), ("model.visual.", "visual.")): + if logical_name.startswith(logical): + name = physical + logical_name.removeprefix(logical) + if name in _PHYSICAL_SPECS and name.endswith((".weight", ".bias")): + return name + raise ValueError(f"Unsupported MiniMax-H3 quantized text logical tensor: {logical_name}") + + +def load_selected_quantized_text_weights( + checkpoint_file: str | Path, logical_names: Iterable[str] +) -> dict[str, np.ndarray]: + """Load requested language/vision tensors and their optional AWQ input scales. + + Only requested payloads are touched. Memory-mapped BF16 vision/norm tensors + retain their file owner; decoded weights use bounded per-row temporaries. + """ + path = Path(checkpoint_file) + identity = validate_quantized_text_checkpoint(path) + requested = tuple(logical_names) + if len(requested) != len(set(requested)): + raise ValueError("MiniMax-H3 text weight request contains duplicates") + names = {logical: _physical_name(logical) for logical in requested} + header, offset = _read_header(path) + + def read(name): + entry = header[name] + return np.memmap(path, mode="r", dtype=_DTYPES[entry["dtype"]], + offset=offset + entry["data_offsets"][0], shape=tuple(entry["shape"])) + + result = {} + for logical, physical in names.items(): + value = read(physical) + base = physical.removesuffix(".weight") + if f"{base}.weight_scale_2" in header: + value = _dequantize_nvfp4(value, read(f"{base}.weight_scale"), read(f"{base}.weight_scale_2")) + prescale = f"{base}.pre_quant_scale" + if prescale in header: + result[f"{logical.removesuffix('.weight')}.pre_quant_scale"] = read(prescale) + elif physical == f"{_EMBEDDING}.weight": + scale = read(f"{base}.weight_scale") + decoded = np.empty(value.shape, dtype=ml_dtypes.bfloat16) + for start in range(0, value.shape[0], 128): + chunk_scale = scale[start:start + 128] + if not np.isfinite(chunk_scale).all() or (chunk_scale < 0).any(): + raise ValueError("MiniMax-H3 embedding scales must be finite and non-negative") + decoded[start:start + 128] = value[start:start + 128].astype(np.float32) * chunk_scale + value = decoded + result[logical] = value + if _source_file_identity(path) != identity.source_file_identity: + raise ValueError("MiniMax-H3 text checkpoint changed during loading") + return result diff --git a/families/minimax_h3/provenance.py b/families/minimax_h3/provenance.py new file mode 100644 index 0000000000..90b23a9de0 --- /dev/null +++ b/families/minimax_h3/provenance.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Path-free public metadata for MiniMax-H3 native bundles.""" + +from __future__ import annotations + +import json +import math +import struct +from pathlib import Path + +from tensorrt_model_connect.bundle_writer import BUNDLE_MAGIC + +from .config import ( + NATIVE_EXPLICIT_CANVAS_SIZES, + SOL_ENGINE_1344X768_124_TO_345F, + native_plan_filenames, +) + + +CHECKPOINT_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc" +CHECKPOINT_REPOSITORY = "MiniMaxAI/MiniMax-H3" + +QUANTIZED_TRANSFORMER_CONFIG = { + "format": "int8_tensorwise_convrot", + "checkpoint_scope": "full_transformer_overlay", + "quantized_weight_scope": ( + "blocks.*.{adaln_proj.linear,attn.qkv_proj,attn.out_proj,mlp.fc1,mlp.fc2}" + ), + "quantized_weight_count": 250, + "activation_precision": "dynamic_int8_rowwise", +} + +SUPER_RESOLUTION_PLAN_FILENAME = "video_super_resolution.plan" +SUPER_RESOLUTION_SECTION = "video_super_resolution_plan" +SUPER_RESOLUTION_MODEL = "realesr-general-x4v3" +SUPER_RESOLUTION_ARCHITECTURE = "SRVGGNetCompact" +SUPER_RESOLUTION_PRIMARY_FILENAME = "realesr-general-x4v3.pth" +SUPER_RESOLUTION_WEAK_FILENAME = "realesr-general-wdn-x4v3.pth" +SUPER_RESOLUTION_SOURCE_SHAPE = [480, 864] +SUPER_RESOLUTION_TARGET_SHAPE = [720, 1296] +SUPER_RESOLUTION_BATCH_PROFILE = [1, 4, 8] +SUPER_RESOLUTION_LEARNED_RESIDUAL_STRENGTH = 0.25 + + +def checkpoint_snapshot_record( + snapshot: Path, + *, + include_transformer_weights: bool = True, +) -> dict[str, object]: + """Describe the selected public checkpoint without recording local paths.""" + + root = Path(snapshot) + required = ["text_encoder", "vae", "audio_vae", "tokenizer"] + if include_transformer_weights: + required.append("transformer") + missing = [name for name in required if not (root / name).exists()] + if missing: + raise FileNotFoundError( + "Incomplete MiniMax-H3 checkpoint; missing: " + ", ".join(missing) + ) + return { + "repository": CHECKPOINT_REPOSITORY, + "revision": CHECKPOINT_REVISION, + "includes_transformer_weights": include_transformer_weights, + } + + +def validate_checkpoint_snapshot_record( + record: object, + *, + include_transformer_weights: bool = True, +) -> dict[str, object]: + expected = { + "repository": CHECKPOINT_REPOSITORY, + "revision": CHECKPOINT_REVISION, + "includes_transformer_weights": include_transformer_weights, + } + if record != expected: + raise ValueError("MiniMax-H3 checkpoint metadata does not match the public model") + return dict(expected) + + +def _checkpoint_record(path: Path, role: str, filename: str) -> dict[str, object]: + value = Path(path) + if not value.is_file(): + raise FileNotFoundError(f"MiniMax-H3 {role} checkpoint is missing: {value}") + if value.name != filename: + raise ValueError(f"MiniMax-H3 {role} checkpoint must be named {filename}") + return {"role": role, "filename": filename, "bytes": value.stat().st_size} + + +def super_resolution_source_identity( + primary_checkpoint: Path, + weak_checkpoint: Path | None, + *, + denoise_strength: float, + learned_residual_strength: float = SUPER_RESOLUTION_LEARNED_RESIDUAL_STRENGTH, +) -> dict[str, object]: + if isinstance(denoise_strength, bool) or not isinstance(denoise_strength, (int, float)): + raise ValueError("MiniMax-H3 super-resolution denoise_strength must be numeric") + if isinstance(learned_residual_strength, bool) or not isinstance( + learned_residual_strength, (int, float) + ): + raise ValueError("MiniMax-H3 learned residual strength must be numeric") + strength = float(denoise_strength) + residual = float(learned_residual_strength) + if not math.isfinite(strength) or not 0.0 <= strength <= 1.0: + raise ValueError("MiniMax-H3 super-resolution denoise_strength must be in [0, 1]") + if not math.isfinite(residual) or not 0.0 <= residual <= 1.0: + raise ValueError("MiniMax-H3 learned residual strength must be in [0, 1]") + if weak_checkpoint is None and strength != 1.0: + raise ValueError("MiniMax-H3 denoise blending requires the weak checkpoint") + sources = [_checkpoint_record(primary_checkpoint, "primary", SUPER_RESOLUTION_PRIMARY_FILENAME)] + if weak_checkpoint is not None: + sources.append(_checkpoint_record(weak_checkpoint, "weak_denoise", SUPER_RESOLUTION_WEAK_FILENAME)) + return { + "model": SUPER_RESOLUTION_MODEL, + "architecture": SUPER_RESOLUTION_ARCHITECTURE, + "denoise_strength": strength, + "learned_residual_strength": residual, + "sources": sources, + } + + +def validate_super_resolution_source_identity(record: object) -> dict[str, object]: + if not isinstance(record, dict): + raise ValueError("MiniMax-H3 super-resolution source metadata must be an object") + if record.get("model") != SUPER_RESOLUTION_MODEL: + raise ValueError("MiniMax-H3 super-resolution model is unsupported") + if record.get("architecture") != SUPER_RESOLUTION_ARCHITECTURE: + raise ValueError("MiniMax-H3 super-resolution architecture is unsupported") + sources = record.get("sources") + if not isinstance(sources, list) or len(sources) not in (1, 2): + raise ValueError("MiniMax-H3 super-resolution sources are invalid") + expected = ( + ("primary", SUPER_RESOLUTION_PRIMARY_FILENAME), + ("weak_denoise", SUPER_RESOLUTION_WEAK_FILENAME), + ) + for source, (role, filename) in zip(sources, expected, strict=False): + if ( + not isinstance(source, dict) + or source.get("role") != role + or source.get("filename") != filename + or not isinstance(source.get("bytes"), int) + or source["bytes"] <= 0 + ): + raise ValueError("MiniMax-H3 super-resolution source metadata is invalid") + strength = record.get("denoise_strength") + residual = record.get("learned_residual_strength") + if not isinstance(strength, (int, float)) or isinstance(strength, bool): + raise ValueError("MiniMax-H3 super-resolution denoise strength is invalid") + if not isinstance(residual, (int, float)) or isinstance(residual, bool): + raise ValueError("MiniMax-H3 super-resolution residual strength is invalid") + if not 0.0 <= float(strength) <= 1.0 or not 0.0 <= float(residual) <= 1.0: + raise ValueError("MiniMax-H3 super-resolution strengths must be in [0, 1]") + if len(sources) == 1 and float(strength) != 1.0: + raise ValueError("MiniMax-H3 denoise blending requires the weak checkpoint") + return dict(record) + + +def super_resolution_bundle_config(source_identity: object) -> dict[str, object]: + source = validate_super_resolution_source_identity(source_identity) + return { + "section": SUPER_RESOLUTION_SECTION, + "mode": "explicit", + "input_name": "frames", + "output_name": "upscaled_frames", + "model": source["model"], + "architecture": source["architecture"], + "source_shape": list(SUPER_RESOLUTION_SOURCE_SHAPE), + "target_shape": list(SUPER_RESOLUTION_TARGET_SHAPE), + "batch_profile": list(SUPER_RESOLUTION_BATCH_PROFILE), + "denoise_strength": source["denoise_strength"], + "learned_residual_strength": source["learned_residual_strength"], + "sources": source["sources"], + } + + +def validate_super_resolution_bundle_config(record: object) -> dict[str, object]: + if not isinstance(record, dict) or record.get("section") != SUPER_RESOLUTION_SECTION: + raise ValueError("MiniMax-H3 super-resolution bundle metadata is invalid") + if record.get("mode") != "explicit": + raise ValueError("MiniMax-H3 legacy or unsupported SR mode; rebuild with super_resolution=true") + if record.get("input_name") != "frames" or record.get("output_name") != "upscaled_frames": + raise ValueError("MiniMax-H3 super-resolution tensor names are invalid") + if record.get("source_shape") != SUPER_RESOLUTION_SOURCE_SHAPE: + raise ValueError("MiniMax-H3 super-resolution source shape is invalid") + if record.get("target_shape") != SUPER_RESOLUTION_TARGET_SHAPE: + raise ValueError("MiniMax-H3 super-resolution target shape is invalid") + if record.get("batch_profile") != SUPER_RESOLUTION_BATCH_PROFILE: + raise ValueError("MiniMax-H3 super-resolution batch profile is invalid") + validate_super_resolution_source_identity( + { + "model": record.get("model"), + "architecture": record.get("architecture"), + "denoise_strength": record.get("denoise_strength"), + "learned_residual_strength": record.get("learned_residual_strength"), + "sources": record.get("sources"), + } + ) + return dict(record) + + +def validate_quantized_transformer_metadata(record: object) -> dict[str, object]: + from .quantized_checkpoint import QUANTIZED_CHECKPOINT_IDENTITY + + expected = QUANTIZED_CHECKPOINT_IDENTITY.bundle_metadata() + if record != expected: + raise ValueError("MiniMax-H3 quantized transformer metadata does not match the public model") + return expected + + +def validate_workspace_limit_bytes( + record: object, + *, + profile=SOL_ENGINE_1344X768_124_TO_345F, + additional_plan_filenames: tuple[str, ...] = (), + excluded_plan_filenames: tuple[str, ...] = (), +) -> dict[str, object]: + expected = tuple( + filename + for filename in native_plan_filenames() + if filename not in excluded_plan_filenames + ) + tuple(additional_plan_filenames) + if not isinstance(record, dict) or set(record) != set(expected): + raise ValueError("MiniMax-H3 workspace metadata must cover every native plan") + for filename, value in record.items(): + if value not in ("tensorrt_default", None) and ( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + ): + raise ValueError(f"MiniMax-H3 workspace value for {filename} is invalid") + del profile + return dict(record) + + +def load_bundle_config(bundle: Path) -> dict[str, object]: + """Load the canonical runtime section from a current ModelConnect bundle.""" + + with Path(bundle).open("rb") as stream: + if stream.read(8) != BUNDLE_MAGIC: + raise ValueError("MiniMax-H3 bundle has invalid magic") + raw_size = stream.read(8) + if len(raw_size) != 8: + raise ValueError("MiniMax-H3 bundle header is truncated") + header_size = struct.unpack(" dict: + config = load_bundle_config(bundle) + if config.get("checkpoint_revision") != CHECKPOINT_REVISION: + raise ValueError("MiniMax-H3 bundle checkpoint revision is invalid") + if config.get("explicit_canvas_sizes") != [list(size) for size in NATIVE_EXPLICIT_CANVAS_SIZES]: + raise ValueError("MiniMax-H3 bundle canvas profile is invalid") + del source_revision + return config diff --git a/families/minimax_h3/quantized_checkpoint.py b/families/minimax_h3/quantized_checkpoint.py new file mode 100644 index 0000000000..7c99d72597 --- /dev/null +++ b/families/minimax_h3/quantized_checkpoint.py @@ -0,0 +1,738 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict loader for the public full MiniMax-H3 INT8 ConvRot transformers. + +The Comfy-Org single-file checkpoint uses the original MiniMax module names, +contiguous Q/K/V thirds, and ``[gate; value]`` SwiGLU rows. The native +TensorRT builders consume the frozen Diffusers names, the same contiguous Q/K/V +layout, and ``[value; gate]`` rows. This module performs the lossless SwiGLU +row permutation while preserving the checkpoint's INT8 data and per-output-row +FP32 scales. + +T2VA/FL2VA and Ref2VA share this tensor layout but use distinct checkpoint +values. The explicit workflow selects the pinned source; it never substitutes +the FL2VA checkpoint for Ref2VA or accepts the pruned AdaLN variant. + +Large tensor payloads are selected through :mod:`safetensors`; validation reads +only the JSON header and the 250 tiny ``comfy_quant`` markers. The released +file size is checked without rereading the 34 GB payload before every build. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +import json +import math +import os +from pathlib import Path +import stat +import struct +from typing import Any, Iterable + +import numpy as np + +from .checkpoint import numpy_state + + +MODEL_ID = "Comfy-Org/MiniMax-H3" +CHECKPOINT_REVISION = "4cc1d817b6184899b41293954329f576cb5ae86b" +CHECKPOINT_FILENAME = "diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors" +CHECKPOINT_BYTES = 34_038_892_334 +REF2VA_CHECKPOINT_FILENAME = "diffusion_models/minimax_h3_ref2va_int8_convrot.safetensors" +REF2VA_CHECKPOINT_BYTES = 34_038_894_550 + +_MAX_HEADER_BYTES = 8 << 20 +_NUM_HEADS = 56 +_HEAD_DIM = 128 +_HIDDEN_SIZE = 5_376 +_INNER_DIM = _NUM_HEADS * _HEAD_DIM +_FFN_DIM = 14_336 +_TIME_DIM = 2_688 +_QUANT_FORMAT = "int8_tensorwise" +_HF_CACHE_REPOSITORY = "models--Comfy-Org--MiniMax-H3" +_HF_LOCAL_DOWNLOAD_METADATA = Path(".cache/huggingface/download") + + +@dataclass(frozen=True) +class ConvRotInt8Weight: + """One Comfy tensorwise-INT8 weight in Diffusers logical row order. + + ``packed_parent`` and ``row_slice`` let the fused-QKV graph path recover a + single packed parent matrix when the loader returned three + logical ``to_q``/``to_k``/``to_v`` entries. No concatenation of those + child views is necessary. + """ + + qweight: np.ndarray + scale: np.ndarray + group_size: int + packed_parent: ConvRotInt8Weight | None = None + row_slice: tuple[int, int] | None = None + is_full_fused_qkv: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.qweight, np.ndarray) or self.qweight.dtype != np.int8: + raise TypeError("MiniMax-H3 ConvRot qweight must be a NumPy int8 array") + if self.qweight.ndim != 2 or not self.qweight.flags.c_contiguous: + raise ValueError("MiniMax-H3 ConvRot qweight must be contiguous [out,in]") + if not isinstance(self.scale, np.ndarray) or self.scale.dtype != np.float32: + raise TypeError("MiniMax-H3 ConvRot scale must be a NumPy float32 array") + if self.scale.shape not in ((self.qweight.shape[0],), (self.qweight.shape[0], 1)): + raise ValueError("MiniMax-H3 ConvRot scale must contain one value per output row") + if not self.scale.flags.c_contiguous: + raise ValueError("MiniMax-H3 ConvRot scale must be contiguous") + _validate_group_size(self.group_size, self.qweight.shape[1]) + if self.is_full_fused_qkv: + if self.packed_parent is not None or self.row_slice is not None: + raise ValueError("A full fused-QKV ConvRot weight cannot itself be a child view") + elif self.packed_parent is not None: + if not self.packed_parent.is_full_fused_qkv or self.row_slice is None: + raise ValueError("A fused-QKV child requires a full parent and row_slice") + start, end = self.row_slice + if not 0 <= start < end <= self.packed_parent.qweight.shape[0]: + raise ValueError("MiniMax-H3 fused-QKV child row_slice is out of bounds") + if end - start != self.qweight.shape[0]: + raise ValueError("MiniMax-H3 fused-QKV child row_slice does not match qweight") + elif self.row_slice is not None: + raise ValueError("MiniMax-H3 ConvRot row_slice requires packed_parent") + + +@dataclass(frozen=True) +class QuantizedSourceFileIdentity: + """Path-free local identity used only to make staged resume fail closed.""" + + device: int + inode: int + size_bytes: int + mtime_ns: int + ctime_ns: int + + def receipt_metadata(self) -> dict[str, int]: + return { + "device": self.device, + "inode": self.inode, + "size_bytes": self.size_bytes, + "mtime_ns": self.mtime_ns, + "ctime_ns": self.ctime_ns, + } + + +@dataclass(frozen=True) +class QuantizedCheckpointIdentity: + """Path-free identity of the pinned public Comfy-Org checkpoint.""" + + model_id: str + revision: str + filename: str + size_bytes: int + tensor_count: int + quantized_weight_count: int + source_file_identity: QuantizedSourceFileIdentity | None = None + + def bundle_metadata(self) -> dict[str, object]: + """Return stable public provenance without a workstation path.""" + + return { + "schema_version": 1, + "model_id": self.model_id, + "revision": self.revision, + "filename": self.filename, + "size_bytes": self.size_bytes, + "tensor_count": self.tensor_count, + "quantized_weight_count": self.quantized_weight_count, + "quantization": "int8_tensorwise_convrot", + "runtime_framework": None, + } + + +QUANTIZED_CHECKPOINT_IDENTITY = QuantizedCheckpointIdentity( + model_id=MODEL_ID, + revision=CHECKPOINT_REVISION, + filename=CHECKPOINT_FILENAME, + size_bytes=CHECKPOINT_BYTES, + tensor_count=1_035, + quantized_weight_count=250, +) +QUANTIZED_REF2VA_CHECKPOINT_IDENTITY = replace( + QUANTIZED_CHECKPOINT_IDENTITY, + filename=REF2VA_CHECKPOINT_FILENAME, + size_bytes=REF2VA_CHECKPOINT_BYTES, +) + + +def _checkpoint_identity(workflow: str) -> QuantizedCheckpointIdentity: + if workflow == "fl2va": + return QUANTIZED_CHECKPOINT_IDENTITY + if workflow == "ref2va": + return QUANTIZED_REF2VA_CHECKPOINT_IDENTITY + raise ValueError("MiniMax-H3 quantized checkpoint workflow must be 'fl2va' or 'ref2va'") + + +_EXPECTED_CONFIG: dict[str, object] = { + "transformer": { + "hidden_size": 5_376, + "num_layers": 50, + "token_refiner_num_layers": 2, + "num_attention_heads": 56, + "attention_head_dim": 128, + "ffn_hidden_size": 14_336, + "latents_dim": 24, + "audio_latents_dim": 32, + "patch_size": [1, 2, 2], + "text_dim": 5_120, + "timestep_input_dim": 256, + "time_embed_hidden_size": 5_376, + "time_embed_dim": 2_688, + "adaln_out_features": 96_768, + "final_adaln_out_features": 10_752, + "rope_inv_freq_len": 16, + "norm_eps": 1.0e-5, + "qk_norm_eps": 1.0e-5, + "final_norm_eps": 1.0e-5, + "image_model": "minimax_h3", + } +} + + +@dataclass(frozen=True) +class _TensorSpec: + dtype: str + shape: tuple[int, ...] + + +def _add_quantized( + specs: dict[str, _TensorSpec], + groups: dict[str, int], + base: str, + shape: tuple[int, int], + group_size: int, +) -> None: + specs[f"{base}.weight"] = _TensorSpec("I8", shape) + specs[f"{base}.weight_scale"] = _TensorSpec("F32", (shape[0], 1)) + marker_bytes = len( + json.dumps( + { + "format": _QUANT_FORMAT, + "convrot": True, + "convrot_groupsize": group_size, + } + ).encode("utf-8") + ) + specs[f"{base}.comfy_quant"] = _TensorSpec("U8", (marker_bytes,)) + groups[base] = group_size + + +def _physical_contract() -> tuple[dict[str, _TensorSpec], dict[str, int]]: + specs: dict[str, _TensorSpec] = { + "video_patch_proj.weight": _TensorSpec("F32", (_HIDDEN_SIZE, 96)), + "video_patch_proj.bias": _TensorSpec("F32", (_HIDDEN_SIZE,)), + "audio_patch_proj.weight": _TensorSpec("F32", (_HIDDEN_SIZE, 32)), + "audio_patch_proj.bias": _TensorSpec("F32", (_HIDDEN_SIZE,)), + "condition_proj.weight": _TensorSpec("BF16", (_HIDDEN_SIZE, 5_120)), + "condition_proj.bias": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + "time_embedder.proj_in.weight": _TensorSpec("F32", (_HIDDEN_SIZE, 256)), + "time_embedder.proj_in.bias": _TensorSpec("F32", (_HIDDEN_SIZE,)), + "time_embedder.proj_out.weight": _TensorSpec("F32", (_TIME_DIM, _HIDDEN_SIZE)), + "time_embedder.proj_out.bias": _TensorSpec("F32", (_TIME_DIM,)), + "token_refiner.final_norm.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + "final_layer.norm.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + "final_layer.adaln_proj.linear.weight": _TensorSpec( + "BF16", (2 * _HIDDEN_SIZE, _TIME_DIM) + ), + "final_layer.adaln_proj.linear.bias": _TensorSpec("BF16", (2 * _HIDDEN_SIZE,)), + "final_layer.video_out.weight": _TensorSpec("F32", (96, _HIDDEN_SIZE)), + "final_layer.video_out.bias": _TensorSpec("F32", (96,)), + "final_layer.audio_out.weight": _TensorSpec("F32", (32, _HIDDEN_SIZE)), + "final_layer.audio_out.bias": _TensorSpec("F32", (32,)), + "rope.inv_freq": _TensorSpec("F32", (16,)), + } + groups: dict[str, int] = {} + for index in range(50): + prefix = f"blocks.{index}" + specs[f"{prefix}.norm1.weight"] = _TensorSpec("BF16", (_HIDDEN_SIZE,)) + specs[f"{prefix}.norm2.weight"] = _TensorSpec("BF16", (_HIDDEN_SIZE,)) + specs[f"{prefix}.attn.q_norm.weight"] = _TensorSpec("BF16", (_HEAD_DIM,)) + specs[f"{prefix}.attn.k_norm.weight"] = _TensorSpec("BF16", (_HEAD_DIM,)) + specs[f"{prefix}.adaln_proj.linear.bias"] = _TensorSpec( + "BF16", (18 * _HIDDEN_SIZE,) + ) + _add_quantized( + specs, + groups, + f"{prefix}.adaln_proj.linear", + (18 * _HIDDEN_SIZE, _TIME_DIM), + 64, + ) + _add_quantized( + specs, + groups, + f"{prefix}.attn.qkv_proj", + (3 * _INNER_DIM, _HIDDEN_SIZE), + 256, + ) + _add_quantized( + specs, + groups, + f"{prefix}.attn.out_proj", + (_HIDDEN_SIZE, _INNER_DIM), + 256, + ) + _add_quantized( + specs, + groups, + f"{prefix}.mlp.fc1", + (2 * _FFN_DIM, _HIDDEN_SIZE), + 256, + ) + _add_quantized( + specs, + groups, + f"{prefix}.mlp.fc2", + (_HIDDEN_SIZE, _FFN_DIM), + 256, + ) + for index in range(2): + prefix = f"token_refiner.blocks.{index}" + specs.update( + { + f"{prefix}.norm1.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + f"{prefix}.norm2.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + f"{prefix}.attn.qkv_proj.weight": _TensorSpec( + "BF16", (3 * _INNER_DIM, _HIDDEN_SIZE) + ), + f"{prefix}.attn.q_norm.weight": _TensorSpec("BF16", (_HEAD_DIM,)), + f"{prefix}.attn.k_norm.weight": _TensorSpec("BF16", (_HEAD_DIM,)), + f"{prefix}.attn.out_proj.weight": _TensorSpec( + "BF16", (_HIDDEN_SIZE, _INNER_DIM) + ), + f"{prefix}.mlp.fc1.weight": _TensorSpec( + "BF16", (2 * _FFN_DIM, _HIDDEN_SIZE) + ), + f"{prefix}.mlp.fc2.weight": _TensorSpec("BF16", (_HIDDEN_SIZE, _FFN_DIM)), + } + ) + if len(specs) != 1_035 or len(groups) != 250: + raise RuntimeError("MiniMax-H3 quantized checkpoint contract is incomplete") + return specs, groups + + +_PHYSICAL_SPECS, _QUANT_GROUPS = _physical_contract() + + +def _logical_contract() -> dict[str, tuple[str, str]]: + direct = { + "proj_in.weight": "video_patch_proj.weight", + "proj_in.bias": "video_patch_proj.bias", + "audio_proj_in.weight": "audio_patch_proj.weight", + "audio_proj_in.bias": "audio_patch_proj.bias", + "context_embedder.weight": "condition_proj.weight", + "context_embedder.bias": "condition_proj.bias", + "time_embedder.linear_1.weight": "time_embedder.proj_in.weight", + "time_embedder.linear_1.bias": "time_embedder.proj_in.bias", + "time_embedder.linear_2.weight": "time_embedder.proj_out.weight", + "time_embedder.linear_2.bias": "time_embedder.proj_out.bias", + "token_refiner.final_norm.weight": "token_refiner.final_norm.weight", + "norm_out.norm.weight": "final_layer.norm.weight", + "norm_out.linear.weight": "final_layer.adaln_proj.linear.weight", + "norm_out.linear.bias": "final_layer.adaln_proj.linear.bias", + "proj_out.weight": "final_layer.video_out.weight", + "proj_out.bias": "final_layer.video_out.bias", + "audio_proj_out.weight": "final_layer.audio_out.weight", + "audio_proj_out.bias": "final_layer.audio_out.bias", + } + result = {logical: (physical, "direct") for logical, physical in direct.items()} + for physical_root, logical_root, count, has_adaln in ( + ("blocks", "transformer_blocks", 50, True), + ("token_refiner.blocks", "token_refiner.refiner_blocks", 2, False), + ): + for index in range(count): + physical = f"{physical_root}.{index}" + logical = f"{logical_root}.{index}" + result[f"{logical}.norm1.weight"] = (f"{physical}.norm1.weight", "direct") + result[f"{logical}.norm2.weight"] = (f"{physical}.norm2.weight", "direct") + for qkv_index, name in enumerate(("q", "k", "v")): + result[f"{logical}.attn.to_{name}.weight"] = ( + f"{physical}.attn.qkv_proj.weight", + f"qkv:{qkv_index}", + ) + result[f"{logical}.attn.norm_q.weight"] = ( + f"{physical}.attn.q_norm.weight", + "direct", + ) + result[f"{logical}.attn.norm_k.weight"] = ( + f"{physical}.attn.k_norm.weight", + "direct", + ) + result[f"{logical}.attn.to_out.0.weight"] = ( + f"{physical}.attn.out_proj.weight", + "direct", + ) + result[f"{logical}.ff.net.0.proj.weight"] = ( + f"{physical}.mlp.fc1.weight", + "swap_swiglu", + ) + result[f"{logical}.ff.net.2.weight"] = ( + f"{physical}.mlp.fc2.weight", + "direct", + ) + if has_adaln: + result[f"{logical}.adaln_proj.linear.weight"] = ( + f"{physical}.adaln_proj.linear.weight", + "direct", + ) + result[f"{logical}.adaln_proj.linear.bias"] = ( + f"{physical}.adaln_proj.linear.bias", + "direct", + ) + return result + + +_LOGICAL_MAP = _logical_contract() + + +def _validate_group_size(group_size: int, in_features: int) -> None: + if not isinstance(group_size, int) or isinstance(group_size, bool) or group_size < 4: + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + value = group_size + while value > 1 and value % 4 == 0: + value //= 4 + if value != 1: + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + if in_features % group_size: + raise ValueError("MiniMax-H3 ConvRot group_size must divide the input width") + + +def _read_header(path: Path) -> tuple[dict[str, Any], int]: + with path.open("rb") as stream: + prefix = stream.read(8) + if len(prefix) != 8: + raise ValueError("MiniMax-H3 quantized checkpoint has an incomplete header") + (header_bytes,) = struct.unpack(" dict[str, dict[str, Any]]: + metadata = header.get("__metadata__") + if not isinstance(metadata, dict) or set(metadata) != {"config"}: + raise ValueError("MiniMax-H3 quantized checkpoint metadata schema is invalid") + try: + config = json.loads(metadata["config"]) + except (TypeError, json.JSONDecodeError) as error: + raise ValueError("MiniMax-H3 quantized checkpoint config metadata is invalid") from error + if config != _EXPECTED_CONFIG: + raise ValueError( + "MiniMax-H3 quantized checkpoint config is not the full released 50-layer architecture" + ) + + tensors = {name: value for name, value in header.items() if name != "__metadata__"} + actual = set(tensors) + expected = set(_PHYSICAL_SPECS) + if actual != expected: + raise ValueError( + "MiniMax-H3 quantized checkpoint tensor inventory mismatch: " + f"missing={sorted(expected - actual)[:8]}, unexpected={sorted(actual - expected)[:8]}" + ) + + dtype_bytes = {"BF16": 2, "F32": 4, "I8": 1, "U8": 1} + regions: list[tuple[int, int, str]] = [] + for name, spec in _PHYSICAL_SPECS.items(): + entry = tensors[name] + if not isinstance(entry, dict) or set(entry) != {"dtype", "shape", "data_offsets"}: + raise ValueError(f"MiniMax-H3 quantized tensor header is invalid: {name}") + shape = entry["shape"] + offsets = entry["data_offsets"] + if entry["dtype"] != spec.dtype or shape != list(spec.shape): + raise ValueError( + f"MiniMax-H3 quantized tensor contract mismatch for {name}: " + f"expected={spec.dtype}{spec.shape}, actual={entry.get('dtype')}{shape}" + ) + if ( + not isinstance(offsets, list) + or len(offsets) != 2 + or any(not isinstance(value, int) or isinstance(value, bool) for value in offsets) + ): + raise ValueError(f"MiniMax-H3 quantized tensor offsets are invalid: {name}") + start, end = offsets + expected_bytes = math.prod(spec.shape) * dtype_bytes[spec.dtype] + if start < 0 or end - start != expected_bytes: + raise ValueError(f"MiniMax-H3 quantized tensor byte range is invalid: {name}") + regions.append((start, end, name)) + cursor = 0 + for start, end, name in sorted(regions): + if start != cursor: + raise ValueError(f"MiniMax-H3 quantized tensor storage is not contiguous before {name}") + cursor = end + if data_offset + cursor != file_size: + raise ValueError("MiniMax-H3 quantized checkpoint payload size is invalid") + return tensors + + +def _read_marker(path: Path, data_offset: int, entry: dict[str, Any]) -> dict[str, Any]: + start, end = entry["data_offsets"] + with path.open("rb") as stream: + stream.seek(data_offset + start) + payload = stream.read(end - start) + try: + marker = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError("MiniMax-H3 comfy_quant marker is invalid JSON") from error + if not isinstance(marker, dict): + raise ValueError("MiniMax-H3 comfy_quant marker must be an object") + return marker + + +def _source_file_identity(path: Path) -> QuantizedSourceFileIdentity: + try: + metadata = path.stat() + except OSError as error: + raise ValueError("MiniMax-H3 quantized checkpoint is unavailable") from error + if not stat.S_ISREG(metadata.st_mode): + raise ValueError("MiniMax-H3 quantized checkpoint must be a regular file") + return QuantizedSourceFileIdentity( + device=int(metadata.st_dev), + inode=int(metadata.st_ino), + size_bytes=int(metadata.st_size), + mtime_ns=int(metadata.st_mtime_ns), + ctime_ns=int(metadata.st_ctime_ns), + ) + + +def _authenticate_quantized_source(path: Path, *, workflow: str = "fl2va") -> None: + """Validate the pinned HF cache/local-dir layout without rereading 34 GB.""" + + identity = _checkpoint_identity(workflow) + relative = Path(identity.filename) + try: + source_root = path.parents[len(relative.parts) - 1] + except IndexError as error: + raise ValueError("MiniMax-H3 quantized checkpoint path is invalid") from error + expected_path = source_root / relative + if os.path.normcase(os.path.abspath(path)) != os.path.normcase(os.path.abspath(expected_path)): + raise ValueError( + "MiniMax-H3 quantized checkpoint must retain its released diffusion_models path" + ) + + if source_root.parent.name == "snapshots": + repository_root = source_root.parent.parent + if source_root.name != identity.revision or repository_root.name != _HF_CACHE_REPOSITORY: + raise ValueError( + "MiniMax-H3 quantized checkpoint is not from the pinned Hugging Face snapshot" + ) + return + + metadata_path = source_root / _HF_LOCAL_DOWNLOAD_METADATA / f"{relative.as_posix()}.metadata" + try: + metadata = metadata_path.read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError) as error: + raise ValueError( + "Cannot validate the MiniMax-H3 quantized checkpoint source. " + "Download it with `hf download Comfy-Org/MiniMax-H3 --revision " + f'{identity.revision} --include "{identity.filename}" --local-dir ` ' + "so Hugging Face writes pinned download metadata." + ) from error + if not metadata or metadata[0] != identity.revision: + raise ValueError( + "MiniMax-H3 quantized checkpoint Hugging Face metadata does not match the pinned " + "revision" + ) + + +def validate_quantized_transformer_checkpoint( + checkpoint_file: str | Path, + *, + workflow: str = "fl2va", +) -> QuantizedCheckpointIdentity: + """Validate the pinned full INT8 ConvRot checkpoint for the selected workflow. + + A Hugging Face cache symlink is accepted, but the lexical filename, exact + byte size, complete safetensors inventory, architecture config, and every + embedded quantization marker are required. + """ + + identity = _checkpoint_identity(workflow) + path = Path(checkpoint_file) + if not path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 quantized checkpoint is missing: {path}") + expected_name = Path(identity.filename).name + if path.name != expected_name: + raise ValueError( + f"MiniMax-H3 quantized checkpoint filename must be {expected_name}, got {path.name}" + ) + size = path.stat().st_size + if size != identity.size_bytes: + raise ValueError( + "MiniMax-H3 quantized checkpoint size mismatch: " + f"expected={identity.size_bytes}, actual={size}" + ) + source_identity = _source_file_identity(path) + _authenticate_quantized_source(path, workflow=workflow) + header, data_offset = _read_header(path) + tensors = _validate_header_inventory(header, data_offset, size) + for base, group_size in sorted(_QUANT_GROUPS.items()): + marker_name = f"{base}.comfy_quant" + marker = _read_marker(path, data_offset, tensors[marker_name]) + expected_marker = { + "format": _QUANT_FORMAT, + "convrot": True, + "convrot_groupsize": group_size, + } + if marker != expected_marker: + raise ValueError( + f"MiniMax-H3 unsupported comfy_quant config for {base}: {marker!r}" + ) + _validate_group_size(group_size, _PHYSICAL_SPECS[f"{base}.weight"].shape[1]) + if _source_file_identity(path) != source_identity: + raise ValueError("MiniMax-H3 quantized checkpoint changed while it was validated") + return replace( + identity, + source_file_identity=source_identity, + ) + + +def _swap_swiglu_rows(array: np.ndarray) -> np.ndarray: + if array.shape[0] % 2: + raise ValueError("MiniMax-H3 fused SwiGLU weight must have an even output width") + half = array.shape[0] // 2 + return np.ascontiguousarray(np.concatenate((array[half:], array[:half]), axis=0)) + + +def _make_quantized_weight( + qweight: np.ndarray, + scale: np.ndarray, + group_size: int, + *, + transform: str, +) -> ConvRotInt8Weight: + qweight = np.asarray(qweight) + scale = np.asarray(scale) + if transform == "swap_swiglu": + qweight = _swap_swiglu_rows(qweight) + scale = _swap_swiglu_rows(scale) + elif transform != "direct": + raise RuntimeError(f"MiniMax-H3 quantized transform is invalid: {transform}") + return ConvRotInt8Weight(qweight=qweight, scale=scale, group_size=group_size) + + +def _make_qkv_values( + qweight: np.ndarray, + scale: np.ndarray | None, + group_size: int | None, +) -> tuple[Any, Any, Any]: + packed_weight = np.ascontiguousarray(qweight) + expected_rows = 3 * _NUM_HEADS * _HEAD_DIM + if packed_weight.shape[0] != expected_rows: + raise ValueError( + f"MiniMax-H3 fused QKV has {packed_weight.shape[0]} rows, expected {expected_rows}" + ) + width = packed_weight.shape[0] // 3 + if scale is None: + return tuple(packed_weight[index * width : (index + 1) * width] for index in range(3)) + if group_size is None: + raise RuntimeError("MiniMax-H3 QKV quantization group is missing") + packed_scale = np.ascontiguousarray(scale) + parent = ConvRotInt8Weight( + qweight=packed_weight, + scale=packed_scale, + group_size=group_size, + is_full_fused_qkv=True, + ) + children = [] + for index in range(3): + start, end = index * width, (index + 1) * width + children.append( + ConvRotInt8Weight( + qweight=parent.qweight[start:end], + scale=parent.scale[start:end], + group_size=group_size, + packed_parent=parent, + row_slice=(start, end), + ) + ) + return tuple(children) + + +def load_selected_quantized_transformer_weights( + checkpoint_file: str | Path, + logical_names: Iterable[str], + *, + workflow: str = "fl2va", +) -> dict[str, Any]: + """Load only requested frozen-Diffusers logical transformer weights.""" + + path = Path(checkpoint_file) + validate_quantized_transformer_checkpoint(path, workflow=workflow) + requested = tuple(logical_names) + if len(requested) != len(set(requested)): + raise ValueError("MiniMax-H3 quantized weight request contains duplicate logical names") + unknown = sorted(set(requested) - set(_LOGICAL_MAP)) + if unknown: + raise ValueError(f"Unsupported MiniMax-H3 quantized logical tensor names: {unknown}") + if not requested: + return {} + + physical_names: set[str] = set() + for logical_name in requested: + physical_weight, _transform = _LOGICAL_MAP[logical_name] + physical_names.add(physical_weight) + base = physical_weight.removesuffix(".weight") + if base in _QUANT_GROUPS: + physical_names.add(f"{base}.weight_scale") + + from safetensors import safe_open + + with safe_open(path, framework="pt", device="cpu") as reader: + state = {name: reader.get_tensor(name) for name in sorted(physical_names)} + arrays = numpy_state(state) + + result: dict[str, Any] = {} + qkv_cache: dict[str, tuple[Any, Any, Any]] = {} + quant_cache: dict[tuple[str, str], ConvRotInt8Weight] = {} + for logical_name in requested: + physical_weight, transform = _LOGICAL_MAP[logical_name] + if transform.startswith("qkv:"): + values = qkv_cache.get(physical_weight) + if values is None: + base = physical_weight.removesuffix(".weight") + scale = arrays.get(f"{base}.weight_scale") + values = _make_qkv_values( + arrays[physical_weight], + scale, + _QUANT_GROUPS.get(base), + ) + qkv_cache[physical_weight] = values + result[logical_name] = values[int(transform.removeprefix("qkv:"))] + continue + base = physical_weight.removesuffix(".weight") + if base in _QUANT_GROUPS: + cache_key = (physical_weight, transform) + value = quant_cache.get(cache_key) + if value is None: + value = _make_quantized_weight( + arrays[physical_weight], + arrays[f"{base}.weight_scale"], + _QUANT_GROUPS[base], + transform=transform, + ) + quant_cache[cache_key] = value + result[logical_name] = value + elif transform == "swap_swiglu": + result[logical_name] = _swap_swiglu_rows(arrays[physical_weight]) + elif transform == "direct": + result[logical_name] = arrays[physical_weight] + else: + raise RuntimeError(f"MiniMax-H3 logical transform is invalid: {transform}") + return result diff --git a/families/minimax_h3/ref2va_audio_encoder_builder.py b/families/minimax_h3/ref2va_audio_encoder_builder.py new file mode 100644 index 0000000000..c6919b799b --- /dev/null +++ b/families/minimax_h3/ref2va_audio_encoder_builder.py @@ -0,0 +1,524 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native FP32 TensorRT encoder for MiniMax-H3 Ref2VA soundtracks. + +The plan consumes stereo as two batch items, follows the released DAC encoder +and causal attention projection, and emits the posterior mean. Ref2VA never +evaluates ``logs_proj`` and never samples this posterior; C++ applies the +checkpoint's 32-channel mean/std normalization and then packs channel-major +rows. No Torch, audio framework, or process sidecar is part of the runtime. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import gc +import math +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .audio_vae_builder import ( + _conv1d, + _fold_checkpoint_weight, + _require_array, +) +from .fl2va_contract import PlanAbi, TensorAbi + + +trt = trt_compat.get_trt() + +AUDIO_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 + + +@dataclass(frozen=True) +class Ref2VAAudioEncoderProfile: + batch_size: int = 2 + sampling_rate: int = 32_000 + min_samples: int = 64_000 + opt_samples: int = 165_600 + max_samples: int = 480_000 + encoder_dim: int = 64 + encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5) + latent_dim: int = 2048 + latent_channels: int = 32 + num_attention_heads: int = 8 + + @property + def hop_length(self) -> int: + return math.prod(self.encoder_rates) + + @property + def sample_profile(self) -> tuple[int, int, int]: + return self.min_samples, self.opt_samples, self.max_samples + + @property + def latent_profile(self) -> tuple[int, int, int]: + return tuple(value // self.hop_length for value in self.sample_profile) + + def validate(self) -> None: + values = ( + self.batch_size, + self.sampling_rate, + self.min_samples, + self.opt_samples, + self.max_samples, + self.encoder_dim, + self.latent_dim, + self.latent_channels, + self.num_attention_heads, + ) + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in values + ): + raise ValueError("MiniMax-H3 Ref2VA audio-encoder dimensions must be positive integers") + if self.batch_size != 2 or self.sampling_rate != 32_000: + raise ValueError("MiniMax-H3 Ref2VA audio encoder requires stereo at 32 kHz") + if not self.min_samples <= self.opt_samples <= self.max_samples: + raise ValueError("MiniMax-H3 Ref2VA audio samples must satisfy min <= opt <= max") + if any(value % self.hop_length for value in self.sample_profile): + raise ValueError( + "MiniMax-H3 Ref2VA audio inputs must be right-padded to the 800-sample hop" + ) + if self.latent_dim % self.num_attention_heads: + raise ValueError("MiniMax-H3 Ref2VA audio attention width must divide its heads") + + +DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE = Ref2VAAudioEncoderProfile() + + +def ref2va_audio_encoder_abi( + profile: Ref2VAAudioEncoderProfile = DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE, +) -> PlanAbi: + """Strict native ABI for explicit audio and attached video soundtracks. + + The C++ caller resamples/interleaves to stereo 32 kHz, separates the two + channels into this batch-of-two representation, and right-pads each clip + to the 800-sample hop. The plan emits the released posterior mode. The + caller then applies the 32 per-channel ``latents_mean``/``latents_std`` + values pinned in bundle metadata and preserves channel-major row order. + """ + + profile.validate() + sample_shapes = tuple((profile.batch_size, 1, samples) for samples in profile.sample_profile) + latent_shapes = tuple( + (profile.batch_size, profile.latent_channels, frames) for frames in profile.latent_profile + ) + return PlanAbi( + filename="ref2va_audio_vae_encoder.plan", + inputs=(TensorAbi("audio_samples", "float32", *sample_shapes),), + outputs=(TensorAbi("posterior_mean", "float32", *latent_shapes),), + ) + + +def checkpoint_keys( + profile: Ref2VAAudioEncoderProfile = DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE, +) -> tuple[str, ...]: + """Return exactly the weights evaluated by ``posterior.mode()``.""" + + profile.validate() + names = [ + "encoder.block.0.weight_g", + "encoder.block.0.weight_v", + "encoder.block.0.bias", + ] + channels = profile.encoder_dim + for stage, _stride in enumerate(profile.encoder_rates, start=1): + prefix = f"encoder.block.{stage}.block" + for residual in range(3): + residual_prefix = f"{prefix}.{residual}.block" + names.extend( + ( + f"{residual_prefix}.0.alpha", + f"{residual_prefix}.1.weight_g", + f"{residual_prefix}.1.weight_v", + f"{residual_prefix}.1.bias", + f"{residual_prefix}.2.alpha", + f"{residual_prefix}.3.weight_g", + f"{residual_prefix}.3.weight_v", + f"{residual_prefix}.3.bias", + ) + ) + names.extend( + ( + f"{prefix}.3.alpha", + f"{prefix}.4.weight_g", + f"{prefix}.4.weight_v", + f"{prefix}.4.bias", + ) + ) + channels *= 2 + names.extend( + ( + f"encoder.block.{len(profile.encoder_rates) + 1}.alpha", + f"encoder.block.{len(profile.encoder_rates) + 2}.weight_g", + f"encoder.block.{len(profile.encoder_rates) + 2}.weight_v", + f"encoder.block.{len(profile.encoder_rates) + 2}.bias", + "pre_block.norm1.weight", + "pre_block.norm1.bias", + "pre_block.attn.qkv.weight", + "pre_block.attn.q_bias", + "pre_block.attn.v_bias", + "pre_block.attn.zero_k_bias", + "pre_block.attn.proj.weight", + "pre_block.attn.proj.bias", + "pre_block.proj.weight", + "pre_block.proj.bias", + "pre_block.norm3.weight", + "pre_block.norm3.bias", + "pre_block.norm2.weight", + "pre_block.norm2.bias", + "pre_block.mlp.norm.weight", + "pre_block.mlp.norm.bias", + "pre_block.mlp.w0.weight", + "pre_block.mlp.w0.bias", + "pre_block.mlp.w1.weight", + "pre_block.mlp.w1.bias", + "pre_block.mlp.w2.weight", + "pre_block.mlp.w2.bias", + "mean_proj.weight", + "mean_proj.bias", + ) + ) + return tuple(names) + + +def _snake(network, hidden, weights, name: str): + channels = int(tuple(hidden.shape)[1]) + alpha = _require_array(weights, f"{name}.alpha", (1, channels, 1)).reshape(1, channels, 1, 1) + alpha_tensor = op.weight_constant(network, alpha) + phase = network.add_elementwise(hidden, alpha_tensor, trt.ElementWiseOperation.PROD).get_output( + 0 + ) + sine = network.add_unary(phase, trt.UnaryOperation.SIN).get_output(0) + square = network.add_elementwise(sine, sine, trt.ElementWiseOperation.PROD).get_output(0) + epsilon = op.constant(network, np.full((1, 1, 1, 1), 1.0e-9, np.float32)) + denominator = network.add_elementwise( + alpha_tensor, epsilon, trt.ElementWiseOperation.SUM + ).get_output(0) + periodic = network.add_elementwise( + square, denominator, trt.ElementWiseOperation.DIV + ).get_output(0) + return network.add_elementwise(hidden, periodic, trt.ElementWiseOperation.SUM).get_output(0) + + +def _residual_unit( + network, + hidden, + weights, + prefix: str, + *, + channels: int, + dilation: int, + owned_weights: list[np.ndarray], +): + update = _snake(network, hidden, weights, f"{prefix}.0") + conv1 = f"{prefix}.1" + update = _conv1d( + network, + update, + _fold_checkpoint_weight(weights, conv1, (channels, channels, 7)), + _require_array(weights, f"{conv1}.bias", (channels,)), + padding=3 * dilation, + dilation=dilation, + name=conv1, + owned_weights=owned_weights, + ) + update = _snake(network, update, weights, f"{prefix}.2") + conv2 = f"{prefix}.3" + update = _conv1d( + network, + update, + _fold_checkpoint_weight(weights, conv2, (channels, channels, 1)), + _require_array(weights, f"{conv2}.bias", (channels,)), + name=conv2, + owned_weights=owned_weights, + ) + return network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + +def _layer_norm(network, hidden, weights, prefix: str, width: int): + rank = len(tuple(hidden.shape)) + shape = (1,) * (rank - 1) + (width,) + gamma = op.weight_constant( + network, _require_array(weights, f"{prefix}.weight", (width,)).reshape(shape) + ) + beta = op.weight_constant( + network, _require_array(weights, f"{prefix}.bias", (width,)).reshape(shape) + ) + layer = network.add_normalization_v2(hidden, gamma, beta, 1 << (rank - 1)) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 audio LayerNorm {prefix}") + layer.name = prefix + layer.epsilon = 1.0e-5 + return layer.get_output(0) + + +def _linear(network, hidden, weights, prefix: str): + weight = _require_array( + weights, + f"{prefix}.weight", + tuple(np.asarray(weights[f"{prefix}.weight"]).shape), + ) + # ``mean_proj`` is a checkpoint Conv1d(k=1). The projection is evaluated + # on [B,T,C] rows here, so its singleton kernel axis is removed without + # changing either the values or their input/output ordering. + if weight.ndim == 3: + if weight.shape[-1] != 1: + raise ValueError(f"MiniMax-H3 audio linear {prefix} must have a unit kernel") + weight = np.ascontiguousarray(weight[..., 0]) + return op.linear( + network, + hidden, + weight, + _require_array( + weights, + f"{prefix}.bias", + tuple(np.asarray(weights[f"{prefix}.bias"]).shape), + ), + bf16=False, + ) + + +def _causal_attention_projection(network, hidden, weights, profile: Ref2VAAudioEncoderProfile): + batch = profile.batch_size + in_width = profile.latent_dim + out_width = profile.latent_channels + heads = profile.num_attention_heads + head_dim = in_width // heads + + normalized = _layer_norm(network, hidden, weights, "pre_block.norm1", in_width) + qkv_weight = _require_array(weights, "pre_block.attn.qkv.weight", (3 * in_width, in_width)) + q_bias = _require_array(weights, "pre_block.attn.q_bias", (in_width,)) + k_bias = _require_array(weights, "pre_block.attn.zero_k_bias", (in_width,)) + if np.any(k_bias != 0.0): + raise ValueError("MiniMax-H3 audio frozen key bias must be exactly zero") + v_bias = _require_array(weights, "pre_block.attn.v_bias", (in_width,)) + qkv = op.linear( + network, + normalized, + qkv_weight, + np.concatenate((q_bias, k_bias, v_bias)), + bf16=False, + ) + + tensors = [] + for index in range(3): + tensor = op.dynamic_slice( + network, + qkv, + (0, 0, index * in_width), + (batch, None, in_width), + ) + reshape = network.add_shuffle(tensor) + reshape.reshape_dims = (batch, -1, heads, head_dim) + reshape.second_transpose = trt.Permutation((0, 2, 1, 3)) + tensors.append(reshape.get_output(0)) + query, key, value = tensors + scale = op.constant( + network, + np.full((1, 1, 1, 1), 1.0 / math.sqrt(head_dim), np.float32), + ) + query = network.add_elementwise(query, scale, trt.ElementWiseOperation.PROD).get_output(0) + attention = network.add_attention( + query, + key, + value, + trt.AttentionNormalizationOp.SOFTMAX, + True, + ) + if attention is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 audio causal attention") + attention.name = "pre_block.attn.native_causal_attention" + attention.metadata = "trtmc.native_op=IAttention;causal=true;decomposable=true" + # TRT-RTX 1.6 has no dedicated causal kernel covering the complete + # dynamic 80..600-row public profile. Allow TRT itself to lower the + # IAttention composite into native layers; this remains inside the plan + # and introduces no plugin, framework, or runtime dependency. + attention.decomposable = True + + # [B,H,T,D] -> mean heads -> [B,T,D] -> average each exact group + # of D/out_width=8 channels, matching adaptive_avg_pool1d. + pooled_heads = network.add_reduce( + attention.get_output(0), + trt.ReduceOperation.AVG, + 1 << 1, + False, + ) + reshape = network.add_shuffle(pooled_heads.get_output(0)) + reshape.reshape_dims = (batch, -1, out_width, head_dim // out_width) + pooled_width = network.add_reduce( + reshape.get_output(0), + trt.ReduceOperation.AVG, + 1 << 3, + False, + ).get_output(0) + attended = _linear(network, pooled_width, weights, "pre_block.attn.proj") + + residual = _layer_norm(network, hidden, weights, "pre_block.norm3", in_width) + residual = _linear(network, residual, weights, "pre_block.proj") + hidden = network.add_elementwise(residual, attended, trt.ElementWiseOperation.SUM).get_output(0) + update = _layer_norm(network, hidden, weights, "pre_block.norm2", out_width) + update = _layer_norm(network, update, weights, "pre_block.mlp.norm", out_width) + left = _linear(network, update, weights, "pre_block.mlp.w0") + gelu = network.add_activation(left, trt.ActivationType.GELU_TANH) + if gelu is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 audio GeGLU activation") + right = _linear(network, update, weights, "pre_block.mlp.w1") + update = network.add_elementwise( + gelu.get_output(0), right, trt.ElementWiseOperation.PROD + ).get_output(0) + update = _linear(network, update, weights, "pre_block.mlp.w2") + return network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + +@op.cleanup_failed_build +def build_ref2va_audio_encoder_engine( + weights: dict, + profile: Ref2VAAudioEncoderProfile = DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE, + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the released stereo soundtrack posterior-mean encoder.""" + + profile.validate() + expected = set(checkpoint_keys(profile)) + missing = sorted(expected - set(weights)) + unexpected = sorted(set(weights) - expected) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 Ref2VA audio encoder checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=AUDIO_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + abi = ref2va_audio_encoder_abi(profile) + samples = network.add_input( + abi.inputs[0].name, + trt.float32, + (profile.batch_size, 1, -1), + ) + optimization = builder.create_optimization_profile() + optimization.set_shape( + "audio_samples", + min=(profile.batch_size, 1, profile.min_samples), + opt=(profile.batch_size, 1, profile.opt_samples), + max=(profile.batch_size, 1, profile.max_samples), + ) + if config.add_optimization_profile(optimization) != 0: + raise RuntimeError("TensorRT rejected MiniMax-H3 Ref2VA audio profile") + expand = network.add_shuffle(samples) + expand.reshape_dims = (profile.batch_size, 1, 1, -1) + hidden = expand.get_output(0) + owned_weights: list[np.ndarray] = [] + + hidden = _conv1d( + network, + hidden, + _fold_checkpoint_weight( + weights, + "encoder.block.0", + (profile.encoder_dim, 1, 7), + ), + _require_array(weights, "encoder.block.0.bias", (profile.encoder_dim,)), + padding=3, + name="encoder.block.0", + owned_weights=owned_weights, + ) + channels = profile.encoder_dim + for stage, stride in enumerate(profile.encoder_rates, start=1): + prefix = f"encoder.block.{stage}.block" + for residual, dilation in enumerate((1, 3, 9)): + hidden = _residual_unit( + network, + hidden, + weights, + f"{prefix}.{residual}.block", + channels=channels, + dilation=dilation, + owned_weights=owned_weights, + ) + hidden = _snake(network, hidden, weights, f"{prefix}.3") + down = f"{prefix}.4" + hidden = _conv1d( + network, + hidden, + _fold_checkpoint_weight( + weights, + down, + (2 * channels, channels, 2 * stride), + ), + _require_array(weights, f"{down}.bias", (2 * channels,)), + stride=stride, + padding=math.ceil(stride / 2), + name=down, + owned_weights=owned_weights, + ) + channels *= 2 + hidden = _snake( + network, + hidden, + weights, + f"encoder.block.{len(profile.encoder_rates) + 1}", + ) + final = f"encoder.block.{len(profile.encoder_rates) + 2}" + hidden = _conv1d( + network, + hidden, + _fold_checkpoint_weight(weights, final, (profile.latent_dim, channels, 3)), + _require_array(weights, f"{final}.bias", (profile.latent_dim,)), + padding=1, + name=final, + owned_weights=owned_weights, + ) + rows = network.add_shuffle(hidden) + rows.first_transpose = trt.Permutation((0, 3, 1, 2)) + rows.reshape_dims = (profile.batch_size, -1, profile.latent_dim) + hidden = _causal_attention_projection(network, rows.get_output(0), weights, profile) + posterior_mean = _linear(network, hidden, weights, "mean_proj") + output = network.add_shuffle(posterior_mean) + output.first_transpose = trt.Permutation((0, 2, 1)) + value = output.get_output(0) + value.name = abi.outputs[0].name + network.mark_output(value) + op.validate_native_network( + network, + expected_attentions=1, + label="Ref2VA audio VAE encoder", + ) + + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 Ref2VA audio encoder") + del network, config, builder, logger, owned_weights + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/ref2va_bundle_contract.py b/families/minimax_h3/ref2va_bundle_contract.py new file mode 100644 index 0000000000..2281885b31 --- /dev/null +++ b/families/minimax_h3/ref2va_bundle_contract.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fail-closed bundle/provenance contract for native MiniMax-H3 Ref2VA.""" + +from __future__ import annotations + +import math + +from .fl2va_contract import PlanAbi, TensorAbi +from .ref2va_checkpoint import ( + CHECKPOINT_REVISION, + COMPONENT_NAME, + MODEL_ID, + TransformerRefIdentity, +) +from .ref2va_contract import ( + MAX_AUDIOS, + MAX_IMAGES, + MAX_REFERENCES, + MAX_REFERENCE_DURATION_SECONDS, + MAX_TOTAL_AUDIO_DURATION_SECONDS, + MAX_TOTAL_VIDEO_DURATION_SECONDS, + MAX_VIDEOS, + MIN_REFERENCE_DURATION_SECONDS, + Ref2VADenoiserProfile, + ref2va_denoiser_abi, + ref2va_denoiser_profiles, + ref2va_first_block_cache_abis, +) +from .ref2va_qwen_contract import ref2va_shared_qwen_profile_metadata +from .quantized_checkpoint import QuantizedCheckpointIdentity, QUANTIZED_REF2VA_CHECKPOINT_IDENTITY + + +# The Qwen language/vision plans are intentionally not listed here: Ref2VA +# binds the same ``text_encoder_plan`` and ``vision_encoder_plan`` sections +# already shared by T2VA/FL2VA. No second Qwen weight copy is permitted. +REF2VA_PLAN_SECTIONS = ( + ("ref2va_denoiser", "ref2va_denoiser.plan", "ref2va_denoiser_plan"), + ( + "ref2va_adaln_precompute", + "ref2va_adaln_precompute.plan", + "ref2va_adaln_precompute_plan", + ), + ( + "ref2va_video_vae_encoder", + "ref2va_video_vae_encoder.plan", + "ref2va_video_vae_encoder_plan", + ), + ( + "ref2va_audio_vae_encoder", + "ref2va_audio_vae_encoder.plan", + "ref2va_audio_vae_encoder_plan", + ), +) + +REF2VA_FIRST_BLOCK_CACHE_SECTIONS = ( + ("ref2va_dit_head", "ref2va_dit_head.plan", "ref2va_dit_head_plan"), + ("ref2va_dit_tail", "ref2va_dit_tail.plan", "ref2va_dit_tail_plan"), + ("ref2va_dit_finish", "ref2va_dit_finish.plan", "ref2va_dit_finish_plan"), + *REF2VA_PLAN_SECTIONS[1:], +) + + +def ref2va_plan_sections(first_block_cache: bool) -> tuple[tuple[str, str, str], ...]: + return REF2VA_FIRST_BLOCK_CACHE_SECTIONS if first_block_cache else REF2VA_PLAN_SECTIONS + + +REF2VA_SHARED_SECTIONS = { + "text_encoder": "text_encoder_plan", + "vision_encoder": "vision_encoder_plan", + "image_vae_encoder": "fl2va_keyframe_vae_encoder_plan", + "video_vae_decoder": "vae_tile_decoder_plan", + "audio_vae_decoder": "audio_vae_decoder_plan", +} + +# The public transformer_ref partition is the guidance-distilled 50-point +# rectified-flow recipe. ``num_inference_steps`` counts the terminal zero, so +# the transformer executes once for each of the first 49 grid points. Keep +# this mode-local instead of inheriting the base T2VA/FL2VA transformer schedule. +REF2VA_SCHEDULER_GRID_POINTS = 50 +REF2VA_TRANSFORMER_FORWARDS = 49 +REF2VA_VIDEO_FLOW_SHIFT = 12.0 +REF2VA_AUDIO_FLOW_SHIFT = 3.0 +REF2VA_GUIDANCE_SCALE = 1.0 + + +def _ref2va_adaln_abi() -> PlanAbi: + outputs = tuple( + TensorAbi( + f"block_modulation_{index}", + "bfloat16", + (12, 6, 5_376), + (12, 6, 5_376), + (12, 6, 5_376), + ) + for index in range(50) + ) + ( + TensorAbi( + "final_modulation", + "bfloat16", + (4, 2, 5_376), + (4, 2, 5_376), + (4, 2, 5_376), + ), + ) + return PlanAbi( + filename="ref2va_adaln_precompute.plan", + inputs=( + TensorAbi( + "timestep_features", + "float32", + (4, 256), + (4, 256), + (4, 256), + ), + ), + outputs=outputs, + ) + + +def _ref2va_video_encoder_abi() -> PlanAbi: + return PlanAbi( + filename="ref2va_video_vae_encoder.plan", + inputs=( + TensorAbi( + "pixel_tile_clip", + "float32", + (1, 3, 17, 256, 256), + (1, 3, 17, 256, 256), + (1, 3, 17, 256, 256), + ), + ), + outputs=( + TensorAbi( + "posterior_parameter_tile_clip", + "float32", + (1, 48, 5, 16, 16), + (1, 48, 5, 16, 16), + (1, 48, 5, 16, 16), + ), + ), + ) + + +def _ref2va_audio_encoder_abi() -> PlanAbi: + return PlanAbi( + filename="ref2va_audio_vae_encoder.plan", + inputs=( + TensorAbi( + "audio_samples", + "float32", + (2, 1, 64_000), + (2, 1, 165_600), + (2, 1, 480_000), + ), + ), + outputs=( + TensorAbi( + "posterior_mean", + "float32", + (2, 32, 80), + (2, 32, 207), + (2, 32, 600), + ), + ), + ) + + +def ref2va_plan_abi_metadata( + profile: Ref2VADenoiserProfile = Ref2VADenoiserProfile(), + *, + first_block_cache: bool = False, +) -> dict[str, object]: + """Serialize every dedicated plan binding without importing TensorRT.""" + + profile.validate() + + def tensor(binding: TensorAbi) -> dict[str, object]: + return { + "name": binding.name, + "dtype": binding.dtype, + "min_shape": list(binding.min_shape), + "opt_shape": list(binding.opt_shape), + "max_shape": list(binding.max_shape), + } + + denoiser = ( + tuple((f"{name}_plan", abi) for name, abi in ref2va_first_block_cache_abis(profile).items()) + if first_block_cache + else (("ref2va_denoiser_plan", ref2va_denoiser_abi(profile)),) + ) + plans = ( + *denoiser, + ("ref2va_adaln_precompute_plan", _ref2va_adaln_abi()), + ("ref2va_video_vae_encoder_plan", _ref2va_video_encoder_abi()), + ("ref2va_audio_vae_encoder_plan", _ref2va_audio_encoder_abi()), + ) + return { + section: { + "filename": abi.filename, + "inputs": [tensor(binding) for binding in abi.inputs], + "outputs": [tensor(binding) for binding in abi.outputs], + } + for section, abi in plans + } + + +def _capacity(profile: Ref2VADenoiserProfile) -> dict[str, list[int]]: + return { + "video_rows": [ + profile.min_video_rows, + profile.opt_video_rows, + profile.max_video_rows, + ], + "audio_rows": [ + profile.min_audio_rows, + profile.opt_audio_rows, + profile.max_audio_rows, + ], + "text_rows": [ + profile.min_text_rows, + profile.opt_text_rows, + profile.max_text_rows, + ], + "packed_rows": [ + profile.min_packed_rows, + profile.opt_packed_rows, + profile.max_packed_rows, + ], + } + + +def ref2va_bundle_metadata( + transformer_ref: TransformerRefIdentity | QuantizedCheckpointIdentity, + profile: Ref2VADenoiserProfile = Ref2VADenoiserProfile(), + *, + first_block_cache: bool = False, + first_block_cache_threshold: float = 0.08, +) -> dict[str, object]: + """Produce path-free metadata only after strict checkpoint validation.""" + + profile.validate() + if not isinstance(first_block_cache, bool): + raise ValueError("MiniMax-H3 ref2va_first_block_cache must be a boolean") + if ( + not isinstance(first_block_cache_threshold, (int, float)) + or isinstance(first_block_cache_threshold, bool) + or not math.isfinite(first_block_cache_threshold) + or first_block_cache_threshold < 0.0 + ): + raise ValueError("MiniMax-H3 Ref2VA cache threshold must be finite and nonnegative") + if isinstance(transformer_ref, QuantizedCheckpointIdentity): + if transformer_ref.bundle_metadata() != QUANTIZED_REF2VA_CHECKPOINT_IDENTITY.bundle_metadata(): + raise ValueError("MiniMax-H3 Ref2VA requires the distinct full Comfy INT8 checkpoint") + elif not isinstance(transformer_ref, TransformerRefIdentity): + raise TypeError("MiniMax-H3 Ref2VA metadata requires validated transformer_ref identity") + elif ( + transformer_ref.model_id != MODEL_ID + or transformer_ref.revision != CHECKPOINT_REVISION + or transformer_ref.component != COMPONENT_NAME + or transformer_ref.tensor_count != 638 + ): + raise ValueError("MiniMax-H3 Ref2VA transformer_ref provenance is incompatible") + denoiser_profiles = ref2va_denoiser_profiles(profile) + denoiser_profile_names = ( + ("five_second_common", "public_dynamic") + if len(denoiser_profiles) == 2 + else ("public_dynamic",) + ) + return { + "ref2va_schema_version": 5 if first_block_cache else 4, + "ref2va_supported": True, + "ref2va_first_block_cache": { + "enabled": first_block_cache, + "threshold": float(first_block_cache_threshold), + }, + "ref2va_scheduler": { + "sigma_grid_points": REF2VA_SCHEDULER_GRID_POINTS, + "transformer_forwards": REF2VA_TRANSFORMER_FORWARDS, + "video_shift": REF2VA_VIDEO_FLOW_SHIFT, + "audio_shift": REF2VA_AUDIO_FLOW_SHIFT, + "guidance_scale": REF2VA_GUIDANCE_SCALE, + "guidance_distilled": True, + }, + "ref2va_transformer_ref": transformer_ref.bundle_metadata(), + "ref2va_plan_sections": { + component: section + for component, _filename, section in ref2va_plan_sections(first_block_cache) + }, + "ref2va_plan_abis": ref2va_plan_abi_metadata(profile, first_block_cache=first_block_cache), + "ref2va_denoiser_profile_count": len(denoiser_profiles), + "ref2va_denoiser_profile_layout": "_then_".join(denoiser_profile_names), + "ref2va_denoiser_profiles": [ + {"name": name, **_capacity(capacity)} + for name, capacity in zip( + denoiser_profile_names, + denoiser_profiles, + strict=True, + ) + ], + "ref2va_shared_sections": dict(REF2VA_SHARED_SECTIONS), + "ref2va_shared_qwen_profiles": ref2va_shared_qwen_profile_metadata(), + "ref2va_limits": { + "max_images": MAX_IMAGES, + "max_videos": MAX_VIDEOS, + "max_explicit_audios": MAX_AUDIOS, + "max_reference_files": MAX_REFERENCES, + "min_seconds_each_video_or_audio": MIN_REFERENCE_DURATION_SECONDS, + "max_seconds_each_video_or_audio": MAX_REFERENCE_DURATION_SECONDS, + "max_total_video_seconds": MAX_TOTAL_VIDEO_DURATION_SECONDS, + "max_total_video_soundtrack_seconds": MAX_TOTAL_VIDEO_DURATION_SECONDS, + "max_total_explicit_audio_seconds": MAX_TOTAL_AUDIO_DURATION_SECONDS, + "audio_can_be_sole_input": True, + "video_soundtrack_stays_attached": True, + }, + "ref2va_capacity": _capacity(profile), + } diff --git a/families/minimax_h3/ref2va_checkpoint.py b/families/minimax_h3/ref2va_checkpoint.py new file mode 100644 index 0000000000..cba5e81ccf --- /dev/null +++ b/families/minimax_h3/ref2va_checkpoint.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict provenance and partitioning for the public H3 ``transformer_ref``. + +Ref2VA is not a mode bit on the T2VA transformer. The released repository has +a second 66.28 GB checkpoint partition with the same architecture and tensor +names but different values. This module deliberately refuses to build from +``transformer/`` or to fall back when ``transformer_ref/`` is absent. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path + + +MODEL_ID = "MiniMaxAI/MiniMax-H3" +CHECKPOINT_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc" +COMPONENT_NAME = "transformer_ref" +TOTAL_TENSOR_BYTES = 66_280_430_080 +CONFIG_BYTES = 546 +INDEX_BYTES = 64_488 + +SHARDS: tuple[tuple[str, int], ...] = ( + ("diffusion_pytorch_model-00001-of-00014.safetensors", 4_825_958_704), + ("diffusion_pytorch_model-00002-of-00014.safetensors", 4_702_158_032), + ("diffusion_pytorch_model-00003-of-00014.safetensors", 4_933_368_192), + ("diffusion_pytorch_model-00004-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00005-of-00014.safetensors", 4_702_158_080), + ("diffusion_pytorch_model-00006-of-00014.safetensors", 4_933_368_232), + ("diffusion_pytorch_model-00007-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00008-of-00014.safetensors", 4_702_158_080), + ("diffusion_pytorch_model-00009-of-00014.safetensors", 4_933_368_232), + ("diffusion_pytorch_model-00010-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00011-of-00014.safetensors", 4_702_158_080), + ("diffusion_pytorch_model-00012-of-00014.safetensors", 4_933_368_232), + ("diffusion_pytorch_model-00013-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00014-of-00014.safetensors", 4_644_161_920), +) + +EXPECTED_CONFIG = { + "_class_name": "MiniMaxH3Transformer3DModel", + "_diffusers_version": "0.36.0.dev0", + "num_attention_heads": 56, + "attention_head_dim": 128, + "hidden_size": 5376, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "in_channels": 24, + "audio_in_channels": 32, + "patch_size": [1, 2, 2], + "text_dim": 5120, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "rope_freq_dim": 16, + "rope_theta": 10_000.0, + "norm_eps": 1.0e-5, + "qk_norm_eps": 1.0e-5, + "final_norm_eps": 1.0e-5, +} + + +def _dit_checkpoint_keys() -> tuple[str, ...]: + names = [ + "proj_in.weight", + "proj_in.bias", + "audio_proj_in.weight", + "audio_proj_in.bias", + "context_embedder.weight", + "context_embedder.bias", + "token_refiner.final_norm.weight", + ] + for index in range(2): + prefix = f"token_refiner.refiner_blocks.{index}" + names.extend( + ( + f"{prefix}.norm1.weight", + f"{prefix}.norm2.weight", + *(f"{prefix}.attn.to_{name}.weight" for name in ("q", "k", "v")), + f"{prefix}.attn.norm_q.weight", + f"{prefix}.attn.norm_k.weight", + f"{prefix}.attn.to_out.0.weight", + f"{prefix}.ff.net.0.proj.weight", + f"{prefix}.ff.net.2.weight", + ) + ) + for index in range(50): + prefix = f"transformer_blocks.{index}" + names.extend( + ( + f"{prefix}.norm1.weight", + f"{prefix}.norm2.weight", + *(f"{prefix}.attn.to_{name}.weight" for name in ("q", "k", "v")), + f"{prefix}.attn.norm_q.weight", + f"{prefix}.attn.norm_k.weight", + f"{prefix}.attn.to_out.0.weight", + f"{prefix}.ff.net.0.proj.weight", + f"{prefix}.ff.net.2.weight", + ) + ) + names.extend( + ( + "norm_out.norm.weight", + "proj_out.weight", + "proj_out.bias", + "audio_proj_out.weight", + "audio_proj_out.bias", + ) + ) + return tuple(names) + + +def _adaln_checkpoint_keys() -> tuple[str, ...]: + names = [ + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + ] + for index in range(50): + prefix = f"transformer_blocks.{index}.adaln_proj.linear" + names.extend((f"{prefix}.weight", f"{prefix}.bias")) + names.extend(("norm_out.linear.weight", "norm_out.linear.bias")) + return tuple(names) + + +REF2VA_DENOISER_KEYS = _dit_checkpoint_keys() +REF2VA_ADALN_KEYS = _adaln_checkpoint_keys() +REF2VA_ALL_KEYS = REF2VA_DENOISER_KEYS + REF2VA_ADALN_KEYS +REF2VA_FINISH_KEYS = tuple( + name + for name in REF2VA_DENOISER_KEYS + if name.startswith(("norm_out.", "proj_out.", "audio_proj_out.")) +) +REF2VA_TAIL_KEYS = tuple( + name + for name in REF2VA_DENOISER_KEYS + if name.startswith("transformer_blocks.") and not name.startswith("transformer_blocks.0.") +) +REF2VA_HEAD_KEYS = tuple( + name + for name in REF2VA_DENOISER_KEYS + if name not in REF2VA_FINISH_KEYS and name not in REF2VA_TAIL_KEYS +) + + +def download_command(local_dir: str = "MiniMax-H3") -> str: + """Return the build-time-only command for the missing public partition.""" + + if not isinstance(local_dir, str) or not local_dir.strip(): + raise ValueError("MiniMax-H3 checkpoint destination must be a non-empty string") + return ( + f"hf download {MODEL_ID} --revision {CHECKPOINT_REVISION} " + f'--include "transformer_ref/*" --local-dir "{local_dir}"' + ) + + +def _require_file(path: Path, *, size: int) -> dict[str, int]: + if not path.is_file() or path.is_symlink(): + raise FileNotFoundError(f"MiniMax-H3 Ref2VA checkpoint file is missing: {path.name}") + actual_size = path.stat().st_size + if actual_size != size: + raise ValueError( + f"MiniMax-H3 Ref2VA checkpoint size mismatch for {path.name}: " + f"expected={size}, actual={actual_size}" + ) + return {"bytes": size} + + +@dataclass(frozen=True) +class TransformerRefIdentity: + model_id: str + revision: str + component: str + tensor_bytes: int + tensor_count: int + files: dict[str, dict[str, int | str]] + + def bundle_metadata(self) -> dict[str, object]: + """Path-free provenance safe to persist in the native bundle.""" + + return { + "schema_version": 1, + "model_id": self.model_id, + "revision": self.revision, + "component": self.component, + "tensor_bytes": self.tensor_bytes, + "tensor_count": self.tensor_count, + "files": self.files, + "runtime_framework": None, + } + + +def validate_transformer_ref_checkpoint(component_dir: str | Path) -> TransformerRefIdentity: + """Validate the released ``transformer_ref`` layout and plan partition. + + The pinned config and index are hashed; large shards are checked against + the exact index names and released sizes without rereading 66 GiB solely + for an integrity pass. + """ + root = Path(component_dir) + if root.name != COMPONENT_NAME: + candidate = root / COMPONENT_NAME + if candidate.is_dir(): + root = candidate + else: + raise FileNotFoundError( + "MiniMax-H3 Ref2VA requires the distinct transformer_ref checkpoint. " + "The T2VA transformer is not a valid fallback. Download it with: " + f"{download_command(str(root))}" + ) + if not root.is_dir() or root.is_symlink(): + raise FileNotFoundError( + "MiniMax-H3 transformer_ref is absent; no T2VA fallback is permitted. " + f"Download it with: {download_command(str(root.parent))}" + ) + + files: dict[str, dict[str, int | str]] = {} + config_path = root / "config.json" + files["config.json"] = _require_file( + config_path, + size=CONFIG_BYTES, + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + if config != EXPECTED_CONFIG: + raise ValueError( + "MiniMax-H3 transformer_ref config does not match the released architecture" + ) + + index_name = "diffusion_pytorch_model.safetensors.index.json" + index_path = root / index_name + files[index_name] = _require_file( + index_path, + size=INDEX_BYTES, + ) + index = json.loads(index_path.read_text(encoding="utf-8")) + if not isinstance(index, dict) or set(index) != {"metadata", "weight_map"}: + raise ValueError("MiniMax-H3 transformer_ref index has an invalid schema") + if index["metadata"] != {"total_size": TOTAL_TENSOR_BYTES}: + raise ValueError("MiniMax-H3 transformer_ref index total_size is invalid") + weight_map = index["weight_map"] + if not isinstance(weight_map, dict): + raise ValueError("MiniMax-H3 transformer_ref weight_map is invalid") + expected_keys = set(REF2VA_ALL_KEYS) + if len(REF2VA_ALL_KEYS) != 638 or len(expected_keys) != 638: + raise RuntimeError("MiniMax-H3 transformer_ref partition is not exhaustive") + actual_keys = set(weight_map) + if actual_keys != expected_keys: + raise ValueError( + "MiniMax-H3 transformer_ref tensor partition mismatch: " + f"missing={sorted(expected_keys - actual_keys)[:8]}, " + f"unexpected={sorted(actual_keys - expected_keys)[:8]}" + ) + expected_shards = {name for name, _size in SHARDS} + if set(weight_map.values()) != expected_shards: + raise ValueError("MiniMax-H3 transformer_ref index does not reference the exact 14 shards") + + for name, size in SHARDS: + files[name] = _require_file( + root / name, + size=size, + ) + allowed = {"config.json", index_name, *expected_shards} + unexpected_files = sorted( + path.name for path in root.iterdir() if path.is_file() and path.name not in allowed + ) + if unexpected_files: + raise ValueError( + f"MiniMax-H3 transformer_ref contains unexpected files: {unexpected_files}" + ) + return TransformerRefIdentity( + model_id=MODEL_ID, + revision=CHECKPOINT_REVISION, + component=COMPONENT_NAME, + tensor_bytes=TOTAL_TENSOR_BYTES, + tensor_count=len(weight_map), + files=files, + ) diff --git a/families/minimax_h3/ref2va_contract.py b/families/minimax_h3/ref2va_contract.py new file mode 100644 index 0000000000..b988226a59 --- /dev/null +++ b/families/minimax_h3/ref2va_contract.py @@ -0,0 +1,1114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native-runtime contracts for the public MiniMax-H3 Ref2VA workflow. + +This module contains no framework runtime. It is the executable specification +used by the bundle builder and by the C++ port for request validation, media +geometry, Qwen3-VL presentation/MRoPE, packed Omni-DiT rows, and per-row noise +levels. Context-IR and Regenerate-2K are intentionally outside this contract. + +The Python is build/test tooling only. A serialized ModelConnect bundle must +implement the same operations in C++/CUDA and TensorRT-RTX; importing this file +at inference time is neither required nor supported. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +import math +from typing import Literal + +import numpy as np + +from .fl2va_contract import PlanAbi, TensorAbi, resolve_canvas_size, validate_canvas + + +ReferenceKind = Literal["image", "video", "audio"] +QwenModality = Literal["text", "image", "video"] + +MAX_IMAGES = 9 +MAX_VIDEOS = 3 +MAX_AUDIOS = 3 +MAX_REFERENCES = 12 +MIN_REFERENCE_DURATION_SECONDS = 2.0 +MAX_REFERENCE_DURATION_SECONDS = 15.0 +MAX_TOTAL_VIDEO_DURATION_SECONDS = 15.0 +MAX_TOTAL_AUDIO_DURATION_SECONDS = 15.0 + +TARGET_FPS = 24.0 +AUDIO_HOP_LENGTH = 800 +AUDIO_CHANNELS = 2 +VIDEO_PATCH_HEIGHT = 2 +VIDEO_PATCH_WIDTH = 2 +VIDEO_FRAMES_PER_CHUNK = 17 +VIDEO_LATENTS_PER_CHUNK = 5 +REFERENCE_IMAGE_SHORT_EDGE = 2048 +CANVAS_MULTIPLE = 32 +QWEN_VISION_PATCH_SIZE = 16 +QWEN_VISION_MERGE_SIZE = 2 +QWEN_VISION_TEMPORAL_PATCH_SIZE = 2 +QWEN_VIDEO_SAMPLE_FPS = 2.0 +QWEN_MAX_POSITION_EMBEDDINGS = 262_144 + +QWEN_VISION_START_TOKEN_ID = 151_652 +QWEN_VISION_END_TOKEN_ID = 151_653 +QWEN_IMAGE_PAD_TOKEN_ID = 151_655 +QWEN_VIDEO_PAD_TOKEN_ID = 151_656 + +H3_VIDEO_TAG = 0 +H3_TEXT_TAG = 1 +H3_AUDIO_TAG = 2 +CONDITION_VIDEO_TIMESTEP = 0.999 +CONDITION_AUDIO_TIMESTEP = 1.0 + +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32.0 + + +@dataclass(frozen=True) +class ReferenceSpec: + """Decoded public reference metadata, in request order. + + ``duration_seconds`` belongs to video and explicit-audio files. A video + soundtrack remains attached to its video and therefore does not increment + the explicit-audio count or the twelve-file count. + """ + + kind: ReferenceKind + duration_seconds: float | None = None + width: int | None = None + height: int | None = None + has_audio: bool = False + + def validate(self) -> None: + if self.kind not in ("image", "video", "audio"): + raise ValueError(f"MiniMax-H3 reference kind is invalid: {self.kind!r}") + if not isinstance(self.has_audio, bool): + raise ValueError("MiniMax-H3 has_audio must be a boolean") + if self.has_audio and self.kind != "video": + raise ValueError("Only a MiniMax-H3 video reference can carry a soundtrack") + + if self.kind in ("image", "video"): + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in (self.width, self.height) + ): + raise ValueError("MiniMax-H3 visual references need positive integer geometry") + assert self.width is not None and self.height is not None + if self.width > 4 * self.height or self.height > 4 * self.width: + raise ValueError( + "MiniMax-H3 visual references must be within the public 1:4..4:1 aspect range" + ) + elif self.width is not None or self.height is not None: + raise ValueError("MiniMax-H3 audio references cannot carry pixel geometry") + + if self.kind == "image": + if self.duration_seconds is not None: + raise ValueError("MiniMax-H3 image references do not have a duration") + return + duration = self.duration_seconds + if ( + isinstance(duration, bool) + or not isinstance(duration, (int, float)) + or not math.isfinite(duration) + or not MIN_REFERENCE_DURATION_SECONDS <= duration <= MAX_REFERENCE_DURATION_SECONDS + ): + raise ValueError("MiniMax-H3 video/audio references must each be 2..15 seconds long") + + +@dataclass(frozen=True) +class ReferenceRequestSummary: + image_count: int + video_count: int + audio_count: int + total_video_seconds: float + total_audio_seconds: float + audio_bearing_count: int + + +def validate_reference_request( + references: Sequence[ReferenceSpec], +) -> ReferenceRequestSummary: + """Fail closed on every public H3-Base Ref2VA file/duration limit.""" + + if not isinstance(references, Sequence) or isinstance(references, (str, bytes)): + raise ValueError("MiniMax-H3 references must be an ordered sequence") + if not references: + raise ValueError("MiniMax-H3 Ref2VA needs at least one reference") + if len(references) > MAX_REFERENCES: + raise ValueError( + f"MiniMax-H3 accepts at most {MAX_REFERENCES} reference files, got {len(references)}" + ) + for reference in references: + if not isinstance(reference, ReferenceSpec): + raise ValueError("MiniMax-H3 references must use ReferenceSpec metadata") + reference.validate() + + image_count = sum(reference.kind == "image" for reference in references) + video_count = sum(reference.kind == "video" for reference in references) + audio_count = sum(reference.kind == "audio" for reference in references) + if image_count > MAX_IMAGES: + raise ValueError(f"MiniMax-H3 accepts at most {MAX_IMAGES} image references") + if video_count > MAX_VIDEOS: + raise ValueError(f"MiniMax-H3 accepts at most {MAX_VIDEOS} video references") + if audio_count > MAX_AUDIOS: + raise ValueError(f"MiniMax-H3 accepts at most {MAX_AUDIOS} explicit audio references") + total_video = math.fsum( + float(reference.duration_seconds) for reference in references if reference.kind == "video" + ) + total_audio = math.fsum( + float(reference.duration_seconds) for reference in references if reference.kind == "audio" + ) + if total_video > MAX_TOTAL_VIDEO_DURATION_SECONDS: + raise ValueError("MiniMax-H3 reference-video duration total exceeds 15 seconds") + if total_audio > MAX_TOTAL_AUDIO_DURATION_SECONDS: + raise ValueError("MiniMax-H3 explicit-audio duration total exceeds 15 seconds") + return ReferenceRequestSummary( + image_count=image_count, + video_count=video_count, + audio_count=audio_count, + total_video_seconds=total_video, + total_audio_seconds=total_audio, + audio_bearing_count=audio_count + + sum(reference.kind == "video" and reference.has_audio for reference in references), + ) + + +def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: + """Return ``(height, width)`` at short-edge 2048 with no area cap. + + The round is Python/IEEE round-half-to-even, matching the released setup + block. The 1:4 and 4:1 endpoints resolve to 2048x8192 and therefore to the + Qwen image processor's exact 16,777,216-pixel ceiling. + """ + + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in (width, height) + ): + raise ValueError("MiniMax-H3 reference image geometry must use positive integers") + if width > 4 * height or height > 4 * width: + raise ValueError("MiniMax-H3 reference image must be within 1:4 and 4:1") + scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + target_height = max(CANVAS_MULTIPLE, round(height * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE) + target_width = max(CANVAS_MULTIPLE, round(width * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE) + return target_height, target_width + + +def resolve_reference_video_size(width: int, height: int) -> tuple[int, int]: + """Put a video reference on the public 768p canvas of its own aspect.""" + + return resolve_canvas_size(width, height) + + +def video_resample_source_indices( + source_frames: int, + source_fps: float, + *, + target_frames: int, + target_fps: float = TARGET_FPS, +) -> tuple[int, ...]: + """Reproduce the reference's constant-frame-rate drop/duplicate pass. + + The returned source row for each output slot is applied before LANCZOS + spatial resize and is truncated to the generated frame count. + """ + + if ( + isinstance(source_frames, bool) + or not isinstance(source_frames, int) + or source_frames <= 0 + or isinstance(target_frames, bool) + or not isinstance(target_frames, int) + or target_frames <= 0 + ): + raise ValueError("MiniMax-H3 video frame counts must be positive integers") + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + for value in (source_fps, target_fps) + ): + raise ValueError("MiniMax-H3 video frame rates must be finite and positive") + if float(source_fps) == float(target_fps): + return tuple(range(min(source_frames, target_frames))) + + scale = float(target_fps) / float(source_fps) + slots = np.floor(np.arange(source_frames, dtype=np.float64) * scale + 0.5).astype(np.int64) + end = math.floor(source_frames * scale + 0.5) + counts = np.diff(slots, append=np.int64(end)) + if np.any(counts < 0): + raise RuntimeError("MiniMax-H3 constant-frame-rate slots are not monotonic") + result = np.repeat(np.arange(source_frames, dtype=np.int64), counts) + return tuple(int(value) for value in result[:target_frames]) + + +def snap_reference_video_frames_down(num_frames: int) -> int: + """Snap a normalized reference down to ``17*n+5`` before VAE encode.""" + + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError("MiniMax-H3 reference video needs a positive frame count") + chunks = max(1, (num_frames - VIDEO_LATENTS_PER_CHUNK) // VIDEO_FRAMES_PER_CHUNK) + return chunks * VIDEO_FRAMES_PER_CHUNK + VIDEO_LATENTS_PER_CHUNK + + +def reference_video_latent_frames(num_frames: int) -> int: + aligned = snap_reference_video_frames_down(num_frames) + return ( + (aligned - VIDEO_LATENTS_PER_CHUNK) // VIDEO_FRAMES_PER_CHUNK + ) * VIDEO_LATENTS_PER_CHUNK + 2 + + +@dataclass(frozen=True) +class ReferenceVideoEncodeSchedule: + """Native clip schedule matching ``AutoencoderKLMiniMaxH3._encode``.""" + + snapped_frames: int + clip_count: int + repeated_tail_frames: int + raw_posterior_frames: int + dropped_tail_latents: int + output_latent_frames: int + + +def reference_video_encode_schedule(num_frames: int) -> ReferenceVideoEncodeSchedule: + """Plan exact 17-frame invocations and the one global three-token drop. + + The final reference clip has five real frames and twelve copies of the + final frame. The VideoVAE produces five posterior frames per clip, then + the released wrapper drops three frames once after concatenating *all* + clips (never once per clip). + """ + + snapped = snap_reference_video_frames_down(num_frames) + if snapped > num_frames: + raise ValueError("MiniMax-H3 reference video is too short for the released 17*n+5 schedule") + clip_count = (snapped + VIDEO_FRAMES_PER_CHUNK - 1) // VIDEO_FRAMES_PER_CHUNK + repeated = clip_count * VIDEO_FRAMES_PER_CHUNK - snapped + raw = clip_count * VIDEO_LATENTS_PER_CHUNK + dropped = 3 + output = raw - dropped + if repeated != VIDEO_FRAMES_PER_CHUNK - VIDEO_LATENTS_PER_CHUNK: + raise RuntimeError("MiniMax-H3 reference video clip padding invariant failed") + if output != reference_video_latent_frames(num_frames): + raise RuntimeError("MiniMax-H3 reference video latent schedule mismatch") + return ReferenceVideoEncodeSchedule( + snapped_frames=snapped, + clip_count=clip_count, + repeated_tail_frames=repeated, + raw_posterior_frames=raw, + dropped_tail_latents=dropped, + output_latent_frames=output, + ) + + +def audio_latent_frames(num_samples: int) -> int: + if isinstance(num_samples, bool) or not isinstance(num_samples, int) or num_samples <= 0: + raise ValueError("MiniMax-H3 reference audio needs a positive sample count") + return (num_samples + AUDIO_HOP_LENGTH - 1) // AUDIO_HOP_LENGTH + + +def qwen_video_condition_sample( + num_frames: int, + *, + fps: float = TARGET_FPS, + sample_fps: float = QWEN_VIDEO_SAMPLE_FPS, + temporal_patch: int = QWEN_VISION_TEMPORAL_PATCH_SIZE, +) -> tuple[tuple[int, ...], tuple[float, ...]]: + """Return the sampled frame indices and one timestamp per merged pair.""" + + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError("MiniMax-H3 Qwen video input needs a positive frame count") + if ( + not math.isfinite(fps) + or not math.isfinite(sample_fps) + or fps <= 0 + or sample_fps <= 0 + or sample_fps > fps + ): + raise ValueError("MiniMax-H3 Qwen video sample rates are invalid") + if ( + isinstance(temporal_patch, bool) + or not isinstance(temporal_patch, int) + or temporal_patch <= 0 + ): + raise ValueError("MiniMax-H3 Qwen temporal patch must be a positive integer") + + stride = fps / sample_fps + indices: list[int] = [] + cursor = 0.0 + while round(cursor) < num_frames: + index = round(cursor) + if not indices or index > indices[-1]: + indices.append(index) + cursor += stride + if len(indices) < temporal_patch: + minimum = round((temporal_patch - 1) * stride) + 1 + raise ValueError( + f"MiniMax-H3 Qwen reads a reference video only when it has at least {minimum} frames" + ) + timestamps = [index / sample_fps for index in range(len(indices))] + timestamps.extend([timestamps[-1]] * (-len(timestamps) % temporal_patch)) + blocks = tuple( + (timestamps[index] + timestamps[index + temporal_patch - 1]) / 2.0 + for index in range(0, len(timestamps), temporal_patch) + ) + return tuple(indices), blocks + + +def qwen_patch_grid(height: int, width: int) -> tuple[int, int]: + validate_canvas(height, width) + if height % QWEN_VISION_PATCH_SIZE or width % QWEN_VISION_PATCH_SIZE: + raise ValueError("MiniMax-H3 Qwen vision input is not patch aligned") + return height // QWEN_VISION_PATCH_SIZE, width // QWEN_VISION_PATCH_SIZE + + +def qwen_merged_rows(height: int, width: int) -> int: + grid_h, grid_w = qwen_patch_grid(height, width) + merge = QWEN_VISION_MERGE_SIZE + if grid_h % merge or grid_w % merge: + raise ValueError("MiniMax-H3 Qwen vision input is not merge aligned") + return (grid_h // merge) * (grid_w // merge) + + +@dataclass(frozen=True) +class PresentationPiece: + """One literal text run or one Qwen visual-pad run.""" + + modality: QwenModality + text: str = "" + height: int = 0 + width: int = 0 + + @property + def vision_rows(self) -> int: + return 0 if self.modality == "text" else qwen_merged_rows(self.height, self.width) + + +def ref2va_presentation_blueprint( + prompt: str, + references: Sequence[ReferenceSpec], + *, + normalized_visual_sizes: Sequence[tuple[int, int]], + normalized_video_frames: Sequence[int], +) -> tuple[PresentationPiece, ...]: + """Build the ordered labels/timestamps/vision blocks before tokenization.""" + + if not isinstance(prompt, str): + raise ValueError("MiniMax-H3 prompt must be one string") + validate_reference_request(references) + if len(normalized_visual_sizes) != sum(ref.kind != "audio" for ref in references): + raise ValueError("MiniMax-H3 visual-size metadata count does not match references") + if len(normalized_video_frames) != sum(ref.kind == "video" for ref in references): + raise ValueError("MiniMax-H3 video-frame metadata count does not match references") + visual_sizes = iter(normalized_visual_sizes) + video_frames = iter(normalized_video_frames) + pieces: list[PresentationPiece] = [] + counts = {"image": 0, "video": 0, "audio": 0} + for reference in references: + if reference.has_audio or reference.kind == "audio": + counts["audio"] += 1 + pieces.append(PresentationPiece("text", f"