diff --git a/cmake/FindJetsonPI.cmake b/cmake/FindJetsonPI.cmake new file mode 100644 index 00000000..c005f82a --- /dev/null +++ b/cmake/FindJetsonPI.cmake @@ -0,0 +1,267 @@ +# FindJetsonPI.cmake — locate an INSTALLED Jetson-PI (the llama.cpp/GGML fork) for +# FlashRT's Jetson-PI provider, as the production alternative to the dev +# add_subdirectory(JETSON_PI_ROOT) path (jetsonpi迁移.txt §5.2). +# +# The narrow C API libs (jetson_pi_pi0 / jetson_pi_llm / jetson_pi_mllm) are the +# FlashRT provider's real link surface; they in turn need mtmd / llama / ggml +# (+ either directly linked ggml-* backend libs or runtime backend modules). +# This module finds them all under one install prefix and exposes the linked +# libraries as imported targets so the provider links JetsonPI::jetson_pi_pi0 +# etc. exactly as it links the in-tree targets in the dev path. +# +# Usage: +# find_package(JetsonPI REQUIRED) +# target_link_libraries(my_provider PRIVATE JetsonPI::jetson_pi_pi0) +# +# Required hint variable: +# JetsonPI_ROOT : install prefix (lib or lib64, include) +# +# Defines imported targets: +# JetsonPI::jetson_pi_pi0 JetsonPI::jetson_pi_llm JetsonPI::jetson_pi_mllm +# JetsonPI::mtmd JetsonPI::llama JetsonPI::ggml JetsonPI::ggml-base +# JetsonPI::ggml-cpu JetsonPI::ggml-cuda JetsonPI::ggml-vulkan +# JetsonPI::ggml-opencl JetsonPI::ggml-sycl +# Backend imported targets exist only for a direct-link GGML package. For a +# GGML_BACKEND_DL package, JetsonPI_BACKEND_DIR and JetsonPI_BACKEND_MODULES +# identify the validated runtime modules and none are linked into the provider. +# +# Module-mode search: Jetson-PI does not ship a config-file package for the +# narrow C API libs (installing an install(EXPORT) graph would require +# mtmd/llama/ggml to be in an export set too). Instead this module locates the +# installed libs via find_library and builds JetsonPI::* imported targets with +# INTERFACE_LINK_LIBRARIES to the sibling installed deps — no export graph +# needed, just a coherent install prefix. +# +# --- module-mode search ------------------------------------------------------- + +if(NOT JetsonPI_ROOT) + set(JetsonPI_FOUND FALSE) + set(JetsonPI_NOT_FOUND_MESSAGE + "JetsonPI_ROOT must name the installed Jetson-PI prefix") + return() +endif() + +set(_JetsonPI_HINT_PATHS ${JetsonPI_ROOT}) + +unset(ggml_DIR CACHE) +find_package(ggml CONFIG QUIET + PATHS ${_JetsonPI_HINT_PATHS} + NO_DEFAULT_PATH) +if(NOT ggml_FOUND) + message(FATAL_ERROR + "JetsonPI_ROOT='${JetsonPI_ROOT}' has no matching ggml-config.cmake. " + "Install Jetson-PI and its bundled GGML into the same prefix.") +endif() +set(JetsonPI_BACKEND_DL ${GGML_BACKEND_DL}) + +# NO_DEFAULT_PATH: same no-system-mix rule as the find_library path below. +# Without it, a stray system header could paper over a bad JetsonPI_ROOT. +unset(JetsonPI_INCLUDE_DIR CACHE) +find_path(JetsonPI_INCLUDE_DIR + NAMES jetson_pi_pi0.h + HINTS ${_JetsonPI_HINT_PATHS} + PATH_SUFFIXES include + DOC "Jetson-PI narrow C API headers" + NO_DEFAULT_PATH) + +# Each lib we need. find_library per-target so a missing backend lib (e.g. +# ggml-vulkan when GGML_VULKAN=OFF) doesn't fail the whole find. We search ONLY +# the hint paths (NO_DEFAULT_PATH): intentionally NOT falling back to system +# default paths. A system /usr/lib libllama.so / libggml.so on a deploy host +# with a distro llama.cpp package would silently mix a different GGML version +# into the narrow libs — exactly the §15.1 "do not mix another llama.cpp/GGML +# version" hazard. A failed find should fail loudly via the FAIL_MESSAGE below, +# not paper over a bad/missing prefix with a stray system lib. +function(_jetsonpi_find_lib var name) + unset(${var} CACHE) + find_library(${var} + NAMES ${name} + HINTS ${_JetsonPI_HINT_PATHS} + PATH_SUFFIXES lib lib64 + DOC "Jetson-PI ${name} library" + NO_DEFAULT_PATH) +endfunction() + +_jetsonpi_find_lib(JetsonPI_pi0_LIBRARY jetson_pi_pi0) +_jetsonpi_find_lib(JetsonPI_llm_LIBRARY jetson_pi_llm) +_jetsonpi_find_lib(JetsonPI_mllm_LIBRARY jetson_pi_mllm) +_jetsonpi_find_lib(JetsonPI_mtmd_LIBRARY mtmd) +_jetsonpi_find_lib(JetsonPI_llama_LIBRARY llama) +_jetsonpi_find_lib(JetsonPI_ggml_LIBRARY ggml) +_jetsonpi_find_lib(JetsonPI_ggml_base_LIBRARY ggml-base) +_jetsonpi_find_lib(JetsonPI_onemath_cublas_LIBRARY onemath_blas_cublas) + +set(JetsonPI_BACKEND_MODULES "") +if(JetsonPI_BACKEND_DL) + if(NOT GGML_BACKEND_DIR OR NOT IS_ABSOLUTE "${GGML_BACKEND_DIR}") + message(FATAL_ERROR + "The Jetson-PI GGML package uses GGML_BACKEND_DL, but its installed " + "ggml-config.cmake does not contain an absolute GGML_BACKEND_DIR. " + "Rebuild Jetson-PI with -DGGML_BACKEND_DIR=/bin so " + "runtime backend discovery is explicit and independent of cwd.") + endif() + get_filename_component(_JetsonPI_root_real "${JetsonPI_ROOT}" REALPATH) + get_filename_component(JetsonPI_BACKEND_DIR "${GGML_BACKEND_DIR}" REALPATH) + string(FIND "${JetsonPI_BACKEND_DIR}/" "${_JetsonPI_root_real}/" + _JetsonPI_backend_prefix_index) + if(NOT _JetsonPI_backend_prefix_index EQUAL 0) + message(FATAL_ERROR + "GGML_BACKEND_DIR='${GGML_BACKEND_DIR}' is outside " + "JetsonPI_ROOT='${JetsonPI_ROOT}'. Install core libraries and runtime " + "backend modules as one coherent deployment package.") + endif() + + foreach(_JetsonPI_backend IN LISTS GGML_AVAILABLE_BACKENDS) + string(REPLACE "-" "_" _JetsonPI_backend_id "${_JetsonPI_backend}") + set(_JetsonPI_backend_var "JetsonPI_${_JetsonPI_backend_id}_MODULE") + unset(${_JetsonPI_backend_var} CACHE) + find_file(${_JetsonPI_backend_var} + NAMES "${CMAKE_SHARED_MODULE_PREFIX}${_JetsonPI_backend}${CMAKE_SHARED_MODULE_SUFFIX}" + HINTS "${JetsonPI_BACKEND_DIR}" + DOC "Jetson-PI ${_JetsonPI_backend} runtime module" + NO_DEFAULT_PATH) + if(NOT ${_JetsonPI_backend_var}) + message(FATAL_ERROR + "Jetson-PI's ggml-config.cmake declares backend " + "'${_JetsonPI_backend}', but its runtime module is missing from " + "GGML_BACKEND_DIR='${JetsonPI_BACKEND_DIR}'.") + endif() + list(APPEND JetsonPI_BACKEND_MODULES "${${_JetsonPI_backend_var}}") + if(_JetsonPI_backend STREQUAL "ggml-cpu") + set(JetsonPI_ggml_cpu_LIBRARY "${${_JetsonPI_backend_var}}") + elseif(_JetsonPI_backend STREQUAL "ggml-cuda") + set(JetsonPI_ggml_cuda_LIBRARY "${${_JetsonPI_backend_var}}") + elseif(_JetsonPI_backend STREQUAL "ggml-vulkan") + set(JetsonPI_ggml_vulkan_LIBRARY "${${_JetsonPI_backend_var}}") + elseif(_JetsonPI_backend STREQUAL "ggml-opencl") + set(JetsonPI_ggml_opencl_LIBRARY "${${_JetsonPI_backend_var}}") + elseif(_JetsonPI_backend STREQUAL "ggml-sycl") + set(JetsonPI_ggml_sycl_LIBRARY "${${_JetsonPI_backend_var}}") + endif() + endforeach() +else() + set(JetsonPI_BACKEND_DIR "") + _jetsonpi_find_lib(JetsonPI_ggml_cpu_LIBRARY ggml-cpu) + _jetsonpi_find_lib(JetsonPI_ggml_cuda_LIBRARY ggml-cuda) + _jetsonpi_find_lib(JetsonPI_ggml_vulkan_LIBRARY ggml-vulkan) + _jetsonpi_find_lib(JetsonPI_ggml_opencl_LIBRARY ggml-opencl) + _jetsonpi_find_lib(JetsonPI_ggml_sycl_LIBRARY ggml-sycl) +endif() + +if (JetsonPI_ggml_opencl_LIBRARY AND NOT JetsonPI_BACKEND_DL) + find_package(OpenCL REQUIRED) +endif() +if (JetsonPI_ggml_sycl_LIBRARY AND NOT JetsonPI_BACKEND_DL) + find_library(JetsonPI_sycl_LIBRARY NAMES sycl HINTS ENV ONEAPI_ROOT PATH_SUFFIXES lib lib64) + find_library(JetsonPI_imf_LIBRARY NAMES imf HINTS ENV ONEAPI_ROOT PATH_SUFFIXES lib lib64) + find_library(JetsonPI_svml_LIBRARY NAMES svml HINTS ENV ONEAPI_ROOT PATH_SUFFIXES lib lib64) + find_library(JetsonPI_intlc_LIBRARY NAMES intlc intlc.so.5 HINTS ENV ONEAPI_ROOT PATH_SUFFIXES lib lib64) + if (NOT JetsonPI_sycl_LIBRARY OR NOT JetsonPI_imf_LIBRARY OR + NOT JetsonPI_svml_LIBRARY OR NOT JetsonPI_intlc_LIBRARY OR + NOT JetsonPI_onemath_cublas_LIBRARY) + message(FATAL_ERROR "JetsonPI ggml-sycl was found, but its DPC++/oneMath runtime libraries were not. Set ONEAPI_ROOT and install libsycl, libimf, libsvml, libintlc, and libonemath_blas_cublas.") + endif() + find_package(CUDAToolkit REQUIRED) +endif() + +include(FindPackageHandleStandardArgs) +# The provider hard-needs the narrow libs + their core transitive deps. Backend +# ggml-* libs are optional (link whichever were built). +find_package_handle_standard_args(JetsonPI + REQUIRED_VARS + JetsonPI_INCLUDE_DIR + JetsonPI_pi0_LIBRARY + JetsonPI_llm_LIBRARY + JetsonPI_mllm_LIBRARY + JetsonPI_mtmd_LIBRARY + JetsonPI_llama_LIBRARY + JetsonPI_ggml_LIBRARY + JetsonPI_ggml_base_LIBRARY + JetsonPI_ggml_cpu_LIBRARY + FAIL_MESSAGE "Jetson-PI not found. Set -DJetsonPI_ROOT= (a prefix containing include/jetson_pi_pi0.h and lib{,64}/libjetson_pi_pi0.so etc.), or use the dev path -DJETSON_PI_ROOT= with add_subdirectory.") + +if (JetsonPI_FOUND) + # Aggregate the ggml backend libs that actually resolved (cuda/vulkan optional). + # Target names use hyphens to match the soname (libggml-cuda.so) and the + # header comment above (JetsonPI::ggml-cuda / ::ggml-vulkan). + set(_JetsonPI_ggml_backends "") + if (JetsonPI_ggml_cuda_LIBRARY AND NOT JetsonPI_BACKEND_DL) + find_package(CUDAToolkit REQUIRED) + list(APPEND _JetsonPI_ggml_backends JetsonPI::ggml-cuda) + endif() + if (JetsonPI_ggml_vulkan_LIBRARY AND NOT JetsonPI_BACKEND_DL) + list(APPEND _JetsonPI_ggml_backends JetsonPI::ggml-vulkan) + endif() + if (JetsonPI_ggml_opencl_LIBRARY AND NOT JetsonPI_BACKEND_DL) + list(APPEND _JetsonPI_ggml_backends JetsonPI::ggml-opencl) + endif() + if (JetsonPI_ggml_sycl_LIBRARY AND NOT JetsonPI_BACKEND_DL) + list(APPEND _JetsonPI_ggml_backends JetsonPI::ggml-sycl) + endif() + + # Helper: declare one imported shared lib target with its transitive deps. + function(_jetsonpi_import_target tgt lib dep) + if (NOT ${lib}) + return() + endif() + if (TARGET ${tgt}) + return() + endif() + add_library(${tgt} SHARED IMPORTED) + set_target_properties(${tgt} PROPERTIES + IMPORTED_LOCATION "${${lib}}" + INTERFACE_INCLUDE_DIRECTORIES "${JetsonPI_INCLUDE_DIR}") + if (dep) + set_target_properties(${tgt} PROPERTIES INTERFACE_LINK_LIBRARIES "${dep}") + endif() + endfunction() + + set(_ggml_deps "JetsonPI::ggml-base") + _jetsonpi_import_target(JetsonPI::ggml-base JetsonPI_ggml_base_LIBRARY "") + if(NOT JetsonPI_BACKEND_DL) + list(APPEND _ggml_deps JetsonPI::ggml-cpu ${_JetsonPI_ggml_backends}) + _jetsonpi_import_target(JetsonPI::ggml-cpu JetsonPI_ggml_cpu_LIBRARY "JetsonPI::ggml-base") + _jetsonpi_import_target(JetsonPI::ggml-cuda JetsonPI_ggml_cuda_LIBRARY "JetsonPI::ggml-base;CUDA::cudart;CUDA::cublas;CUDA::cuda_driver") + _jetsonpi_import_target(JetsonPI::ggml-vulkan JetsonPI_ggml_vulkan_LIBRARY "JetsonPI::ggml-base") + _jetsonpi_import_target(JetsonPI::ggml-opencl JetsonPI_ggml_opencl_LIBRARY "JetsonPI::ggml-base;OpenCL::OpenCL") + _jetsonpi_import_target(JetsonPI::ggml-sycl JetsonPI_ggml_sycl_LIBRARY "JetsonPI::ggml-base;${JetsonPI_onemath_cublas_LIBRARY};${JetsonPI_sycl_LIBRARY};${JetsonPI_imf_LIBRARY};${JetsonPI_svml_LIBRARY};${JetsonPI_intlc_LIBRARY};CUDA::cublas;CUDA::cuda_driver") + endif() + _jetsonpi_import_target(JetsonPI::ggml JetsonPI_ggml_LIBRARY "${_ggml_deps}") + _jetsonpi_import_target(JetsonPI::llama JetsonPI_llama_LIBRARY "JetsonPI::ggml") + _jetsonpi_import_target(JetsonPI::mtmd JetsonPI_mtmd_LIBRARY "JetsonPI::llama;JetsonPI::ggml") + _jetsonpi_import_target(JetsonPI::jetson_pi_pi0 JetsonPI_pi0_LIBRARY "JetsonPI::mtmd;JetsonPI::llama;JetsonPI::ggml") + _jetsonpi_import_target(JetsonPI::jetson_pi_llm JetsonPI_llm_LIBRARY "JetsonPI::llama;JetsonPI::ggml") + _jetsonpi_import_target(JetsonPI::jetson_pi_mllm JetsonPI_mllm_LIBRARY "JetsonPI::mtmd;JetsonPI::llama;JetsonPI::ggml") + + set(JetsonPI_LIBRARY_DIRS "") + foreach(_JetsonPI_library + JetsonPI_pi0_LIBRARY JetsonPI_llm_LIBRARY JetsonPI_mllm_LIBRARY + JetsonPI_mtmd_LIBRARY JetsonPI_llama_LIBRARY JetsonPI_ggml_LIBRARY + JetsonPI_ggml_base_LIBRARY JetsonPI_ggml_cpu_LIBRARY + JetsonPI_ggml_cuda_LIBRARY JetsonPI_ggml_vulkan_LIBRARY + JetsonPI_ggml_opencl_LIBRARY JetsonPI_ggml_sycl_LIBRARY + JetsonPI_onemath_cublas_LIBRARY JetsonPI_sycl_LIBRARY + JetsonPI_imf_LIBRARY JetsonPI_svml_LIBRARY JetsonPI_intlc_LIBRARY) + if(${_JetsonPI_library}) + get_filename_component(_JetsonPI_library_dir "${${_JetsonPI_library}}" DIRECTORY) + list(APPEND JetsonPI_LIBRARY_DIRS "${_JetsonPI_library_dir}") + endif() + endforeach() + list(REMOVE_DUPLICATES JetsonPI_LIBRARY_DIRS) + + if(JetsonPI_BACKEND_DL) + message(STATUS + "Jetson-PI found (module mode, dynamic GGML backends in ${JetsonPI_BACKEND_DIR}): ${JetsonPI_INCLUDE_DIR}") + else() + message(STATUS "Jetson-PI found (module mode): ${JetsonPI_INCLUDE_DIR}") + endif() +endif() + +mark_as_advanced(JetsonPI_INCLUDE_DIR JetsonPI_BACKEND_DIR + JetsonPI_pi0_LIBRARY JetsonPI_llm_LIBRARY JetsonPI_mllm_LIBRARY + JetsonPI_mtmd_LIBRARY JetsonPI_llama_LIBRARY + JetsonPI_ggml_LIBRARY JetsonPI_ggml_base_LIBRARY + JetsonPI_ggml_cpu_LIBRARY JetsonPI_ggml_cuda_LIBRARY JetsonPI_ggml_vulkan_LIBRARY + JetsonPI_ggml_opencl_LIBRARY JetsonPI_ggml_sycl_LIBRARY + JetsonPI_onemath_cublas_LIBRARY JetsonPI_sycl_LIBRARY + JetsonPI_imf_LIBRARY JetsonPI_svml_LIBRARY JetsonPI_intlc_LIBRARY) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a4ed087d..930423a2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -33,6 +33,14 @@ option(FLASHRT_CPP_WITH_PI05_SM120_TARGET "Build the PI0.5 SM120 frontend target" OFF) option(FLASHRT_CPP_WITH_PI05_SM110_TARGET "Build the PI0.5 SM110 frontend target" OFF) +option(FLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER + "Build the optional llama.cpp provider frontend" OFF) +option(FLASHRT_CPP_WITH_JETSON_PI + "Back the llama.cpp provider with Jetson-PI" OFF) +set(JETSON_PI_ROOT "" CACHE PATH + "Jetson-PI source root for an in-tree provider build") +set(JetsonPI_ROOT "" CACHE PATH + "Jetson-PI install prefix for a packaged provider build") if((FLASHRT_CPP_WITH_PI05 OR FLASHRT_CPP_WITH_SENTENCEPIECE) AND NOT FLASHRT_ENABLE_NATIVE_CPP) message(FATAL_ERROR @@ -45,6 +53,12 @@ if((FLASHRT_CPP_WITH_PI05_SM120_TARGET OR message(FATAL_ERROR "PI0.5 hardware targets require FLASHRT_CPP_WITH_PI05=ON") endif() +if(FLASHRT_CPP_WITH_JETSON_PI AND + NOT FLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_JETSON_PI requires " + "FLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER=ON") +endif() if(FLASHRT_ENABLE_NATIVE_CPP AND NOT FLASHRT_CPP_WITH_EXEC) message(FATAL_ERROR "FLASHRT_ENABLE_NATIVE_CPP requires FLASHRT_CPP_WITH_EXEC=ON") @@ -72,6 +86,46 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) endif() +if(FLASHRT_CPP_WITH_JETSON_PI) + if(JETSON_PI_ROOT AND JetsonPI_ROOT) + message(FATAL_ERROR + "Set only one Jetson-PI integration root: JETSON_PI_ROOT for a " + "source build or JetsonPI_ROOT for an installed package") + endif() + if(NOT JETSON_PI_ROOT AND NOT JetsonPI_ROOT) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_JETSON_PI requires JETSON_PI_ROOT or JetsonPI_ROOT") + endif() + + if(JETSON_PI_ROOT) + if(NOT EXISTS "${JETSON_PI_ROOT}/CMakeLists.txt") + message(FATAL_ERROR + "JETSON_PI_ROOT does not identify a Jetson-PI source tree") + endif() + set(LLAMA_BUILD_COMMON ON CACHE BOOL "" FORCE) + set(LLAMA_CURL OFF CACHE BOOL "" FORCE) + set(LLAMA_HTTPLIB OFF CACHE BOOL "" FORCE) + set(LLAMA_BUILD_TESTS OFF CACHE BOOL "" FORCE) + set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "" FORCE) + set(JETSON_PI_BUILD_MTMD ON CACHE BOOL "" FORCE) + set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) + set(LLAMA_BUILD_SERVER OFF CACHE BOOL "" FORCE) + set(LLAMA_TOOLS_INSTALL OFF CACHE BOOL "" FORCE) + add_subdirectory("${JETSON_PI_ROOT}" + "${CMAKE_CURRENT_BINARY_DIR}/jetson-pi" + EXCLUDE_FROM_ALL) + set(_flashrt_jetson_pi_pi0 jetson_pi_pi0) + set(_flashrt_jetson_pi_llm jetson_pi_llm) + set(_flashrt_jetson_pi_mllm jetson_pi_mllm) + else() + list(PREPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../cmake) + find_package(JetsonPI REQUIRED) + list(POP_FRONT CMAKE_MODULE_PATH) + set(_flashrt_jetson_pi_pi0 JetsonPI::jetson_pi_pi0) + set(_flashrt_jetson_pi_llm JetsonPI::jetson_pi_llm) + set(_flashrt_jetson_pi_mllm JetsonPI::jetson_pi_mllm) + endif() +endif() if(FLASHRT_CPP_WITH_PI05_SM120_TARGET) if(NOT FLASHRT_CPP_WITH_CUDA_KERNELS) message(FATAL_ERROR @@ -227,6 +281,62 @@ target_include_directories(flashrt_cpp_vla target_link_libraries(flashrt_cpp_vla INTERFACE flashrt_cpp_modalities flashrt_cpp) +if(FLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER) + add_library(flashrt_cpp_llama_cpp_provider STATIC + providers/llama_cpp/src/checkpoint_identity.cpp + providers/llama_cpp/src/pi0_runtime.cpp + providers/llama_cpp/src/llm_runtime.cpp + providers/llama_cpp/src/mllm_runtime.cpp) + target_include_directories(flashrt_cpp_llama_cpp_provider + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/providers/llama_cpp/include) + target_link_libraries(flashrt_cpp_llama_cpp_provider + PUBLIC flashrt_runtime) + + if(FLASHRT_CPP_WITH_JETSON_PI) + target_sources(flashrt_cpp_llama_cpp_provider PRIVATE + providers/llama_cpp/src/jetson_pi_engine.cpp + providers/llama_cpp/src/jetson_pi_link_check.cpp) + target_compile_definitions(flashrt_cpp_llama_cpp_provider PRIVATE + FLASHRT_CPP_WITH_JETSON_PI=1) + target_link_libraries(flashrt_cpp_llama_cpp_provider PRIVATE + ${_flashrt_jetson_pi_pi0} + ${_flashrt_jetson_pi_llm} + ${_flashrt_jetson_pi_mllm}) + + add_library(flashrt_cpp_llama_cpp_provider_c SHARED + providers/llama_cpp/src/checkpoint_identity.cpp + providers/llama_cpp/src/pi0_runtime.cpp + providers/llama_cpp/src/llm_runtime.cpp + providers/llama_cpp/src/mllm_runtime.cpp + providers/llama_cpp/src/jetson_pi_engine.cpp + providers/llama_cpp/src/provider_open.cpp) + target_include_directories(flashrt_cpp_llama_cpp_provider_c PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/providers/llama_cpp/include) + target_compile_definitions(flashrt_cpp_llama_cpp_provider_c PRIVATE + FLASHRT_CPP_WITH_JETSON_PI=1) + target_compile_options(flashrt_cpp_llama_cpp_provider_c PRIVATE + $<$:-fvisibility=hidden>) + target_link_libraries(flashrt_cpp_llama_cpp_provider_c PRIVATE + flashrt_runtime + ${_flashrt_jetson_pi_pi0} + ${_flashrt_jetson_pi_llm} + ${_flashrt_jetson_pi_mllm}) + if(UNIX AND NOT APPLE) + target_link_options(flashrt_cpp_llama_cpp_provider_c PRIVATE + "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/providers/llama_cpp/provider_exports.map" + "LINKER:--disable-new-dtags") + endif() + set_target_properties(flashrt_cpp_llama_cpp_provider_c PROPERTIES + INSTALL_RPATH "$ORIGIN") + install(TARGETS flashrt_cpp_llama_cpp_provider_c flashrt_runtime + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}) + install(FILES + ${CMAKE_CURRENT_SOURCE_DIR}/../runtime/include/flashrt/model_runtime.h + ${CMAKE_CURRENT_SOURCE_DIR}/../runtime/include/flashrt/runtime.h + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/flashrt) + endif() +endif() + if(FLASHRT_ENABLE_NATIVE_CPP) add_library(flashrt_cpp_native STATIC native/src/config_object.cpp) @@ -532,6 +642,90 @@ if(FLASHRT_CPP_WITH_PI05) endif() if(BUILD_TESTING) + if(FLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER) + add_executable(test_llama_cpp_provider + tests/test_llama_cpp_provider.cpp) + target_link_libraries(test_llama_cpp_provider + PRIVATE flashrt_cpp_llama_cpp_provider flashrt_runtime) + add_test(NAME llama_cpp_provider COMMAND test_llama_cpp_provider) + + if(FLASHRT_CPP_WITH_JETSON_PI) + add_executable(test_llama_cpp_jetson_pi_link + tests/test_llama_cpp_jetson_pi_link.cpp) + target_link_libraries(test_llama_cpp_jetson_pi_link + PRIVATE flashrt_cpp_llama_cpp_provider) + add_test(NAME llama_cpp_jetson_pi_link + COMMAND test_llama_cpp_jetson_pi_link) + + if(UNIX AND NOT APPLE) + add_test(NAME llama_cpp_provider_exports + COMMAND ${CMAKE_COMMAND} + -DPROVIDER_DSO=$ + -DCMAKE_NM_TOOL=${CMAKE_NM} + -P ${CMAKE_CURRENT_SOURCE_DIR}/tests/check_llama_cpp_provider_exports.cmake) + endif() + + add_executable(test_llama_cpp_jetson_pi_llm + tests/test_llama_cpp_jetson_pi_llm.cpp) + target_link_libraries(test_llama_cpp_jetson_pi_llm + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_llm}) + add_test(NAME llama_cpp_jetson_pi_llm + COMMAND test_llama_cpp_jetson_pi_llm) + + add_executable(test_llama_cpp_jetson_pi_mllm + tests/test_llama_cpp_jetson_pi_mllm.cpp) + target_link_libraries(test_llama_cpp_jetson_pi_mllm + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_mllm}) + add_test(NAME llama_cpp_jetson_pi_mllm + COMMAND test_llama_cpp_jetson_pi_mllm) + + add_executable(test_llama_cpp_jetson_pi_llm_parity + tests/test_llama_cpp_jetson_pi_llm_parity.cpp) + target_link_libraries(test_llama_cpp_jetson_pi_llm_parity + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_llm}) + add_test(NAME llama_cpp_jetson_pi_llm_parity + COMMAND test_llama_cpp_jetson_pi_llm_parity) + + add_executable(test_llama_cpp_jetson_pi_mllm_parity + tests/test_llama_cpp_jetson_pi_mllm_parity.cpp) + target_link_libraries(test_llama_cpp_jetson_pi_mllm_parity + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_mllm}) + add_test(NAME llama_cpp_jetson_pi_mllm_parity + COMMAND test_llama_cpp_jetson_pi_mllm_parity) + + if(JETSON_PI_ROOT AND EXISTS "${JETSON_PI_ROOT}/vendor") + add_executable(test_llama_cpp_jetson_pi_engine + tests/test_llama_cpp_jetson_pi_engine.cpp) + target_include_directories(test_llama_cpp_jetson_pi_engine PRIVATE + ${JETSON_PI_ROOT}/vendor) + target_link_libraries(test_llama_cpp_jetson_pi_engine + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_pi0}) + add_test(NAME llama_cpp_jetson_pi_engine + COMMAND test_llama_cpp_jetson_pi_engine) + + add_executable(test_llama_cpp_jetson_pi_parity + tests/test_llama_cpp_jetson_pi_parity.cpp) + target_include_directories(test_llama_cpp_jetson_pi_parity PRIVATE + ${JETSON_PI_ROOT}/vendor) + target_link_libraries(test_llama_cpp_jetson_pi_parity + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_pi0}) + add_test(NAME llama_cpp_jetson_pi_parity + COMMAND test_llama_cpp_jetson_pi_parity) + + # Opt-in PI0.5 reference-policy state-parity test. CI skips it when the + # external checkpoint and fixture environment are unavailable. + add_executable(test_llama_cpp_jetson_pi_state_parity + tests/test_llama_cpp_jetson_pi_state_parity.cpp) + target_include_directories(test_llama_cpp_jetson_pi_state_parity PRIVATE + ${JETSON_PI_ROOT}/vendor) + target_link_libraries(test_llama_cpp_jetson_pi_state_parity + PRIVATE flashrt_cpp_llama_cpp_provider ${_flashrt_jetson_pi_pi0}) + add_test(NAME llama_cpp_jetson_pi_state_parity + COMMAND test_llama_cpp_jetson_pi_state_parity) + endif() + endif() + endif() + if(TARGET flashrt_native_cpp_cuda_support) add_executable(test_csrc_rope_table tests/test_csrc_rope_table.cu) target_link_libraries(test_csrc_rope_table PRIVATE diff --git a/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/c_api.h b/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/c_api.h new file mode 100644 index 00000000..18aefa74 --- /dev/null +++ b/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/c_api.h @@ -0,0 +1,248 @@ +#ifndef FLASHRT_PROVIDERS_LLAMA_CPP_C_API_H +#define FLASHRT_PROVIDERS_LLAMA_CPP_C_API_H + +#include "flashrt/model_runtime.h" + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +enum frt_llama_cpp_model_family { + FRT_LLAMA_CPP_MODEL_PI0 = 1, + FRT_LLAMA_CPP_MODEL_LLM = 2, + FRT_LLAMA_CPP_MODEL_MLLM = 3, +}; + +enum frt_llama_cpp_pi0_port { + FRT_LLAMA_CPP_PI0_PORT_IMAGES = 0, + FRT_LLAMA_CPP_PI0_PORT_PROMPT = 1, + FRT_LLAMA_CPP_PI0_PORT_STATE = 2, + FRT_LLAMA_CPP_PI0_PORT_ACTIONS = 3, +}; + +enum frt_llama_cpp_pi0_stage_index { + FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER = 0, + FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT = 1, + FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION = 2, +}; + +enum frt_llama_cpp_llm_port { + FRT_LLAMA_CPP_LLM_PORT_PROMPT = 0, + FRT_LLAMA_CPP_LLM_PORT_TEXT = 1, + FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN = 2, + FRT_LLAMA_CPP_LLM_PORT_LOGITS = 3, + FRT_LLAMA_CPP_LLM_PORT_IS_EOG = 4, + FRT_LLAMA_CPP_LLM_PORT_TOKENS = 5, +}; + +enum frt_llama_cpp_llm_stage_index { + FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER = 0, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET = 1, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL = 2, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE = 3, +}; + +enum frt_llama_cpp_mllm_port { + FRT_LLAMA_CPP_MLLM_PORT_IMAGES = 0, + FRT_LLAMA_CPP_MLLM_PORT_PROMPT = 1, + FRT_LLAMA_CPP_MLLM_PORT_TEXT = 2, + FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN = 3, + FRT_LLAMA_CPP_MLLM_PORT_LOGITS = 4, + FRT_LLAMA_CPP_MLLM_PORT_IS_EOG = 5, +}; + +enum frt_llama_cpp_mllm_stage_index { + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER = 0, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_RESET = 1, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL = 2, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE = 3, +}; + +typedef struct frt_llama_cpp_pi0_config { + uint32_t struct_size; + + const char* model_path; + const char* mmproj_path; + const char* backend; + + uint32_t n_views; + uint32_t image_height; + uint32_t image_width; + uint32_t image_channels; + uint32_t action_steps; + uint32_t action_dim; + + const char* model_identity; /* append-only checkpoint content identity */ + const char* mmproj_identity; /* append-only checkpoint content identity */ +} frt_llama_cpp_pi0_config; + +#define FRT_LLAMA_CPP_PI0_CONFIG_BASE_SIZE \ + (offsetof(frt_llama_cpp_pi0_config, model_identity)) + +typedef struct frt_llama_cpp_llm_config { + uint32_t struct_size; + + const char* model_path; + const char* backend; + + uint32_t n_ctx; /* KV context size; 0 = from model */ + int32_t n_threads; /* CPU threads; 0 = hardware_concurrency */ + + float temp; /* sampler temperature (<=0 = greedy) */ + int32_t top_k; /* 0 = disabled */ + float top_p; /* 0 = disabled */ + uint32_t seed; /* RNG seed */ + uint32_t max_tokens; /* cap on generated tokens per infer */ + + const char* model_identity; /* append-only checkpoint content identity */ +} frt_llama_cpp_llm_config; + +#define FRT_LLAMA_CPP_LLM_CONFIG_BASE_SIZE \ + (offsetof(frt_llama_cpp_llm_config, model_identity)) + +typedef struct frt_llama_cpp_mllm_config { + uint32_t struct_size; + + const char* model_path; + const char* mmproj_path; + const char* backend; + + uint32_t n_ctx; /* KV context size; 0 = 4096 fallback */ + int32_t n_threads; /* CPU threads; 0 = hardware_concurrency */ + + float temp; /* sampler temperature (<=0 = greedy) */ + int32_t top_k; /* 0 = disabled */ + float top_p; /* 0 = disabled */ + uint32_t seed; /* RNG seed */ + uint32_t max_tokens; /* cap on generated tokens per infer */ + + const char* model_identity; /* append-only checkpoint content identity */ + const char* mmproj_identity; /* append-only checkpoint content identity */ +} frt_llama_cpp_mllm_config; + +#define FRT_LLAMA_CPP_MLLM_CONFIG_BASE_SIZE \ + (offsetof(frt_llama_cpp_mllm_config, model_identity)) + +typedef struct frt_llama_cpp_engine_v1 { + uint32_t struct_size; + /* Capability bits. Older engines initialize this reserved word to zero. */ + uint32_t reserved; + + void* self; + /* retain/release must be both null for a borrowed engine, or both set for + * a reference-counted engine. Asymmetric ownership hooks are rejected. */ + void (*retain)(void* self); + void (*release)(void* self); + + int (*set_input)(void* self, uint32_t port, + const void* data, uint64_t bytes, int stream); + int (*run_infer)(void* self); + int (*get_output)(void* self, uint32_t port, + void* out, uint64_t capacity, uint64_t* written, + int stream); + /* Should return a non-null borrowed string. If it violates that contract, + * the wrapper reports a stable boundary error string. */ + const char* (*last_error)(void* self); + + /* Optional append-only tail. Engines exposing this callback support the + * provider's model-family-specific stage indices (e.g. LLM reset/prefill/ + * decode). Older engines end at last_error and remain valid. */ + int (*run_stage)(void* self, uint32_t stage); +} frt_llama_cpp_engine_v1; + +#define FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE \ + (offsetof(frt_llama_cpp_engine_v1, run_stage)) +#define FRT_LLAMA_CPP_ENGINE_V1_RUN_STAGE_SIZE \ + (offsetof(frt_llama_cpp_engine_v1, run_stage) + \ + sizeof(((frt_llama_cpp_engine_v1*)0)->run_stage)) + +/* Stored in frt_llama_cpp_engine_v1::reserved. A callback alone is + * insufficient to distinguish a real boundary from cached-result shims. */ +#define FRT_LLAMA_CPP_ENGINE_CAP_PI0_REAL_CONTEXT_ACTION UINT32_C(1) + +typedef struct frt_llama_cpp_engine_factory_v1 { + uint32_t struct_size; + uint32_t reserved; + + void* self; + /* Returns one owned engine reference in out_engine. The config pointer + * and its string fields are borrowed and valid only for the duration of + * the callback; factories must copy anything they keep. The runtime-open + * wrapper consumes the returned reference after retaining the model + * runtime's own engine reference. Factory-created engines must provide + * symmetric retain/release hooks; borrowed engines are only accepted by + * the lower create_with_engine entry point. On nonzero return, no engine + * ownership is transferred and out_engine is ignored. */ + int (*create_pi0)(void* self, const frt_llama_cpp_pi0_config* config, + frt_llama_cpp_engine_v1* out_engine); + /* Same contract as create_pi0 but for a generic GGUF LLM (text in -> + * text out). May be NULL if the factory only serves Pi0. */ + int (*create_llm)(void* self, const frt_llama_cpp_llm_config* config, + frt_llama_cpp_engine_v1* out_engine); + /* Same contract as create_pi0 but for a multimodal LLM (images + text in + * -> text out). May be NULL if the factory does not serve VLMs. */ + int (*create_mllm)(void* self, const frt_llama_cpp_mllm_config* config, + frt_llama_cpp_engine_v1* out_engine); + const char* (*last_error)(void* self); +} frt_llama_cpp_engine_factory_v1; + +/* Description of the most recent runtime-open failure on the calling thread. + * Always returns a non-null borrowed string, valid until the next + * frt_llama_cpp_*_runtime_open_with_engine_factory call on that thread. */ +const char* frt_llama_cpp_runtime_open_error(void); + +int frt_llama_cpp_pi0_runtime_create_with_engine( + const frt_llama_cpp_pi0_config* config, + const frt_llama_cpp_engine_v1* engine, + frt_model_runtime_v1** out); + +/* Provider-specific JSON open path for the current dependency-injection + * boundary. Required JSON fields: + * model_family="pi0", model_path, mmproj_path, backend, + * n_views, image_height, image_width, image_channels, + * action_steps, action_dim. + * No field has a default; missing or mismatched fields fail hard. */ +int frt_llama_cpp_pi0_runtime_open_with_engine_factory( + const char* config_json, + const frt_llama_cpp_engine_factory_v1* factory, + frt_model_runtime_v1** out); + +int frt_llama_cpp_llm_runtime_create_with_engine( + const frt_llama_cpp_llm_config* config, + const frt_llama_cpp_engine_v1* engine, + frt_model_runtime_v1** out); + +/* Provider-specific JSON open path for generic GGUF LLM. Required JSON + * fields: + * model_family="llm", model_path, backend, + * n_ctx, n_threads, temp, top_k, top_p, seed, max_tokens. + * No field has a default; missing or mismatched fields fail hard. The factory + * must provide create_llm. */ +int frt_llama_cpp_llm_runtime_open_with_engine_factory( + const char* config_json, + const frt_llama_cpp_engine_factory_v1* factory, + frt_model_runtime_v1** out); + +int frt_llama_cpp_mllm_runtime_create_with_engine( + const frt_llama_cpp_mllm_config* config, + const frt_llama_cpp_engine_v1* engine, + frt_model_runtime_v1** out); + +/* Provider-specific JSON open path for multimodal LLM. Required JSON fields: + * model_family="mllm", model_path, mmproj_path, backend, + * n_ctx, n_threads, temp, top_k, top_p, seed, max_tokens. + * No field has a default; missing or mismatched fields fail hard. The factory + * must provide create_mllm. */ +int frt_llama_cpp_mllm_runtime_open_with_engine_factory( + const char* config_json, + const frt_llama_cpp_engine_factory_v1* factory, + frt_model_runtime_v1** out); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* FLASHRT_PROVIDERS_LLAMA_CPP_C_API_H */ diff --git a/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/jetson_pi_engine.h b/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/jetson_pi_engine.h new file mode 100644 index 00000000..113926ee --- /dev/null +++ b/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/jetson_pi_engine.h @@ -0,0 +1,40 @@ +#ifndef FLASHRT_PROVIDERS_LLAMA_CPP_JETSON_PI_ENGINE_H +#define FLASHRT_PROVIDERS_LLAMA_CPP_JETSON_PI_ENGINE_H + +// +// Default Jetson-PI Pi0 engine factory. +// +// Returns a borrowed pointer to a process-global factory vtable backed by the +// Jetson-PI jetson_pi_pi0 policy library. Pass it to +// frt_llama_cpp_pi0_runtime_open_with_engine_factory to open a Pi0 runtime +// that drives a real Jetson-PI whole-graph infer. +// +// Returns NULL when FlashRT was built without FLASHRT_CPP_WITH_JETSON_PI. +// The pointer is valid for the process lifetime; do not release it. +// +// factory->create_pi0 is thread-safe (errors are reported through a +// thread-local sink queried via factory->last_error). Each successful +// create_pi0 yields one owned engine reference; the runtime-open wrapper +// transfers it to the v1 runtime, which releases it on drop. +// +// Engine error codes (returned by set_input/run_infer/get_output, surfaced +// through v1 verbs): 0 ok; -1 null self / precondition; -2 invalid config or +// model; -5 action output buffer too small; -7 actions not ready (run_infer +// did not complete since the last set_input); -8 generic infer failure. +// last_error returns a non-null string describing the most recent failure. +// + +#include "flashrt/providers/llama_cpp/c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +const frt_llama_cpp_engine_factory_v1* +frt_llama_cpp_default_engine_factory(void); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* FLASHRT_PROVIDERS_LLAMA_CPP_JETSON_PI_ENGINE_H */ diff --git a/cpp/providers/llama_cpp/provider_exports.map b/cpp/providers/llama_cpp/provider_exports.map new file mode 100644 index 00000000..cf9de216 --- /dev/null +++ b/cpp/providers/llama_cpp/provider_exports.map @@ -0,0 +1,6 @@ +FLASHRT_LLAMA_CPP_PROVIDER_1.0 { + global: + frt_model_runtime_open_v1; + local: + *; +}; diff --git a/cpp/providers/llama_cpp/src/checkpoint_identity.cpp b/cpp/providers/llama_cpp/src/checkpoint_identity.cpp new file mode 100644 index 00000000..e8b68761 --- /dev/null +++ b/cpp/providers/llama_cpp/src/checkpoint_identity.cpp @@ -0,0 +1,172 @@ +#include "checkpoint_identity.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace flashrt::providers::llama_cpp { +namespace { + +thread_local std::string g_runtime_open_error; + +constexpr std::array kSha256Round = { + 0x428a2f98u, 0x71374491u, 0xb5c0fbcfu, 0xe9b5dba5u, + 0x3956c25bu, 0x59f111f1u, 0x923f82a4u, 0xab1c5ed5u, + 0xd807aa98u, 0x12835b01u, 0x243185beu, 0x550c7dc3u, + 0x72be5d74u, 0x80deb1feu, 0x9bdc06a7u, 0xc19bf174u, + 0xe49b69c1u, 0xefbe4786u, 0x0fc19dc6u, 0x240ca1ccu, + 0x2de92c6fu, 0x4a7484aau, 0x5cb0a9dcu, 0x76f988dau, + 0x983e5152u, 0xa831c66du, 0xb00327c8u, 0xbf597fc7u, + 0xc6e00bf3u, 0xd5a79147u, 0x06ca6351u, 0x14292967u, + 0x27b70a85u, 0x2e1b2138u, 0x4d2c6dfcu, 0x53380d13u, + 0x650a7354u, 0x766a0abbu, 0x81c2c92eu, 0x92722c85u, + 0xa2bfe8a1u, 0xa81a664bu, 0xc24b8b70u, 0xc76c51a3u, + 0xd192e819u, 0xd6990624u, 0xf40e3585u, 0x106aa070u, + 0x19a4c116u, 0x1e376c08u, 0x2748774cu, 0x34b0bcb5u, + 0x391c0cb3u, 0x4ed8aa4au, 0x5b9cca4fu, 0x682e6ff3u, + 0x748f82eeu, 0x78a5636fu, 0x84c87814u, 0x8cc70208u, + 0x90befffau, 0xa4506cebu, 0xbef9a3f7u, 0xc67178f2u, +}; + +uint32_t rotate_right(uint32_t value, uint32_t bits) { + return (value >> bits) | (value << (32u - bits)); +} + +struct Sha256 { + std::array state = { + 0x6a09e667u, 0xbb67ae85u, 0x3c6ef372u, 0xa54ff53au, + 0x510e527fu, 0x9b05688cu, 0x1f83d9abu, 0x5be0cd19u, + }; + std::array block{}; + uint64_t bytes = 0; + size_t used = 0; + + void transform(const uint8_t* input) { + uint32_t words[64]; + for (size_t i = 0; i < 16; ++i) { + words[i] = (static_cast(input[i * 4]) << 24) | + (static_cast(input[i * 4 + 1]) << 16) | + (static_cast(input[i * 4 + 2]) << 8) | + static_cast(input[i * 4 + 3]); + } + for (size_t i = 16; i < 64; ++i) { + const uint32_t s0 = rotate_right(words[i - 15], 7) ^ + rotate_right(words[i - 15], 18) ^ + (words[i - 15] >> 3); + const uint32_t s1 = rotate_right(words[i - 2], 17) ^ + rotate_right(words[i - 2], 19) ^ + (words[i - 2] >> 10); + words[i] = words[i - 16] + s0 + words[i - 7] + s1; + } + uint32_t a = state[0], b = state[1], c = state[2], d = state[3]; + uint32_t e = state[4], f = state[5], g = state[6], h = state[7]; + for (size_t i = 0; i < 64; ++i) { + const uint32_t sum1 = rotate_right(e, 6) ^ rotate_right(e, 11) ^ + rotate_right(e, 25); + const uint32_t choice = (e & f) ^ (~e & g); + const uint32_t temp1 = h + sum1 + choice + kSha256Round[i] + words[i]; + const uint32_t sum0 = rotate_right(a, 2) ^ rotate_right(a, 13) ^ + rotate_right(a, 22); + const uint32_t majority = (a & b) ^ (a & c) ^ (b & c); + const uint32_t temp2 = sum0 + majority; + h = g; g = f; f = e; e = d + temp1; + d = c; c = b; b = a; a = temp1 + temp2; + } + state[0] += a; state[1] += b; state[2] += c; state[3] += d; + state[4] += e; state[5] += f; state[6] += g; state[7] += h; + } + + void update(const void* data, size_t size) { + const auto* input = static_cast(data); + bytes += size; + while (size > 0) { + const size_t take = std::min(size, block.size() - used); + std::copy(input, input + take, block.begin() + used); + input += take; + size -= take; + used += take; + if (used == block.size()) { + transform(block.data()); + used = 0; + } + } + } + + std::array finish() { + const uint64_t bit_count = bytes * 8u; + const uint8_t one = 0x80; + update(&one, 1); + const uint8_t zero = 0; + while (used != 56) update(&zero, 1); + uint8_t length[8]; + for (size_t i = 0; i < 8; ++i) { + length[7 - i] = static_cast(bit_count >> (i * 8)); + } + update(length, sizeof(length)); + std::array digest{}; + for (size_t i = 0; i < state.size(); ++i) { + digest[i * 4] = static_cast(state[i] >> 24); + digest[i * 4 + 1] = static_cast(state[i] >> 16); + digest[i * 4 + 2] = static_cast(state[i] >> 8); + digest[i * 4 + 3] = static_cast(state[i]); + } + return digest; + } +}; + +} // namespace + +bool checkpoint_identity(const char* path, std::string* identity, + std::string* error) { + if (!path || !*path || !identity) { + if (error) *error = "invalid checkpoint identity arguments"; + return false; + } + std::ifstream file(path, std::ios::binary); + if (!file) { + if (error) *error = std::string("failed to open checkpoint for identity: ") + path; + return false; + } + Sha256 hash; + std::vector chunk(1024 * 1024); + while (file) { + file.read(chunk.data(), static_cast(chunk.size())); + const std::streamsize read = file.gcount(); + if (read > 0) hash.update(chunk.data(), static_cast(read)); + } + if (!file.eof()) { + if (error) *error = std::string("failed to read checkpoint for identity: ") + path; + return false; + } + const auto digest = hash.finish(); + char hex[65]; + for (size_t i = 0; i < digest.size(); ++i) { + std::snprintf(hex + i * 2, 3, "%02x", digest[i]); + } + hex[64] = '\0'; + *identity = hex; + if (error) error->clear(); + return true; +} + +void clear_runtime_open_error() { + g_runtime_open_error.clear(); +} + +void set_runtime_open_error(const std::string& error) { + g_runtime_open_error = error; +} + +const char* runtime_open_error() { + return g_runtime_open_error.c_str(); +} + +} // namespace flashrt::providers::llama_cpp + +extern "C" const char* frt_llama_cpp_runtime_open_error(void) { + return flashrt::providers::llama_cpp::runtime_open_error(); +} diff --git a/cpp/providers/llama_cpp/src/checkpoint_identity.h b/cpp/providers/llama_cpp/src/checkpoint_identity.h new file mode 100644 index 00000000..5833fb5d --- /dev/null +++ b/cpp/providers/llama_cpp/src/checkpoint_identity.h @@ -0,0 +1,16 @@ +#ifndef FLASHRT_PROVIDERS_LLAMA_CPP_CHECKPOINT_IDENTITY_H +#define FLASHRT_PROVIDERS_LLAMA_CPP_CHECKPOINT_IDENTITY_H + +#include + +namespace flashrt::providers::llama_cpp { + +bool checkpoint_identity(const char* path, std::string* identity, + std::string* error); +void clear_runtime_open_error(); +void set_runtime_open_error(const std::string& error); +const char* runtime_open_error(); + +} // namespace flashrt::providers::llama_cpp + +#endif diff --git a/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp b/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp new file mode 100644 index 00000000..c225b9fe --- /dev/null +++ b/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp @@ -0,0 +1,1214 @@ +// Jetson-PI Pi0 engine: a thin mapping from the FlashRT frt_llama_cpp_engine_v1 +// vtable onto the Jetson-PI jetson_pi_pi0 policy C API. The Pi0 infer glue +// (marker injection, tokenize, encode+decode, KV reset, state padding) lives +// inside jetson_pi_pi0; this file only translates frt_image_view/prompt/state +// into jetson_pi_pi0 inputs and the action chunk back out. +// +// Built only when FLASHRT_CPP_WITH_JETSON_PI is on. No GGML types are exposed +// through the FlashRT public header (c_api.h / jetson_pi_engine.h stay clean). + +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" + +#if defined(FLASHRT_CPP_WITH_JETSON_PI) + +#include "flashrt/model_runtime.h" +#include "jetson_pi_pi0.h" +#include "jetson_pi_llm.h" +#include "jetson_pi_mllm.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct Engine { + jetson_pi_pi0 * pi0 = nullptr; + + // Config snapshot (owned strings). + std::string model_path; + std::string mmproj_path; + std::string backend; + uint32_t n_views = 0; + uint32_t image_height = 0; + uint32_t image_width = 0; + uint32_t action_steps = 0; + uint32_t action_dim = 0; + + // Per-tick transient state, fed by set_input and consumed by run_infer. + std::vector rgb_scratch; // packed RGB per view, concatenated + std::vector image_ptrs; // n_views pointers into rgb_scratch + std::string prompt; + std::vector state; // input proprioception (any width) + std::vector actions_buf; // output action chunk, size == action_steps*action_dim + bool images_set = false; + bool prompt_set = false; + bool state_set = false; + + std::string last_error; + std::atomic refs{1}; + + void set_error(const std::string & m) { last_error = m; } + void clear_error() { last_error.clear(); } +}; + +// Swizzle one frt_image_view into packed RGB (mtmd_bitmap's format). Honors +// stride_bytes and the declared pixel_format. Returns false on a bad view. +bool view_to_rgb(const frt_image_view & v, uint32_t expected_w, + uint32_t expected_h, std::vector & out) { + if (v.struct_size < sizeof(frt_image_view) || !v.data) return false; + if (v.width <= 0 || v.height <= 0) return false; + if (static_cast(v.width) != expected_w || + static_cast(v.height) != expected_h) { + return false; + } + const int w = v.width; + const int h = v.height; + int ch_src; + switch (v.pixel_format) { + case FRT_RT_PIXEL_RGB8: + case FRT_RT_PIXEL_BGR8: + ch_src = 3; break; + case FRT_RT_PIXEL_RGBA8: + case FRT_RT_PIXEL_BGRA8: + ch_src = 4; break; + case FRT_RT_PIXEL_GRAY8: + ch_src = 1; break; + default: + return false; // reject unknown pixel formats (no silent fallback) + } + if (v.stride_bytes < 0) return false; + const uint64_t row_bytes = static_cast(w) * ch_src; + const uint64_t stride = v.stride_bytes == 0 + ? row_bytes : static_cast(v.stride_bytes); + if (stride < row_bytes) return false; + const uint64_t preceding_rows = static_cast(h - 1); + if (preceding_rows != 0 && + stride > (std::numeric_limits::max() - row_bytes) / + preceding_rows) { + return false; + } + const uint64_t required_bytes = preceding_rows * stride + row_bytes; + const uint64_t rgb_bytes = static_cast(w) * h * 3; + if (v.bytes < required_bytes || + required_bytes > static_cast( + std::numeric_limits::max()) || + rgb_bytes > std::numeric_limits::max()) { + return false; + } + + out.resize(static_cast(rgb_bytes)); + const uint8_t * src = static_cast(v.data); + for (int y = 0; y < h; ++y) { + const uint8_t * row = src + static_cast(y) * stride; + for (int x = 0; x < w; ++x) { + const uint8_t * px = row + x * ch_src; + uint8_t * dst = &out[(static_cast(y) * w + x) * 3]; + switch (v.pixel_format) { + case FRT_RT_PIXEL_RGB8: + dst[0] = px[0]; dst[1] = px[1]; dst[2] = px[2]; + break; + case FRT_RT_PIXEL_BGR8: + dst[0] = px[2]; dst[1] = px[1]; dst[2] = px[0]; + break; + case FRT_RT_PIXEL_RGBA8: + dst[0] = px[0]; dst[1] = px[1]; dst[2] = px[2]; + break; + case FRT_RT_PIXEL_BGRA8: + dst[0] = px[2]; dst[1] = px[1]; dst[2] = px[0]; + break; + case FRT_RT_PIXEL_GRAY8: + dst[0] = dst[1] = dst[2] = px[0]; + break; + default: + return false; // unreachable; ch_src gate above + } + } + } + return true; +} + +// Open errors are reported through the factory's thread_local sink because +// the contract is the returned engine is zeroed on failure (no handle to +// query last_error from). This covers both jetson_pi_pi0_open failures and +// the engine's own create_pi0 validation (config/shape mismatch, OOM). +static thread_local std::string g_create_error; +static void set_create_error(const std::string & m) { g_create_error = m; } + +int32_t pi0_status_to_engine(int32_t s) { + switch (s) { + case JETSON_PI_PI0_OK: return 0; + case JETSON_PI_PI0_ACTION_NOT_READY: return -7; + case JETSON_PI_PI0_BUFFER_TOO_SMALL: return -5; + case JETSON_PI_PI0_INVALID: return -2; + default: return -8; + } +} + +// ---- engine vtable ---------------------------------------------------------- + +void engine_retain(void * self) { + static_cast(self)->refs.fetch_add(1, std::memory_order_relaxed); +} + +void engine_release(void * self) { + Engine * e = static_cast(self); + if (e->refs.fetch_sub(1, std::memory_order_acq_rel) == 1) { + if (e->pi0) jetson_pi_pi0_close(e->pi0); + delete e; + } +} + +int engine_set_input(void * self, uint32_t port, const void * data, + uint64_t bytes, int /*stream*/) { + Engine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + int32_t discard_status = jetson_pi_pi0_discard_context(e->pi0); + if (discard_status != JETSON_PI_PI0_OK) { + e->set_error(std::string("jetson_pi_pi0_discard_context failed: ") + + jetson_pi_pi0_last_error(e->pi0)); + return pi0_status_to_engine(discard_status); + } + // Any new input invalidates the previous tick's actions so get_output + // cannot return stale data without a fresh run_infer. + e->actions_buf.clear(); + switch (port) { + case FRT_LLAMA_CPP_PI0_PORT_IMAGES: { + e->images_set = false; + e->image_ptrs.clear(); + e->rgb_scratch.clear(); + if (!data || bytes % sizeof(frt_image_view) != 0) { + e->set_error("images payload must be frt_image_view[]"); + return -1; + } + const uint64_t n = bytes / sizeof(frt_image_view); + if (n != e->n_views) { + e->set_error("images view count != config.n_views"); + return -1; + } + const auto * views = static_cast(data); + std::vector packed; + packed.reserve(static_cast(e->n_views) * e->image_width * + e->image_height * 3); + for (uint32_t i = 0; i < e->n_views; ++i) { + std::vector rgb; + if (!view_to_rgb(views[i], e->image_width, e->image_height, + rgb)) { + e->set_error("invalid frt_image_view at index " + + std::to_string(i)); + return -1; + } + packed.insert(packed.end(), rgb.begin(), rgb.end()); + } + // Now that packed is fully grown (no more reallocs), compute the + // per-view pointers into it. Storing them earlier would risk + // dangling pointers if a later insert reallocated packed. + e->rgb_scratch = std::move(packed); + const size_t per_view = + static_cast(e->image_width) * e->image_height * 3; + e->image_ptrs.resize(e->n_views); + for (uint32_t i = 0; i < e->n_views; ++i) { + e->image_ptrs[i] = e->rgb_scratch.data() + i * per_view; + } + e->images_set = true; + return 0; + } + case FRT_LLAMA_CPP_PI0_PORT_PROMPT: { + e->prompt_set = false; + e->prompt.clear(); + if (!data || bytes == 0) { + e->set_error("empty prompt"); + return -1; + } + e->prompt.assign(static_cast(data), bytes); + e->prompt_set = true; + return 0; + } + case FRT_LLAMA_CPP_PI0_PORT_STATE: { + e->state_set = false; + e->state.clear(); + if (!data) { + e->set_error("null state"); + return -1; + } + if (bytes == 0 || bytes % sizeof(float) != 0) { + e->set_error("state bytes is not a nonzero multiple of sizeof(float)"); + return -1; + } + // State width is policy-specific and independent of action_dim: + // PI0.5 discretizes up to 8 proprioception values into the prompt, + // while legacy Pi0 consumes an action_dim tensor. Accept any + // float count here; the backend enforces the per-model bound. + const size_t n = bytes / sizeof(float); + e->state.assign(static_cast(data), + static_cast(data) + n); + e->state_set = true; + return 0; + } + case FRT_LLAMA_CPP_PI0_PORT_ACTIONS: + default: + e->set_error("unknown/invalid pi0 input port"); + return -1; + } +} + +int engine_run_infer(void * self) { + Engine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (!e->images_set || !e->prompt_set || !e->state_set) { + e->set_error("infer requires images, prompt, and state to be set"); + return -1; + } + // jetson_pi_pi0_infer writes the action chunk into a buffer it sizes from + // the model; allocate action_steps*action_dim floats. + std::vector actions( + static_cast(e->action_steps) * e->action_dim); + size_t written = 0; + int32_t s = jetson_pi_pi0_infer(e->pi0, + e->image_ptrs.data(), e->image_ptrs.size(), + e->prompt.data(), e->prompt.size(), + e->state.data(), e->state.size(), + actions.data(), actions.size(), + &written); + if (s != JETSON_PI_PI0_OK) { + e->set_error(std::string("jetson_pi_pi0_infer failed: ") + + jetson_pi_pi0_last_error(e->pi0)); + return pi0_status_to_engine(s); + } + e->actions_buf.assign(actions.begin(), actions.end()); + return 0; +} + +int engine_run_stage(void * self, uint32_t stage) { + Engine * e = static_cast(self); + if (!e) return -1; + if (stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER) { + return engine_run_infer(self); + } + e->clear_error(); + if (!e->images_set || !e->prompt_set || !e->state_set) { + e->set_error("Pi0 stage requires images, prompt, and state to be set"); + return -1; + } + if (stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT) { + e->actions_buf.clear(); + int32_t s = jetson_pi_pi0_context( + e->pi0, e->image_ptrs.data(), e->image_ptrs.size(), + e->prompt.data(), e->prompt.size(), e->state.data(), + e->state.size()); + if (s != JETSON_PI_PI0_OK) { + e->set_error(std::string("jetson_pi_pi0_context failed: ") + + jetson_pi_pi0_last_error(e->pi0)); + return pi0_status_to_engine(s); + } + return 0; + } + if (stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION) { + std::vector actions( + static_cast(e->action_steps) * e->action_dim); + size_t written = 0; + int32_t s = jetson_pi_pi0_action( + e->pi0, actions.data(), actions.size(), &written); + if (s != JETSON_PI_PI0_OK) { + e->set_error(std::string("jetson_pi_pi0_action failed: ") + + jetson_pi_pi0_last_error(e->pi0)); + return pi0_status_to_engine(s); + } + e->actions_buf.assign(actions.begin(), actions.end()); + return 0; + } + e->set_error("unknown Pi0 engine stage"); + return -1; +} + +const char * engine_last_error(void * self) { + Engine * e = static_cast(self); + if (!e) return "null jetson_pi engine"; + if (e->last_error.empty()) return "ok"; + return e->last_error.c_str(); +} + +int engine_get_output(void * self, uint32_t port, void * out, + uint64_t capacity, uint64_t * written, int /*stream*/) { + Engine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (port != FRT_LLAMA_CPP_PI0_PORT_ACTIONS) { + e->set_error("pi0 only exports the actions port"); + return -1; + } + if (!out || !written) { + e->set_error("get_output requires out and written"); + return -1; + } + const size_t need_elems = + static_cast(e->action_steps) * e->action_dim; + const uint64_t need_bytes = need_elems * sizeof(float); + *written = need_bytes; + if (capacity < need_bytes) { + e->set_error("action output buffer too small"); + return -5; + } + if (e->actions_buf.size() != need_elems) { + e->set_error("actions not ready; run_infer did not complete"); + return -7; + } + std::memcpy(out, e->actions_buf.data(), need_bytes); + return 0; +} + +} // namespace + +// --------------------------------------------------------------------------- +// LLM engine (generic GGUF text completion). Distinct IO shape from Pi0 +// (1 prompt in -> 1 text out), so it has its own Engine struct + verbs, but +// reuses the frt_llama_cpp_engine_v1 vtable shape and the default factory. +// --------------------------------------------------------------------------- + +namespace { + +struct LlmEngine { + jetson_pi_llm * llm = nullptr; + + std::string prompt; // set_input(PROMPT) stash + std::vector tokens; + std::string text_buf; // run_infer output, get_output(TEXT) source + std::vector logits_buf; + int32_t next_token = 0; + int32_t is_eog = 0; + bool next_token_ready = false; + bool prompt_set = false; + + std::string last_error; + std::atomic refs{1}; + + void set_error(const std::string & m) { last_error = m; } + void clear_error() { last_error.clear(); } +}; + +int llm_engine_run_infer(void * self); + +void llm_engine_retain(void * self) { + static_cast(self)->refs.fetch_add(1, std::memory_order_relaxed); +} + +void llm_engine_release(void * self) { + LlmEngine * e = static_cast(self); + if (e->refs.fetch_sub(1, std::memory_order_acq_rel) == 1) { + if (e->llm) jetson_pi_llm_close(e->llm); + delete e; + } +} + +int llm_engine_set_input(void * self, uint32_t port, const void * data, + uint64_t bytes, int /*stream*/) { + LlmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (port != FRT_LLAMA_CPP_LLM_PORT_PROMPT && + port != FRT_LLAMA_CPP_LLM_PORT_TOKENS) { + e->set_error("unknown llm input port"); + return -1; + } + e->prompt_set = false; + e->prompt.clear(); + e->tokens.clear(); + e->text_buf.clear(); + e->logits_buf.clear(); + e->next_token_ready = false; + if (port == FRT_LLAMA_CPP_LLM_PORT_PROMPT) { + if (!data || bytes == 0) { + e->set_error("empty prompt"); + return -1; + } + e->prompt.assign(static_cast(data), bytes); + e->prompt_set = true; + } else { + if (!data || bytes == 0 || bytes % sizeof(int32_t) != 0) { + e->set_error("tokens payload must be a non-empty int32 array"); + return -1; + } + const auto * begin = static_cast(data); + e->tokens.assign(begin, begin + bytes / sizeof(int32_t)); + e->prompt_set = true; + } + return 0; +} + +int llm_engine_refresh_logits(LlmEngine * e) { + size_t n_logits = 0; + int32_t s = jetson_pi_llm_get_logits(e->llm, nullptr, 0, &n_logits); + if (s != JETSON_PI_LLM_BUFFER_TOO_SMALL || n_logits == 0) { + e->set_error(std::string("jetson_pi_llm_get_logits size query failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + e->logits_buf.resize(n_logits); + s = jetson_pi_llm_get_logits(e->llm, e->logits_buf.data(), + e->logits_buf.size(), &n_logits); + if (s != JETSON_PI_LLM_OK) { + e->logits_buf.clear(); + e->set_error(std::string("jetson_pi_llm_get_logits failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + return 0; +} + +int llm_engine_run_stage(void * self, uint32_t stage) { + LlmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER) { + return llm_engine_run_infer(self); + } + if (stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET) { + const int32_t s = jetson_pi_llm_reset(e->llm); + if (s != JETSON_PI_LLM_OK) { + e->set_error(std::string("jetson_pi_llm_reset failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + e->text_buf.clear(); + e->logits_buf.clear(); + e->next_token_ready = false; + return 0; + } + if (stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL) { + if (!e->prompt_set) { + e->set_error("prefill requires prompt to be set"); + return -1; + } + const int32_t s = e->tokens.empty() + ? jetson_pi_llm_prefill(e->llm, e->prompt.data(), e->prompt.size()) + : jetson_pi_llm_prefill_tokens(e->llm, e->tokens.data(), + e->tokens.size()); + if (s != JETSON_PI_LLM_OK) { + e->set_error(std::string("jetson_pi_llm_prefill failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + e->text_buf.clear(); + e->next_token_ready = false; + return llm_engine_refresh_logits(e); + } + if (stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE) { + e->next_token_ready = false; + e->logits_buf.clear(); + int32_t is_eog = 0; + const int32_t s = jetson_pi_llm_decode_step( + e->llm, &e->next_token, &is_eog); + if (s != JETSON_PI_LLM_OK) { + e->set_error(std::string("jetson_pi_llm_decode_step failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + e->next_token_ready = true; + e->is_eog = is_eog; + e->logits_buf.clear(); + if (!is_eog) { + size_t piece_size = 0; + int32_t piece_status = jetson_pi_llm_token_to_piece( + e->llm, e->next_token, nullptr, 0, &piece_size); + if (piece_status != JETSON_PI_LLM_BUFFER_TOO_SMALL) { + e->set_error(std::string("token piece size query failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + if (piece_size == 0) return llm_engine_refresh_logits(e); + const size_t old_size = e->text_buf.size(); + e->text_buf.resize(old_size + piece_size); + piece_status = jetson_pi_llm_token_to_piece( + e->llm, e->next_token, e->text_buf.data() + old_size, + piece_size, &piece_size); + if (piece_status != JETSON_PI_LLM_OK) { + e->text_buf.resize(old_size); + e->set_error(std::string("token_to_piece failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + return llm_engine_refresh_logits(e); + } + return 0; + } + e->set_error("unknown llm engine stage"); + return -1; +} + +int llm_engine_run_infer(void * self) { + LlmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (!e->prompt_set) { + e->set_error("infer requires prompt to be set"); + return -1; + } + size_t cap = 0; + int32_t s = jetson_pi_llm_generate(e->llm, e->prompt.data(), + e->prompt.size(), nullptr, 0, &cap); + if (s != JETSON_PI_LLM_BUFFER_TOO_SMALL || cap == 0) { + e->set_error(std::string("jetson_pi_llm_generate size query failed: ") + + jetson_pi_llm_last_error(e->llm)); + return -8; + } + e->text_buf.assign(cap, '\0'); + size_t written = 0; + s = jetson_pi_llm_generate(e->llm, e->prompt.data(), e->prompt.size(), + e->text_buf.data(), cap, &written); + if (s != JETSON_PI_LLM_OK) { + e->set_error(std::string("jetson_pi_llm_generate failed: ") + + jetson_pi_llm_last_error(e->llm)); + e->text_buf.clear(); + return -8; + } + e->text_buf.resize(written); + e->logits_buf.clear(); + e->next_token_ready = false; + return 0; +} + +int llm_engine_get_output(void * self, uint32_t port, void * out, + uint64_t capacity, uint64_t * written, + int /*stream*/) { + LlmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (!written) { + e->set_error("get_output requires written"); + return -1; + } + if (port == FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN) { + *written = sizeof(e->next_token); + if (!e->next_token_ready) { + e->set_error("next token not ready"); + return -7; + } + if (!out || capacity < sizeof(e->next_token)) { + e->set_error("next token output buffer too small"); + return -5; + } + std::memcpy(out, &e->next_token, sizeof(e->next_token)); + return 0; + } + if (port == FRT_LLAMA_CPP_LLM_PORT_IS_EOG) { + *written = sizeof(e->is_eog); + if (!e->next_token_ready) { + e->set_error("is_eog not ready"); + return -7; + } + if (!out || capacity < sizeof(e->is_eog)) { + e->set_error("is_eog output buffer too small"); + return -5; + } + std::memcpy(out, &e->is_eog, sizeof(e->is_eog)); + return 0; + } + if (port == FRT_LLAMA_CPP_LLM_PORT_LOGITS) { + const uint64_t need = e->logits_buf.size() * sizeof(float); + *written = need; + if (need == 0) { + e->set_error("logits not ready"); + return -7; + } + if (!out || capacity < need) { + e->set_error("logits output buffer too small"); + return -5; + } + std::memcpy(out, e->logits_buf.data(), need); + return 0; + } + if (port != FRT_LLAMA_CPP_LLM_PORT_TEXT) { + e->set_error("unknown llm output port"); + return -1; + } + const uint64_t need = e->text_buf.size(); + *written = need; + // Size-query (out=NULL) or too-small buffer: report need, return nonzero. + if (!out || capacity < need) { + if (need == 0) { + e->set_error("text not ready; run_infer did not complete"); + return -7; + } + e->set_error("text output buffer too small"); + return -5; + } + std::memcpy(out, e->text_buf.data(), need); + return 0; +} + +const char * llm_engine_last_error(void * self) { + LlmEngine * e = static_cast(self); + if (!e) return "null jetson_pi llm engine"; + if (e->last_error.empty()) return "ok"; + return e->last_error.c_str(); +} + +// --------------------------------------------------------------------------- +// MLLM engine (multimodal LLM: images + prompt -> text). IO shape = Pi0's +// image input + LLM's text output. Reuses view_to_rgb for images and the +// LLM generate contract for text. +// --------------------------------------------------------------------------- + +struct MllmEngine { + jetson_pi_mllm * mllm = nullptr; + + // Per-tick transient state. + std::vector rgb_scratch; + std::vector image_ptrs; + uint32_t n_images = 0; + uint32_t image_height = 0; + uint32_t image_width = 0; + std::string prompt; + std::string text_buf; + std::vector logits_buf; + int32_t next_token = 0; + int32_t is_eog = 0; + bool next_token_ready = false; + bool images_set = false; + bool prompt_set = false; + + std::string last_error; + std::atomic refs{1}; + + void set_error(const std::string & m) { last_error = m; } + void clear_error() { last_error.clear(); } +}; + +int mllm_engine_run_infer(void * self); + +void mllm_engine_retain(void * self) { + static_cast(self)->refs.fetch_add(1, std::memory_order_relaxed); +} + +void mllm_engine_release(void * self) { + MllmEngine * e = static_cast(self); + if (e->refs.fetch_sub(1, std::memory_order_acq_rel) == 1) { + if (e->mllm) jetson_pi_mllm_close(e->mllm); + delete e; + } +} + +int mllm_engine_set_input(void * self, uint32_t port, const void * data, + uint64_t bytes, int /*stream*/) { + MllmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + switch (port) { + case FRT_LLAMA_CPP_MLLM_PORT_IMAGES: { + e->images_set = false; + e->image_ptrs.clear(); + e->rgb_scratch.clear(); + e->n_images = 0; + e->image_height = 0; + e->image_width = 0; + e->text_buf.clear(); + e->logits_buf.clear(); + e->next_token_ready = false; + if (bytes % sizeof(frt_image_view) != 0 || + (bytes != 0 && !data)) { + e->set_error("images payload must be frt_image_view[]"); + return -1; + } + const uint64_t n = bytes / sizeof(frt_image_view); + if (n == 0) { + // n_images == 0 means text-only; allow but record zero. + e->images_set = true; + return 0; + } + if (n > std::numeric_limits::max()) { + e->set_error("images view count exceeds uint32_t"); + return -1; + } + const auto * views = static_cast(data); + // All images must share the first view's H/W (jetson_pi_mllm + // takes a single image_height/image_width for all images). + if (views[0].width <= 0 || views[0].height <= 0) { + e->set_error("invalid frt_image_view at index 0"); + return -1; + } + const uint32_t w = static_cast(views[0].width); + const uint32_t h = static_cast(views[0].height); + std::vector packed; + for (uint64_t i = 0; i < n; ++i) { + std::vector rgb; + if (!view_to_rgb(views[i], w, h, rgb)) { + e->set_error("invalid frt_image_view at index " + + std::to_string(i)); + return -1; + } + packed.insert(packed.end(), rgb.begin(), rgb.end()); + } + e->rgb_scratch = std::move(packed); + const size_t per_view = static_cast(w) * h * 3; + e->image_ptrs.resize(n); + for (uint64_t i = 0; i < n; ++i) { + e->image_ptrs[i] = e->rgb_scratch.data() + i * per_view; + } + e->n_images = static_cast(n); + e->image_height = h; + e->image_width = w; + e->images_set = true; + return 0; + } + case FRT_LLAMA_CPP_MLLM_PORT_PROMPT: { + e->prompt_set = false; + e->prompt.clear(); + e->text_buf.clear(); + e->logits_buf.clear(); + e->next_token_ready = false; + if (!data || bytes == 0) { + e->set_error("empty prompt"); + return -1; + } + e->prompt.assign(static_cast(data), bytes); + e->prompt_set = true; + return 0; + } + case FRT_LLAMA_CPP_MLLM_PORT_TEXT: + default: + e->set_error("unknown/invalid mllm input port"); + return -1; + } +} + +int mllm_engine_refresh_logits(MllmEngine * e) { + size_t n_logits = 0; + int32_t s = jetson_pi_mllm_get_logits(e->mllm, nullptr, 0, &n_logits); + if (s != JETSON_PI_MLLM_BUFFER_TOO_SMALL || n_logits == 0) { + e->set_error(std::string("jetson_pi_mllm_get_logits size query failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + e->logits_buf.resize(n_logits); + s = jetson_pi_mllm_get_logits(e->mllm, e->logits_buf.data(), + e->logits_buf.size(), &n_logits); + if (s != JETSON_PI_MLLM_OK) { + e->logits_buf.clear(); + e->set_error(std::string("jetson_pi_mllm_get_logits failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + return 0; +} + +int mllm_engine_run_stage(void * self, uint32_t stage) { + MllmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (stage == FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER) { + return mllm_engine_run_infer(self); + } + if (stage == FRT_LLAMA_CPP_MLLM_STAGE_INDEX_RESET) { + const int32_t s = jetson_pi_mllm_reset(e->mllm); + if (s != JETSON_PI_MLLM_OK) { + e->set_error(std::string("jetson_pi_mllm_reset failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + e->text_buf.clear(); + e->logits_buf.clear(); + e->next_token_ready = false; + return 0; + } + if (stage == FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL) { + if (!e->images_set || !e->prompt_set) { + e->set_error("prefill requires images and prompt to be set"); + return -1; + } + const int32_t s = jetson_pi_mllm_prefill( + e->mllm, e->image_ptrs.data(), e->n_images, + e->image_height, e->image_width, + e->prompt.data(), e->prompt.size()); + if (s != JETSON_PI_MLLM_OK) { + e->set_error(std::string("jetson_pi_mllm_prefill failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + e->text_buf.clear(); + e->next_token_ready = false; + return mllm_engine_refresh_logits(e); + } + if (stage == FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE) { + e->next_token_ready = false; + e->logits_buf.clear(); + int32_t is_eog = 0; + const int32_t s = jetson_pi_mllm_decode_step( + e->mllm, &e->next_token, &is_eog); + if (s != JETSON_PI_MLLM_OK) { + e->set_error(std::string("jetson_pi_mllm_decode_step failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + e->next_token_ready = true; + e->is_eog = is_eog; + e->logits_buf.clear(); + if (is_eog) return 0; + size_t piece_size = 0; + int32_t piece_status = jetson_pi_mllm_token_to_piece( + e->mllm, e->next_token, nullptr, 0, &piece_size); + if (piece_status != JETSON_PI_MLLM_BUFFER_TOO_SMALL) { + e->set_error(std::string("token piece size query failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + if (piece_size > 0) { + const size_t old_size = e->text_buf.size(); + e->text_buf.resize(old_size + piece_size); + piece_status = jetson_pi_mllm_token_to_piece( + e->mllm, e->next_token, e->text_buf.data() + old_size, + piece_size, &piece_size); + if (piece_status != JETSON_PI_MLLM_OK) { + e->text_buf.resize(old_size); + e->set_error(std::string("token_to_piece failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + } + return mllm_engine_refresh_logits(e); + } + e->set_error("unknown mllm engine stage"); + return -1; +} + +int mllm_engine_run_infer(void * self) { + MllmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (!e->images_set || !e->prompt_set) { + e->set_error("infer requires images and prompt to be set"); + return -1; + } + size_t cap = 0; + int32_t s = jetson_pi_mllm_infer( + e->mllm, e->image_ptrs.data(), e->n_images, + e->image_height, e->image_width, + e->prompt.data(), e->prompt.size(), nullptr, 0, &cap); + if (s != JETSON_PI_MLLM_BUFFER_TOO_SMALL || cap == 0) { + e->set_error(std::string("jetson_pi_mllm_infer size query failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + return -8; + } + e->text_buf.assign(cap, '\0'); + size_t written = 0; + s = jetson_pi_mllm_infer(e->mllm, + e->image_ptrs.data(), e->n_images, + e->image_height, e->image_width, + e->prompt.data(), e->prompt.size(), + e->text_buf.data(), cap, &written); + if (s != JETSON_PI_MLLM_OK) { + e->set_error(std::string("jetson_pi_mllm_infer failed: ") + + jetson_pi_mllm_last_error(e->mllm)); + e->text_buf.clear(); + return -8; + } + e->text_buf.resize(written); + e->logits_buf.clear(); + e->next_token_ready = false; + return 0; +} + +int mllm_engine_get_output(void * self, uint32_t port, void * out, + uint64_t capacity, uint64_t * written, + int /*stream*/) { + MllmEngine * e = static_cast(self); + if (!e) return -1; + e->clear_error(); + if (!written) { + e->set_error("get_output requires written"); + return -1; + } + if (port == FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN || + port == FRT_LLAMA_CPP_MLLM_PORT_IS_EOG) { + if (!e->next_token_ready) { + e->set_error("decode output not ready"); + return -7; + } + const uint64_t need = sizeof(int32_t); + *written = need; + if (!out || capacity < need) { + e->set_error("decode scalar output buffer too small"); + return -5; + } + const int32_t value = port == FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN + ? e->next_token : e->is_eog; + std::memcpy(out, &value, need); + return 0; + } + if (port == FRT_LLAMA_CPP_MLLM_PORT_LOGITS) { + const uint64_t need = e->logits_buf.size() * sizeof(float); + *written = need; + if (need == 0) { + e->set_error("logits not ready"); + return -7; + } + if (!out || capacity < need) { + e->set_error("logits output buffer too small"); + return -5; + } + std::memcpy(out, e->logits_buf.data(), need); + return 0; + } + if (port != FRT_LLAMA_CPP_MLLM_PORT_TEXT) { + e->set_error("unknown mllm output port"); + return -1; + } + const uint64_t need = e->text_buf.size(); + *written = need; + if (!out || capacity < need) { + if (need == 0) { + e->set_error("text not ready; run_infer did not complete"); + return -7; + } + e->set_error("text output buffer too small"); + return -5; + } + std::memcpy(out, e->text_buf.data(), need); + return 0; +} + +const char * mllm_engine_last_error(void * self) { + MllmEngine * e = static_cast(self); + if (!e) return "null jetson_pi mllm engine"; + if (e->last_error.empty()) return "ok"; + return e->last_error.c_str(); +} + +} // namespace + +extern "C" const frt_llama_cpp_engine_factory_v1* +frt_llama_cpp_default_engine_factory(void) { + static const frt_llama_cpp_engine_factory_v1 factory = []{ + frt_llama_cpp_engine_factory_v1 f{}; + f.struct_size = sizeof(f); + f.self = nullptr; + f.create_pi0 = [](void * /*self*/, + const frt_llama_cpp_pi0_config * config, + frt_llama_cpp_engine_v1 * out) -> int { + g_create_error.clear(); + if (!config || config->struct_size < FRT_LLAMA_CPP_PI0_CONFIG_BASE_SIZE || !out) { + set_create_error("invalid create_pi0 arguments"); + return -1; + } + if (!config->model_path || !config->model_path[0] || + !config->mmproj_path || !config->mmproj_path[0] || + !config->backend || !config->backend[0] || + !config->n_views || config->n_views > 3 || + !config->image_height || + !config->image_width || !config->image_channels || + !config->action_steps || !config->action_dim) { + set_create_error("create_pi0 config has invalid fields"); + return -1; + } +#if defined(JETSON_PI_PI0_HAS_CAPABILITIES_API) && \ + JETSON_PI_PI0_HAS_CAPABILITIES_API + const uint32_t pi0_capabilities = jetson_pi_pi0_capabilities(); + if ((pi0_capabilities & + JETSON_PI_PI0_CAP_REAL_CONTEXT_ACTION) == 0) { + set_create_error( + "loaded Jetson-PI runtime does not provide real Pi0 context/action split"); + return -1; + } +#endif + jetson_pi_pi0_config jc{}; + jc.struct_size = sizeof(jc); + jc.model_path = config->model_path; + jc.mmproj_path = config->mmproj_path; + jc.backend = config->backend; + jc.n_views = config->n_views; + jc.image_height = config->image_height; + jc.image_width = config->image_width; + jc.n_threads = 0; // let jetson_pi_pi0 pick hardware_concurrency + + jetson_pi_pi0 * pi0 = nullptr; + int32_t s = jetson_pi_pi0_open(&jc, &pi0); + if (s != JETSON_PI_PI0_OK || !pi0) { + set_create_error(std::string("jetson_pi_pi0_open failed: ") + + jetson_pi_pi0_open_error()); + return -1; + } + + // Validate the model's real action shape against the config. + uint32_t steps = 0, dim = 0; + if (jetson_pi_pi0_action_shape(pi0, &steps, &dim) != + JETSON_PI_PI0_OK) { + set_create_error("jetson_pi_pi0_action_shape failed"); + jetson_pi_pi0_close(pi0); + return -1; + } + if (steps != config->action_steps || dim != config->action_dim) { + set_create_error( + "action shape mismatch: config=" + + std::to_string(config->action_steps) + "x" + + std::to_string(config->action_dim) + " model=" + + std::to_string(steps) + "x" + std::to_string(dim)); + jetson_pi_pi0_close(pi0); + return -1; + } + + Engine * e = new (std::nothrow) Engine(); + if (!e) { + set_create_error("engine allocation failed"); + jetson_pi_pi0_close(pi0); + return -5; + } + e->pi0 = pi0; + e->model_path = config->model_path; + e->mmproj_path = config->mmproj_path; + e->backend = config->backend; + e->n_views = config->n_views; + e->image_height = config->image_height; + e->image_width = config->image_width; + e->action_steps = config->action_steps; + e->action_dim = config->action_dim; + + out->struct_size = sizeof(*out); + out->reserved = 0; +#if defined(JETSON_PI_PI0_HAS_REAL_CONTEXT_ACTION) && \ + JETSON_PI_PI0_HAS_REAL_CONTEXT_ACTION && \ + defined(JETSON_PI_PI0_HAS_CAPABILITIES_API) && \ + JETSON_PI_PI0_HAS_CAPABILITIES_API + out->reserved |= + FRT_LLAMA_CPP_ENGINE_CAP_PI0_REAL_CONTEXT_ACTION; +#endif + out->self = e; + out->retain = engine_retain; + out->release = engine_release; + out->set_input = engine_set_input; + out->run_infer = engine_run_infer; + out->get_output = engine_get_output; + out->last_error = engine_last_error; + out->run_stage = engine_run_stage; + return 0; + }; + f.create_llm = [](void * /*self*/, + const frt_llama_cpp_llm_config * config, + frt_llama_cpp_engine_v1 * out) -> int { + g_create_error.clear(); + if (!config || config->struct_size < FRT_LLAMA_CPP_LLM_CONFIG_BASE_SIZE || !out) { + set_create_error("invalid create_llm arguments"); + return -1; + } + if (!config->model_path || !config->model_path[0] || + !config->backend || !config->backend[0]) { + set_create_error("create_llm config has empty/zero fields"); + return -1; + } + jetson_pi_llm_config jc{}; + jc.struct_size = sizeof(jc); + jc.model_path = config->model_path; + jc.backend = config->backend; + jc.n_ctx = config->n_ctx; + jc.n_threads = config->n_threads; + jc.temp = config->temp; + jc.top_k = config->top_k; + jc.top_p = config->top_p; + jc.seed = config->seed; + jc.max_tokens = config->max_tokens; + + jetson_pi_llm * llm = nullptr; + int32_t s = jetson_pi_llm_open(&jc, &llm); + if (s != JETSON_PI_LLM_OK || !llm) { + set_create_error(std::string("jetson_pi_llm_open failed: ") + + jetson_pi_llm_open_error()); + return -1; + } + + LlmEngine * e = new (std::nothrow) LlmEngine(); + if (!e) { + set_create_error("llm engine allocation failed"); + jetson_pi_llm_close(llm); + return -5; + } + e->llm = llm; + + out->struct_size = sizeof(*out); + out->reserved = 0; + out->self = e; + out->retain = llm_engine_retain; + out->release = llm_engine_release; + out->set_input = llm_engine_set_input; + out->run_infer = llm_engine_run_infer; + out->get_output = llm_engine_get_output; + out->last_error = llm_engine_last_error; + out->run_stage = llm_engine_run_stage; + return 0; + }; + f.create_mllm = [](void * /*self*/, + const frt_llama_cpp_mllm_config * config, + frt_llama_cpp_engine_v1 * out) -> int { + g_create_error.clear(); + if (!config || config->struct_size < FRT_LLAMA_CPP_MLLM_CONFIG_BASE_SIZE || !out) { + set_create_error("invalid create_mllm arguments"); + return -1; + } + if (!config->model_path || !config->model_path[0] || + !config->mmproj_path || !config->mmproj_path[0] || + !config->backend || !config->backend[0]) { + set_create_error("create_mllm config has empty/zero fields"); + return -1; + } + jetson_pi_mllm_config jc{}; + jc.struct_size = sizeof(jc); + jc.model_path = config->model_path; + jc.mmproj_path = config->mmproj_path; + jc.backend = config->backend; + jc.n_ctx = config->n_ctx; + jc.n_threads = config->n_threads; + jc.temp = config->temp; + jc.top_k = config->top_k; + jc.top_p = config->top_p; + jc.seed = config->seed; + jc.max_tokens = config->max_tokens; + + jetson_pi_mllm * mllm = nullptr; + int32_t s = jetson_pi_mllm_open(&jc, &mllm); + if (s != JETSON_PI_MLLM_OK || !mllm) { + set_create_error(std::string("jetson_pi_mllm_open failed: ") + + jetson_pi_mllm_open_error()); + return -1; + } + + MllmEngine * e = new (std::nothrow) MllmEngine(); + if (!e) { + set_create_error("mllm engine allocation failed"); + jetson_pi_mllm_close(mllm); + return -5; + } + e->mllm = mllm; + + out->struct_size = sizeof(*out); + out->reserved = 0; + out->self = e; + out->retain = mllm_engine_retain; + out->release = mllm_engine_release; + out->set_input = mllm_engine_set_input; + out->run_infer = mllm_engine_run_infer; + out->get_output = mllm_engine_get_output; + out->last_error = mllm_engine_last_error; + out->run_stage = mllm_engine_run_stage; + return 0; + }; + f.last_error = [](void * /*self*/) -> const char * { + return g_create_error.c_str(); + }; + return f; + }(); + return &factory; +} + +#else /* !FLASHRT_CPP_WITH_JETSON_PI */ + +extern "C" const frt_llama_cpp_engine_factory_v1* +frt_llama_cpp_default_engine_factory(void) { + return nullptr; +} + +#endif /* FLASHRT_CPP_WITH_JETSON_PI */ diff --git a/cpp/providers/llama_cpp/src/jetson_pi_link_check.cpp b/cpp/providers/llama_cpp/src/jetson_pi_link_check.cpp new file mode 100644 index 00000000..8426df18 --- /dev/null +++ b/cpp/providers/llama_cpp/src/jetson_pi_link_check.cpp @@ -0,0 +1,47 @@ +#include "flashrt/providers/llama_cpp/c_api.h" + +#include "llama.h" +#include "mtmd.h" +#include "mtmd-helper.h" + +#include + +extern "C" int frt_llama_cpp_jetson_pi_link_check(void) { + mtmd_context_params mtmd_params = mtmd_context_params_default(); + llama_model_params model_params = llama_model_default_params(); + llama_context_params context_params = llama_context_default_params(); + (void)mtmd_params; + (void)model_params; + (void)context_params; + + const char* marker = mtmd_default_marker(); + if (!marker || std::strlen(marker) == 0) return -1; + + /* Force a link-time reference to the Pi0 whole-request entry point. The + * default-parameter symbols alone do not exercise this translation unit's + * dependency on the Pi0 inference path. */ + volatile auto eval_pi0 = &mtmd_helper_eval_chunks_pi0; + (void)eval_pi0; + + frt_llama_cpp_pi0_config cfg{}; + cfg.struct_size = sizeof(cfg); + cfg.model_path = "model.gguf"; + cfg.mmproj_path = "mmproj.gguf"; + cfg.backend = "cpu"; + cfg.n_views = 1; + cfg.image_height = 224; + cfg.image_width = 224; + cfg.image_channels = 3; + cfg.action_steps = 2; + cfg.action_dim = 4; + + frt_llama_cpp_engine_v1 engine{}; + engine.struct_size = sizeof(engine); + frt_model_runtime_v1* model = nullptr; + if (frt_llama_cpp_pi0_runtime_create_with_engine(&cfg, &engine, &model) != + -1) { + if (model) model->release(model->owner); + return -1; + } + return 0; +} diff --git a/cpp/providers/llama_cpp/src/llm_runtime.cpp b/cpp/providers/llama_cpp/src/llm_runtime.cpp new file mode 100644 index 00000000..129b85e9 --- /dev/null +++ b/cpp/providers/llama_cpp/src/llm_runtime.cpp @@ -0,0 +1,536 @@ +// frt_llama_cpp_llm_runtime — v1 runtime wrapper for a generic GGUF LLM +// engine. Mirrors pi0_runtime.cpp's shape: 2 STAGED ports (prompt TEXT in, +// text TEXT out) + 1 callback "infer" stage. Strict JSON open path. +// +// The engine vtable (frt_llama_cpp_engine_v1) is reused: set_input(PROMPT), +// run_infer(), get_output(TEXT). No GGML types appear here. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "checkpoint_identity.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +struct RuntimeOwner { + frt_llama_cpp_engine_v1 engine{}; + std::string last_error; + int64_t prompt_shape[1] = {-1}; // variable-length UTF-8 + int64_t text_shape[1] = {-1}; // variable-length UTF-8 output + int64_t token_shape[1] = {1}; + int64_t eog_shape[1] = {1}; + int64_t logits_shape[1] = {-1}; + int64_t tokens_shape[1] = {-1}; + bool staged_decode = false; +}; + +constexpr uint32_t kInferExecutor = 37; +constexpr uint32_t kResetExecutor = 4; +constexpr uint32_t kPrefillExecutor = 91; +constexpr uint32_t kDecodeExecutor = 113; + +int unsupported_prepare(void* self, uint32_t, frt_shape_key) { + auto* owner = static_cast(self); + if (owner) owner->last_error = "llama_cpp LLM provider has no graph variants"; + return -3; +} + +const char* engine_error(RuntimeOwner* owner) { + const char* err = owner->engine.last_error(owner->engine.self); + return err ? err : "llama_cpp LLM engine returned a null error"; +} + +int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, + int stream) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->engine.set_input(owner->engine.self, port, data, + bytes, stream); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int get_output(void* self, uint32_t port, void* out, uint64_t capacity, + uint64_t* written, int stream) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->engine.get_output(owner->engine.self, port, out, + capacity, written, stream); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int run_engine_stage(void* self, uint32_t stage) { + auto* owner = static_cast(self); + if (!owner) return -1; + int rc = 0; + if (stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER) { + rc = owner->engine.run_infer(owner->engine.self); + } else if (owner->staged_decode && owner->engine.run_stage && + (stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET || + stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL || + stage == FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE)) { + rc = owner->engine.run_stage(owner->engine.self, stage); + } else { + owner->last_error = "unknown or unsupported llama_cpp LLM stage"; + return -1; + } + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int step(void* self) { + return run_engine_stage(self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER); +} + +int run_opaque(void* self, uint32_t executor_ref) { + switch (executor_ref) { + case kInferExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER); + case kResetExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET); + case kPrefillExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL); + case kDecodeExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE); + default: + static_cast(self)->last_error = + "unknown llama_cpp LLM executor ref"; + return -2; + } +} + +const char* last_error(void* self) { + auto* owner = static_cast(self); + if (!owner) return "null llama_cpp LLM runtime"; + if (!owner->last_error.empty()) return owner->last_error.c_str(); + return engine_error(owner); +} + +void destroy_owner(void* self) { + auto* owner = static_cast(self); + if (owner->engine.release) owner->engine.release(owner->engine.self); + delete owner; +} + +} // namespace + +extern "C" int frt_llama_cpp_llm_runtime_create_with_engine( + const frt_llama_cpp_llm_config* config, + const frt_llama_cpp_engine_v1* engine, + frt_model_runtime_v1** out) { + if (!out) return -1; + *out = nullptr; + if (!config || config->struct_size < FRT_LLAMA_CPP_LLM_CONFIG_BASE_SIZE || + !config->model_path || !config->model_path[0] || + !config->backend || !config->backend[0]) { + return -1; + } + if (!engine || engine->struct_size < FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE || + !engine->self || !engine->set_input || !engine->run_infer || + !engine->get_output || !engine->last_error || + static_cast(engine->retain) != + static_cast(engine->release)) { + return -1; + } + + auto* owner = new (std::nothrow) RuntimeOwner(); + if (!owner) return -5; + std::memcpy(&owner->engine, engine, + std::min(engine->struct_size, sizeof(owner->engine))); + owner->staged_decode = + engine->struct_size >= FRT_LLAMA_CPP_ENGINE_V1_RUN_STAGE_SIZE && + owner->engine.run_stage; + if (owner->engine.retain) owner->engine.retain(owner->engine.self); + + frt_runtime_builder b = frt_model_runtime_builder_create_metadata(); + if (!b) { + destroy_owner(owner); + return -5; + } + + int rc = 0; + rc |= frt_runtime_builder_add_port( + b, "prompt", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, owner->prompt_shape, 1, 0, + nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "text", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, owner->text_shape, 1, 0, + nullptr, 0, 0); + if (owner->staged_decode) { + rc |= frt_runtime_builder_add_port( + b, "next_token", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_I32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + owner->token_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "logits", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + owner->logits_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "is_eog", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_I32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + owner->eog_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "tokens", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_I32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 0, + owner->tokens_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_generic_stage( + b, "reset", FRT_GENERIC_STAGE_OPAQUE, kResetExecutor, + nullptr, 0); + const uint32_t prefill_after[1] = {0}; + rc |= frt_runtime_builder_add_generic_stage( + b, "prefill", FRT_GENERIC_STAGE_OPAQUE, kPrefillExecutor, + prefill_after, 1); + const uint32_t decode_after[1] = {1}; + rc |= frt_runtime_builder_add_generic_stage( + b, "decode", FRT_GENERIC_STAGE_OPAQUE, kDecodeExecutor, + decode_after, 1); + } else { + rc |= frt_runtime_builder_add_generic_stage( + b, "infer", FRT_GENERIC_STAGE_OPAQUE, kInferExecutor, + nullptr, 0); + } + rc |= frt_runtime_builder_set_generic_stage_runner(b, owner, run_opaque); + rc |= frt_runtime_builder_add_identity(b, "provider", "llama_cpp"); + rc |= frt_runtime_builder_add_identity(b, "model_family", "llm"); + rc |= frt_runtime_builder_add_identity(b, "model_path", config->model_path); + if (config->struct_size >= sizeof(*config) && config->model_identity && + config->model_identity[0]) { + rc |= frt_runtime_builder_add_identity( + b, "weights_sha256", config->model_identity); + } + rc |= frt_runtime_builder_add_identity(b, "backend", config->backend); + rc |= frt_runtime_builder_add_identity( + b, "n_ctx", std::to_string(config->n_ctx).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "n_threads", std::to_string(config->n_threads).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "temp", std::to_string(config->temp).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "top_k", std::to_string(config->top_k).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "top_p", std::to_string(config->top_p).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "seed", std::to_string(config->seed).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "max_tokens", std::to_string(config->max_tokens).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "stage_plan", owner->staged_decode ? "staged_decode" : "full"); + if (rc != 0) { + (void)frt_runtime_builder_finish(b, nullptr, nullptr, nullptr); + destroy_owner(owner); + return -1; + } + + frt_model_runtime_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.set_input = set_input; + verbs.get_output = get_output; + verbs.prepare = unsupported_prepare; + verbs.step = step; + verbs.last_error = last_error; + + frt_model_runtime_v1* model = frt_runtime_builder_finish_model( + b, &verbs, owner, owner, nullptr, destroy_owner); + if (!model) { + (void)frt_runtime_builder_finish(b, nullptr, nullptr, nullptr); + destroy_owner(owner); + return -1; + } + *out = model; + return 0; +} + +extern "C" int frt_llama_cpp_llm_runtime_open_with_engine_factory( + const char* config_json, + const frt_llama_cpp_engine_factory_v1* factory, + frt_model_runtime_v1** out) { + flashrt::providers::llama_cpp::clear_runtime_open_error(); + flashrt::providers::llama_cpp::set_runtime_open_error( + "invalid LLM runtime open arguments or JSON config"); + if (!out) return -1; + *out = nullptr; + if (!factory || + factory->struct_size < sizeof(frt_llama_cpp_engine_factory_v1) || + !factory->create_llm || !factory->last_error) { + return -1; + } + if (!config_json) { + return -1; + } + + frt_llama_cpp_llm_config config{}; + config.struct_size = sizeof(config); + std::string model_family; + std::string model_path; + std::string backend; + std::string model_identity; + bool seen_model_family = false; + bool seen_model_path = false; + bool seen_backend = false; + bool seen_n_ctx = false; + bool seen_n_threads = false; + bool seen_temp = false; + bool seen_top_k = false; + bool seen_top_p = false; + bool seen_seed = false; + bool seen_max_tokens = false; + const char* p = config_json; + auto skip_ws = [&p]() { + while (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r') ++p; + }; + auto parse_string = [&p](std::string* out) { + if (*p != '"') return false; + ++p; + out->clear(); + while (*p && *p != '"') { + if (*p == '\\') { + ++p; + switch (*p) { + case '"': + case '\\': + case '/': + out->push_back(*p++); + break; + case 'b': + out->push_back('\b'); + ++p; + break; + case 'f': + out->push_back('\f'); + ++p; + break; + case 'n': + out->push_back('\n'); + ++p; + break; + case 'r': + out->push_back('\r'); + ++p; + break; + case 't': + out->push_back('\t'); + ++p; + break; + default: + return false; + } + } else { + const unsigned char ch = static_cast(*p); + if (ch < 0x20) return false; + out->push_back(*p++); + } + } + if (*p != '"') return false; + ++p; + return true; + }; + auto parse_u32 = [&p](uint32_t* out) { + uint32_t value = 0; + if (*p < '0' || *p > '9') return false; + if (*p == '0') { + ++p; + } else { + while (*p >= '0' && *p <= '9') { + const uint32_t digit = static_cast(*p - '0'); + if (value > (UINT32_MAX - digit) / 10u) return false; + value = value * 10u + digit; + ++p; + } + } + *out = value; + return true; + }; + auto parse_i32 = [&p](int32_t* out) { + bool neg = false; + if (*p == '-') { neg = true; ++p; } + int32_t value = 0; + if (*p < '0' || *p > '9') return false; + if (*p == '0') { + ++p; + } else { + while (*p >= '0' && *p <= '9') { + const int32_t digit = static_cast(*p - '0'); + if (value > (INT32_MAX - digit) / 10) return false; + value = value * 10 + digit; + ++p; + } + } + *out = neg ? -value : value; + return true; + }; + auto parse_f32 = [&p](float* out) { + // Accepts an optional sign, integer part, optional fraction, optional + // decimal exponent (e.g. "0.8", "-0.5", "1e-3", ".25"). Matches what + // Python's json.dumps can emit for float fields. Does NOT accept + // inf/nan (rejected as parse failure). + bool neg = false; + if (*p == '-' || *p == '+') { neg = (*p == '-'); ++p; } + float mant = 0.0f; + bool saw_digit = false; + while (*p >= '0' && *p <= '9') { + mant = mant * 10.0f + static_cast(*p - '0'); + ++p; + saw_digit = true; + } + if (*p == '.') { + ++p; + float scale = 0.1f; + while (*p >= '0' && *p <= '9') { + mant += static_cast(*p - '0') * scale; + scale *= 0.1f; + ++p; + saw_digit = true; + } + } + if (!saw_digit) return false; + if (*p == 'e' || *p == 'E') { + ++p; + bool exp_neg = false; + if (*p == '-' || *p == '+') { exp_neg = (*p == '-'); ++p; } + int exp = 0; + bool saw_exp_digit = false; + while (*p >= '0' && *p <= '9') { + exp = exp * 10 + (*p - '0'); + ++p; + saw_exp_digit = true; + } + if (!saw_exp_digit) return false; + float scale = 1.0f; + for (int i = 0; i < exp; ++i) scale *= 10.0f; + if (exp_neg) mant /= scale; else mant *= scale; + } + *out = neg ? -mant : mant; + return true; + }; + skip_ws(); + if (*p != '{') return -1; + ++p; + skip_ws(); + if (*p == '}') return -1; + for (;;) { + std::string key; + if (!parse_string(&key)) return -1; + skip_ws(); + if (*p != ':') return -1; + ++p; + skip_ws(); + if (key == "model_family") { + if (seen_model_family || !parse_string(&model_family) || + model_family.empty()) { + return -1; + } + seen_model_family = true; + } else if (key == "model_path") { + if (seen_model_path || !parse_string(&model_path) || + model_path.empty()) { + return -1; + } + seen_model_path = true; + } else if (key == "backend") { + if (seen_backend || !parse_string(&backend) || + backend.empty()) { + return -1; + } + seen_backend = true; + } else if (key == "n_ctx") { + if (seen_n_ctx || !parse_u32(&config.n_ctx)) return -1; + seen_n_ctx = true; + } else if (key == "n_threads") { + if (seen_n_threads || !parse_i32(&config.n_threads)) return -1; + seen_n_threads = true; + } else if (key == "temp") { + if (seen_temp || !parse_f32(&config.temp)) return -1; + seen_temp = true; + } else if (key == "top_k") { + if (seen_top_k || !parse_i32(&config.top_k)) return -1; + seen_top_k = true; + } else if (key == "top_p") { + if (seen_top_p || !parse_f32(&config.top_p)) return -1; + seen_top_p = true; + } else if (key == "seed") { + if (seen_seed || !parse_u32(&config.seed)) return -1; + seen_seed = true; + } else if (key == "max_tokens") { + if (seen_max_tokens || !parse_u32(&config.max_tokens)) return -1; + seen_max_tokens = true; + } else { + return -1; + } + skip_ws(); + if (*p == '}') { + ++p; + break; + } + if (*p != ',') return -1; + ++p; + skip_ws(); + } + skip_ws(); + if (*p != '\0') return -1; + if (!seen_model_family || model_family != "llm" || !seen_model_path || + !seen_backend || !seen_n_ctx || !seen_n_threads || !seen_temp || + !seen_top_k || !seen_top_p || !seen_seed || !seen_max_tokens) { + return -1; + } + config.model_path = model_path.c_str(); + config.backend = backend.c_str(); + std::string identity_error; + if (!flashrt::providers::llama_cpp::checkpoint_identity( + config.model_path, &model_identity, &identity_error)) { + flashrt::providers::llama_cpp::set_runtime_open_error(identity_error); + return -1; + } + config.model_identity = model_identity.c_str(); + + frt_llama_cpp_engine_v1 engine{}; + const int rc = factory->create_llm(factory->self, &config, &engine); + if (rc != 0) { + const char* error = factory->last_error(factory->self); + flashrt::providers::llama_cpp::set_runtime_open_error( + error ? error : "LLM engine factory failed without an error"); + return rc; + } + if (engine.struct_size < FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE || + !engine.self || !engine.retain || !engine.release || + !engine.set_input || !engine.run_infer || !engine.get_output || + !engine.last_error) { + flashrt::providers::llama_cpp::set_runtime_open_error( + "factory returned an invalid LLM engine"); + return -1; + } + frt_model_runtime_v1* model = nullptr; + const int create_rc = + frt_llama_cpp_llm_runtime_create_with_engine(&config, &engine, &model); + engine.release(engine.self); + if (create_rc != 0) { + flashrt::providers::llama_cpp::set_runtime_open_error( + "failed to create LLM model runtime"); + return create_rc; + } + *out = model; + flashrt::providers::llama_cpp::clear_runtime_open_error(); + return 0; +} diff --git a/cpp/providers/llama_cpp/src/mllm_runtime.cpp b/cpp/providers/llama_cpp/src/mllm_runtime.cpp new file mode 100644 index 00000000..ce54622a --- /dev/null +++ b/cpp/providers/llama_cpp/src/mllm_runtime.cpp @@ -0,0 +1,544 @@ +// frt_llama_cpp_mllm_runtime — v1 runtime wrapper for a multimodal LLM +// engine. Mirrors llm_runtime.cpp's shape but with 3 STAGED ports (images +// IMAGE in, prompt TEXT in, text TEXT out) + 1 callback "infer" stage. +// Strict JSON open path. No GGML types appear here. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "checkpoint_identity.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +struct RuntimeOwner { + frt_llama_cpp_engine_v1 engine{}; + std::string last_error; + // images port: NHWC [n_images, H, W, 3]; -1 = variable n_images. + int64_t images_shape[4] = {-1, -1, -1, 3}; + int64_t prompt_shape[1] = {-1}; + int64_t text_shape[1] = {-1}; + int64_t token_shape[1] = {1}; + int64_t logits_shape[1] = {-1}; + int64_t eog_shape[1] = {1}; + bool staged_decode = false; +}; + +constexpr uint32_t kInferExecutor = 37; +constexpr uint32_t kResetExecutor = 4; +constexpr uint32_t kPrefillExecutor = 91; +constexpr uint32_t kDecodeExecutor = 113; + +int unsupported_prepare(void* self, uint32_t, frt_shape_key) { + auto* owner = static_cast(self); + if (owner) owner->last_error = "llama_cpp MLLM provider has no graph variants"; + return -3; +} + +const char* engine_error(RuntimeOwner* owner) { + const char* err = owner->engine.last_error(owner->engine.self); + return err ? err : "llama_cpp MLLM engine returned a null error"; +} + +int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, + int stream) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->engine.set_input(owner->engine.self, port, data, + bytes, stream); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int get_output(void* self, uint32_t port, void* out, uint64_t capacity, + uint64_t* written, int stream) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->engine.get_output(owner->engine.self, port, out, + capacity, written, stream); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int run_engine_stage(void* self, uint32_t stage) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->staged_decode + ? owner->engine.run_stage(owner->engine.self, stage) + : (stage == FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER + ? owner->engine.run_infer(owner->engine.self) : -1); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int step(void* self) { + return run_engine_stage(self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER); +} + +int run_opaque(void* self, uint32_t executor_ref) { + switch (executor_ref) { + case kInferExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER); + case kResetExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_RESET); + case kPrefillExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL); + case kDecodeExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE); + default: + static_cast(self)->last_error = + "unknown llama_cpp MLLM executor ref"; + return -2; + } +} + +const char* last_error(void* self) { + auto* owner = static_cast(self); + if (!owner) return "null llama_cpp MLLM runtime"; + if (!owner->last_error.empty()) return owner->last_error.c_str(); + return engine_error(owner); +} + +void destroy_owner(void* self) { + auto* owner = static_cast(self); + if (!owner) return; + if (owner->engine.release) owner->engine.release(owner->engine.self); + delete owner; +} + +} // namespace + +extern "C" int frt_llama_cpp_mllm_runtime_create_with_engine( + const frt_llama_cpp_mllm_config* config, + const frt_llama_cpp_engine_v1* engine, + frt_model_runtime_v1** out) { + if (!out) return -1; + *out = nullptr; + if (!config || config->struct_size < FRT_LLAMA_CPP_MLLM_CONFIG_BASE_SIZE || + !config->model_path || !config->model_path[0] || + !config->mmproj_path || !config->mmproj_path[0] || + !config->backend || !config->backend[0]) { + return -1; + } + if (!engine || engine->struct_size < FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE || + !engine->self || !engine->set_input || !engine->run_infer || + !engine->get_output || !engine->last_error || + static_cast(engine->retain) != + static_cast(engine->release)) { + return -1; + } + + auto* owner = new (std::nothrow) RuntimeOwner(); + if (!owner) return -5; + std::memcpy(&owner->engine, engine, + std::min(engine->struct_size, sizeof(owner->engine))); + owner->staged_decode = + engine->struct_size >= FRT_LLAMA_CPP_ENGINE_V1_RUN_STAGE_SIZE && + owner->engine.run_stage; + if (owner->engine.retain) owner->engine.retain(owner->engine.self); + + frt_runtime_builder b = frt_model_runtime_builder_create_metadata(); + if (!b) { + destroy_owner(owner); + return -5; + } + + int rc = 0; + rc |= frt_runtime_builder_add_port( + b, "images", FRT_RT_MOD_IMAGE, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_NHWC, + FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, owner->images_shape, 4, 0, + nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "prompt", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, owner->prompt_shape, 1, 0, + nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "text", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, owner->text_shape, 1, 0, + nullptr, 0, 0); + if (owner->staged_decode) { + rc |= frt_runtime_builder_add_port( + b, "next_token", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_I32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + owner->token_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "logits", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + owner->logits_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "is_eog", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_I32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + owner->eog_shape, 1, 0, nullptr, 0, 0); + rc |= frt_runtime_builder_add_generic_stage( + b, "reset", FRT_GENERIC_STAGE_OPAQUE, kResetExecutor, + nullptr, 0); + const uint32_t prefill_after[1] = {0}; + rc |= frt_runtime_builder_add_generic_stage( + b, "prefill", FRT_GENERIC_STAGE_OPAQUE, kPrefillExecutor, + prefill_after, 1); + const uint32_t decode_after[1] = {1}; + rc |= frt_runtime_builder_add_generic_stage( + b, "decode", FRT_GENERIC_STAGE_OPAQUE, kDecodeExecutor, + decode_after, 1); + } else { + rc |= frt_runtime_builder_add_generic_stage( + b, "infer", FRT_GENERIC_STAGE_OPAQUE, kInferExecutor, + nullptr, 0); + } + rc |= frt_runtime_builder_set_generic_stage_runner(b, owner, run_opaque); + rc |= frt_runtime_builder_add_identity(b, "provider", "llama_cpp"); + rc |= frt_runtime_builder_add_identity(b, "model_family", "mllm"); + rc |= frt_runtime_builder_add_identity(b, "model_path", config->model_path); + rc |= frt_runtime_builder_add_identity(b, "mmproj_path", config->mmproj_path); + if (config->struct_size >= sizeof(*config) && config->model_identity && + config->model_identity[0] && config->mmproj_identity && + config->mmproj_identity[0]) { + rc |= frt_runtime_builder_add_identity( + b, "weights_sha256", config->model_identity); + rc |= frt_runtime_builder_add_identity( + b, "mmproj_sha256", config->mmproj_identity); + } + rc |= frt_runtime_builder_add_identity(b, "backend", config->backend); + rc |= frt_runtime_builder_add_identity( + b, "n_ctx", std::to_string(config->n_ctx).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "n_threads", std::to_string(config->n_threads).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "temp", std::to_string(config->temp).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "top_k", std::to_string(config->top_k).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "top_p", std::to_string(config->top_p).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "seed", std::to_string(config->seed).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "max_tokens", std::to_string(config->max_tokens).c_str()); + rc |= frt_runtime_builder_add_identity( + b, "stage_plan", owner->staged_decode ? "staged_decode" : "full"); + if (rc != 0) { + (void)frt_runtime_builder_finish(b, nullptr, nullptr, nullptr); + destroy_owner(owner); + return -1; + } + + frt_model_runtime_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.set_input = set_input; + verbs.get_output = get_output; + verbs.prepare = unsupported_prepare; + verbs.step = step; + verbs.last_error = last_error; + + frt_model_runtime_v1* model = frt_runtime_builder_finish_model( + b, &verbs, owner, owner, nullptr, destroy_owner); + if (!model) { + (void)frt_runtime_builder_finish(b, nullptr, nullptr, nullptr); + destroy_owner(owner); + return -1; + } + *out = model; + return 0; +} + +extern "C" int frt_llama_cpp_mllm_runtime_open_with_engine_factory( + const char* config_json, + const frt_llama_cpp_engine_factory_v1* factory, + frt_model_runtime_v1** out) { + flashrt::providers::llama_cpp::clear_runtime_open_error(); + flashrt::providers::llama_cpp::set_runtime_open_error( + "invalid MLLM runtime open arguments or JSON config"); + if (!out) return -1; + *out = nullptr; + if (!factory || + factory->struct_size < sizeof(frt_llama_cpp_engine_factory_v1) || + !factory->create_mllm || !factory->last_error) { + return -1; + } + if (!config_json) { + return -1; + } + + frt_llama_cpp_mllm_config config{}; + config.struct_size = sizeof(config); + std::string model_family; + std::string model_path; + std::string mmproj_path; + std::string backend; + std::string model_identity; + std::string mmproj_identity; + bool seen_model_family = false; + bool seen_model_path = false; + bool seen_mmproj_path = false; + bool seen_backend = false; + bool seen_n_ctx = false; + bool seen_n_threads = false; + bool seen_temp = false; + bool seen_top_k = false; + bool seen_top_p = false; + bool seen_seed = false; + bool seen_max_tokens = false; + const char* p = config_json; + auto skip_ws = [&p]() { + while (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r') ++p; + }; + auto parse_string = [&p](std::string* out) { + if (*p != '"') return false; + ++p; + out->clear(); + while (*p && *p != '"') { + if (*p == '\\') { + ++p; + switch (*p) { + case '"': + case '\\': + case '/': + out->push_back(*p++); + break; + case 'b': + out->push_back('\b'); + ++p; + break; + case 'f': + out->push_back('\f'); + ++p; + break; + case 'n': + out->push_back('\n'); + ++p; + break; + case 'r': + out->push_back('\r'); + ++p; + break; + case 't': + out->push_back('\t'); + ++p; + break; + default: + return false; + } + } else { + const unsigned char ch = static_cast(*p); + if (ch < 0x20) return false; + out->push_back(*p++); + } + } + if (*p != '"') return false; + ++p; + return true; + }; + auto parse_u32 = [&p](uint32_t* out) { + uint32_t value = 0; + if (*p < '0' || *p > '9') return false; + if (*p == '0') { + ++p; + } else { + while (*p >= '0' && *p <= '9') { + const uint32_t digit = static_cast(*p - '0'); + if (value > (UINT32_MAX - digit) / 10u) return false; + value = value * 10u + digit; + ++p; + } + } + *out = value; + return true; + }; + auto parse_i32 = [&p](int32_t* out) { + bool neg = false; + if (*p == '-') { neg = true; ++p; } + int32_t value = 0; + if (*p < '0' || *p > '9') return false; + if (*p == '0') { + ++p; + } else { + while (*p >= '0' && *p <= '9') { + const int32_t digit = static_cast(*p - '0'); + if (value > (INT32_MAX - digit) / 10) return false; + value = value * 10 + digit; + ++p; + } + } + *out = neg ? -value : value; + return true; + }; + auto parse_f32 = [&p](float* out) { + bool neg = false; + if (*p == '-' || *p == '+') { neg = (*p == '-'); ++p; } + float mant = 0.0f; + bool saw_digit = false; + while (*p >= '0' && *p <= '9') { + mant = mant * 10.0f + static_cast(*p - '0'); + ++p; + saw_digit = true; + } + if (*p == '.') { + ++p; + float scale = 0.1f; + while (*p >= '0' && *p <= '9') { + mant += static_cast(*p - '0') * scale; + scale *= 0.1f; + ++p; + saw_digit = true; + } + } + if (!saw_digit) return false; + if (*p == 'e' || *p == 'E') { + ++p; + bool exp_neg = false; + if (*p == '-' || *p == '+') { exp_neg = (*p == '-'); ++p; } + int exp = 0; + bool saw_exp_digit = false; + while (*p >= '0' && *p <= '9') { + exp = exp * 10 + (*p - '0'); + ++p; + saw_exp_digit = true; + } + if (!saw_exp_digit) return false; + float scale = 1.0f; + for (int i = 0; i < exp; ++i) scale *= 10.0f; + if (exp_neg) mant /= scale; else mant *= scale; + } + *out = neg ? -mant : mant; + return true; + }; + skip_ws(); + if (*p != '{') return -1; + ++p; + skip_ws(); + if (*p == '}') return -1; + for (;;) { + std::string key; + if (!parse_string(&key)) return -1; + skip_ws(); + if (*p != ':') return -1; + ++p; + skip_ws(); + if (key == "model_family") { + if (seen_model_family || !parse_string(&model_family) || + model_family.empty()) { + return -1; + } + seen_model_family = true; + } else if (key == "model_path") { + if (seen_model_path || !parse_string(&model_path) || + model_path.empty()) { + return -1; + } + seen_model_path = true; + } else if (key == "mmproj_path") { + if (seen_mmproj_path || !parse_string(&mmproj_path) || + mmproj_path.empty()) { + return -1; + } + seen_mmproj_path = true; + } else if (key == "backend") { + if (seen_backend || !parse_string(&backend) || + backend.empty()) { + return -1; + } + seen_backend = true; + } else if (key == "n_ctx") { + if (seen_n_ctx || !parse_u32(&config.n_ctx)) return -1; + seen_n_ctx = true; + } else if (key == "n_threads") { + if (seen_n_threads || !parse_i32(&config.n_threads)) return -1; + seen_n_threads = true; + } else if (key == "temp") { + if (seen_temp || !parse_f32(&config.temp)) return -1; + seen_temp = true; + } else if (key == "top_k") { + if (seen_top_k || !parse_i32(&config.top_k)) return -1; + seen_top_k = true; + } else if (key == "top_p") { + if (seen_top_p || !parse_f32(&config.top_p)) return -1; + seen_top_p = true; + } else if (key == "seed") { + if (seen_seed || !parse_u32(&config.seed)) return -1; + seen_seed = true; + } else if (key == "max_tokens") { + if (seen_max_tokens || !parse_u32(&config.max_tokens)) return -1; + seen_max_tokens = true; + } else { + return -1; + } + skip_ws(); + if (*p == '}') { + ++p; + break; + } + if (*p != ',') return -1; + ++p; + skip_ws(); + } + skip_ws(); + if (*p != '\0') return -1; + if (!seen_model_family || model_family != "mllm" || !seen_model_path || + !seen_mmproj_path || !seen_backend || !seen_n_ctx || !seen_n_threads || + !seen_temp || !seen_top_k || !seen_top_p || !seen_seed || + !seen_max_tokens) { + return -1; + } + config.model_path = model_path.c_str(); + config.mmproj_path = mmproj_path.c_str(); + config.backend = backend.c_str(); + std::string identity_error; + if (!flashrt::providers::llama_cpp::checkpoint_identity( + config.model_path, &model_identity, &identity_error) || + !flashrt::providers::llama_cpp::checkpoint_identity( + config.mmproj_path, &mmproj_identity, &identity_error)) { + flashrt::providers::llama_cpp::set_runtime_open_error(identity_error); + return -1; + } + config.model_identity = model_identity.c_str(); + config.mmproj_identity = mmproj_identity.c_str(); + + frt_llama_cpp_engine_v1 engine{}; + const int rc = factory->create_mllm(factory->self, &config, &engine); + if (rc != 0) { + const char* error = factory->last_error(factory->self); + flashrt::providers::llama_cpp::set_runtime_open_error( + error ? error : "MLLM engine factory failed without an error"); + return rc; + } + if (engine.struct_size < FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE || + !engine.self || !engine.retain || !engine.release || + !engine.set_input || !engine.run_infer || !engine.get_output || + !engine.last_error) { + if (engine.release) engine.release(engine.self); + flashrt::providers::llama_cpp::set_runtime_open_error( + "factory returned an invalid MLLM engine"); + return -1; + } + frt_model_runtime_v1* model = nullptr; + const int create_rc = + frt_llama_cpp_mllm_runtime_create_with_engine(&config, &engine, &model); + engine.release(engine.self); + if (create_rc != 0) { + flashrt::providers::llama_cpp::set_runtime_open_error( + "failed to create MLLM model runtime"); + return create_rc; + } + *out = model; + flashrt::providers::llama_cpp::clear_runtime_open_error(); + return 0; +} diff --git a/cpp/providers/llama_cpp/src/pi0_runtime.cpp b/cpp/providers/llama_cpp/src/pi0_runtime.cpp new file mode 100644 index 00000000..f7fd2c14 --- /dev/null +++ b/cpp/providers/llama_cpp/src/pi0_runtime.cpp @@ -0,0 +1,473 @@ +#include "flashrt/providers/llama_cpp/c_api.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "checkpoint_identity.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +struct RuntimeOwner { + frt_llama_cpp_engine_v1 engine{}; + std::string last_error; + int64_t image_shape[4] = {}; + int64_t state_shape[1] = {}; + int64_t action_shape[2] = {}; + bool staged_context_action = false; +}; + +constexpr uint32_t kInferExecutor = 37; +constexpr uint32_t kContextExecutor = 4; +constexpr uint32_t kActionExecutor = 91; + +int unsupported_prepare(void* self, uint32_t, frt_shape_key) { + auto* owner = static_cast(self); + if (owner) owner->last_error = "llama_cpp Pi0 provider has no graph variants"; + return -3; +} + +const char* engine_error(RuntimeOwner* owner) { + const char* err = owner->engine.last_error(owner->engine.self); + return err ? err : "llama_cpp Pi0 engine returned a null error"; +} + +int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, + int stream) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->engine.set_input(owner->engine.self, port, data, + bytes, stream); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int get_output(void* self, uint32_t port, void* out, uint64_t capacity, + uint64_t* written, int stream) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->engine.get_output(owner->engine.self, port, out, + capacity, written, stream); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int run_engine_stage(void* self, uint32_t stage) { + auto* owner = static_cast(self); + if (!owner) return -1; + const int rc = owner->staged_context_action + ? owner->engine.run_stage(owner->engine.self, stage) + : (stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER + ? owner->engine.run_infer(owner->engine.self) : -1); + if (rc != 0) { + owner->last_error = engine_error(owner); + } else { + owner->last_error.clear(); + } + return rc; +} + +int step(void* self) { + return run_engine_stage(self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER); +} + +int run_opaque(void* self, uint32_t executor_ref) { + switch (executor_ref) { + case kInferExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER); + case kContextExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT); + case kActionExecutor: + return run_engine_stage( + self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION); + default: + static_cast(self)->last_error = + "unknown llama_cpp Pi0 executor ref"; + return -2; + } +} + +const char* last_error(void* self) { + auto* owner = static_cast(self); + if (!owner) return "null llama_cpp Pi0 runtime"; + if (!owner->last_error.empty()) return owner->last_error.c_str(); + return engine_error(owner); +} + +void destroy_owner(void* self) { + auto* owner = static_cast(self); + if (owner->engine.release) owner->engine.release(owner->engine.self); + delete owner; +} + +} // namespace + +extern "C" int frt_llama_cpp_pi0_runtime_create_with_engine( + const frt_llama_cpp_pi0_config* config, + const frt_llama_cpp_engine_v1* engine, + frt_model_runtime_v1** out) { + if (!out) return -1; + *out = nullptr; + if (!config || config->struct_size < FRT_LLAMA_CPP_PI0_CONFIG_BASE_SIZE || + !config->model_path || !config->model_path[0] || + !config->mmproj_path || !config->mmproj_path[0] || + !config->backend || !config->backend[0] || !config->n_views || + config->n_views > 3 || + !config->image_height || !config->image_width || + !config->image_channels || + !config->action_steps || !config->action_dim) { + return -1; + } + if (!engine || engine->struct_size < FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE || + !engine->self || !engine->set_input || !engine->run_infer || + !engine->get_output || !engine->last_error || + static_cast(engine->retain) != + static_cast(engine->release)) { + return -1; + } + + auto* owner = new (std::nothrow) RuntimeOwner(); + if (!owner) return -5; + std::memcpy(&owner->engine, engine, + std::min(engine->struct_size, sizeof(owner->engine))); + owner->staged_context_action = + engine->struct_size >= FRT_LLAMA_CPP_ENGINE_V1_RUN_STAGE_SIZE && + owner->engine.run_stage && + (owner->engine.reserved & + FRT_LLAMA_CPP_ENGINE_CAP_PI0_REAL_CONTEXT_ACTION) != 0; + if (owner->engine.retain) owner->engine.retain(owner->engine.self); + + owner->image_shape[0] = static_cast(config->n_views); + owner->image_shape[1] = static_cast(config->image_height); + owner->image_shape[2] = static_cast(config->image_width); + owner->image_shape[3] = static_cast(config->image_channels); + // State width is model-specific: PI0.5 accepts 1..8 values, while legacy + // Pi0 accepts 1..action_dim. Publish a bucket-variable dimension and let + // the backend enforce the selected model's bound. + owner->state_shape[0] = -1; + owner->action_shape[0] = static_cast(config->action_steps); + owner->action_shape[1] = static_cast(config->action_dim); + + frt_runtime_builder b = frt_model_runtime_builder_create_metadata(); + if (!b) { + destroy_owner(owner); + return -5; + } + + int rc = 0; + const int64_t prompt_shape[1] = {-1}; + rc |= frt_runtime_builder_add_port( + b, "images", FRT_RT_MOD_IMAGE, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_NHWC, + FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, owner->image_shape, 4, 0, + nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "prompt", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, prompt_shape, 1, 0, + nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "state", FRT_RT_MOD_STATE, FRT_RT_DTYPE_F32, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, owner->state_shape, 1, 0, + nullptr, 0, 0); + rc |= frt_runtime_builder_add_port( + b, "actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, owner->action_shape, 2, 0, + nullptr, 0, 0); + if (owner->staged_context_action) { + rc |= frt_runtime_builder_add_generic_stage( + b, "context", FRT_GENERIC_STAGE_OPAQUE, kContextExecutor, + nullptr, 0); + const uint32_t action_after[1] = {0}; + rc |= frt_runtime_builder_add_generic_stage( + b, "action", FRT_GENERIC_STAGE_OPAQUE, kActionExecutor, + action_after, 1); + } else { + rc |= frt_runtime_builder_add_generic_stage( + b, "infer", FRT_GENERIC_STAGE_OPAQUE, kInferExecutor, + nullptr, 0); + } + rc |= frt_runtime_builder_set_generic_stage_runner(b, owner, run_opaque); + rc |= frt_runtime_builder_add_identity(b, "provider", "llama_cpp"); + rc |= frt_runtime_builder_add_identity(b, "model_family", "pi0"); + rc |= frt_runtime_builder_add_identity(b, "model_path", config->model_path); + rc |= frt_runtime_builder_add_identity(b, "mmproj_path", config->mmproj_path); + if (config->struct_size >= sizeof(*config) && config->model_identity && + config->model_identity[0] && config->mmproj_identity && + config->mmproj_identity[0]) { + rc |= frt_runtime_builder_add_identity( + b, "weights_sha256", config->model_identity); + rc |= frt_runtime_builder_add_identity( + b, "mmproj_sha256", config->mmproj_identity); + } + rc |= frt_runtime_builder_add_identity(b, "backend", config->backend); + rc |= frt_runtime_builder_add_identity( + b, "stage_plan", + owner->staged_context_action ? "context_action" : "full"); + if (rc != 0) { + (void)frt_runtime_builder_finish(b, nullptr, nullptr, nullptr); + destroy_owner(owner); + return -1; + } + + frt_model_runtime_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.set_input = set_input; + verbs.get_output = get_output; + verbs.prepare = unsupported_prepare; + verbs.step = step; + verbs.last_error = last_error; + + frt_model_runtime_v1* model = frt_runtime_builder_finish_model( + b, &verbs, owner, owner, nullptr, destroy_owner); + if (!model) { + (void)frt_runtime_builder_finish(b, nullptr, nullptr, nullptr); + destroy_owner(owner); + return -1; + } + *out = model; + return 0; +} + +extern "C" int frt_llama_cpp_pi0_runtime_open_with_engine_factory( + const char* config_json, + const frt_llama_cpp_engine_factory_v1* factory, + frt_model_runtime_v1** out) { + flashrt::providers::llama_cpp::clear_runtime_open_error(); + flashrt::providers::llama_cpp::set_runtime_open_error( + "invalid Pi0 runtime open arguments or JSON config"); + if (!out) return -1; + *out = nullptr; + if (!factory || + factory->struct_size < sizeof(frt_llama_cpp_engine_factory_v1) || + !factory->create_pi0 || !factory->last_error) { + return -1; + } + if (!config_json) { + return -1; + } + + frt_llama_cpp_pi0_config config{}; + config.struct_size = sizeof(config); + std::string model_family; + std::string model_path; + std::string mmproj_path; + std::string backend; + std::string model_identity; + std::string mmproj_identity; + bool seen_model_family = false; + bool seen_model_path = false; + bool seen_mmproj_path = false; + bool seen_backend = false; + bool seen_n_views = false; + bool seen_image_height = false; + bool seen_image_width = false; + bool seen_image_channels = false; + bool seen_action_steps = false; + bool seen_action_dim = false; + const char* p = config_json; + auto skip_ws = [&p]() { + while (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r') ++p; + }; + auto parse_string = [&p](std::string* out) { + if (*p != '"') return false; + ++p; + out->clear(); + while (*p && *p != '"') { + if (*p == '\\') { + ++p; + switch (*p) { + case '"': + case '\\': + case '/': + out->push_back(*p++); + break; + case 'b': + out->push_back('\b'); + ++p; + break; + case 'f': + out->push_back('\f'); + ++p; + break; + case 'n': + out->push_back('\n'); + ++p; + break; + case 'r': + out->push_back('\r'); + ++p; + break; + case 't': + out->push_back('\t'); + ++p; + break; + default: + return false; + } + } else { + const unsigned char ch = static_cast(*p); + if (ch < 0x20) return false; + out->push_back(*p++); + } + } + if (*p != '"') return false; + ++p; + return true; + }; + auto parse_u32 = [&p](uint32_t* out) { + uint32_t value = 0; + if (*p < '0' || *p > '9') return false; + if (*p == '0') { + ++p; + } else { + while (*p >= '0' && *p <= '9') { + const uint32_t digit = static_cast(*p - '0'); + if (value > (UINT32_MAX - digit) / 10u) return false; + value = value * 10u + digit; + ++p; + } + } + if (value == 0) return false; + *out = value; + return true; + }; + skip_ws(); + if (*p != '{') return -1; + ++p; + skip_ws(); + if (*p == '}') return -1; + for (;;) { + std::string key; + if (!parse_string(&key)) return -1; + skip_ws(); + if (*p != ':') return -1; + ++p; + skip_ws(); + if (key == "model_family") { + if (seen_model_family || !parse_string(&model_family) || + model_family.empty()) { + return -1; + } + seen_model_family = true; + } else if (key == "model_path") { + if (seen_model_path || !parse_string(&model_path) || + model_path.empty()) { + return -1; + } + seen_model_path = true; + } else if (key == "mmproj_path") { + if (seen_mmproj_path || !parse_string(&mmproj_path) || + mmproj_path.empty()) { + return -1; + } + seen_mmproj_path = true; + } else if (key == "backend") { + if (seen_backend || !parse_string(&backend) || backend.empty()) { + return -1; + } + seen_backend = true; + } else if (key == "n_views") { + if (seen_n_views || !parse_u32(&config.n_views)) return -1; + seen_n_views = true; + } else if (key == "image_height") { + if (seen_image_height || !parse_u32(&config.image_height)) { + return -1; + } + seen_image_height = true; + } else if (key == "image_width") { + if (seen_image_width || !parse_u32(&config.image_width)) return -1; + seen_image_width = true; + } else if (key == "image_channels") { + if (seen_image_channels || !parse_u32(&config.image_channels)) { + return -1; + } + seen_image_channels = true; + } else if (key == "action_steps") { + if (seen_action_steps || !parse_u32(&config.action_steps)) { + return -1; + } + seen_action_steps = true; + } else if (key == "action_dim") { + if (seen_action_dim || !parse_u32(&config.action_dim)) return -1; + seen_action_dim = true; + } else { + return -1; + } + skip_ws(); + if (*p == '}') { + ++p; + break; + } + if (*p != ',') return -1; + ++p; + skip_ws(); + } + skip_ws(); + if (*p != '\0') return -1; + if (!seen_model_family || model_family != "pi0" || !seen_model_path || + !seen_mmproj_path || !seen_backend || !seen_n_views || + !seen_image_height || !seen_image_width || !seen_image_channels || + !seen_action_steps || !seen_action_dim || config.n_views > 3) { + return -1; + } + config.model_path = model_path.c_str(); + config.mmproj_path = mmproj_path.c_str(); + config.backend = backend.c_str(); + std::string identity_error; + if (!flashrt::providers::llama_cpp::checkpoint_identity( + config.model_path, &model_identity, &identity_error) || + !flashrt::providers::llama_cpp::checkpoint_identity( + config.mmproj_path, &mmproj_identity, &identity_error)) { + flashrt::providers::llama_cpp::set_runtime_open_error(identity_error); + return -1; + } + config.model_identity = model_identity.c_str(); + config.mmproj_identity = mmproj_identity.c_str(); + + frt_llama_cpp_engine_v1 engine{}; + const int rc = factory->create_pi0(factory->self, &config, &engine); + if (rc != 0) { + const char* error = factory->last_error(factory->self); + flashrt::providers::llama_cpp::set_runtime_open_error( + error ? error : "Pi0 engine factory failed without an error"); + return rc; + } + if (engine.struct_size < FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE || + !engine.self || !engine.retain || !engine.release || + !engine.set_input || !engine.run_infer || !engine.get_output || + !engine.last_error) { + if (engine.struct_size >= + offsetof(frt_llama_cpp_engine_v1, release) + sizeof(engine.release) && + engine.self && engine.release) { + engine.release(engine.self); + } + flashrt::providers::llama_cpp::set_runtime_open_error( + "factory returned an invalid Pi0 engine"); + return -1; + } + frt_model_runtime_v1* model = nullptr; + const int create_rc = + frt_llama_cpp_pi0_runtime_create_with_engine(&config, &engine, + &model); + engine.release(engine.self); + if (create_rc != 0) { + flashrt::providers::llama_cpp::set_runtime_open_error( + "failed to create Pi0 model runtime"); + return create_rc; + } + *out = model; + flashrt::providers::llama_cpp::clear_runtime_open_error(); + return 0; +} diff --git a/cpp/providers/llama_cpp/src/provider_open.cpp b/cpp/providers/llama_cpp/src/provider_open.cpp new file mode 100644 index 00000000..42aab5af --- /dev/null +++ b/cpp/providers/llama_cpp/src/provider_open.cpp @@ -0,0 +1,36 @@ +#include "flashrt/providers/llama_cpp/c_api.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "checkpoint_identity.h" + +#if defined(_WIN32) +#define FRT_PROVIDER_EXPORT __declspec(dllexport) +#else +#define FRT_PROVIDER_EXPORT __attribute__((visibility("default"))) +#endif + +extern "C" FRT_PROVIDER_EXPORT int frt_model_runtime_open_v1( + const char* config_json, frt_model_runtime_v1** out) { + if (!out) return -1; + *out = nullptr; + if (!config_json) return -1; + + const frt_llama_cpp_engine_factory_v1* factory = + frt_llama_cpp_default_engine_factory(); + if (!factory) return -3; + + int rc = frt_llama_cpp_pi0_runtime_open_with_engine_factory( + config_json, factory, out); + if (rc == 0) return 0; + rc = frt_llama_cpp_llm_runtime_open_with_engine_factory( + config_json, factory, out); + if (rc == 0) return 0; + rc = frt_llama_cpp_mllm_runtime_open_with_engine_factory( + config_json, factory, out); + if (rc == 0) return 0; + + flashrt::providers::llama_cpp::set_runtime_open_error( + "configuration does not describe a supported provider model family"); + return rc; +} + +#undef FRT_PROVIDER_EXPORT diff --git a/cpp/tests/check_llama_cpp_provider_exports.cmake b/cpp/tests/check_llama_cpp_provider_exports.cmake new file mode 100644 index 00000000..eb99e0d8 --- /dev/null +++ b/cpp/tests/check_llama_cpp_provider_exports.cmake @@ -0,0 +1,25 @@ +if(NOT DEFINED PROVIDER_DSO OR NOT DEFINED CMAKE_NM_TOOL) + message(FATAL_ERROR "PROVIDER_DSO and CMAKE_NM_TOOL are required") +endif() + +execute_process( + COMMAND "${CMAKE_NM_TOOL}" -D --defined-only "${PROVIDER_DSO}" + RESULT_VARIABLE nm_result + OUTPUT_VARIABLE nm_output + ERROR_VARIABLE nm_error) +if(NOT nm_result EQUAL 0) + message(FATAL_ERROR "dynamic symbol audit failed: ${nm_error}") +endif() + +string(REGEX MATCHALL "[^\n]+" symbol_lines "${nm_output}") +set(found_open FALSE) +foreach(line IN LISTS symbol_lines) + if(line MATCHES "[ ]frt_model_runtime_open_v1(@@[^ ]+)?$") + set(found_open TRUE) + elseif(NOT line MATCHES "[ ]FLASHRT_LLAMA_CPP_PROVIDER_[^ ]+$") + message(FATAL_ERROR "unexpected provider export: ${line}") + endif() +endforeach() +if(NOT found_open) + message(FATAL_ERROR "frt_model_runtime_open_v1 is not exported") +endif() diff --git a/cpp/tests/fixtures/prepare_pi0_fixture.py b/cpp/tests/fixtures/prepare_pi0_fixture.py new file mode 100644 index 00000000..8f93c120 --- /dev/null +++ b/cpp/tests/fixtures/prepare_pi0_fixture.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Convert a LIBERO Pi0 .npz fixture into the PNG/bin/txt files the C++ engine +test consumes. + +Output files (in --out-dir): + image.png primary camera, 224x224x3 uint8 + wrist_image.png wrist camera, 224x224x3 uint8 + state.bin action_dim float32 (real state zero-padded to action_dim) + prompt.txt task prompt, trailing newline trimmed + +The .npz key names are probed defensively; pass --image-key / --wrist-key / +--state-key to override. +""" +import argparse +import os +import sys + +import numpy as np + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--npz", required=True) + ap.add_argument("--prompt", required=True) + ap.add_argument("--out-dir", required=True) + ap.add_argument("--action-dim", type=int, default=32) + ap.add_argument("--image-key", default=None) + ap.add_argument("--wrist-key", default=None) + ap.add_argument("--state-key", default=None) + args = ap.parse_args() + + try: + from PIL import Image + except ImportError: + sys.exit("Pillow is required: pip install pillow") + + d = np.load(args.npz) + keys = list(d.files) + print(f"npz keys: {keys}") + + def pick(candidates, override): + if override: + if override not in keys: + sys.exit(f"override key {override!r} not in npz; have {keys}") + return d[override] + for c in candidates: + if c in keys: + return d[c] + sys.exit(f"none of {candidates} found in npz; have {keys}") + + img = pick(["image", "agentview_image", "observation.images.image"], + args.image_key) + wrist = pick(["wrist_image", "wristview_image", + "observation.images.wrist_image"], args.wrist_key) + state = pick(["state", "robot_state", "observation.state", "qpos"], + args.state_key) + + def to_hwc_uint8(a): + a = np.asarray(a) + if a.ndim == 3 and a.shape[0] in (3, 4) and a.shape[-1] not in (3, 4): + a = np.transpose(a, (1, 2, 0)) # CHW -> HWC + a = a.astype(np.uint8, copy=False) + if a.ndim == 2: # grayscale -> RGB + a = np.stack([a, a, a], axis=-1) + if a.shape[-1] == 4: + a = a[..., :3] + return a + + img = to_hwc_uint8(img) + wrist = to_hwc_uint8(wrist) + print(f"image {img.shape} {img.dtype}; wrist {wrist.shape} {wrist.dtype}; " + f"state {np.asarray(state).shape} {np.asarray(state).dtype}") + + os.makedirs(args.out_dir, exist_ok=True) + Image.fromarray(img).save(os.path.join(args.out_dir, "image.png")) + Image.fromarray(wrist).save(os.path.join(args.out_dir, "wrist_image.png")) + + state = np.asarray(state).astype(np.float32).reshape(-1) + if state.size > args.action_dim: + sys.exit(f"state has {state.size} values, more than action_dim=" + f"{args.action_dim}") + padded = np.zeros(args.action_dim, dtype=np.float32) + padded[:state.size] = state + padded.tofile(os.path.join(args.out_dir, "state.bin")) + + with open(args.prompt) as f: + prompt = f.read().rstrip("\n") + with open(os.path.join(args.out_dir, "prompt.txt"), "w") as f: + f.write(prompt) + + print(f"wrote image.png, wrist_image.png, state.bin ({padded.size} f32), " + f"prompt.txt to {args.out_dir}") + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/llama_cpp_generic_plan.h b/cpp/tests/llama_cpp_generic_plan.h new file mode 100644 index 00000000..5ba42950 --- /dev/null +++ b/cpp/tests/llama_cpp_generic_plan.h @@ -0,0 +1,38 @@ +#ifndef FLASHRT_TESTS_LLAMA_CPP_GENERIC_PLAN_H +#define FLASHRT_TESTS_LLAMA_CPP_GENERIC_PLAN_H + +#include "flashrt/model_runtime.h" + +#include + +inline const frt_generic_stage_plan_ext_v1* llama_cpp_generic_plan( + const frt_model_runtime_v1* model) { + if (!model || + model->struct_size < FRT_MODEL_RUNTIME_V1_QUERY_EXTENSION_SIZE || + !model->query_extension) { + return nullptr; + } + const void* extension = nullptr; + if (model->query_extension( + model, FRT_EXT_GENERIC_STAGE_PLAN_V1, + FRT_GENERIC_STAGE_PLAN_ABI_VERSION, &extension) != 0) { + return nullptr; + } + return static_cast(extension); +} + +inline int llama_cpp_run_generic_stage( + const frt_model_runtime_v1* model, const char* name) { + const auto* plan = llama_cpp_generic_plan(model); + if (!plan || !name || !plan->run_opaque) return -3; + for (uint64_t i = 0; i < plan->n_stages; ++i) { + const auto& stage = plan->stages[i]; + if (stage.name && std::strcmp(stage.name, name) == 0) { + if (stage.executor_kind != FRT_GENERIC_STAGE_OPAQUE) return -3; + return plan->run_opaque(plan->stage_self, stage.executor_ref); + } + } + return -3; +} + +#endif diff --git a/cpp/tests/test_llama_cpp_jetson_pi_engine.cpp b/cpp/tests/test_llama_cpp_jetson_pi_engine.cpp new file mode 100644 index 00000000..d8550916 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_engine.cpp @@ -0,0 +1,308 @@ +// End-to-end Pi0 engine test: drives the real Jetson-PI engine factory with +// Pi0 weights (when available) and asserts a sane action chunk. Skips (returns +// 0) when the weights/fixture env vars are unset, so CI without weights still +// passes. +// +// Coverage: +// A. bogus model path fails cleanly (no-weights contract) +// B. end-to-end single Pi0 tick: set_input images/prompt/state -> run_stage +// infer -> get_output actions; asserts shape, no NaN/Inf, non-zero. +// C. config vs model action_dim mismatch: open must reject (engine freed). +// D. multi-tick: a second infer with fresh inputs must (1) return -7 from +// get_output after set_input but before run_infer (staleness guard), +// and (2) reproduce the first tick's actions (KV reset, no leak). +// +// Env: +// FLASHRT_PI0_MODEL path to Pi0 policy GGUF +// FLASHRT_PI0_MMPROJ path to VIT mmproj GGUF +// FLASHRT_PI0_FIXTURE_DIR dir containing image.png, wrist_image.png, +// state.bin (action_dim float32), prompt.txt +// FLASHRT_PI0_ACTION_STEPS required model-specific action horizon. +// FLASHRT_PI0_ACTION_DIM required model-specific action width. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#define STB_IMAGE_IMPLEMENTATION +#define STBI_ONLY_PNG +#include "stb/stb_image.h" + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +namespace { + +bool file_exists(const char * p) { + std::ifstream f(p); + return f.good(); +} + +std::string read_file(const std::string & path, bool * ok) { + std::ifstream f(path, std::ios::binary); + if (!f) { if (ok) *ok = false; return {}; } + std::string s((std::istreambuf_iterator(f)), + std::istreambuf_iterator()); + if (ok) *ok = true; + return s; +} + +} // namespace + +int main() { + const char * model_env = std::getenv("FLASHRT_PI0_MODEL"); + const char * mmproj_env = std::getenv("FLASHRT_PI0_MMPROJ"); + const char * fixture_env = std::getenv("FLASHRT_PI0_FIXTURE_DIR"); + if (!model_env || !mmproj_env || !fixture_env || + !file_exists(model_env) || !file_exists(mmproj_env)) { + std::printf("SKIP - FLASHRT_PI0_MODEL / FLASHRT_PI0_MMPROJ / " + "FLASHRT_PI0_FIXTURE_DIR not set or files missing\n"); + return 0; + } + + const std::string fixture_dir = fixture_env; + const std::string img_path = fixture_dir + "/image.png"; + const std::string wrist_path = fixture_dir + "/wrist_image.png"; + const std::string state_path = fixture_dir + "/state.bin"; + const std::string prompt_path = fixture_dir + "/prompt.txt"; + if (!file_exists(img_path.c_str()) || + !file_exists(wrist_path.c_str()) || + !file_exists(state_path.c_str()) || + !file_exists(prompt_path.c_str())) { + std::printf("SKIP - fixture files missing in %s\n", fixture_dir.c_str()); + return 0; + } + + // ---- sub-test A: bogus model path fails (no-weights contract) ---------- + const frt_llama_cpp_engine_factory_v1 * factory = + frt_llama_cpp_default_engine_factory(); + CHECK(factory != nullptr, "default engine factory is non-null"); + CHECK(factory->create_pi0 != nullptr && factory->last_error != nullptr, + "factory vtable complete"); + + const char * bogus_json = + "{\"model_family\":\"pi0\"," + "\"model_path\":\"/nonexistent/bogus.gguf\"," + "\"mmproj_path\":\"/nonexistent/bogus-mmproj.gguf\"," + "\"backend\":\"cpu\",\"n_views\":2,\"image_height\":224," + "\"image_width\":224,\"image_channels\":3," + "\"action_steps\":10,\"action_dim\":32}"; + frt_model_runtime_v1 * bogus = nullptr; + int rc = frt_llama_cpp_pi0_runtime_open_with_engine_factory( + bogus_json, factory, &bogus); + CHECK(rc != 0 && bogus == nullptr, + "open with bogus model path fails without crashing"); + + // ---- sub-test B: end-to-end Pi0 tick ----------------------------------- + // Action dims come from env (verified pi0_base is 50x32; other checkpoints + // may differ and must provide their model-specific shape explicitly). + const char * steps_env = std::getenv("FLASHRT_PI0_ACTION_STEPS"); + const char * dim_env = std::getenv("FLASHRT_PI0_ACTION_DIM"); + if (!steps_env || !dim_env) { + std::printf("SKIP - FLASHRT_PI0_ACTION_STEPS/ACTION_DIM not set\n"); + return 0; + } + // Backend is "cpu" by default (byte-identical to the original test). Set + // FLASHRT_PI0_BACKEND=cuda to run the real forward pass on the GPU (the + // Jetson-PI engine maps backend=="cuda" to n_gpu_layers=9999 and use_gpu + // for the mmproj). CUDA_VISIBLE_DEVICES selects the physical card. + // Applied to the real-Pi0 JSON only (sub-test B); the bogus-path and + // action_dim-mismatch sub-tests keep "cpu" hardcoded so they stay cheap + // and deterministic. + const char * be_env = std::getenv("FLASHRT_PI0_BACKEND"); + const std::string backend = (be_env && be_env[0]) ? std::string(be_env) + : std::string("cpu"); + long action_steps = std::atol(steps_env); + long action_dim = std::atol(dim_env); + if (action_steps <= 0 || action_dim <= 0 || action_steps > 10000 || + action_dim > 10000) { + std::printf("SKIP - bad FLASHRT_PI0_ACTION_STEPS/DIM\n"); + return 0; + } + + // ---- sub-test C: config vs model action_dim mismatch fails cleanly ---- + // The selected model uses the configured shape; claim a wrong action_dim + // and expect create_pi0 to reject without leaking the opened engine. + // Backend is hardcoded "cpu" here on purpose: this path only exercises the + // action_dim rejection, so paying for a full GPU model load + 37-layer + // offload (which FLASHRT_PI0_BACKEND=cuda would trigger) is wasteful and + // nondeterministic. The env-driven backend below is for the real forward + // pass only. + { + std::string mismatch_json = + std::string("{") + + "\"model_family\":\"pi0\"," + "\"model_path\":\"" + model_env + "\"," + "\"mmproj_path\":\"" + mmproj_env + "\"," + "\"backend\":\"cpu\"," + "\"n_views\":2,\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":" + + std::to_string(action_steps) + + ",\"action_dim\":" + std::to_string(action_dim + 1) + "}"; + frt_model_runtime_v1 * m = nullptr; + int mrc = frt_llama_cpp_pi0_runtime_open_with_engine_factory( + mismatch_json.c_str(), factory, &m); + CHECK(mrc != 0 && m == nullptr, + "open rejects action_dim mismatch with model"); + } + + std::string json = + std::string("{") + + "\"model_family\":\"pi0\"," + "\"model_path\":\"" + model_env + "\"," + "\"mmproj_path\":\"" + mmproj_env + "\"," + "\"backend\":\"" + backend + "\"," + "\"n_views\":2,\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":" + std::to_string(action_steps) + + ",\"action_dim\":" + std::to_string(action_dim) + "}"; + frt_model_runtime_v1 * model = nullptr; + rc = frt_llama_cpp_pi0_runtime_open_with_engine_factory( + json.c_str(), factory, &model); + if (rc != 0 || !model) { + std::printf("FAIL: open real Pi0 runtime (rc=%d): %s\n", rc, + factory->last_error(factory->self)); + return 1; + } + CHECK(rc == 0 && model != nullptr, "open real Pi0 runtime from JSON"); + + // Load fixture. + int iw = 0, ih = 0, ic = 0; + unsigned char * img = stbi_load(img_path.c_str(), &iw, &ih, &ic, 3); + CHECK(img != nullptr && iw == 224 && ih == 224, "load image.png 224x224"); + int ww = 0, wh = 0, wc = 0; + unsigned char * wrist = stbi_load(wrist_path.c_str(), &ww, &wh, &wc, 3); + CHECK(wrist != nullptr && ww == 224 && wh == 224, "load wrist_image.png 224x224"); + bool state_ok = false; + std::string state_bytes = read_file(state_path, &state_ok); + CHECK(state_ok && state_bytes.size() == + static_cast(action_dim) * sizeof(float), + "state.bin matches action_dim float32"); + bool prompt_ok = false; + std::string prompt = read_file(prompt_path, &prompt_ok); + CHECK(prompt_ok && !prompt.empty(), "prompt.txt non-empty"); + if (!prompt.empty() && prompt.back() == '\n') prompt.pop_back(); + + if (img && wrist && state_ok && prompt_ok) { + frt_image_view views[2]; + views[0].struct_size = sizeof(frt_image_view); + views[0].pixel_format = FRT_RT_PIXEL_RGB8; + views[0].data = img; + views[0].bytes = static_cast(224) * 224 * 3; + views[0].width = 224; views[0].height = 224; views[0].stride_bytes = 0; + views[0].reserved = 0; views[0].timestamp_ns = 0; + views[1] = views[0]; + views[1].data = wrist; + + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, + &views[0], sizeof(views), -1) == 0, + "set_input images"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "set_input prompt"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), + state_bytes.size(), -1) == 0, + "set_input state"); + + rc = model->verbs.step(model->self); + if (rc != 0) { + std::printf("FAIL: run_stage infer (rc=%d): %s\n", rc, + model->verbs.last_error(model->self)); + g_fail = 1; + } else { + std::printf("ok : run_stage infer\n"); + } + + std::vector actions( + static_cast(action_steps) * action_dim); + uint64_t written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_PI0_PORT_ACTIONS, + actions.data(), actions.size() * sizeof(float), &written, -1); + CHECK(rc == 0 && + written == static_cast(action_steps) * action_dim * + sizeof(float), + "get_output actions shape matches config"); + if (rc == 0) { + bool nan_inf = false, all_zero = true; + for (float v : actions) { + if (std::isnan(v) || std::isinf(v)) { nan_inf = true; break; } + if (v != 0.0f) all_zero = false; + } + CHECK(!nan_inf, "actions contain no NaN/Inf"); + CHECK(!all_zero, "actions are not all zero"); + } + + // ---- multi-tick: a second infer with fresh inputs must not leak the + // first tick's KV (KV reset) nor return the first tick's actions + // (staleness guard). We reuse the same inputs; the action should + // reproduce within tolerance, and a get_output before run_infer (after + // set_input) must return -7 (actions not ready). ---- + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, + &views[0], sizeof(views), -1) == 0, + "tick2 set_input images"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "tick2 set_input prompt"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), + state_bytes.size(), -1) == 0, + "tick2 set_input state"); + // set_input must have invalidated actions_buf: get_output now -7. + { + uint64_t w2 = 0; + int rc2 = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_PI0_PORT_ACTIONS, + actions.data(), actions.size() * sizeof(float), &w2, -1); + CHECK(rc2 != 0, + "get_output after set_input (before run_infer) fails"); + } + CHECK(model->verbs.step(model->self) == 0, + "tick2 run_stage infer"); + std::vector actions2( + static_cast(action_steps) * action_dim); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_PI0_PORT_ACTIONS, + actions2.data(), actions2.size() * sizeof(float), &written, -1); + CHECK(rc == 0 && + written == static_cast(action_steps) * action_dim * + sizeof(float), + "tick2 get_output actions shape matches config"); + if (rc == 0) { + // Same inputs -> deterministically reproducible action chunk + // (Pi0 with fixed seed noise). Verify the two ticks match. + bool match = actions.size() == actions2.size(); + for (size_t i = 0; match && i < actions.size(); ++i) { + if (std::fabs(actions[i] - actions2[i]) > 1e-5f) match = false; + } + CHECK(match, "tick2 actions reproduce tick1 (KV reset, no leak)"); + } + } + + model->release(model->owner); + if (img) stbi_image_free(img); + if (wrist) stbi_image_free(wrist); + + std::printf(g_fail ? "\n== JETSON_PI ENGINE FAILED ==\n" + : "\n== JETSON_PI ENGINE PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_link.cpp b/cpp/tests/test_llama_cpp_jetson_pi_link.cpp new file mode 100644 index 00000000..8e6b195d --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_link.cpp @@ -0,0 +1,15 @@ +#include "flashrt/providers/llama_cpp/c_api.h" + +#include + +extern "C" int frt_llama_cpp_jetson_pi_link_check(void); + +int main() { + if (frt_llama_cpp_jetson_pi_link_check() != 0) { + std::printf("FAIL: FlashRT llama_cpp provider Jetson-PI link check\n"); + return 1; + } + + std::printf("ok : FlashRT llama_cpp provider links Jetson-PI llama/mtmd APIs\n"); + return 0; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_llm.cpp b/cpp/tests/test_llama_cpp_jetson_pi_llm.cpp new file mode 100644 index 00000000..ed00e2a8 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_llm.cpp @@ -0,0 +1,377 @@ +// End-to-end LLM engine test: drives the real Jetson-PI LLM engine factory +// with a GGUF LLM (when available) and asserts a sane text completion. +// Skips (returns 0) when FLASHRT_LLM_MODEL is unset. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +int main() { + const char * model_env = std::getenv("FLASHRT_LLM_MODEL"); + if (!model_env || !*model_env || !std::fopen(model_env, "rb")) { + std::printf("SKIP - FLASHRT_LLM_MODEL not set or file missing\n"); + return 0; + } + const char * backend_env = std::getenv("FLASHRT_LLM_BACKEND"); + const std::string backend = backend_env && *backend_env ? backend_env : "cpu"; + + const frt_llama_cpp_engine_factory_v1 * factory = + frt_llama_cpp_default_engine_factory(); + CHECK(factory != nullptr && factory->create_llm != nullptr, + "default engine factory exposes create_llm"); + + // Bogus model path must fail cleanly. + const char * bogus_json = + "{\"model_family\":\"llm\"," + "\"model_path\":\"/nonexistent/bogus.gguf\"," + "\"backend\":\"cpu\"," + "\"n_ctx\":2048,\"n_threads\":0," + "\"temp\":0.8,\"top_k\":40,\"top_p\":0.9,\"seed\":1,\"max_tokens\":64}"; + frt_model_runtime_v1 * bogus = nullptr; + CHECK(frt_llama_cpp_llm_runtime_open_with_engine_factory( + bogus_json, factory, &bogus) != 0 && bogus == nullptr, + "open LLM with bogus model path fails without crashing"); + CHECK(std::strstr(frt_llama_cpp_runtime_open_error(), + "failed to open checkpoint for identity"), + "bogus model path reports checkpoint identity error"); + + // Real model. + const char * prompt = "What is 2 plus 2? The answer is"; + std::string json = + std::string("{") + + "\"model_family\":\"llm\"," + "\"model_path\":\"" + model_env + "\"," + "\"backend\":\"" + backend + "\"," + "\"n_ctx\":2048,\"n_threads\":0," + "\"temp\":0.0,\"top_k\":0,\"top_p\":0.0,\"seed\":1,\"max_tokens\":16}"; + frt_model_runtime_v1 * model = nullptr; + int rc = frt_llama_cpp_llm_runtime_open_with_engine_factory( + json.c_str(), factory, &model); + if (rc != 0 || !model) { + std::printf("FAIL: open LLM runtime (rc=%d): %s\n", rc, + factory->last_error(factory->self)); + return 1; + } + CHECK(rc == 0 && model != nullptr, "open LLM runtime from JSON"); + const auto* plan = llama_cpp_generic_plan(model); + std::printf(" runtime stages (%llu):", + static_cast(plan ? plan->n_stages : 0)); + for (uint64_t i = 0; plan && i < plan->n_stages; ++i) { + std::printf(" %s", plan->stages[i].name); + } + std::printf("\n"); + CHECK(model->n_stages == 0 && plan && plan->n_stages == 3, + "LLM runtime exposes selected reset/prefill/decode plan"); + CHECK(model->n_ports == 6, + "LLM runtime exposes prompt/text/token/logits/eog/tokens ports"); + CHECK(model->exp && std::strstr(model->exp->identity, + "weights_sha256="), + "deployment identity includes checkpoint content digest"); + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_LOGITS, + nullptr, 0, nullptr, -1) == -1, + "get_output rejects null written without crashing"); + + rc = model->verbs.set_input(model->self, + FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1); + CHECK(rc == 0, "set_input prompt"); + + rc = model->verbs.step(model->self); + if (rc != 0) { + std::printf("FAIL: run_stage infer (rc=%d): %s\n", rc, + model->verbs.last_error(model->self)); + g_fail = 1; + } else { + std::printf("ok : run_stage infer\n"); + } + + uint64_t written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, nullptr, 0, &written, -1); + CHECK(rc == -5 && written > 0, "query one-shot text size"); + std::string text(written, '\0'); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, &text[0], text.size(), + &written, -1); + CHECK(rc == 0 && written > 0, "get_output text non-empty"); + if (rc == 0) { + text.resize(written); + std::printf(" generated: %s\n", text.c_str()); + // Sanity: contains at least one printable character. + bool printable = false; + for (char c : text) { + if (c >= 0x20 && c < 0x7f) { printable = true; break; } + } + CHECK(printable, "generated text contains printable chars"); + } + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + nullptr, 0, -1) != 0, + "invalid prompt replacement is rejected"); + CHECK(model->verbs.step(model->self) != 0, + "failed prompt replacement invalidates the previous prompt"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0, + "restore valid prompt after failed replacement"); + + CHECK(llama_cpp_run_generic_stage(model, "reset") == 0, + "run_stage reset"); + const auto prefill_start = std::chrono::steady_clock::now(); + rc = llama_cpp_run_generic_stage(model, "prefill"); + const auto prefill_end = std::chrono::steady_clock::now(); + CHECK(rc == 0, "run_stage prefill"); + std::printf(" staged prefill latency: %.3f ms\n", + std::chrono::duration( + prefill_end - prefill_start).count()); + uint64_t logits_bytes = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_LOGITS, nullptr, 0, + &logits_bytes, -1); + CHECK(rc != 0 && logits_bytes > 1000 * sizeof(float), + "prefill exposes logits size"); + std::vector logits(logits_bytes / sizeof(float)); + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_LOGITS, logits.data(), + logits_bytes, &logits_bytes, -1); + CHECK(rc == 0, "prefill get_output logits"); + bool finite_logits = true; + for (float value : logits) finite_logits &= std::isfinite(value); + CHECK(finite_logits, "prefill logits contain no NaN/Inf"); + + std::string staged_text; + std::vector baseline_tokens; + std::vector baseline_eog; + double decode_ms = 0.0; + for (uint32_t i = 0; i < 16; ++i) { + const auto decode_start = std::chrono::steady_clock::now(); + rc = llama_cpp_run_generic_stage(model, "decode"); + const auto decode_end = std::chrono::steady_clock::now(); + decode_ms += std::chrono::duration( + decode_end - decode_start).count(); + CHECK(rc == 0, "run_stage decode"); + int32_t token = 0; + int32_t is_eog = 0; + uint64_t scalar_bytes = 0; + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, &token, + sizeof(token), &scalar_bytes, -1) == 0 && + scalar_bytes == sizeof(token), + "decode exposes next_token"); + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_IS_EOG, &is_eog, + sizeof(is_eog), &scalar_bytes, -1) == 0 && + scalar_bytes == sizeof(is_eog), + "decode exposes is_eog"); + baseline_tokens.push_back(token); + baseline_eog.push_back(is_eog); + if (is_eog) break; + } + std::printf(" staged decode throughput: %.2f token/s " + "(%zu tokens in %.3f ms)\n", + baseline_tokens.size() * 1000.0 / decode_ms, + baseline_tokens.size(), decode_ms); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, nullptr, 0, &written, -1); + CHECK(rc == -5 && written > 0, "query staged text size"); + staged_text.assign(written, '\0'); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, staged_text.data(), + staged_text.size(), &written, -1); + CHECK(rc == 0 && written > 0, "staged decode exposes accumulated text"); + staged_text.resize(written); + CHECK(staged_text == text, "staged decode text matches one-shot infer"); + rc = model->verbs.step(model->self); + CHECK(rc == 0, "one-shot infer after staged decode"); + uint64_t stale_written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, nullptr, 0, + &stale_written, -1); + CHECK(rc == -7, "one-shot infer invalidates staged token output"); + + frt_model_runtime_v1* peer = nullptr; + rc = frt_llama_cpp_llm_runtime_open_with_engine_factory( + json.c_str(), factory, &peer); + CHECK(rc == 0 && peer, "open independent peer LLM session"); + if (peer) { + const char* peer_prompt = + "Complete this sequence with one number: 3, 6, 9,"; + CHECK(peer->verbs.set_input( + peer->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + peer_prompt, std::strlen(peer_prompt), -1) == 0, + "peer set_input distinct prompt"); + CHECK(llama_cpp_run_generic_stage(peer, "reset") == 0 && + llama_cpp_run_generic_stage(peer, "prefill") == 0, + "prepare peer standalone baseline"); + std::vector peer_baseline_tokens; + std::vector peer_baseline_eog; + for (uint32_t i = 0; i < 16; ++i) { + CHECK(llama_cpp_run_generic_stage(peer, "decode") == 0, + "decode peer standalone baseline"); + int32_t token = 0; + int32_t is_eog = 0; + uint64_t scalar_bytes = 0; + CHECK(peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, &token, + sizeof(token), &scalar_bytes, -1) == 0 && + scalar_bytes == sizeof(token) && + peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_IS_EOG, &is_eog, + sizeof(is_eog), &scalar_bytes, -1) == 0, + "capture peer standalone token sequence"); + peer_baseline_tokens.push_back(token); + peer_baseline_eog.push_back(is_eog); + if (is_eog) break; + } + uint64_t peer_baseline_written = 0; + CHECK(peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, + nullptr, 0, &peer_baseline_written, -1) == -5 && + peer_baseline_written > 0, + "query peer standalone baseline text size"); + std::string peer_baseline_text(peer_baseline_written, '\0'); + peer_baseline_written = 0; + CHECK(peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, + peer_baseline_text.data(), peer_baseline_text.size(), + &peer_baseline_written, -1) == 0, + "read peer standalone baseline text"); + peer_baseline_text.resize(peer_baseline_written); + CHECK(llama_cpp_run_generic_stage(model, "reset") == 0 && + llama_cpp_run_generic_stage(peer, "reset") == 0, + "reset two independent sessions"); + CHECK(llama_cpp_run_generic_stage(model, "prefill") == 0 && + llama_cpp_run_generic_stage(peer, "prefill") == 0, + "prefill two independent sessions"); + const size_t interleaved_steps = + std::max(baseline_tokens.size(), peer_baseline_tokens.size()); + for (size_t i = 0; i < interleaved_steps; ++i) { + uint64_t scalar_bytes = 0; + if (i < baseline_tokens.size()) { + int32_t token = 0; + int32_t is_eog = 0; + CHECK(llama_cpp_run_generic_stage(model, "decode") == 0 && + model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &token, sizeof(token), &scalar_bytes, -1) == 0 && + model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_IS_EOG, + &is_eog, sizeof(is_eog), &scalar_bytes, -1) == 0 && + token == baseline_tokens[i] && + is_eog == baseline_eog[i], + "interleaved primary token matches standalone baseline"); + } + if (i < peer_baseline_tokens.size()) { + int32_t token = 0; + int32_t is_eog = 0; + CHECK(llama_cpp_run_generic_stage(peer, "decode") == 0 && + peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &token, sizeof(token), &scalar_bytes, -1) == 0 && + peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_IS_EOG, + &is_eog, sizeof(is_eog), &scalar_bytes, -1) == 0 && + token == peer_baseline_tokens[i] && + is_eog == peer_baseline_eog[i], + "interleaved peer token matches standalone baseline"); + } + } + uint64_t model_written = 0; + uint64_t peer_written = 0; + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, + nullptr, 0, &model_written, -1) == -5 && + model_written > 0 && + peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, + nullptr, 0, &peer_written, -1) == -5 && + peer_written > 0, + "query interleaved session output sizes"); + std::string model_text(model_written, '\0'); + std::string peer_text(peer_written, '\0'); + model_written = 0; + peer_written = 0; + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, model_text.data(), + model_text.size(), &model_written, -1) == 0 && + peer->verbs.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, peer_text.data(), + peer_text.size(), &peer_written, -1) == 0, + "read interleaved session outputs"); + model_text.resize(model_written); + peer_text.resize(peer_written); + CHECK(model_text == staged_text && peer_text == peer_baseline_text, + "distinct sessions reproduce baselines without KV leakage"); + peer->release(peer->owner); + } + + model->release(model->owner); + + std::string budget_json = + std::string("{") + + "\"model_family\":\"llm\"," + + "\"model_path\":\"" + model_env + "\"," + + "\"backend\":\"" + backend + "\"," + + "\"n_ctx\":2048,\"n_threads\":0," + "\"temp\":0.0,\"top_k\":0,\"top_p\":0.0,\"seed\":1," + "\"max_tokens\":1}"; + frt_model_runtime_v1 * budget_model = nullptr; + rc = frt_llama_cpp_llm_runtime_open_with_engine_factory( + budget_json.c_str(), factory, &budget_model); + CHECK(rc == 0 && budget_model, "open max_tokens=1 LLM runtime"); + if (budget_model) { + CHECK(budget_model->verbs.set_input( + budget_model->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0 && + llama_cpp_run_generic_stage(budget_model, "prefill") == 0 && + llama_cpp_run_generic_stage(budget_model, "decode") == 0, + "max_tokens=1 session permits exactly one decode"); + int32_t first_is_eog = 1; + uint64_t scalar_written = 0; + CHECK(budget_model->verbs.get_output( + budget_model->self, FRT_LLAMA_CPP_LLM_PORT_IS_EOG, + &first_is_eog, sizeof(first_is_eog), &scalar_written, -1) == 0 && + first_is_eog == 0, + "budget test prompt first token is not EOG"); + CHECK(llama_cpp_run_generic_stage(budget_model, "decode") != 0 && + std::strstr(budget_model->verbs.last_error( + budget_model->self), + "max_tokens"), + "staged decode rejects calls beyond max_tokens"); + int32_t stale_scalar = 0; + CHECK(budget_model->verbs.get_output( + budget_model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &stale_scalar, sizeof(stale_scalar), &scalar_written, -1) == -7 && + budget_model->verbs.get_output( + budget_model->self, FRT_LLAMA_CPP_LLM_PORT_IS_EOG, + &stale_scalar, sizeof(stale_scalar), &scalar_written, -1) == -7, + "failed decode invalidates staged scalar outputs"); + CHECK(llama_cpp_run_generic_stage(budget_model, "prefill") == 0 && + llama_cpp_run_generic_stage(budget_model, "decode") == 0, + "prefill restores staged decode budget"); + budget_model->release(budget_model->owner); + } + + std::printf(g_fail ? "\n== JETSON_PI LLM FAILED ==\n" + : "\n== JETSON_PI LLM PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_llm_parity.cpp b/cpp/tests/test_llama_cpp_jetson_pi_llm_parity.cpp new file mode 100644 index 00000000..b64a8e3e --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_llm_parity.cpp @@ -0,0 +1,243 @@ +// Numerical-parity test: FlashRT's Jetson-PI LLM provider (via +// frt_model_runtime_v1) vs a DIRECT jetson_pi_llm call, same prompt / sampling +// / backend. Proves the FlashRT port/stage/verb plumbing does not perturb the +// generated text relative to the native narrow C API. +// +// Closes jetsonpi迁移.txt §14 (parity vs native). Skips (returns 0) when +// FLASHRT_LLM_MODEL is unset. +// +// Determinism foundation: greedy sampling (temp<=0) builds a +// llama_sampler_init_greedy()-only chain (argmax, no RNG) — fully +// deterministic regardless of seed. KV is cleared per generate. So both paths +// must produce the EXACT same token sequence / text. (temp>0 uses an RNG +// sampler and is not deterministic across handles; parity uses greedy only.) +// +// Env: +// FLASHRT_LLM_MODEL path to a GGUF LLM +// FLASHRT_LLM_BACKEND (optional) "cpu" (default) or "cuda". +// +// Note: both paths load the model independently, so a CPU run pays two full +// loads (several minutes for a 0.6B model); CUDA is the practical default +// when a GPU is available. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "jetson_pi_llm.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +int main() { + const char * model_env = std::getenv("FLASHRT_LLM_MODEL"); + bool model_ok = false; + if (model_env && *model_env) { + std::ifstream f(model_env, std::ios::binary); + model_ok = f.good(); + } + if (!model_ok) { + std::printf("SKIP - FLASHRT_LLM_MODEL not set or file missing\n"); + return 0; + } + const char * be_env = std::getenv("FLASHRT_LLM_BACKEND"); + const std::string backend = (be_env && be_env[0]) ? std::string(be_env) + : std::string("cpu"); + + // Shared prompt + sampling. Greedy (temp=0) -> deterministic argmax. + const char * prompt = "What is 2 plus 2? The answer is"; + const uint32_t n_ctx = 2048; + const int32_t n_threads = 0; + const float temp = 0.0f; + const int32_t top_k = 0; + const float top_p = 0.0f; + const uint32_t seed = 1; + const uint32_t max_tokens = 64; + + // ---- PATH A: FlashRT frt_model_runtime_v1 wrapper ---------------------- + std::string text_flashrt; + std::vector logits_flashrt; + int32_t token_flashrt = 0; + { + const frt_llama_cpp_engine_factory_v1 * factory = + frt_llama_cpp_default_engine_factory(); + CHECK(factory != nullptr && factory->create_llm != nullptr, + "default engine factory exposes create_llm"); + + std::string json = + std::string("{") + + "\"model_family\":\"llm\"," + "\"model_path\":\"" + model_env + "\"," + "\"backend\":\"" + backend + "\"," + "\"n_ctx\":" + std::to_string(n_ctx) + "," + "\"n_threads\":" + std::to_string(n_threads) + "," + "\"temp\":" + std::to_string(temp) + "," + "\"top_k\":" + std::to_string(top_k) + "," + "\"top_p\":" + std::to_string(top_p) + "," + "\"seed\":" + std::to_string(seed) + "," + "\"max_tokens\":" + std::to_string(max_tokens) + "}"; + frt_model_runtime_v1 * model = nullptr; + int rc = frt_llama_cpp_llm_runtime_open_with_engine_factory( + json.c_str(), factory, &model); + if (rc != 0 || !model) { + std::printf("FAIL: open FlashRT LLM runtime (rc=%d): %s\n", rc, + factory->last_error(factory->self)); + g_fail = 1; + } else { + CHECK(true, "open FlashRT LLM runtime from JSON"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0, + "FlashRT set_input prompt"); + CHECK(model->verbs.step(model->self) == 0, + "FlashRT run_stage infer"); + uint64_t written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, + nullptr, 0, &written, -1); + CHECK(rc == -5 && written > 0, + "FlashRT query generated text size"); + text_flashrt.assign(written, '\0'); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, &text_flashrt[0], + text_flashrt.size(), &written, -1); + CHECK(rc == 0 && written > 0, "FlashRT get_output text non-empty"); + text_flashrt.resize(written); + std::printf(" FlashRT text: %s\n", text_flashrt.c_str()); + CHECK(llama_cpp_run_generic_stage(model, "reset") == 0 && + llama_cpp_run_generic_stage(model, "prefill") == 0, + "FlashRT staged prefill"); + uint64_t logits_bytes = 0; + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_LOGITS, + nullptr, 0, &logits_bytes, -1) != 0 && logits_bytes > 0, + "FlashRT query logits size"); + logits_flashrt.resize(logits_bytes / sizeof(float)); + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_LOGITS, + logits_flashrt.data(), logits_bytes, &logits_bytes, + -1) == 0 && + llama_cpp_run_generic_stage(model, "decode") == 0, + "FlashRT read logits and decode first token"); + uint64_t token_bytes = 0; + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &token_flashrt, sizeof(token_flashrt), &token_bytes, + -1) == 0, + "FlashRT read first token"); + model->release(model->owner); + } + } + + // ---- PATH B: direct jetson_pi_llm call --------------------------------- + std::string text_native; + std::vector logits_native; + int32_t token_native = 0; + { + jetson_pi_llm_config jc{}; + jc.struct_size = sizeof(jc); + jc.model_path = model_env; + jc.backend = backend.c_str(); + jc.n_ctx = n_ctx; + jc.n_threads = n_threads; + jc.temp = temp; + jc.top_k = top_k; + jc.top_p = top_p; + jc.seed = seed; + jc.max_tokens = max_tokens; + + jetson_pi_llm * llm = nullptr; + int32_t s = jetson_pi_llm_open(&jc, &llm); + if (s != JETSON_PI_LLM_OK || !llm) { + std::printf("FAIL: jetson_pi_llm_open (rc=%d): %s\n", s, + jetson_pi_llm_open_error()); + g_fail = 1; + } else { + CHECK(true, "direct jetson_pi_llm_open"); + size_t written = 0; + s = jetson_pi_llm_generate(llm, prompt, std::strlen(prompt), + nullptr, 0, &written); + CHECK(s == JETSON_PI_LLM_BUFFER_TOO_SMALL && written > 0, + "direct query generated text bound"); + text_native.assign(written, '\0'); + written = 0; + s = jetson_pi_llm_generate(llm, prompt, std::strlen(prompt), + text_native.data(), text_native.size(), + &written); + CHECK(s == JETSON_PI_LLM_OK && written > 0, + "direct jetson_pi_llm_generate produced text"); + if (s != JETSON_PI_LLM_OK) { + std::printf(" generate error: %s\n", jetson_pi_llm_last_error(llm)); + } + text_native.resize(written); + std::printf(" native text: %s\n", text_native.c_str()); + CHECK(jetson_pi_llm_reset(llm) == JETSON_PI_LLM_OK && + jetson_pi_llm_prefill(llm, prompt, std::strlen(prompt)) == + JETSON_PI_LLM_OK, + "direct staged prefill"); + size_t logits_count = 0; + CHECK(jetson_pi_llm_get_logits( + llm, nullptr, 0, &logits_count) == + JETSON_PI_LLM_BUFFER_TOO_SMALL && logits_count > 0, + "direct query logits size"); + logits_native.resize(logits_count); + CHECK(jetson_pi_llm_get_logits( + llm, logits_native.data(), logits_native.size(), + &logits_count) == JETSON_PI_LLM_OK, + "direct read logits"); + int32_t is_eog = 0; + CHECK(jetson_pi_llm_decode_step( + llm, &token_native, &is_eog) == JETSON_PI_LLM_OK, + "direct decode first token"); + jetson_pi_llm_close(llm); + } + } + + if (g_fail) { + std::printf("\n== LLM PARITY FAILED ==\n"); + return g_fail; + } + + // Greedy = argmax, deterministic -> EXACT equality (no tolerance). + if (text_flashrt == text_native) { + CHECK(true, "FlashRT text == direct jetson_pi_llm text (exact)"); + } else { + std::printf("FAIL: text mismatch (FlashRT %zu bytes, native %zu bytes)\n", + text_flashrt.size(), text_native.size()); + size_t cmn = 0; + while (cmn < text_flashrt.size() && cmn < text_native.size() && + text_flashrt[cmn] == text_native[cmn]) ++cmn; + std::printf(" first divergence at byte %zu: FlashRT=0x%02x native=0x%02x\n", + cmn, + cmn < text_flashrt.size() ? (unsigned char)text_flashrt[cmn] : 0, + cmn < text_native.size() ? (unsigned char)text_native[cmn] : 0); + g_fail = 1; + } + CHECK(token_flashrt == token_native, + "FlashRT first token matches direct narrow API"); + bool logits_match = logits_flashrt.size() == logits_native.size(); + float logits_max_diff = 0.0f; + for (size_t i = 0; logits_match && i < logits_flashrt.size(); ++i) { + const float diff = std::fabs(logits_flashrt[i] - logits_native[i]); + if (diff > logits_max_diff) logits_max_diff = diff; + if (diff > 1e-6f) logits_match = false; + } + std::printf(" logits max abs diff = %.9g\n", logits_max_diff); + CHECK(logits_match, "FlashRT prefill logits match direct narrow API"); + + std::printf(g_fail ? "\n== LLM PARITY FAILED ==\n" + : "\n== LLM PARITY PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_mllm.cpp b/cpp/tests/test_llama_cpp_jetson_pi_mllm.cpp new file mode 100644 index 00000000..edcd3f68 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_mllm.cpp @@ -0,0 +1,297 @@ +// End-to-end MLLM engine test: drives the real Jetson-PI MLLM engine factory +// with a GGUF VLM (when available) + a programmatically generated test image, +// and asserts a sane text completion. Skips (returns 0) when env vars unset. +// +// Env: +// FLASHRT_MLLM_MODEL path to VLM GGUF (e.g. Qwen2.5-VL-3B-Instruct-q4_0.gguf) +// FLASHRT_MLLM_MMPROJ path to VIT mmproj GGUF + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" + +#include +#include +#include +#include +#include +#include +#include + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +int main() { + const char * model_env = std::getenv("FLASHRT_MLLM_MODEL"); + const char * mmproj_env = std::getenv("FLASHRT_MLLM_MMPROJ"); + const char * backend_env = std::getenv("FLASHRT_MLLM_BACKEND"); + const char * backend = (backend_env && *backend_env) ? backend_env : "cpu"; + auto file_exists = [](const char * path) { + if (!path || !*path) return false; + FILE * f = std::fopen(path, "rb"); + if (!f) return false; + std::fclose(f); + return true; + }; + if (!file_exists(model_env) || !file_exists(mmproj_env)) { + std::printf("SKIP - FLASHRT_MLLM_MODEL / FLASHRT_MLLM_MMPROJ not set or missing\n"); + return 0; + } + + const frt_llama_cpp_engine_factory_v1 * factory = + frt_llama_cpp_default_engine_factory(); + CHECK(factory != nullptr && factory->create_mllm != nullptr, + "default engine factory exposes create_mllm"); + + // Bogus model path must fail cleanly. + const char * bogus_json = + "{\"model_family\":\"mllm\"," + "\"model_path\":\"/nonexistent/bogus.gguf\"," + "\"mmproj_path\":\"/nonexistent/bogus-mmproj.gguf\"," + "\"backend\":\"cpu\"," + "\"n_ctx\":2048,\"n_threads\":0," + "\"temp\":0.0,\"top_k\":0,\"top_p\":0.0,\"seed\":1,\"max_tokens\":32}"; + frt_model_runtime_v1 * bogus = nullptr; + CHECK(frt_llama_cpp_mllm_runtime_open_with_engine_factory( + bogus_json, factory, &bogus) != 0 && bogus == nullptr, + "open MLLM with bogus model path fails without crashing"); + + // Real VLM. + std::string json = + std::string("{") + + "\"model_family\":\"mllm\"," + "\"model_path\":\"" + model_env + "\"," + "\"mmproj_path\":\"" + mmproj_env + "\"," + "\"backend\":\"" + std::string(backend) + "\"," + "\"n_ctx\":2048,\"n_threads\":0," + "\"temp\":0.0,\"top_k\":0,\"top_p\":0.0,\"seed\":1,\"max_tokens\":16}"; + frt_model_runtime_v1 * model = nullptr; + int rc = frt_llama_cpp_mllm_runtime_open_with_engine_factory( + json.c_str(), factory, &model); + if (rc != 0 || !model) { + std::printf("FAIL: open MLLM runtime (rc=%d): %s\n", rc, + factory->last_error(factory->self)); + return 1; + } + CHECK(rc == 0 && model != nullptr, "open MLLM runtime from JSON"); + const auto* plan = llama_cpp_generic_plan(model); + CHECK(model->n_stages == 0 && plan && plan->n_stages == 3, + "MLLM runtime exposes selected reset/prefill/decode plan"); + CHECK(model->n_ports == 6, + "MLLM runtime exposes images/prompt/text/token/logits/eog ports"); + CHECK(model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_LOGITS, + nullptr, 0, nullptr, -1) == -1, + "get_output rejects null written without crashing"); + + // Generate a 224x224 solid-red RGB image inline (no stb_image_write dep). + const int W = 224, H = 224; + std::vector rgb(static_cast(W) * H * 3); + for (size_t i = 0; i < rgb.size(); i += 3) { + rgb[i] = 255; rgb[i + 1] = 0; rgb[i + 2] = 0; + } + + frt_image_view view; + view.struct_size = sizeof(frt_image_view); + view.pixel_format = FRT_RT_PIXEL_RGB8; + view.data = rgb.data(); + view.bytes = rgb.size(); + view.width = W; view.height = H; view.stride_bytes = 0; + view.reserved = 0; view.timestamp_ns = 0; + + rc = model->verbs.set_input(model->self, + FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &view, sizeof(view), -1); + CHECK(rc == 0, "set_input images"); + + const char * prompt = "Describe this image in one sentence."; + rc = model->verbs.set_input(model->self, + FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1); + CHECK(rc == 0, "set_input prompt"); + + rc = model->verbs.step(model->self); + if (rc != 0) { + std::printf("FAIL: run_stage infer (rc=%d): %s\n", rc, + model->verbs.last_error(model->self)); + g_fail = 1; + } else { + std::printf("ok : run_stage infer\n"); + } + + uint64_t written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_TEXT, nullptr, 0, &written, -1); + CHECK(rc == -5 && written > 0, "query one-shot text size"); + std::string text(written, '\0'); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_TEXT, &text[0], text.size(), + &written, -1); + CHECK(rc == 0 && written > 0, "get_output text non-empty"); + if (rc == 0) { + text.resize(written); + std::printf(" generated: %s\n", text.c_str()); + bool printable = false; + for (char c : text) { + if (c >= 0x20 && c < 0x7f) { printable = true; break; } + } + CHECK(printable, "generated text contains printable chars"); + } + frt_image_view bad_view = view; + bad_view.stride_bytes = 1; + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &bad_view, sizeof(bad_view), -1) != 0, + "undersized image stride is rejected"); + CHECK(model->verbs.step(model->self) != 0, + "failed image replacement invalidates previous images"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &view, sizeof(view), -1) == 0, + "restore valid images after failed replacement"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + nullptr, 0, -1) != 0, + "invalid prompt replacement is rejected"); + CHECK(model->verbs.step(model->self) != 0, + "failed prompt replacement invalidates previous prompt"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0, + "restore valid prompt after failed replacement"); + + rc = llama_cpp_run_generic_stage(model, "reset"); + CHECK(rc == 0, "run_stage reset"); + const auto prefill_start = std::chrono::steady_clock::now(); + rc = llama_cpp_run_generic_stage(model, "prefill"); + const auto prefill_end = std::chrono::steady_clock::now(); + CHECK(rc == 0, "run_stage prefill"); + std::printf(" staged MLLM prefill latency: %.3f ms\n", + std::chrono::duration( + prefill_end - prefill_start).count()); + + uint64_t logits_bytes = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_LOGITS, nullptr, 0, + &logits_bytes, -1); + CHECK(rc == -5 && logits_bytes > 0 && logits_bytes % sizeof(float) == 0, + "prefill exposes logits size"); + std::vector logits(logits_bytes / sizeof(float)); + uint64_t logits_written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_LOGITS, logits.data(), + logits_bytes, &logits_written, -1); + CHECK(rc == 0 && logits_written == logits_bytes, + "prefill get_output logits"); + bool finite = true; + for (float value : logits) finite = finite && std::isfinite(value); + CHECK(finite, "prefill logits contain no NaN/Inf"); + + double decode_ms = 0.0; + uint32_t decoded_tokens = 0; + for (uint32_t i = 0; i < 16; ++i) { + const auto decode_start = std::chrono::steady_clock::now(); + rc = llama_cpp_run_generic_stage(model, "decode"); + const auto decode_end = std::chrono::steady_clock::now(); + decode_ms += std::chrono::duration( + decode_end - decode_start).count(); + ++decoded_tokens; + CHECK(rc == 0, "run_stage decode"); + int32_t token = 0; + int32_t is_eog = 0; + uint64_t scalar_written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN, &token, + sizeof(token), &scalar_written, -1); + CHECK(rc == 0 && scalar_written == sizeof(token), + "decode exposes next_token"); + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_IS_EOG, &is_eog, + sizeof(is_eog), &scalar_written, -1); + CHECK(rc == 0 && scalar_written == sizeof(is_eog), + "decode exposes is_eog"); + if (is_eog) break; + } + std::printf(" staged MLLM decode throughput: %.2f token/s " + "(%u tokens in %.3f ms)\n", + decoded_tokens * 1000.0 / decode_ms, + decoded_tokens, decode_ms); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_TEXT, nullptr, 0, &written, -1); + CHECK(rc == -5 && written > 0, "query staged text size"); + std::string staged(written, '\0'); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_TEXT, staged.data(), staged.size(), + &written, -1); + CHECK(rc == 0 && written > 0, "staged decode exposes accumulated text"); + if (rc == 0) staged.resize(written); + CHECK(staged == text, "staged decode text matches one-shot infer"); + rc = model->verbs.step(model->self); + CHECK(rc == 0, "one-shot infer after staged decode"); + uint64_t stale_written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN, nullptr, 0, + &stale_written, -1); + CHECK(rc == -7, "one-shot infer invalidates staged token output"); + + model->release(model->owner); + + std::string budget_json = + std::string("{") + + "\"model_family\":\"mllm\"," + + "\"model_path\":\"" + model_env + "\"," + + "\"mmproj_path\":\"" + mmproj_env + "\"," + + "\"backend\":\"" + std::string(backend) + "\"," + + "\"n_ctx\":2048,\"n_threads\":0," + "\"temp\":0.0,\"top_k\":0,\"top_p\":0.0,\"seed\":1," + "\"max_tokens\":1}"; + frt_model_runtime_v1 * budget_model = nullptr; + rc = frt_llama_cpp_mllm_runtime_open_with_engine_factory( + budget_json.c_str(), factory, &budget_model); + CHECK(rc == 0 && budget_model, "open max_tokens=1 MLLM runtime"); + if (budget_model) { + CHECK(budget_model->verbs.set_input( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &view, sizeof(view), -1) == 0 && + budget_model->verbs.set_input( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0 && + llama_cpp_run_generic_stage(budget_model, "prefill") == 0 && + llama_cpp_run_generic_stage(budget_model, "decode") == 0, + "max_tokens=1 MLLM session permits exactly one decode"); + int32_t first_is_eog = 1; + uint64_t scalar_written = 0; + CHECK(budget_model->verbs.get_output( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_IS_EOG, + &first_is_eog, sizeof(first_is_eog), &scalar_written, -1) == 0 && + first_is_eog == 0, + "budget test image prompt first token is not EOG"); + CHECK(llama_cpp_run_generic_stage(budget_model, "decode") != 0 && + std::strstr(budget_model->verbs.last_error( + budget_model->self), + "max_tokens"), + "staged decode rejects calls beyond max_tokens"); + int32_t stale_scalar = 0; + CHECK(budget_model->verbs.get_output( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN, + &stale_scalar, sizeof(stale_scalar), &scalar_written, -1) == -7 && + budget_model->verbs.get_output( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_IS_EOG, + &stale_scalar, sizeof(stale_scalar), &scalar_written, -1) == -7, + "failed decode invalidates staged scalar outputs"); + CHECK(llama_cpp_run_generic_stage(budget_model, "prefill") == 0 && + llama_cpp_run_generic_stage(budget_model, "decode") == 0, + "prefill restores staged decode budget"); + budget_model->release(budget_model->owner); + } + + std::printf(g_fail ? "\n== JETSON_PI MLLM FAILED ==\n" + : "\n== JETSON_PI MLLM PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_mllm_parity.cpp b/cpp/tests/test_llama_cpp_jetson_pi_mllm_parity.cpp new file mode 100644 index 00000000..ee3de72f --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_mllm_parity.cpp @@ -0,0 +1,128 @@ +// Numerical parity: FlashRT MLLM provider versus direct jetson_pi_mllm. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "jetson_pi_mllm.h" + +#include +#include +#include +#include +#include +#include + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +int main() { + const char* model_path = std::getenv("FLASHRT_MLLM_MODEL"); + const char* mmproj_path = std::getenv("FLASHRT_MLLM_MMPROJ"); + const char* backend_env = std::getenv("FLASHRT_MLLM_BACKEND"); + const std::string backend = + backend_env && *backend_env ? backend_env : "cpu"; + if (!model_path || !mmproj_path || + !std::ifstream(model_path, std::ios::binary).good() || + !std::ifstream(mmproj_path, std::ios::binary).good()) { + std::printf("SKIP - MLLM model/mmproj missing\n"); + return 0; + } + + const uint32_t width = 224; + const uint32_t height = 224; + const uint32_t max_tokens = 16; + const char* prompt = "Describe this image in one sentence."; + std::vector rgb(static_cast(width) * height * 3); + for (size_t i = 0; i < rgb.size(); i += 3) rgb[i] = 255; + + std::string flashrt_text; + { + const auto* factory = frt_llama_cpp_default_engine_factory(); + std::string json = + std::string("{") + + "\"model_family\":\"mllm\",\"model_path\":\"" + model_path + + "\",\"mmproj_path\":\"" + mmproj_path + + "\",\"backend\":\"" + backend + + "\",\"n_ctx\":2048,\"n_threads\":0,\"temp\":0.0," + "\"top_k\":0,\"top_p\":0.0,\"seed\":1,\"max_tokens\":16}"; + frt_model_runtime_v1* model = nullptr; + int rc = frt_llama_cpp_mllm_runtime_open_with_engine_factory( + json.c_str(), factory, &model); + CHECK(rc == 0 && model, "open FlashRT MLLM runtime"); + if (model) { + frt_image_view image{}; + image.struct_size = sizeof(image); + image.pixel_format = FRT_RT_PIXEL_RGB8; + image.data = rgb.data(); + image.bytes = rgb.size(); + image.width = width; + image.height = height; + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &image, sizeof(image), -1) == 0 && + model->verbs.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0 && + model->verbs.step(model->self) == 0, + "run FlashRT MLLM image+text infer"); + uint64_t written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_TEXT, + nullptr, 0, &written, -1); + CHECK(rc == -5 && written > 0, "query FlashRT MLLM text size"); + flashrt_text.assign(written, '\0'); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_MLLM_PORT_TEXT, + flashrt_text.data(), flashrt_text.size(), &written, -1); + CHECK(rc == 0 && written > 0, "read FlashRT MLLM text"); + flashrt_text.resize(written); + model->release(model->owner); + } + } + + std::string direct_text; + { + jetson_pi_mllm_config config{}; + config.struct_size = sizeof(config); + config.model_path = model_path; + config.mmproj_path = mmproj_path; + config.backend = backend.c_str(); + config.n_ctx = 2048; + config.temp = 0.0f; + config.seed = 1; + config.max_tokens = max_tokens; + jetson_pi_mllm* handle = nullptr; + int32_t rc = jetson_pi_mllm_open(&config, &handle); + CHECK(rc == JETSON_PI_MLLM_OK && handle, + "open direct jetson_pi_mllm"); + if (handle) { + const uint8_t* images[] = {rgb.data()}; + size_t written = 0; + rc = jetson_pi_mllm_infer( + handle, images, 1, height, width, prompt, + std::strlen(prompt), nullptr, 0, &written); + CHECK(rc == JETSON_PI_MLLM_BUFFER_TOO_SMALL && written > 0, + "query direct MLLM text bound"); + direct_text.assign(written, '\0'); + written = 0; + rc = jetson_pi_mllm_infer( + handle, images, 1, height, width, prompt, + std::strlen(prompt), direct_text.data(), direct_text.size(), + &written); + CHECK(rc == JETSON_PI_MLLM_OK && written > 0, + "run direct MLLM image+text infer"); + direct_text.resize(written); + jetson_pi_mllm_close(handle); + } + } + + CHECK(flashrt_text == direct_text, + "FlashRT MLLM text matches direct narrow API exactly"); + std::printf(g_fail ? "\n== MLLM PARITY FAILED ==\n" + : "\n== MLLM PARITY PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_parity.cpp b/cpp/tests/test_llama_cpp_jetson_pi_parity.cpp new file mode 100644 index 00000000..efa27d95 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_parity.cpp @@ -0,0 +1,421 @@ +// Numerical-parity test: FlashRT's Jetson-PI Pi0 provider (via +// frt_model_runtime_v1) vs a DIRECT jetson_pi_pi0 call, same inputs / backend +// / thread count. Proves the FlashRT port/stage/verb plumbing does not +// perturb the action chunk relative to the native narrow C API. +// +// Closes jetsonpi迁移.txt §14 (parity vs native). Skips (returns 0) when the +// weights/fixture env vars are unset, so CI without weights still passes. +// +// Determinism foundation: Pi0's action-chunk noise latent is +// create_normal_noise_cpp11(..., seed=41) (Jetson-PI/src/llama-context.cpp), +// reset to fresh noise + action_ready=false per tick; KV/position are cleared +// before every infer. So given fixed model+inputs+backend+n_threads, both +// paths must produce identical actions (modulo float-arith ULP). The existing +// tick-repro test already relies on this (test_llama_cpp_jetson_pi_engine.cpp). +// +// FlashRT adds NO numerical perturbation vs the direct call (verified): image +// swizzle is an identity copy for RGB8 fixtures (load via stbi_load(...,3)), +// prompt/state are copied verbatim, and marker injection plus per-model state +// handling happen INSIDE jetson_pi_pi0_infer for both (PI0.5 serializes state +// into the prompt; legacy Pi0 zero-pads it to action_dim). Backend +// n_gpu_layers/use_gpu mapping is identical, and n_threads=0 on both resolves +// to the same hardware_concurrency(). +// A parity failure means a real FlashRT-port bug. +// +// Env: +// FLASHRT_PI0_MODEL path to Pi0 policy GGUF +// FLASHRT_PI0_MMPROJ path to VIT mmproj GGUF +// FLASHRT_PI0_FIXTURE_DIR dir with image.png, wrist_image.png, state.bin, prompt.txt +// FLASHRT_PI0_ACTION_STEPS required model-specific action horizon +// (the verified pi0_base checkpoint uses 50). +// FLASHRT_PI0_ACTION_DIM required model-specific action width +// (the verified pi0_base checkpoint uses 32). +// FLASHRT_PI0_BACKEND required explicit backend, e.g. "cpu" or "cuda". +// CUDA may have ULP drift from atomics/warp reductions; +// the 1e-5 tolerance is headroom. CPU expects ~0. +// FLASHRT_PI0_CLI_ACTION_LOG (optional) log emitted by llama-mtmd-cli for +// the same fixture and effective prompt. When set, +// the `Pi0 action: [...]` line must be bit-identical +// to FlashRT's action chunk. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "jetson_pi_pi0.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define STB_IMAGE_IMPLEMENTATION +#define STBI_ONLY_PNG +#include "stb/stb_image.h" + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +namespace { +bool file_exists(const char * p) { + std::ifstream f(p); + return f.good(); +} +std::string read_file(const std::string & path, bool * ok) { + std::ifstream f(path, std::ios::binary); + if (!f) { if (ok) *ok = false; return {}; } + std::string s((std::istreambuf_iterator(f)), + std::istreambuf_iterator()); + if (ok) *ok = true; + return s; +} +} // namespace + +int main() { + const char * model_env = std::getenv("FLASHRT_PI0_MODEL"); + const char * mmproj_env = std::getenv("FLASHRT_PI0_MMPROJ"); + const char * fixture_env = std::getenv("FLASHRT_PI0_FIXTURE_DIR"); + const char * be_env = std::getenv("FLASHRT_PI0_BACKEND"); + if (!model_env || !mmproj_env || !fixture_env || !be_env || !be_env[0] || + !file_exists(model_env) || !file_exists(mmproj_env)) { + std::printf("SKIP - FLASHRT_PI0_MODEL / FLASHRT_PI0_MMPROJ / " + "FLASHRT_PI0_FIXTURE_DIR / FLASHRT_PI0_BACKEND not set " + "or files missing\n"); + return 0; + } + const std::string fixture_dir = fixture_env; + const std::string img_path = fixture_dir + "/image.png"; + const std::string wrist_path = fixture_dir + "/wrist_image.png"; + const std::string state_path = fixture_dir + "/state.bin"; + const std::string prompt_path = fixture_dir + "/prompt.txt"; + if (!file_exists(img_path.c_str()) || !file_exists(wrist_path.c_str()) || + !file_exists(state_path.c_str()) || !file_exists(prompt_path.c_str())) { + std::printf("SKIP - fixture files missing in %s\n", fixture_dir.c_str()); + return 0; + } + + const char * steps_env = std::getenv("FLASHRT_PI0_ACTION_STEPS"); + const char * dim_env = std::getenv("FLASHRT_PI0_ACTION_DIM"); + if (!steps_env || !dim_env) { + std::printf("SKIP - FLASHRT_PI0_ACTION_STEPS/ACTION_DIM not set\n"); + return 0; + } + const std::string backend = be_env; + long action_steps = std::atol(steps_env); + long action_dim = std::atol(dim_env); + if (action_steps <= 0 || action_dim <= 0 || action_steps > 10000 || + action_dim > 10000) { + std::printf("SKIP - bad FLASHRT_PI0_ACTION_STEPS/DIM\n"); + return 0; + } + + // Load fixtures ONCE; both paths consume the same bytes. + int iw = 0, ih = 0, ic = 0; + unsigned char * img = stbi_load(img_path.c_str(), &iw, &ih, &ic, 3); + CHECK(img != nullptr && iw == 224 && ih == 224, "load image.png 224x224"); + int ww = 0, wh = 0, wc = 0; + unsigned char * wrist = stbi_load(wrist_path.c_str(), &ww, &wh, &wc, 3); + CHECK(wrist != nullptr && ww == 224 && wh == 224, "load wrist_image.png 224x224"); + bool state_ok = false; + std::string state_bytes = read_file(state_path, &state_ok); + // State width is policy-specific and independent of action_dim. Accept any + // non-empty float fixture; the backend enforces the selected model's bound. + const size_t state_count = + state_bytes.size() / sizeof(float); + CHECK(state_ok && state_bytes.size() % sizeof(float) == 0 && + state_count > 0, + "state.bin contains non-empty float32 policy state"); + bool prompt_ok = false; + std::string prompt = read_file(prompt_path, &prompt_ok); + CHECK(prompt_ok && !prompt.empty(), "prompt.txt non-empty"); + if (!prompt.empty() && prompt.back() == '\n') prompt.pop_back(); + + if (!img || !wrist || !state_ok || !prompt_ok) { + if (img) stbi_image_free(img); + if (wrist) stbi_image_free(wrist); + std::printf(g_fail ? "\n== PI0 PARITY FAILED ==\n" + : "\n== PI0 PARITY SKIPPED ==\n"); + return g_fail; + } + + const size_t n_elems = static_cast(action_steps) * action_dim; + const size_t n_bytes = n_elems * sizeof(float); + + // ---- PATH A: FlashRT frt_model_runtime_v1 wrapper ---------------------- + std::vector actions_flashrt(n_elems, 0.0f); + std::vector actions_split(n_elems, 0.0f); + { + const frt_llama_cpp_engine_factory_v1 * factory = + frt_llama_cpp_default_engine_factory(); + CHECK(factory != nullptr, "default engine factory is non-null"); + + std::string json = + std::string("{") + + "\"model_family\":\"pi0\"," + "\"model_path\":\"" + model_env + "\"," + "\"mmproj_path\":\"" + mmproj_env + "\"," + "\"backend\":\"" + backend + "\"," + "\"n_views\":2,\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":" + std::to_string(action_steps) + + ",\"action_dim\":" + std::to_string(action_dim) + "}"; + frt_model_runtime_v1 * model = nullptr; + int rc = frt_llama_cpp_pi0_runtime_open_with_engine_factory( + json.c_str(), factory, &model); + if (rc != 0 || !model) { + std::printf("FAIL: open FlashRT Pi0 runtime (rc=%d): %s\n", rc, + factory->last_error(factory->self)); + g_fail = 1; + } else { + CHECK(true, "open FlashRT Pi0 runtime from JSON"); + const auto* plan = llama_cpp_generic_plan(model); + CHECK(model->n_stages == 0 && plan && plan->n_stages == 2, + "FlashRT Pi0 runtime exposes selected context/action plan"); + + frt_image_view views[2]; + views[0].struct_size = sizeof(frt_image_view); + views[0].pixel_format = FRT_RT_PIXEL_RGB8; + views[0].data = img; + views[0].bytes = static_cast(224) * 224 * 3; + views[0].width = 224; views[0].height = 224; views[0].stride_bytes = 0; + views[0].reserved = 0; views[0].timestamp_ns = 0; + views[1] = views[0]; + views[1].data = wrist; + + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, + &views[0], sizeof(views), -1) == 0, + "FlashRT set_input images"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "FlashRT set_input prompt"); + CHECK(model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), + state_bytes.size(), -1) == 0, + "FlashRT set_input state"); + CHECK(model->verbs.step(model->self) == 0, + "FlashRT run_stage infer"); + + uint64_t written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_PI0_PORT_ACTIONS, + actions_flashrt.data(), n_bytes, &written, -1); + CHECK(rc == 0 && written == n_bytes, + "FlashRT get_output actions shape matches config"); + CHECK(llama_cpp_run_generic_stage(model, "context") == 0, + "FlashRT run_stage context"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "new Pi0 input discards pending context"); + CHECK(llama_cpp_run_generic_stage(model, "action") == -7, + "action after replacement input is not ready"); + CHECK(llama_cpp_run_generic_stage(model, "context") == 0, + "FlashRT rebuilds context after replacement input"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), 0, -1) != 0, + "empty replacement state is rejected"); + CHECK(llama_cpp_run_generic_stage(model, "context") != 0, + "failed state replacement invalidates the previous state"); + CHECK(llama_cpp_run_generic_stage(model, "action") != 0, + "failed replacement input blocks action execution"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), state_bytes.size(), -1) == 0, + "restore valid Pi0 state after invalid replacement"); + frt_image_view bad_views[2] = {views[0], views[1]}; + bad_views[0].stride_bytes = 1; + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_IMAGES, + bad_views, sizeof(bad_views), -1) != 0, + "undersized Pi0 image stride is rejected"); + CHECK(llama_cpp_run_generic_stage(model, "context") != 0, + "failed image replacement invalidates previous images"); + CHECK(model->verbs.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_IMAGES, + views, sizeof(views), -1) == 0, + "restore valid Pi0 images after failed replacement"); + CHECK(llama_cpp_run_generic_stage(model, "context") == 0, + "FlashRT rebuilds context after failed replacement"); + CHECK(llama_cpp_run_generic_stage(model, "action") == 0, + "FlashRT run_stage action"); + written = 0; + rc = model->verbs.get_output( + model->self, FRT_LLAMA_CPP_PI0_PORT_ACTIONS, + actions_split.data(), n_bytes, &written, -1); + CHECK(rc == 0 && written == n_bytes, + "FlashRT split stages produced action chunk"); + CHECK(llama_cpp_run_generic_stage(model, "action") == -7, + "Pi0 action stage consumes pending context exactly once"); + model->release(model->owner); + } + } + + // ---- PATH B: direct jetson_pi_pi0 call --------------------------------- + std::vector actions_native(n_elems, 0.0f); + { + jetson_pi_pi0_config jc{}; + jc.struct_size = sizeof(jc); + jc.model_path = model_env; + jc.mmproj_path = mmproj_env; + jc.backend = backend.c_str(); + jc.n_views = 2; + jc.image_height = 224; + jc.image_width = 224; + jc.n_threads = 0; // match FlashRT's hardcoded 0 (jetson_pi_engine.cpp) + + jetson_pi_pi0 * pi0 = nullptr; + int32_t s = jetson_pi_pi0_open(&jc, &pi0); + if (s != JETSON_PI_PI0_OK || !pi0) { + std::printf("FAIL: jetson_pi_pi0_open (rc=%d): %s\n", s, + jetson_pi_pi0_open_error()); + g_fail = 1; + } else { + CHECK(true, "direct jetson_pi_pi0_open"); + + uint32_t steps = 0, dim = 0; + CHECK(jetson_pi_pi0_action_shape(pi0, &steps, &dim) == + JETSON_PI_PI0_OK && + steps == static_cast(action_steps) && + dim == static_cast(action_dim), + "direct action_shape matches config"); + + const uint8_t * imgs[2] = { img, wrist }; + size_t written = 0; + CHECK(jetson_pi_pi0_action( + pi0, actions_native.data(), n_elems, &written) == + JETSON_PI_PI0_ACTION_NOT_READY && written == 0, + "direct action-not-ready reports zero elements written"); + written = 0; + CHECK(jetson_pi_pi0_infer( + pi0, imgs, 2, prompt.data(), prompt.size(), + reinterpret_cast(state_bytes.data()), + state_bytes.size() / sizeof(float), actions_native.data(), + n_elems - 1, &written) == + JETSON_PI_PI0_BUFFER_TOO_SMALL && written == n_elems, + "too-small whole infer reports required size before context"); + written = 123; + CHECK(jetson_pi_pi0_action( + pi0, actions_native.data(), n_elems, &written) == + JETSON_PI_PI0_ACTION_NOT_READY && written == 0, + "invalid whole infer leaves no consumable context"); + CHECK(jetson_pi_pi0_context( + pi0, imgs, 2, prompt.data(), prompt.size(), + reinterpret_cast(state_bytes.data()), + state_bytes.size() / sizeof(float)) == JETSON_PI_PI0_OK, + "direct context prepares pending action"); + CHECK(jetson_pi_pi0_context( + pi0, imgs, 2, nullptr, 0, + reinterpret_cast(state_bytes.data()), + state_bytes.size() / sizeof(float)) == JETSON_PI_PI0_INVALID, + "failed direct context replacement is rejected"); + written = 123; + CHECK(jetson_pi_pi0_action( + pi0, actions_native.data(), n_elems, &written) == + JETSON_PI_PI0_ACTION_NOT_READY && written == 0, + "failed direct context discards prior pending context"); + s = jetson_pi_pi0_infer(pi0, imgs, 2, + prompt.data(), prompt.size(), + reinterpret_cast(state_bytes.data()), + state_bytes.size() / sizeof(float), + actions_native.data(), n_elems, &written); + CHECK(s == JETSON_PI_PI0_OK && written == n_elems, + "direct jetson_pi_pi0_infer produced action chunk"); + if (s != JETSON_PI_PI0_OK) { + std::printf(" infer error: %s\n", jetson_pi_pi0_last_error(pi0)); + } + jetson_pi_pi0_close(pi0); + } + } + + if (img) stbi_image_free(img); + if (wrist) stbi_image_free(wrist); + + if (g_fail) { + std::printf("\n== PI0 PARITY FAILED ==\n"); + return g_fail; + } + + // ---- compare ----------------------------------------------------------- + bool nan_inf = false; + for (float v : actions_flashrt) { + if (std::isnan(v) || std::isinf(v)) { nan_inf = true; break; } + } + CHECK(!nan_inf, "FlashRT actions contain no NaN/Inf"); + + float split_max_diff = 0.0f; + for (size_t i = 0; i < n_elems; ++i) { + split_max_diff = std::max( + split_max_diff, std::fabs(actions_flashrt[i] - actions_split[i])); + } + std::printf(" whole-vs-split max abs diff = %.9g\n", split_max_diff); + CHECK(split_max_diff == 0.0f, + "FlashRT context/action is bit-identical to whole infer"); + + float max_diff = 0.0f; + size_t diverge_idx = 0; + for (size_t i = 0; i < n_elems; ++i) { + float d = std::fabs(actions_flashrt[i] - actions_native[i]); + if (d > max_diff) { max_diff = d; diverge_idx = i; } + } + std::printf(" max abs diff = %.3g (n_elems=%zu, first diverge idx=%zu)\n", + max_diff, n_elems, diverge_idx); + CHECK(max_diff <= 1e-5f, + "FlashRT actions match direct jetson_pi_pi0 (max_diff <= 1e-5)"); + + const char * cli_action_log = std::getenv("FLASHRT_PI0_CLI_ACTION_LOG"); + if (cli_action_log && cli_action_log[0]) { + std::ifstream log(cli_action_log); + CHECK(log.good(), "open native Pi0 CLI action log"); + std::vector actions_cli; + std::string line; + const std::string prefix = "Pi0 action: ["; + while (std::getline(log, line)) { + if (line.rfind(prefix, 0) != 0 || line.back() != ']') { + continue; + } + std::string payload = line.substr( + prefix.size(), line.size() - prefix.size() - 1); + for (char & ch : payload) { + if (ch == ',') ch = ' '; + } + std::istringstream values(payload); + float value = 0.0f; + while (values >> value) actions_cli.push_back(value); + if (!values.eof()) actions_cli.clear(); + break; + } + CHECK(actions_cli.size() == n_elems, + "native Pi0 CLI log contains the complete action chunk"); + if (actions_cli.size() == n_elems) { + float cli_max_diff = 0.0f; + size_t cli_diverge_idx = 0; + for (size_t i = 0; i < n_elems; ++i) { + const float d = std::fabs(actions_flashrt[i] - actions_cli[i]); + if (d > cli_max_diff) { + cli_max_diff = d; + cli_diverge_idx = i; + } + } + std::printf(" CLI max abs diff = %.9g (n_elems=%zu, first diverge idx=%zu)\n", + cli_max_diff, n_elems, cli_diverge_idx); + CHECK(std::memcmp(actions_flashrt.data(), actions_cli.data(), + n_bytes) == 0, + "FlashRT actions are bit-identical to native Pi0 CLI"); + } + } + + std::printf(g_fail ? "\n== PI0 PARITY FAILED ==\n" + : "\n== PI0 PARITY PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_jetson_pi_state_parity.cpp b/cpp/tests/test_llama_cpp_jetson_pi_state_parity.cpp new file mode 100644 index 00000000..ad663bb9 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_state_parity.cpp @@ -0,0 +1,225 @@ +// PI0.5 reference-policy state-parity scaffold. +// +// Merged PKU-SEC-Lab/Jetson-PI-Edge#1 implements the PI0.5 contract: the policy +// consumes discretized state inside the `Task: ..., State: ...;\nAction:` +// prompt rather than the legacy llama_set_pi0_state tensor. This real-model +// test remains opt-in because CI does not carry the multi-gigabyte checkpoint; +// it also verifies the fixture is detected as PI0.5 before inference. Follow-up +// #2 completes the capability contract across legacy Pi0 and PI0.5. +// +// Boundary-region coverage matters: the openpi discretizer maps x<-1 to a +// literal -1 token (distinct from bin 0), x in [-1,1) to floor((x+1)*128), +// and x>=1 to 255. A wrong binning formula still passes a naive "changes when +// state changes" check, so each region is exercised explicitly. +// +// Production callers pass finite policy-normalized state. The test also covers +// finite outliers because the server contract defines explicit bins below -1 +// and at or above 1. +// +// Env: +// FLASHRT_PI0_MODEL path to PI0.5 policy GGUF +// FLASHRT_PI0_MMPROJ path to VIT mmproj GGUF +// FLASHRT_PI0_FIXTURE_DIR dir with image.png, wrist_image.png, prompt.txt +// FLASHRT_PI0_ACTION_STEPS required model-specific action horizon. +// FLASHRT_PI0_ACTION_DIM required model-specific action width. +// FLASHRT_PI0_BACKEND required explicit backend, e.g. "cpu" or "cuda". +// FLASHRT_PI0_STATE_PARITY_READY set to "1" only when the backend build +// includes PI0.5 prompt-state serialization. + +#include "flashrt/providers/llama_cpp/c_api.h" +#include "llama_cpp_generic_plan.h" +#include "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "pi-model-detect.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#define STB_IMAGE_IMPLEMENTATION +#include "stb/stb_image.h" + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + const std::string check_msg = (msg); \ + if (!(cond)) { std::printf("FAIL: %s\n", check_msg.c_str()); g_fail = 1; } \ + else { std::printf("ok : %s\n", check_msg.c_str()); } \ +} while (0) + +namespace { +bool file_exists(const char * p) { + std::ifstream f(p); + return f.good(); +} +} // namespace + +int main() { + const char * ready = std::getenv("FLASHRT_PI0_STATE_PARITY_READY"); + if (!ready || ready[0] != '1') { + std::printf("SKIP - PI0.5 state parity requires an explicit real-model " + "run (set FLASHRT_PI0_STATE_PARITY_READY=1)\n"); + return 0; + } + const char * model_env = std::getenv("FLASHRT_PI0_MODEL"); + const char * mmproj_env = std::getenv("FLASHRT_PI0_MMPROJ"); + const char * fixture_env = std::getenv("FLASHRT_PI0_FIXTURE_DIR"); + const char * be_env = std::getenv("FLASHRT_PI0_BACKEND"); + if (!model_env || !mmproj_env || !fixture_env || !be_env || !be_env[0] || + !file_exists(model_env) || !file_exists(mmproj_env)) { + std::printf("SKIP - FLASHRT_PI0_MODEL / FLASHRT_PI0_MMPROJ / " + "FLASHRT_PI0_FIXTURE_DIR / FLASHRT_PI0_BACKEND not set " + "or files missing\n"); + return 0; + } + const pi_model_detect_result detected = + pi_model_detect_gguf_pair(model_env, mmproj_env); + CHECK(detected.kind == PI_MODEL_PI05, + "state parity GGUF is detected as PI0.5"); + if (detected.kind != PI_MODEL_PI05) { + std::printf(" detected=%s reason=%s\n", + pi_model_kind_name(detected.kind), detected.reason.c_str()); + return 1; + } + const std::string fixture_dir = fixture_env; + const std::string img_path = fixture_dir + "/image.png"; + const std::string wrist_path = fixture_dir + "/wrist_image.png"; + const std::string prompt_path = fixture_dir + "/prompt.txt"; + if (!file_exists(img_path.c_str()) || !file_exists(wrist_path.c_str()) || + !file_exists(prompt_path.c_str())) { + std::printf("SKIP - fixture files missing in %s\n", fixture_dir.c_str()); + return 0; + } + + const char * steps_env = std::getenv("FLASHRT_PI0_ACTION_STEPS"); + const char * dim_env = std::getenv("FLASHRT_PI0_ACTION_DIM"); + if (!steps_env || !dim_env) { + std::printf("SKIP - FLASHRT_PI0_ACTION_STEPS/ACTION_DIM not set\n"); + return 0; + } + const std::string backend = be_env; + long action_steps = std::atol(steps_env); + long action_dim = std::atol(dim_env); + if (action_steps <= 0 || action_dim <= 0) { + std::printf("SKIP - bad FLASHRT_PI0_ACTION_STEPS/DIM\n"); + return 0; + } + + int iw = 0, ih = 0, ic = 0; + unsigned char * img = stbi_load(img_path.c_str(), &iw, &ih, &ic, 3); + CHECK(img != nullptr && iw == 224 && ih == 224, "load image.png 224x224"); + int ww = 0, wh = 0, wc = 0; + unsigned char * wrist = stbi_load(wrist_path.c_str(), &ww, &wh, &wc, 3); + CHECK(wrist != nullptr && ww == 224 && wh == 224, "load wrist_image.png 224x224"); + std::ifstream prompt_file(prompt_path, std::ios::binary); + std::string prompt((std::istreambuf_iterator(prompt_file)), + std::istreambuf_iterator()); + const bool prompt_ok = prompt_file.good() || prompt_file.eof(); + CHECK(prompt_ok && !prompt.empty(), "prompt.txt non-empty"); + if (!prompt.empty() && prompt.back() == '\n') prompt.pop_back(); + + if (!img || !wrist || !prompt_ok) { + if (img) stbi_image_free(img); + if (wrist) stbi_image_free(wrist); + std::printf(g_fail ? "\n== PI0 STATE PARITY FAILED ==\n" + : "\n== PI0 STATE PARITY SKIPPED ==\n"); + return g_fail; + } + + const size_t n_elems = static_cast(action_steps) * action_dim; + const size_t n_bytes = n_elems * sizeof(float); + // PI0.5 proprioception width is 8 (openpi libero), independent of + // action_dim. The backend discretizes up to 8 state values into the + // prompt and rejects wider state. The state port is bucket-variable. + const long state_dim = 8; + + const frt_llama_cpp_engine_factory_v1 * factory = + frt_llama_cpp_default_engine_factory(); + CHECK(factory != nullptr, "default engine factory is non-null"); + + std::string json = + std::string("{") + + "\"model_family\":\"pi0\"," + "\"model_path\":\"" + model_env + "\"," + "\"mmproj_path\":\"" + mmproj_env + "\"," + "\"backend\":\"" + backend + "\"," + "\"n_views\":2,\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":" + std::to_string(action_steps) + + ",\"action_dim\":" + std::to_string(action_dim) + "}"; + + auto run_with_state = [&](const std::vector & state, + std::vector & out) -> bool { + frt_model_runtime_v1 * model = nullptr; + if (frt_llama_cpp_pi0_runtime_open_with_engine_factory( + json.c_str(), factory, &model) != 0 || !model) { + std::printf("FAIL: open FlashRT Pi0 runtime: %s\n", + factory->last_error(factory->self)); + g_fail = 1; + return false; + } + frt_image_view views[2]; + views[0].struct_size = sizeof(frt_image_view); + views[0].pixel_format = FRT_RT_PIXEL_RGB8; + views[0].data = img; + views[0].bytes = static_cast(224) * 224 * 3; + views[0].width = 224; views[0].height = 224; views[0].stride_bytes = 0; + views[0].reserved = 0; views[0].timestamp_ns = 0; + views[1] = views[0]; + views[1].data = wrist; + bool ok = true; + ok &= model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, &views[0], sizeof(views), -1) == 0; + ok &= model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, prompt.data(), prompt.size(), -1) == 0; + ok &= model->verbs.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_STATE, state.data(), state.size() * sizeof(float), + -1) == 0; + ok &= model->verbs.step(model->self) == 0; + uint64_t written = 0; + ok &= model->verbs.get_output(model->self, + FRT_LLAMA_CPP_PI0_PORT_ACTIONS, out.data(), n_bytes, &written, -1) == 0; + ok &= written == n_bytes; + model->release(model->owner); + return ok; + }; + + // Baseline: zero state (all bins map to the -1..0 region). + std::vector zero_state(state_dim, 0.0f); + std::vector actions_zero(n_elems, 0.0f); + CHECK(run_with_state(zero_state, actions_zero), "infer with zero state"); + + // Boundary-region states exercising every openpi bin branch. Each must + // differ from the zero-state baseline once PI0.5 state is serialized into + // the prompt. The outer cases intentionally exercise defined outlier bins. + struct Case { const char * name; float value; }; + const Case cases[] = { + {"x<-1 (literal -1 token)", -2.0f}, + {"x in (-1,0)", -0.5f}, + {"x in (0,1)", 0.5f}, + {"x>=1 (bin 255)", 1.5f}, + }; + for (const Case & c : cases) { + std::vector s(state_dim, c.value); + std::vector a(n_elems, 0.0f); + CHECK(run_with_state(s, a), std::string("infer with ") + c.name); + float max_diff = 0.0f; + for (size_t i = 0; i < n_elems; ++i) { + max_diff = std::max(max_diff, std::fabs(a[i] - actions_zero[i])); + } + std::printf(" %s: max abs diff vs zero-state = %.9g\n", c.name, max_diff); + // This assertion is reached only after the READY gate above confirms + // that the build includes PI0.5 prompt-state serialization. + CHECK(max_diff > 0.0f, + std::string("state in prompt changes actions: ") + c.name); + } + + if (img) stbi_image_free(img); + if (wrist) stbi_image_free(wrist); + + std::printf(g_fail ? "\n== PI0 STATE PARITY FAILED ==\n" + : "\n== PI0 STATE PARITY PASSED ==\n"); + return g_fail; +} diff --git a/cpp/tests/test_llama_cpp_provider.cpp b/cpp/tests/test_llama_cpp_provider.cpp new file mode 100644 index 00000000..047f55a0 --- /dev/null +++ b/cpp/tests/test_llama_cpp_provider.cpp @@ -0,0 +1,655 @@ +#include "flashrt/providers/llama_cpp/c_api.h" + +#include +#include +#include +#include +#include +#include + +static int g_fail = 0; +#define CHECK(cond, msg) do { \ + if (!(cond)) { std::printf("FAIL: %s\n", (msg)); g_fail = 1; } \ + else { std::printf("ok : %s\n", (msg)); } \ +} while (0) + +namespace { + +struct FakeEngine { + int retains = 0; + int releases = 0; + int infer = 0; + int stage_calls = 0; + uint32_t last_stage = UINT32_MAX; + int set_input_calls = 0; + uint32_t last_port = 999; + float actions[4] = {1.0f, 2.0f, 3.0f, 4.0f}; + const char* last_error = ""; +}; + +struct FakeFactory { + FakeEngine* engine = nullptr; + int creates = 0; + frt_llama_cpp_pi0_config seen{}; + std::string seen_model_path; + std::string seen_mmproj_path; + std::string seen_backend; + const char* last_error = ""; + bool return_borrowed = false; + bool return_retain_only = false; + bool return_undersized = false; + bool return_null_self = false; + bool return_missing_set_input = false; + bool fail_after_engine = false; +}; + +void retain_engine(void* p) { + static_cast(p)->retains += 1; +} + +void release_engine(void* p) { + static_cast(p)->releases += 1; +} + +int set_input(void* p, uint32_t port, const void* data, uint64_t bytes, + int stream) { + (void)data; + (void)bytes; + (void)stream; + auto* engine = static_cast(p); + engine->set_input_calls += 1; + engine->last_port = port; + return 0; +} + +int run_infer(void* p) { + auto* engine = static_cast(p); + engine->infer += 1; + engine->actions[0] += 10.0f; + return 0; +} + +int run_stage(void* p, uint32_t stage) { + auto* engine = static_cast(p); + engine->stage_calls += 1; + engine->last_stage = stage; + return 0; +} + +int get_output(void* p, uint32_t port, void* out, uint64_t capacity, + uint64_t* written, int stream) { + (void)stream; + auto* engine = static_cast(p); + if (port != FRT_LLAMA_CPP_PI0_PORT_ACTIONS) { + engine->last_error = "unknown output port"; + return -1; + } + const uint64_t need = sizeof(engine->actions); + if (written) *written = need; + if (capacity < need) { + engine->last_error = "action output buffer is too small"; + return -5; + } + std::memcpy(out, engine->actions, need); + engine->last_error = ""; + return 0; +} + +const char* last_error(void* p) { + return static_cast(p)->last_error; +} + +const char* null_last_error(void*) { + return nullptr; +} + +int create_pi0_engine(void* p, const frt_llama_cpp_pi0_config* config, + frt_llama_cpp_engine_v1* out) { + auto* factory = static_cast(p); + factory->creates += 1; + if (!config || !out || !factory->engine) { + factory->last_error = "invalid factory input"; + return -1; + } + if (factory->fail_after_engine) { + out->struct_size = sizeof(*out); + out->self = factory->engine; + out->retain = retain_engine; + out->release = release_engine; + factory->last_error = "factory failed"; + return -7; + } + factory->seen = *config; + factory->seen_model_path = config->model_path ? config->model_path : ""; + factory->seen_mmproj_path = config->mmproj_path ? config->mmproj_path : ""; + factory->seen_backend = config->backend ? config->backend : ""; + out->struct_size = factory->return_undersized ? 8u : sizeof(*out); + out->self = factory->return_null_self ? nullptr : factory->engine; + out->retain = (factory->return_borrowed ? nullptr : retain_engine); + out->release = (factory->return_borrowed || factory->return_retain_only) + ? nullptr + : release_engine; + out->set_input = factory->return_missing_set_input ? nullptr : set_input; + out->run_infer = run_infer; + out->get_output = get_output; + out->last_error = last_error; + return 0; +} + +const char* factory_last_error(void* p) { + return static_cast(p)->last_error; +} + +const frt_generic_stage_plan_ext_v1* generic_plan( + const frt_model_runtime_v1* model) { + if (!model || + model->struct_size < FRT_MODEL_RUNTIME_V1_QUERY_EXTENSION_SIZE || + !model->query_extension) return nullptr; + const void* extension = nullptr; + if (model->query_extension( + model, FRT_EXT_GENERIC_STAGE_PLAN_V1, 1, &extension) != 0) { + return nullptr; + } + return static_cast(extension); +} + +} // namespace + +int main() { + frt_model_runtime_v1* bad = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine(nullptr, nullptr, &bad) + == -1 && bad == nullptr, + "create rejects missing config and engine"); + + FakeEngine engine; + frt_llama_cpp_engine_v1 engine_api{}; + engine_api.struct_size = sizeof(engine_api); + engine_api.self = &engine; + engine_api.retain = retain_engine; + engine_api.release = release_engine; + engine_api.set_input = set_input; + engine_api.run_infer = run_infer; + engine_api.get_output = get_output; + engine_api.last_error = last_error; + + FakeEngine old_engine; + frt_llama_cpp_engine_v1 old_engine_api = engine_api; + old_engine_api.self = &old_engine; + old_engine_api.struct_size = FRT_LLAMA_CPP_ENGINE_V1_BASE_SIZE; + old_engine_api.run_stage = nullptr; + + frt_model_runtime_v1* old_pi0 = nullptr; + + frt_llama_cpp_llm_config llm_cfg{}; + llm_cfg.struct_size = FRT_LLAMA_CPP_LLM_CONFIG_BASE_SIZE; + llm_cfg.model_path = "/models/llm.gguf"; + llm_cfg.backend = "cpu"; + llm_cfg.n_ctx = 2048; + llm_cfg.max_tokens = 16; + frt_model_runtime_v1* old_llm = nullptr; + CHECK(frt_llama_cpp_llm_runtime_create_with_engine( + &llm_cfg, &old_engine_api, &old_llm) == 0 && old_llm && + old_llm->n_ports == 2 && generic_plan(old_llm) && + generic_plan(old_llm)->n_stages == 1, + "old-prefix LLM engine exposes only infer schema"); + if (old_llm) old_llm->release(old_llm->owner); + + frt_llama_cpp_mllm_config mllm_cfg{}; + mllm_cfg.struct_size = FRT_LLAMA_CPP_MLLM_CONFIG_BASE_SIZE; + mllm_cfg.model_path = "/models/mllm.gguf"; + mllm_cfg.mmproj_path = "/models/mllm-mmproj.gguf"; + mllm_cfg.backend = "cpu"; + mllm_cfg.n_ctx = 2048; + mllm_cfg.max_tokens = 16; + frt_model_runtime_v1* old_mllm = nullptr; + CHECK(frt_llama_cpp_mllm_runtime_create_with_engine( + &mllm_cfg, &old_engine_api, &old_mllm) == 0 && old_mllm && + old_mllm->n_ports == 3 && generic_plan(old_mllm) && + generic_plan(old_mllm)->n_stages == 1, + "old-prefix MLLM engine exposes only infer schema"); + if (old_mllm) old_mllm->release(old_mllm->owner); + + frt_llama_cpp_pi0_config cfg{}; + cfg.struct_size = FRT_LLAMA_CPP_PI0_CONFIG_BASE_SIZE; + cfg.model_path = "/models/pi0.gguf"; + cfg.mmproj_path = "/models/pi0-mmproj.gguf"; + cfg.backend = "cpu"; + cfg.n_views = 2; + cfg.image_height = 224; + cfg.image_width = 224; + cfg.image_channels = 3; + cfg.action_steps = 2; + cfg.action_dim = 2; + + cfg.n_views = 4; + bad = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine( + &cfg, &old_engine_api, &bad) == -1 && bad == nullptr, + "create rejects more than three Pi0 views"); + cfg.n_views = 2; + + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine( + &cfg, &old_engine_api, &old_pi0) == 0 && old_pi0 && + generic_plan(old_pi0) && + generic_plan(old_pi0)->n_stages == 1, + "old-prefix Pi0 engine exposes only infer schema"); + if (old_pi0) old_pi0->release(old_pi0->owner); + + frt_llama_cpp_engine_v1 asymmetric = engine_api; + asymmetric.retain = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine(&cfg, &asymmetric, + &bad) == -1 && + bad == nullptr, + "create rejects release without retain"); + + asymmetric = engine_api; + asymmetric.release = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine(&cfg, &asymmetric, + &bad) == -1 && + bad == nullptr, + "create rejects retain without release"); + + frt_llama_cpp_engine_v1 borrowed = engine_api; + borrowed.retain = nullptr; + borrowed.release = nullptr; + frt_model_runtime_v1* borrowed_model = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine(&cfg, &borrowed, + &borrowed_model) == 0 && + borrowed_model, + "create accepts a borrowed explicit engine"); + borrowed_model->release(borrowed_model->owner); + + frt_model_runtime_v1* model = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine(&cfg, &engine_api, + &model) == 0 && model, + "create llama_cpp Pi0 runtime with explicit engine"); + CHECK(engine.retains == 1, "engine retained once"); + const frt_generic_stage_plan_ext_v1* plan = generic_plan(model); + CHECK(model->abi_version == FRT_MODEL_RUNTIME_ABI_VERSION && + model->n_ports == 4 && model->n_stages == 0 && plan && + plan->n_stages == 1, + "runtime exposes v1 metadata and one generic authority"); + CHECK(model->exp && model->exp->ctx == nullptr && + model->exp->n_graphs == 0, + "provider runtime has no FlashRT exec graph"); + CHECK(std::strcmp(model->ports[FRT_LLAMA_CPP_PI0_PORT_IMAGES].name, + "images") == 0 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_IMAGES].update == + FRT_RT_PORT_STAGED && + model->ports[FRT_LLAMA_CPP_PI0_PORT_IMAGES].shape[0] == 2 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_IMAGES].shape[3] == 3, + "images port schema"); + CHECK(std::strcmp(model->ports[FRT_LLAMA_CPP_PI0_PORT_PROMPT].name, + "prompt") == 0 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_PROMPT].modality == + FRT_RT_MOD_TEXT && + model->ports[FRT_LLAMA_CPP_PI0_PORT_PROMPT].shape[0] == -1, + "prompt port schema"); + CHECK(std::strcmp(model->ports[FRT_LLAMA_CPP_PI0_PORT_STATE].name, + "state") == 0 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_STATE].dtype == + FRT_RT_DTYPE_F32 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_STATE].shape[0] == -1, + "state port schema is model-variable"); + CHECK(std::strcmp(model->ports[FRT_LLAMA_CPP_PI0_PORT_ACTIONS].name, + "actions") == 0 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_ACTIONS].direction == + FRT_RT_PORT_OUT && + model->ports[FRT_LLAMA_CPP_PI0_PORT_ACTIONS].shape[0] == 2 && + model->ports[FRT_LLAMA_CPP_PI0_PORT_ACTIONS].shape[1] == 2, + "actions port schema"); + CHECK(std::strcmp(plan->stages[0].name, "infer") == 0 && + plan->stages[0].executor_kind == FRT_GENERIC_STAGE_OPAQUE && + plan->stages[0].executor_ref == 37, + "infer generic stage schema"); + CHECK(std::strstr(model->exp->identity, "provider=llama_cpp\n") && + std::strstr(model->exp->identity, "model_family=pi0\n") && + std::strstr(model->exp->identity, + "gstage-v1:0:5:infer:1:37:0:"), + "identity carries provider, model family, and generic stage"); + + CHECK(model->verbs.set_input(model->self, FRT_LLAMA_CPP_PI0_PORT_PROMPT, + "hello", 5, -1) == 0 && + engine.set_input_calls == 1 && + engine.last_port == FRT_LLAMA_CPP_PI0_PORT_PROMPT, + "set_input delegates to engine"); + CHECK(plan->run_opaque(plan->stage_self, + plan->stages[0].executor_ref) == 0 && + engine.infer == 1, + "generic runner delegates executor ref to the engine"); + float out[4] = {}; + uint64_t written = 0; + CHECK(model->verbs.get_output(model->self, + FRT_LLAMA_CPP_PI0_PORT_ACTIONS, + out, sizeof(out), &written, -1) == 0 && + written == sizeof(out) && std::fabs(out[0] - 11.0f) < 0.01f, + "get_output delegates to engine"); + CHECK(model->verbs.prepare(model->self, 0, 0) == -3 && + std::strstr(model->verbs.last_error(model->self), + "no graph variants"), + "prepare hard-errors instead of fabricating graph variants"); + model->release(model->owner); + CHECK(engine.releases == 1, "engine released once"); + + // A callback without the explicit capability may be the old cached-action + // shim, so it must not publish a staged plan. + FakeEngine staged_engine; + frt_llama_cpp_engine_v1 staged_engine_api = engine_api; + staged_engine_api.self = &staged_engine; + staged_engine_api.run_stage = run_stage; + staged_engine_api.struct_size = FRT_LLAMA_CPP_ENGINE_V1_RUN_STAGE_SIZE; + frt_model_runtime_v1* staged_model = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine( + &cfg, &staged_engine_api, &staged_model) == 0 && staged_model, + "create staged Pi0 runtime"); + const auto* staged_plan = generic_plan(staged_model); + CHECK(staged_plan && staged_plan->n_stages == 1 && + std::strcmp(staged_plan->stages[0].name, "infer") == 0 && + staged_plan->stages[0].executor_ref == 37, + "uncapable Pi0 runtime publishes a single infer authority"); + CHECK(staged_model->verbs.step(staged_model->self) == 0 && + staged_engine.infer == 1, + "whole-model step delegates to infer (run_stage not used)"); + if (staged_model) staged_model->release(staged_model->owner); + + FakeEngine capable_engine; + frt_llama_cpp_engine_v1 capable_engine_api = staged_engine_api; + capable_engine_api.self = &capable_engine; + capable_engine_api.reserved = + FRT_LLAMA_CPP_ENGINE_CAP_PI0_REAL_CONTEXT_ACTION; + frt_model_runtime_v1* capable_model = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine( + &cfg, &capable_engine_api, &capable_model) == 0 && capable_model, + "create real-split Pi0 runtime"); + const auto* capable_plan = generic_plan(capable_model); + CHECK(capable_plan && capable_plan->n_stages == 2 && + std::strcmp(capable_plan->stages[0].name, "context") == 0 && + std::strcmp(capable_plan->stages[1].name, "action") == 0 && + capable_plan->stages[1].n_after == 1 && + capable_plan->stages[1].after[0] == 0, + "capable Pi0 runtime publishes context then action"); + CHECK(capable_plan->run_opaque( + capable_plan->stage_self, + capable_plan->stages[0].executor_ref) == 0 && + capable_engine.last_stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT, + "context authority delegates to the real split callback"); + CHECK(capable_plan->run_opaque( + capable_plan->stage_self, + capable_plan->stages[1].executor_ref) == 0 && + capable_engine.last_stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION, + "action authority delegates to the real split callback"); + CHECK(capable_model->verbs.step(capable_model->self) == 0 && + capable_engine.last_stage == FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, + "whole-model step remains available on a split-capable engine"); + if (capable_model) capable_model->release(capable_model->owner); + + FakeEngine null_error_engine; + frt_llama_cpp_engine_v1 null_error_api = engine_api; + null_error_api.self = &null_error_engine; + null_error_api.last_error = null_last_error; + frt_model_runtime_v1* null_error_model = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine( + &cfg, &null_error_api, &null_error_model) == 0 && + null_error_model, + "create accepts engine with nullable error implementation"); + CHECK(null_error_model->verbs.get_output( + null_error_model->self, FRT_LLAMA_CPP_PI0_PORT_PROMPT, + out, sizeof(out), &written, -1) == -1 && + std::strstr(null_error_model->verbs.last_error( + null_error_model->self), + "null error"), + "null engine errors are reported without crashing"); + null_error_model->release(null_error_model->owner); + + FakeEngine factory_engine; + FakeFactory factory; + factory.engine = &factory_engine; + frt_llama_cpp_engine_factory_v1 factory_api{}; + factory_api.struct_size = sizeof(factory_api); + factory_api.self = &factory; + factory_api.create_pi0 = create_pi0_engine; + factory_api.last_error = factory_last_error; + + const std::string identity_prefix = + "/tmp/flashrt-identity-" + std::to_string(getpid()); + const std::string identity_model = identity_prefix + "-model.gguf"; + const std::string identity_mmproj = identity_prefix + "-mmproj.gguf"; + { + std::string model_bytes(64 * 1024 + 16, 'a'); + model_bytes.back() = 'x'; + std::ofstream(identity_model, std::ios::binary) + .write(model_bytes.data(), model_bytes.size()); + std::ofstream(identity_mmproj, std::ios::binary) << "mmproj-a"; + } + const std::string open_json = + "{\"model_family\":\"pi0\",\"model_path\":\"" + + identity_model + "\",\"mmproj_path\":\"" + identity_mmproj + + "\",\"backend\":\"cpu\",\"n_views\":2," + "\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":2,\"action_dim\":2}"; + frt_model_runtime_v1* opened = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &factory_api, &opened) == 0 && + opened && factory.creates == 1, + "open_with_engine_factory creates a Pi0 runtime from JSON"); + CHECK(factory.seen_model_path == identity_model && + factory.seen_mmproj_path == identity_mmproj && + factory.seen_backend == "cpu" && + factory.seen.n_views == 2 && + factory.seen.action_steps == 2 && + factory.seen.action_dim == 2, + "factory receives parsed Pi0 config"); + CHECK(factory_engine.retains == 1 && factory_engine.releases == 1, + "factory engine reference is transferred to the runtime"); + CHECK(opened->verbs.step(opened->self) == 0 && + factory_engine.infer == 1, + "opened runtime step delegates whole-model infer"); + const uint64_t first_fingerprint = opened->exp->fingerprint; + opened->release(opened->owner); + CHECK(factory_engine.releases == 2, + "opened runtime releases retained factory engine"); + { + std::string model_bytes(64 * 1024 + 16, 'a'); + model_bytes.back() = 'y'; + std::ofstream(identity_model, std::ios::binary | std::ios::trunc) + .write(model_bytes.data(), model_bytes.size()); + } + frt_model_runtime_v1* changed_checkpoint = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &factory_api, &changed_checkpoint) == 0 && + changed_checkpoint && + changed_checkpoint->exp->fingerprint != first_fingerprint && + std::strstr(changed_checkpoint->exp->identity, + "weights_sha256=ee5c5dda1cef40957d80fafab7b0eeddeea052e221a311ee80ec4972cbdc5d5f"), + "same-size checkpoint tail change alters deployment fingerprint"); + const uint64_t model_changed_fingerprint = + changed_checkpoint ? changed_checkpoint->exp->fingerprint : 0; + if (changed_checkpoint) changed_checkpoint->release(changed_checkpoint->owner); + { + std::ofstream(identity_mmproj, std::ios::binary | std::ios::trunc) + << "mmproj-b"; + } + frt_model_runtime_v1* changed_mmproj = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &factory_api, &changed_mmproj) == 0 && + changed_mmproj && + changed_mmproj->exp->fingerprint != model_changed_fingerprint && + std::strstr(changed_mmproj->exp->identity, + "mmproj_sha256=d772e18bbe6501374cfab6a5d76a93f9288b8f2bcb4ec9694356b32e3bcf4c1c"), + "mmproj prefix changes deployment fingerprint"); + + const std::string cuda_json = + "{\"model_family\":\"pi0\",\"model_path\":\"" + + identity_model + "\",\"mmproj_path\":\"" + identity_mmproj + + "\",\"backend\":\"cuda\",\"n_views\":2," + "\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":2,\"action_dim\":2}"; + frt_model_runtime_v1* cuda_identity = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + cuda_json.c_str(), &factory_api, &cuda_identity) == 0 && + cuda_identity && + changed_mmproj && + cuda_identity->exp->fingerprint != + changed_mmproj->exp->fingerprint, + "backend change alters deployment fingerprint"); + bool same_port_schema = cuda_identity && changed_mmproj && + cuda_identity->n_ports == changed_mmproj->n_ports; + if (same_port_schema) { + for (uint64_t i = 0; i < cuda_identity->n_ports; ++i) { + const frt_runtime_port_desc& lhs = cuda_identity->ports[i]; + const frt_runtime_port_desc& rhs = changed_mmproj->ports[i]; + same_port_schema &= std::strcmp(lhs.name, rhs.name) == 0 && + lhs.modality == rhs.modality && + lhs.dtype == rhs.dtype && + lhs.layout == rhs.layout && + lhs.direction == rhs.direction && + lhs.update == rhs.update && + lhs.required == rhs.required && + lhs.rank == rhs.rank; + for (uint32_t dim = 0; same_port_schema && dim < lhs.rank; ++dim) { + same_port_schema &= lhs.shape[dim] == rhs.shape[dim]; + } + } + } + CHECK(same_port_schema, "backend switch preserves port schema"); + if (cuda_identity) cuda_identity->release(cuda_identity->owner); + if (changed_mmproj) changed_mmproj->release(changed_mmproj->owner); + + const int creates_before_missing_checkpoint = factory.creates; + const std::string missing_checkpoint_json = + "{\"model_family\":\"pi0\",\"model_path\":\"" + + identity_prefix + "-missing.gguf\",\"mmproj_path\":\"" + + identity_mmproj + + "\",\"backend\":\"cpu\",\"n_views\":2," + "\"image_height\":224,\"image_width\":224," + "\"image_channels\":3,\"action_steps\":2,\"action_dim\":2}"; + frt_model_runtime_v1* missing_checkpoint = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + missing_checkpoint_json.c_str(), &factory_api, + &missing_checkpoint) == -1 && + missing_checkpoint == nullptr && + factory.creates == creates_before_missing_checkpoint && + std::strstr(frt_llama_cpp_runtime_open_error(), + "failed to open checkpoint for identity"), + "missing checkpoint reports identity error before factory creation"); + + frt_model_runtime_v1* missing = nullptr; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + "{\"model_family\":\"pi0\",\"model_path\":\"/models/pi0.gguf\"}", + &factory_api, &missing) == -1 && + missing == nullptr, + "open rejects incomplete JSON config"); + CHECK(std::strstr(frt_llama_cpp_runtime_open_error(), + "invalid Pi0 runtime open arguments or JSON config"), + "incomplete JSON reports runtime-open validation error"); + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + "{\"model_family\":\"pi0\",\"model_path\":\"/models/pi0.gguf\"," + "\"mmproj_path\":\"/models/pi0-mmproj.gguf\"," + "\"backend\":\"cpu\",\"n_views\":2,\"image_height\":224," + "\"image_width\":224,\"image_channels\":3," + "\"action_steps\":2,\"action_dim\":2,\"unexpected\":1}", + &factory_api, &missing) == -1 && + missing == nullptr, + "open rejects unknown JSON fields"); + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + "{\"model_family\":\"pi0\",\"model_family\":\"pi0\"," + "\"model_path\":\"/models/pi0.gguf\"," + "\"mmproj_path\":\"/models/pi0-mmproj.gguf\"," + "\"backend\":\"cpu\",\"n_views\":2,\"image_height\":224," + "\"image_width\":224,\"image_channels\":3," + "\"action_steps\":2,\"action_dim\":2}", + &factory_api, &missing) == -1 && + missing == nullptr, + "open rejects duplicate JSON fields"); + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + "{\"model_family\":\"llm\",\"model_path\":\"/models/pi0.gguf\"," + "\"mmproj_path\":\"/models/pi0-mmproj.gguf\"," + "\"backend\":\"cpu\",\"n_views\":2,\"image_height\":224," + "\"image_width\":224,\"image_channels\":3," + "\"action_steps\":2,\"action_dim\":2}", + &factory_api, &missing) == -1 && + missing == nullptr, + "open rejects non-Pi0 model family"); + + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + "{\"model_family\":\"pi0\",\"model_path\":\"/models/pi0.gguf\"," + "\"mmproj_path\":\"/models/pi0-mmproj.gguf\"," + "\"backend\":\"cpu\",\"n_views\":0,\"image_height\":224," + "\"image_width\":224,\"image_channels\":3," + "\"action_steps\":2,\"action_dim\":2}", + &factory_api, &missing) == -1 && + missing == nullptr, + "open rejects zero-sized numeric config fields"); + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + "{\"model_family\":\"pi0\",\"model_path\":\"/models/pi0.gguf\"," + "\"mmproj_path\":\"/models/pi0-mmproj.gguf\"," + "\"backend\":\"cpu\",\"n_views\":4,\"image_height\":224," + "\"image_width\":224,\"image_channels\":3," + "\"action_steps\":2,\"action_dim\":2}", + &factory_api, &missing) == -1 && + missing == nullptr, + "open rejects more than three Pi0 views"); + FakeFactory borrowed_factory = factory; + borrowed_factory.engine = &factory_engine; + borrowed_factory.return_borrowed = true; + frt_llama_cpp_engine_factory_v1 borrowed_factory_api = factory_api; + borrowed_factory_api.self = &borrowed_factory; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &borrowed_factory_api, &missing) == -1 && + missing == nullptr, + "open rejects borrowed engines from factories"); + const int releases_before_invalid_engine = factory_engine.releases; + FakeFactory invalid_factory = factory; + invalid_factory.engine = &factory_engine; + invalid_factory.return_undersized = true; + frt_llama_cpp_engine_factory_v1 invalid_factory_api = factory_api; + invalid_factory_api.self = &invalid_factory; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &invalid_factory_api, &missing) == -1 && + missing == nullptr && + factory_engine.releases == releases_before_invalid_engine, + "open rejects undersized factory engines without calling release"); + invalid_factory = factory; + invalid_factory.engine = &factory_engine; + invalid_factory.return_null_self = true; + invalid_factory_api.self = &invalid_factory; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &invalid_factory_api, &missing) == -1 && + missing == nullptr && + factory_engine.releases == releases_before_invalid_engine, + "open rejects null factory engine self without calling release"); + invalid_factory = factory; + invalid_factory.engine = &factory_engine; + invalid_factory.return_retain_only = true; + invalid_factory_api.self = &invalid_factory; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &invalid_factory_api, &missing) == -1 && + missing == nullptr && + factory_engine.releases == releases_before_invalid_engine, + "open rejects asymmetric factory engines without calling release"); + invalid_factory = factory; + invalid_factory.engine = &factory_engine; + invalid_factory.return_missing_set_input = true; + invalid_factory_api.self = &invalid_factory; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &invalid_factory_api, &missing) == -1 && + missing == nullptr && + factory_engine.releases == releases_before_invalid_engine + 1, + "open releases owned factory engines missing hot-path hooks"); + invalid_factory = factory; + invalid_factory.engine = &factory_engine; + invalid_factory.fail_after_engine = true; + invalid_factory_api.self = &invalid_factory; + CHECK(frt_llama_cpp_pi0_runtime_open_with_engine_factory( + open_json.c_str(), &invalid_factory_api, &missing) == -7 && + missing == nullptr && + factory_engine.releases == releases_before_invalid_engine + 1, + "open ignores out_engine when factory create fails"); + + std::remove(identity_model.c_str()); + std::remove(identity_mmproj.c_str()); + + std::printf(g_fail ? "\n== LLAMA_CPP PROVIDER FAILED ==\n" + : "\n== LLAMA_CPP PROVIDER PASSED ==\n"); + return g_fail; +} diff --git a/docs/jetson_pi_usage.md b/docs/jetson_pi_usage.md new file mode 100644 index 00000000..923b5795 --- /dev/null +++ b/docs/jetson_pi_usage.md @@ -0,0 +1,168 @@ +# Jetson-PI Provider + +The Jetson-PI integration is an optional llama.cpp/GGML model provider. It +publishes the standard `frt_model_runtime_v1` contract and the selected generic +stage-plan extension. FlashRT runtime code contains no Jetson-PI or model-family +branches. + +## Build + +The provider is disabled by default. Install Jetson-PI first, then install +FlashRT into the same prefix so the provider and its dependency closure are +co-located: + +```bash +cmake -S cpp -B build/provider \ + -DFLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER=ON \ + -DFLASHRT_CPP_WITH_JETSON_PI=ON \ + -DJetsonPI_ROOT= +cmake --build build/provider --target flashrt_cpp_llama_cpp_provider_c +cmake --install build/provider --prefix +``` + +`` and `` should identify the same +deployment prefix. FlashRT does not copy or take ownership of third-party +libraries. + +For development against a Jetson-PI source tree, use +`-DJETSON_PI_ROOT=` instead. Set exactly one of +`JETSON_PI_ROOT` and `JetsonPI_ROOT`. + +`FLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER=ON` without +`FLASHRT_CPP_WITH_JETSON_PI=ON` builds the provider contract and fake-engine +tests without linking Jetson-PI. This mode is suitable for CPU-only contract +CI. + +The installed provider DSO exports one entry point: + +```c +int frt_model_runtime_open_v1( + const char* config_json, frt_model_runtime_v1** out); +``` + +All provider and backend implementation symbols remain hidden. + +## Runtime Contract + +The returned object always uses model-runtime v1: + +- ports describe model inputs and outputs; +- `step()` performs one whole-model request; +- `FRT_EXT_GENERIC_STAGE_PLAN_V1` publishes one selected stage plan; +- Jetson-PI state, KV caches, samplers, and intermediate tensors remain + provider-private; +- outputs use ordinary STAGED `get_output`; no provider-specific buffer ABI is + required. + +VLA (Pi0) engines that declare the real context/action capability publish +`context -> action`; `step()` remains the whole-request `infer` entry point. +Older engine vtables that do not declare the capability publish their original +single `infer` plan. Autoregressive LLM/MLLM engines with staged capability +publish `reset -> prefill -> decode`; engines without that capability publish a +single `infer` stage. All plans use OPAQUE executors, so a generic host never +sees llama.cpp or GGML types. + +[PKU-SEC-Lab/Jetson-PI-Edge#1](https://github.com/PKU-SEC-Lab/Jetson-PI-Edge/pull/1) +is merged and provides the real PI0.5 prepare/decode boundary plus +`Task: ..., State: ...;\nAction:` serialization. The follow-up +[Jetson-PI-Edge#2](https://github.com/PKU-SEC-Lab/Jetson-PI-Edge/pull/2) +adds the legacy Pi0 prepare path with owned VIT embeddings, fixes persistent +GPU encoded-KV lifetime, and declares the capability only when the boundary is +valid for both detected Pi0 model kinds. FlashRT does not infer capability from +the presence of a callback alone. + +## Python + +### VLA + +```python +import flash_rt + +model = flash_rt.load_model( + "policy.gguf", + framework="jetson_pi", + config="pi0", + mmproj_path="vision.gguf", + backend="cuda", + num_views=2, + action_steps=10, + action_dim=32, +) + +actions = model.predict(images, prompt=prompt, state=state) +``` + +`images` is an ordered list of contiguous RGB `uint8` HWC arrays. `state` is +the non-empty, finite, policy-normalized proprioception vector, passed without +padding. The provider does not load checkpoint normalization statistics or +normalize raw physical observations. PI0.5 accepts at most 8 `float32` values, +while legacy Pi0 accepts at most `action_dim` and lets the backend zero-pad +shorter input. The result is a copied +`float32[action_steps, action_dim]` array. + +`model.predict(...)` runs one whole request. With a split-capable backend, +`model._pipe.context(...)` prepares and retains one encoded context and +`model._pipe.action()` consumes it exactly once. Replacing or rejecting an +input invalidates any pending context. A backend without the explicit +capability exposes only `predict(...)`. + +### Text LLM + +```python +llm = flash_rt.load_model( + "model.gguf", + framework="jetson_pi", + config="llm", + backend="cuda", + max_tokens=128, +) + +text = llm.generate("Write a short summary.") +``` + +For host-driven decoding, call `reset()`, `prefill()`, and `decode()`. The host +may stop by not issuing another decode call. + +### Multimodal LLM + +```python +mllm = flash_rt.load_model( + "model.gguf", + framework="jetson_pi", + config="mllm", + mmproj_path="vision.gguf", + backend="cuda", + max_tokens=128, +) + +text = mllm.generate(images, "Describe the scene.") +``` + +## C++ Host + +Load the provider DSO, resolve `FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL`, and probe the +returned object using `FRT_MODEL_RUNTIME_V1_BASE_SIZE`. Probe +`query_extension` only when `struct_size` reaches +`FRT_MODEL_RUNTIME_V1_QUERY_EXTENSION_SIZE`. + +When `FRT_EXT_GENERIC_STAGE_PLAN_V1` is present: + +1. Validate the extension version and size. +2. Validate the dependency DAG. +3. Stage input ports through `set_input`. +4. Execute the selected plan in dependency order. +5. Read staged outputs through `get_output`. +6. Release the model runtime once. + +Nexus provides the generic host-side adopter for this flow. Provider code owns +model execution; Nexus owns lifecycle and scheduling policy. + +## Failure Rules + +- Unknown or duplicate JSON fields are rejected. +- Missing checkpoints fail before engine creation. +- Unsupported backends fail instead of falling back silently. +- Failed input replacement invalidates the old ready state. +- Stage ordering and decode budgets are enforced by the provider. +- Provider identity includes checkpoint digests, backend, port schema, and the + selected stage plan. diff --git a/docs/jetson_pi_validation_matrix.md b/docs/jetson_pi_validation_matrix.md new file mode 100644 index 00000000..3b034072 --- /dev/null +++ b/docs/jetson_pi_validation_matrix.md @@ -0,0 +1,91 @@ +# Jetson-PI Validation + +This matrix defines the merge gates for the optional provider. Commands use +placeholders intentionally; validation records must not publish workstation +paths, device assignments, model storage locations, or private environment +details. + +## Required Gates + +| Area | Gate | +|---|---| +| Default isolation | Provider options OFF builds no provider target or DSO and changes no existing model route. | +| Provider-only CI | Contract tests build with CUDA, exec, and Jetson-PI disabled. | +| Public ABI | Provider returns `frt_model_runtime_v1`; no parallel runtime ABI is defined. | +| Extension | Exactly one valid `GENERIC_STAGE_PLAN_V1` authority is published. | +| DSO surface | Dynamic exports equal `frt_model_runtime_open_v1` plus the version node. | +| Lifecycle | Open, retain/release, repeated instances, and all failure paths are leak-free. | +| Identity | Checkpoint bytes, backend, port schema, and selected plan affect the fingerprint. | +| VLA | Whole-step `infer`, split `context -> action`, and the direct provider path produce identical action buffers. The split is published only when the backend declares a real context/action capability. | +| LLM | Whole generation and host-driven `reset -> prefill -> decode` match exactly under deterministic sampling. | +| MLLM | Whole generation and staged decode match exactly under deterministic sampling. | +| Install tree | An out-of-tree consumer loads the installed DSO without source-tree headers or build-tree RPATH. | +| Nexus | Generic OPAQUE execution is synchronous, dependency ordered, and does not alter native graph adoption. | + +## Contract Tests + +The repository keeps focused tests for: + +- old-prefix and current engine-vtable compatibility; +- strict JSON parsing and checkpoint identity; +- port schema, selected-plan records, stable executor references, and DAG + dependencies; +- missing staged verbs, stale output, failed replacement, decode budget, and + multi-instance isolation; +- direct narrow-API parity for VLA, LLM, and MLLM paths; +- Python whole-request and staged APIs; +- exact provider DSO exports. + +Tests requiring model files skip explicitly when fixtures are unavailable. +Skipping is not parity evidence; release qualification must provide the model +matrix below. + +## Release Matrix + +At least one supported device for each advertised backend must run: + +| Model face | Whole request | Selected plan | Direct parity | Repeated session | +|---|---:|---:|---:|---:| +| VLA | required | `context -> action` | required | required | +| Text LLM | required | required | required | required | +| Multimodal LLM | required | required | required | required | + +VLA qualification uses a backend that explicitly declares the real +context/action capability. The merged backend change +[Jetson-PI-Edge#1](https://github.com/PKU-SEC-Lab/Jetson-PI-Edge/pull/1) +provides the PI0.5 encode/decode boundary and prompt-state serialization. The +follow-up [Jetson-PI-Edge#2](https://github.com/PKU-SEC-Lab/Jetson-PI-Edge/pull/2) +extends the real boundary to legacy Pi0 with owned VIT embeddings and publishes +the cross-model capability contract. The opt-in parity tests also require +actual GGUF model-kind detection, one-shot pending-context semantics, and +whole/split byte equality. + +For deterministic paths, text/token outputs are exact and VLA action buffers +are byte-identical unless the provider documents a backend-specific numerical +contract. Shape, dtype, finite-value, stale-state, and teardown checks are +always required. + +## Portable Build Check + +```bash +cmake -S cpp -B build/provider-test \ + -DBUILD_TESTING=ON \ + -DFLASHRT_CPP_WITH_LLAMA_CPP_PROVIDER=ON \ + -DFLASHRT_CPP_WITH_JETSON_PI=ON \ + -DJetsonPI_ROOT= +cmake --build build/provider-test +ctest --test-dir build/provider-test --output-on-failure +``` + +After installation, audit the DSO using the platform's dynamic-symbol and +runtime-dependency tools. Linux acceptance requires: + +```text +defined public symbol: frt_model_runtime_open_v1 +RPATH/RUNPATH: $ORIGIN or no embedded path +build/source paths: absent +``` + +The final cross-repository gate installs FlashRT provider and Nexus into clean +prefixes, then runs both the native graph producer and this generic provider +through the same Nexus host lifecycle. diff --git a/flash_rt/api.py b/flash_rt/api.py index d6fcf0c5..4e4a64d7 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -299,7 +299,20 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, use_fp16=False, use_fp8=True, state_prompt_mode="exact", - state_prompt_fixed_max_len=None): + state_prompt_fixed_max_len=None, + *, + mmproj_path=None, + backend="cpu", + action_steps=None, + action_dim=None, + lib_path=None, + n_ctx=0, + n_threads=0, + temp=0.8, + top_k=40, + top_p=0.9, + seed=1, + max_tokens=512): """Load a FlashRT model. Args: @@ -436,16 +449,64 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, "Qwen3VlTorchFrontendRtx\n" "See docs/qwen3_vl_fp8_sm89.md and docs/qwen3_vl_nvfp4.md.") - if config not in ("pi05", "groot", "groot_n17", "pi0", "pi0fast", + if framework == "jetson_pi": + if config not in ("pi0", "pi05", "llm", "mllm"): + raise ValueError( + f"Unknown Jetson-PI config: {config}. " + "Supported: pi0, pi05, llm, mllm") + elif config not in ("pi05", "groot", "groot_n17", "pi0", "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video", "cosmos3_edge", "nexn2"): raise ValueError( f"Unknown config: {config}. " f"Supported: pi05, groot, groot_n17, pi0, pi0fast, motus, " f"wan22_ti2v_5b, cosmos3_video, cosmos3_edge, nexn2") - if framework not in ("torch", "jax"): + if framework not in ("torch", "jax", "jetson_pi"): raise ValueError( - f"Unknown framework: {framework}. Supported: torch, jax") + f"Unknown framework: {framework}. Supported: torch, jax, jetson_pi") + + # Drives the Jetson-PI provider through frt_model_runtime_v1 via ctypes. + # No torch/jax or GPU architecture detection is involved. The action chunk + # shape is passed explicitly by the caller (the verified pi0_base is 50x32). + if framework == "jetson_pi": + if config == "llm": + from flash_rt.frontends.jetson_pi.llm import LlmJetsonPiFrontend + return LlmJetsonPiFrontend( + checkpoint, + backend=backend, + n_ctx=n_ctx, + n_threads=n_threads, + temp=temp, + top_k=top_k, + top_p=top_p, + seed=seed, + max_tokens=max_tokens, + lib_path=lib_path) + if config == "mllm": + from flash_rt.frontends.jetson_pi.mllm import MllmJetsonPiFrontend + return MllmJetsonPiFrontend( + checkpoint, + mmproj_path=mmproj_path, + backend=backend, + n_ctx=n_ctx, + n_threads=n_threads, + temp=temp, + top_k=top_k, + top_p=top_p, + seed=seed, + max_tokens=max_tokens, + lib_path=lib_path) + # default / "pi0": VLA path + from flash_rt.frontends.jetson_pi.pi0 import Pi0JetsonPiFrontend + pipe = Pi0JetsonPiFrontend( + checkpoint, + mmproj_path=mmproj_path, + backend=backend, + num_views=num_views, + action_steps=action_steps, + action_dim=action_dim, + lib_path=lib_path) + return VLAModel(pipe, framework) # When use_fp4=True, the default resolves to the best-known production # FP4 config (full 18 encoder FFN layers + AWQ + P1 split-GU). Passing diff --git a/flash_rt/frontends/jetson_pi/__init__.py b/flash_rt/frontends/jetson_pi/__init__.py new file mode 100644 index 00000000..0549ec6a --- /dev/null +++ b/flash_rt/frontends/jetson_pi/__init__.py @@ -0,0 +1,12 @@ +"""Jetson-PI (llama.cpp/GGML) frontends for FlashRT. + +- :class:`Pi0JetsonPiFrontend` — Pi0 VLA (vision-language-action) provider. +- :class:`LlmJetsonPiFrontend` — generic GGUF LLM (text completion) provider. +- :class:`MllmJetsonPiFrontend` — multimodal LLM (vision+text) provider. +""" + +from .pi0 import Pi0JetsonPiFrontend +from .llm import LlmJetsonPiFrontend +from .mllm import MllmJetsonPiFrontend + +__all__ = ["Pi0JetsonPiFrontend", "LlmJetsonPiFrontend", "MllmJetsonPiFrontend"] diff --git a/flash_rt/frontends/jetson_pi/llm.py b/flash_rt/frontends/jetson_pi/llm.py new file mode 100644 index 00000000..cc78064c --- /dev/null +++ b/flash_rt/frontends/jetson_pi/llm.py @@ -0,0 +1,261 @@ +"""Jetson-PI generic GGUF LLM frontend — drives the Jetson-PI llama.cpp LLM +provider through the FlashRT ``frt_model_runtime_v1`` C ABI via ctypes. + +This is the Python entry for plain text completion (Pi0 is a VLA and +goes through :mod:`flash_rt.frontends.jetson_pi.pi0`). One raw prompt in, +one generated text blob out per :meth:`generate` call. The caller is +responsible for applying the chat template; the engine only does raw +prompt -> text. + +``flash_rt.load_model(framework="jetson_pi", config="llm", ...)`` returns a +:class:`LlmJetsonPiFrontend` directly (it is NOT wrapped in ``VLAModel`` — +LLMs are not VLA and do not take images). +""" + +from __future__ import annotations + +import ctypes +import json +import os + +import numpy as np + +from .pi0 import ( # reuse the ctypes mirrors + lib finder + FrtModelRuntimeV1, + _FRT_MODEL_RUNTIME_ABI_VERSION, + _find_lib, + _run_generic_stage, +) + +# Port / stage indices (c_api.h: FRT_LLAMA_CPP_LLM_*) +PORT_PROMPT = 0 +PORT_TEXT = 1 +PORT_NEXT_TOKEN = 2 +PORT_LOGITS = 3 +PORT_IS_EOG = 4 +PORT_TOKENS = 5 +STAGE_INFER = 0 +STAGE_RESET = 1 +STAGE_PREFILL = 2 +STAGE_DECODE = 3 + + +class LlmJetsonPiFrontend: + """Generic GGUF LLM completion frontend backed by the Jetson-PI provider.""" + + def __init__(self, checkpoint, *, backend="cpu", n_ctx=0, n_threads=0, + temp=0.8, top_k=40, top_p=0.9, seed=1, max_tokens=512, + lib_path=None, **_unused): + if max_tokens <= 0: + raise ValueError("max_tokens must be > 0") + self.max_tokens = int(max_tokens) + self._lib_path = _find_lib(lib_path, env_var="FLASHRT_LLM_LIB") + self._lib = ctypes.CDLL(self._lib_path) + + self._lib.frt_model_runtime_open_v1.argtypes = [ + ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)] + self._lib.frt_model_runtime_open_v1.restype = ctypes.c_int + + config = { + "model_family": "llm", + "model_path": str(checkpoint), + "backend": backend, + "n_ctx": int(n_ctx), + "n_threads": int(n_threads), + "temp": float(temp), + "top_k": int(top_k), + "top_p": float(top_p), + "seed": int(seed), + "max_tokens": int(max_tokens), + } + config_json = json.dumps(config, ensure_ascii=False).encode("utf-8") + + model_ptr = ctypes.c_void_p(0) + rc = self._lib.frt_model_runtime_open_v1( + config_json, ctypes.byref(model_ptr)) + if rc != 0 or not model_ptr.value: + raise RuntimeError( + f"frt_model_runtime_open_v1 failed for LLM (rc={rc})") + + self._model = FrtModelRuntimeV1.from_address(model_ptr.value) + if self._model.abi_version != _FRT_MODEL_RUNTIME_ABI_VERSION: + raise RuntimeError( + f"frt_model_runtime_v1 abi_version={self._model.abi_version}, " + f"expected {_FRT_MODEL_RUNTIME_ABI_VERSION}.") + if self._model.struct_size < ctypes.sizeof(FrtModelRuntimeV1): + raise RuntimeError( + f"frt_model_runtime_v1 struct_size={self._model.struct_size} " + f"< ctypes sizeof {ctypes.sizeof(FrtModelRuntimeV1)}.") + self._model_ptr = model_ptr + + def generate(self, prompt): + """Run one whole-prompt completion. Returns the generated text (str).""" + if self._model is None: + raise RuntimeError("LlmJetsonPiFrontend is closed") + if isinstance(prompt, str): + prompt_bytes = prompt.encode("utf-8") + else: + prompt_bytes = bytes(prompt) + verbs = self._model.verbs + self_ = self._model.self + + rc = verbs.set_input(self_, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + rc = verbs.step(self_) + self._check(rc, "step infer") + + written = ctypes.c_uint64(0) + rc = verbs.get_output(self_, PORT_TEXT, None, 0, + ctypes.byref(written), -1) + if rc != -5 or written.value == 0: + self._check(rc, "query text output size") + raise RuntimeError("query text output size returned zero bytes") + cap = written.value + out = (ctypes.c_char * cap)() + written.value = 0 + rc = verbs.get_output(self_, PORT_TEXT, out, cap, + ctypes.byref(written), -1) + self._check(rc, "get_output text") + return out.raw[:written.value].decode("utf-8", errors="replace") + + def reset(self): + """Clear the current KV-cache and sampler state.""" + if self._model is None: + raise RuntimeError("LlmJetsonPiFrontend is closed") + rc = _run_generic_stage(self._model, "reset") + self._check(rc, "generic stage reset") + + def prefill(self, prompt=None, *, tokens=None): + """Start a session from either raw ``prompt`` or int32 ``tokens``.""" + if self._model is None: + raise RuntimeError("LlmJetsonPiFrontend is closed") + if (prompt is None) == (tokens is None): + raise ValueError("exactly one of prompt or tokens is required") + verbs = self._model.verbs + if tokens is not None: + token_array = np.ascontiguousarray(tokens, dtype=np.int32) + if token_array.ndim != 1 or token_array.size == 0: + raise ValueError("tokens must be a non-empty 1-D int32 array") + rc = verbs.set_input( + self._model.self, PORT_TOKENS, token_array.ctypes.data, + token_array.nbytes, -1) + self._check(rc, "set_input tokens") + else: + if isinstance(prompt, str): + prompt_bytes = prompt.encode("utf-8") + else: + prompt_bytes = bytes(prompt) + rc = verbs.set_input(self._model.self, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + rc = _run_generic_stage(self._model, "prefill") + self._check(rc, "generic stage prefill") + return self.get_logits() + + def decode(self): + """Sample and decode one token from the current session.""" + if self._model is None: + raise RuntimeError("LlmJetsonPiFrontend is closed") + verbs = self._model.verbs + rc = _run_generic_stage(self._model, "decode") + self._check(rc, "generic stage decode") + + next_token = ctypes.c_int32() + written = ctypes.c_uint64(0) + rc = verbs.get_output( + self._model.self, PORT_NEXT_TOKEN, ctypes.byref(next_token), + ctypes.sizeof(next_token), ctypes.byref(written), -1) + self._check(rc, "get_output next_token") + if written.value != ctypes.sizeof(next_token): + raise RuntimeError( + f"get_output next_token wrote {written.value} bytes, expected " + f"{ctypes.sizeof(next_token)}") + + is_eog = ctypes.c_int32() + written = ctypes.c_uint64(0) + rc = verbs.get_output( + self._model.self, PORT_IS_EOG, ctypes.byref(is_eog), + ctypes.sizeof(is_eog), ctypes.byref(written), -1) + self._check(rc, "get_output is_eog") + if written.value != ctypes.sizeof(is_eog): + raise RuntimeError( + f"get_output is_eog wrote {written.value} bytes, expected " + f"{ctypes.sizeof(is_eog)}") + + written = ctypes.c_uint64(0) + rc = verbs.get_output(self._model.self, PORT_TEXT, None, 0, + ctypes.byref(written), -1) + if rc != -5 or written.value == 0: + self._check(rc, "query text output size") + raise RuntimeError("query text output size returned zero bytes") + cap = written.value + out = (ctypes.c_char * cap)() + written.value = 0 + rc = verbs.get_output(self._model.self, PORT_TEXT, out, cap, + ctypes.byref(written), -1) + self._check(rc, "get_output text") + return { + "token": int(next_token.value), + "is_eog": bool(is_eog.value), + "text": out.raw[:written.value].decode("utf-8", errors="replace"), + } + + def get_logits(self): + """Copy the current next-token logits into a NumPy float32 array.""" + if self._model is None: + raise RuntimeError("LlmJetsonPiFrontend is closed") + verbs = self._model.verbs + required = ctypes.c_uint64(0) + rc = verbs.get_output( + self._model.self, PORT_LOGITS, None, 0, + ctypes.byref(required), -1) + if rc != -5 or required.value == 0: + self._check(rc, "query logits size") + raise RuntimeError("query logits size returned no bytes") + if required.value % np.dtype(np.float32).itemsize != 0: + raise RuntimeError( + f"logits byte size {required.value} is not float32-aligned") + logits = np.empty( + required.value // np.dtype(np.float32).itemsize, dtype=np.float32) + written = ctypes.c_uint64(0) + rc = verbs.get_output( + self._model.self, PORT_LOGITS, logits.ctypes.data, + logits.nbytes, ctypes.byref(written), -1) + self._check(rc, "get_output logits") + if written.value != logits.nbytes: + raise RuntimeError( + f"get_output logits wrote {written.value} bytes, expected " + f"{logits.nbytes}") + return logits + + def infer(self, observation, debug=False): + """VLAModel-predict-parity entry: observation['prompt'] -> {'text': ...}.""" + _ = debug + prompt = observation.get("prompt") if isinstance(observation, dict) else None + if prompt is None: + raise ValueError("observation['prompt'] is required") + return {"text": self.generate(prompt)} + + def close(self): + if getattr(self, "_model", None) is not None: + if self._model.release: + self._model.release(self._model.owner) + self._model = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + def _check(self, rc, what): + if rc != 0: + err = b"" + try: + err = self._model.verbs.last_error(self._model.self) or b"" + except Exception: + pass + raise RuntimeError( + f"{what} failed (rc={rc}): " + f"{(err.decode(errors='replace') if err else 'no error')}") diff --git a/flash_rt/frontends/jetson_pi/mllm.py b/flash_rt/frontends/jetson_pi/mllm.py new file mode 100644 index 00000000..673765e3 --- /dev/null +++ b/flash_rt/frontends/jetson_pi/mllm.py @@ -0,0 +1,289 @@ +"""Jetson-PI multimodal LLM frontend — drives the Jetson-PI llama.cpp MLLM +provider through the FlashRT ``frt_model_runtime_v1`` C ABI via ctypes. + +This is the Python entry for vision-language models (images + prompt -> +text). The caller is responsible for applying the chat template; the engine +only does raw prompt + media markers → text. + +``flash_rt.load_model(framework="jetson_pi", config="mllm", ...)`` returns a +:class:`MllmJetsonPiFrontend` directly (it is NOT wrapped in ``VLAModel`` — +MLLMs output text, not actions). +""" + +from __future__ import annotations + +import ctypes +import json +import os + +import numpy as np + +from .pi0 import ( + FrtImageView, + FrtModelRuntimeV1, + FRT_RT_PIXEL_RGB8, + _FRT_MODEL_RUNTIME_ABI_VERSION, + _find_lib, + _run_generic_stage, +) + +PORT_IMAGES = 0 +PORT_PROMPT = 1 +PORT_TEXT = 2 +PORT_NEXT_TOKEN = 3 +PORT_LOGITS = 4 +PORT_IS_EOG = 5 +STAGE_INFER = 0 +STAGE_RESET = 1 +STAGE_PREFILL = 2 +STAGE_DECODE = 3 + + +class MllmJetsonPiFrontend: + """Multimodal LLM frontend backed by the Jetson-PI provider.""" + + def __init__(self, checkpoint, *, mmproj_path, backend="cpu", + n_ctx=0, n_threads=0, temp=0.8, top_k=40, top_p=0.9, + seed=1, max_tokens=512, lib_path=None, **_unused): + if max_tokens <= 0: + raise ValueError("max_tokens must be > 0") + if not mmproj_path: + raise ValueError("mmproj_path is required for MLLM") + self.max_tokens = int(max_tokens) + self._lib_path = _find_lib(lib_path, env_var="FLASHRT_MLLM_LIB") + self._lib = ctypes.CDLL(self._lib_path) + + self._lib.frt_model_runtime_open_v1.argtypes = [ + ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)] + self._lib.frt_model_runtime_open_v1.restype = ctypes.c_int + + config = { + "model_family": "mllm", + "model_path": str(checkpoint), + "mmproj_path": str(mmproj_path), + "backend": backend, + "n_ctx": int(n_ctx), + "n_threads": int(n_threads), + "temp": float(temp), + "top_k": int(top_k), + "top_p": float(top_p), + "seed": int(seed), + "max_tokens": int(max_tokens), + } + config_json = json.dumps(config, ensure_ascii=False).encode("utf-8") + + model_ptr = ctypes.c_void_p(0) + rc = self._lib.frt_model_runtime_open_v1( + config_json, ctypes.byref(model_ptr)) + if rc != 0 or not model_ptr.value: + raise RuntimeError( + f"frt_model_runtime_open_v1 failed for MLLM (rc={rc})") + + self._model = FrtModelRuntimeV1.from_address(model_ptr.value) + if self._model.abi_version != _FRT_MODEL_RUNTIME_ABI_VERSION: + raise RuntimeError( + f"frt_model_runtime_v1 abi_version={self._model.abi_version}, " + f"expected {_FRT_MODEL_RUNTIME_ABI_VERSION}.") + if self._model.struct_size < ctypes.sizeof(FrtModelRuntimeV1): + raise RuntimeError( + f"frt_model_runtime_v1 struct_size={self._model.struct_size} " + f"< ctypes sizeof {ctypes.sizeof(FrtModelRuntimeV1)}.") + self._model_ptr = model_ptr + + def _make_image_views(self, images): + """Convert a list of numpy RGB uint8 arrays to frt_image_view[].""" + n = len(images) + Views = FrtImageView * n + views = Views() + self._rgb_bufs = [] + for i, img in enumerate(images): + arr = np.ascontiguousarray(img, dtype=np.uint8) + if arr.ndim != 3 or arr.shape[2] != 3: + raise ValueError( + f"Image {i} must be HWC RGB uint8, got shape {arr.shape}") + self._rgb_bufs.append(arr) + views[i].struct_size = ctypes.sizeof(FrtImageView) + views[i].pixel_format = FRT_RT_PIXEL_RGB8 + views[i].data = arr.ctypes.data + views[i].bytes = arr.nbytes + views[i].width = arr.shape[1] + views[i].height = arr.shape[0] + views[i].stride_bytes = 0 + views[i].reserved = 0 + views[i].timestamp_ns = 0 + return views + + def generate(self, images, prompt): + """Run one multimodal completion. + + Args: + images: list of numpy arrays (H,W,3) uint8 RGB. May be empty for + text-only (though a pure LLM frontend is more efficient). + prompt: text prompt (str or bytes). Raw — caller applies chat + template. The engine injects media markers for each image. + + Returns: + Generated text (str). + """ + if self._model is None: + raise RuntimeError("MllmJetsonPiFrontend is closed") + verbs = self._model.verbs + self_ = self._model.self + + views = self._make_image_views(images) + rc = verbs.set_input(self_, PORT_IMAGES, ctypes.byref(views), + ctypes.sizeof(views), -1) + self._check(rc, "set_input images") + + if isinstance(prompt, str): + prompt_bytes = prompt.encode("utf-8") + else: + prompt_bytes = bytes(prompt) + rc = verbs.set_input(self_, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + + rc = verbs.step(self_) + self._check(rc, "step infer") + + written = ctypes.c_uint64(0) + rc = verbs.get_output(self_, PORT_TEXT, None, 0, + ctypes.byref(written), -1) + if rc != -5 or written.value == 0: + self._check(rc, "query text output size") + raise RuntimeError("query text output size returned zero bytes") + cap = written.value + out = (ctypes.c_char * cap)() + written.value = 0 + rc = verbs.get_output(self_, PORT_TEXT, out, cap, + ctypes.byref(written), -1) + self._check(rc, "get_output text") + return out.raw[:written.value].decode("utf-8", errors="replace") + + def reset(self): + """Clear the current multimodal KV-cache and sampler state.""" + if self._model is None: + raise RuntimeError("MllmJetsonPiFrontend is closed") + rc = _run_generic_stage(self._model, "reset") + self._check(rc, "generic stage reset") + + def prefill(self, images, prompt): + """Encode images and prompt, returning first next-token logits.""" + if self._model is None: + raise RuntimeError("MllmJetsonPiFrontend is closed") + verbs = self._model.verbs + views = self._make_image_views(images) + rc = verbs.set_input(self._model.self, PORT_IMAGES, ctypes.byref(views), + ctypes.sizeof(views), -1) + self._check(rc, "set_input images") + if isinstance(prompt, str): + prompt_bytes = prompt.encode("utf-8") + else: + prompt_bytes = bytes(prompt) + rc = verbs.set_input(self._model.self, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + rc = _run_generic_stage(self._model, "prefill") + self._check(rc, "generic stage prefill") + return self.get_logits() + + def decode(self): + """Sample and decode one token from the current multimodal session.""" + if self._model is None: + raise RuntimeError("MllmJetsonPiFrontend is closed") + verbs = self._model.verbs + rc = _run_generic_stage(self._model, "decode") + self._check(rc, "generic stage decode") + next_token = ctypes.c_int32() + written = ctypes.c_uint64(0) + rc = verbs.get_output( + self._model.self, PORT_NEXT_TOKEN, ctypes.byref(next_token), + ctypes.sizeof(next_token), ctypes.byref(written), -1) + self._check(rc, "get_output next_token") + is_eog = ctypes.c_int32() + written = ctypes.c_uint64(0) + rc = verbs.get_output( + self._model.self, PORT_IS_EOG, ctypes.byref(is_eog), + ctypes.sizeof(is_eog), ctypes.byref(written), -1) + self._check(rc, "get_output is_eog") + written = ctypes.c_uint64(0) + rc = verbs.get_output(self._model.self, PORT_TEXT, None, 0, + ctypes.byref(written), -1) + if rc != -5 or written.value == 0: + self._check(rc, "query text output size") + raise RuntimeError("query text output size returned zero bytes") + cap = written.value + out = (ctypes.c_char * cap)() + written.value = 0 + rc = verbs.get_output(self._model.self, PORT_TEXT, out, cap, + ctypes.byref(written), -1) + self._check(rc, "get_output text") + return { + "token": int(next_token.value), + "is_eog": bool(is_eog.value), + "text": out.raw[:written.value].decode("utf-8", errors="replace"), + } + + def get_logits(self): + """Copy current next-token logits into a NumPy float32 array.""" + if self._model is None: + raise RuntimeError("MllmJetsonPiFrontend is closed") + verbs = self._model.verbs + required = ctypes.c_uint64(0) + rc = verbs.get_output(self._model.self, PORT_LOGITS, None, 0, + ctypes.byref(required), -1) + if rc != -5 or required.value == 0: + self._check(rc, "query logits size") + raise RuntimeError("query logits size returned no bytes") + if required.value % np.dtype(np.float32).itemsize != 0: + raise RuntimeError( + f"logits byte size {required.value} is not float32-aligned") + logits = np.empty( + required.value // np.dtype(np.float32).itemsize, dtype=np.float32) + written = ctypes.c_uint64(0) + rc = verbs.get_output(self._model.self, PORT_LOGITS, logits.ctypes.data, + logits.nbytes, ctypes.byref(written), -1) + self._check(rc, "get_output logits") + if written.value != logits.nbytes: + raise RuntimeError( + f"get_output logits wrote {written.value} bytes, expected " + f"{logits.nbytes}") + return logits + + def infer(self, observation, debug=False): + """VLAModel-predict-parity entry. + + observation keys: 'images' (list of HWC uint8 arrays), 'prompt' (str). + Returns: {'text': str}. + """ + _ = debug + if not isinstance(observation, dict): + raise ValueError("observation must be a dict") + images = observation.get("images", []) + prompt = observation.get("prompt") + if prompt is None: + raise ValueError("observation['prompt'] is required") + return {"text": self.generate(images, prompt)} + + def close(self): + if getattr(self, "_model", None) is not None: + if self._model.release: + self._model.release(self._model.owner) + self._model = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + def _check(self, rc, what): + if rc != 0: + err = b"" + try: + err = self._model.verbs.last_error(self._model.self) or b"" + except Exception: + pass + raise RuntimeError( + f"{what} failed (rc={rc}): " + f"{(err.decode(errors='replace') if err else 'no error')}") diff --git a/flash_rt/frontends/jetson_pi/pi0.py b/flash_rt/frontends/jetson_pi/pi0.py new file mode 100644 index 00000000..43b90086 --- /dev/null +++ b/flash_rt/frontends/jetson_pi/pi0.py @@ -0,0 +1,420 @@ +"""Jetson-PI Pi0 frontend — drives the Jetson-PI llama.cpp/GGML Pi0 provider +through the FlashRT ``frt_model_runtime_v1`` C ABI via ctypes. + +It dlopens the SHARED provider library +(``libflashrt_cpp_llama_cpp_provider_c.so``) +built under ``FLASHRT_CPP_WITH_JETSON_PI``, opens a Pi0 runtime through +``frt_model_runtime_open_v1``, and drives one whole-model Pi0 infer per +``infer(observation)`` call. + +The frontend intentionally does no action unnormalization / LIBERO slicing: +it returns the raw ``action_steps × action_dim`` action chunk the model +produces. Higher layers (or the caller) post-process as needed. + +Memory: the Jetson-PI engine copies all inputs on ``set_input`` (see +``jetson_pi_engine.cpp``), so numpy arrays need not be kept alive past the +``infer`` call. +""" + +from __future__ import annotations + +import ctypes +import ctypes.util +import json +import os + +import numpy as np + +# ---- frt_image_view ctypes mirror (matches runtime/include/flashrt/model_runtime.h) ---- + +FRT_RT_PIXEL_RGB8 = 0 + + +class FrtImageView(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("pixel_format", ctypes.c_uint32), + ("data", ctypes.c_void_p), + ("bytes", ctypes.c_uint64), + ("width", ctypes.c_int32), + ("height", ctypes.c_int32), + ("stride_bytes", ctypes.c_int32), + ("reserved", ctypes.c_uint32), + ("timestamp_ns", ctypes.c_uint64), + ] + + +# ---- frt_model_runtime_v1 + GENERIC_STAGE_PLAN_V1 ctypes mirrors ----------- + +_SetInputFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_uint32, + ctypes.c_void_p, ctypes.c_uint64, ctypes.c_int) +_GetOutputFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_uint32, + ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint64), + ctypes.c_int) +_PrepareFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint64) +_StepFn = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p) +_LastErrorFn = ctypes.CFUNCTYPE(ctypes.c_char_p, ctypes.c_void_p) +_RetainReleaseFn = ctypes.CFUNCTYPE(None, ctypes.c_void_p) + + +class _FrtVerbsV1(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("reserved", ctypes.c_uint32), + ("set_input", _SetInputFn), + ("get_output", _GetOutputFn), + ("prepare", _PrepareFn), + ("step", _StepFn), + ("last_error", _LastErrorFn), + ] + + +class FrtModelRuntimeV1(ctypes.Structure): + pass + + +_QueryExtensionFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.POINTER(FrtModelRuntimeV1), ctypes.c_uint64, + ctypes.c_uint32, ctypes.POINTER(ctypes.c_void_p)) +_RunOpaqueFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_uint32) + + +class _GenericStageDescV1(ctypes.Structure): + _fields_ = [ + ("name", ctypes.c_char_p), + ("executor_kind", ctypes.c_uint32), + ("executor_ref", ctypes.c_uint32), + ("n_after", ctypes.c_uint32), + ("after", ctypes.POINTER(ctypes.c_uint32)), + ] + + +class _GenericStagePlanV1(ctypes.Structure): + _fields_ = [ + ("abi_version", ctypes.c_uint32), + ("struct_size", ctypes.c_uint32), + ("stages", ctypes.POINTER(_GenericStageDescV1)), + ("n_stages", ctypes.c_uint64), + ("stage_self", ctypes.c_void_p), + ("run_opaque", _RunOpaqueFn), + ] + + +FrtModelRuntimeV1._fields_ = [ + ("abi_version", ctypes.c_uint32), + ("struct_size", ctypes.c_uint32), + ("exp", ctypes.c_void_p), + ("ports", ctypes.c_void_p), + ("n_ports", ctypes.c_uint64), + ("stages", ctypes.c_void_p), + ("n_stages", ctypes.c_uint64), + ("self", ctypes.c_void_p), + ("verbs", _FrtVerbsV1), + ("owner", ctypes.c_void_p), + ("retain", _RetainReleaseFn), + ("release", _RetainReleaseFn), + ("query_extension", _QueryExtensionFn), +] + + +# Port indices (c_api.h: FRT_LLAMA_CPP_PI0_PORT_*) +PORT_IMAGES = 0 +PORT_PROMPT = 1 +PORT_STATE = 2 +PORT_ACTIONS = 3 +STAGE_INFER = 0 +STAGE_CONTEXT = 1 +STAGE_ACTION = 2 + +_FRT_MODEL_RUNTIME_ABI_VERSION = 1 +_FRT_EXT_GENERIC_STAGE_PLAN_V1 = 1 +_F32_BYTES = np.dtype(np.float32).itemsize + + +def _generic_plan(model): + extension = ctypes.c_void_p() + rc = model.query_extension( + ctypes.byref(model), _FRT_EXT_GENERIC_STAGE_PLAN_V1, 1, + ctypes.byref(extension)) + if rc != 0 or not extension.value: + raise RuntimeError(f"generic stage plan query failed (rc={rc})") + plan = _GenericStagePlanV1.from_address(extension.value) + if plan.abi_version < 1 or not plan.stages or not plan.run_opaque: + raise RuntimeError("provider returned an invalid generic stage plan") + return plan + + +def _run_generic_stage(model, name): + plan = _generic_plan(model) + encoded = name.encode("utf-8") + for index in range(plan.n_stages): + stage = plan.stages[index] + if stage.name == encoded: + return plan.run_opaque(plan.stage_self, stage.executor_ref) + raise RuntimeError(f"generic stage {name!r} is not in the selected plan") + + +def _find_lib(lib_path, env_var="FLASHRT_PI0_LIB"): + """Resolve the SHARED provider .so path. + + Priority: explicit ``lib_path`` kwarg, environment override, then the + platform dynamic-loader search path. An explicit path is a hard contract. + """ + if lib_path is not None: + if not os.path.exists(lib_path): + raise RuntimeError( + f"lib_path does not exist: {lib_path}") + return lib_path + env = os.environ.get(env_var) + if env and os.path.exists(env): + return env + return (ctypes.util.find_library("flashrt_cpp_llama_cpp_provider_c") or + "libflashrt_cpp_llama_cpp_provider_c.so") + + +class Pi0JetsonPiFrontend: + """Pi0 VLA frontend backed by the Jetson-PI llama.cpp/GGML provider.""" + + def __init__(self, checkpoint, *, mmproj_path=None, backend="cpu", + num_views=2, image_height=224, image_width=224, + action_steps=None, action_dim=None, + lib_path=None, **_unused): + if mmproj_path is None: + raise ValueError("mmproj_path is required for the Jetson-PI Pi0 frontend") + if action_steps is None or action_dim is None: + raise ValueError( + "action_steps and action_dim must be set explicitly (the " + "verified pi0_base checkpoint is 50x32; use the selected " + "checkpoint's model-specific shape)") + self.num_views = int(num_views) + if self.num_views < 1 or self.num_views > 3: + raise ValueError("num_views must be in [1, 3] for Jetson-PI Pi0") + self.image_height = int(image_height) + self.image_width = int(image_width) + self.action_steps = int(action_steps) + self.action_dim = int(action_dim) + self._prompt = b"" + self._lib_path = _find_lib(lib_path) + self._lib = ctypes.CDLL(self._lib_path) + + self._lib.frt_model_runtime_open_v1.argtypes = [ + ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)] + self._lib.frt_model_runtime_open_v1.restype = ctypes.c_int + + config = { + "model_family": "pi0", + "model_path": str(checkpoint), + "mmproj_path": str(mmproj_path), + "backend": backend, + "n_views": self.num_views, + "image_height": int(image_height), + "image_width": int(image_width), + "image_channels": 3, + "action_steps": int(action_steps), + "action_dim": int(action_dim), + } + # ensure_ascii=False: the C-side JSON parser (pi0_runtime.cpp) does + # not handle \uXXXX escapes, only raw UTF-8 bytes. + config_json = json.dumps(config, ensure_ascii=False).encode("utf-8") + + model_ptr = ctypes.c_void_p(0) + rc = self._lib.frt_model_runtime_open_v1( + config_json, ctypes.byref(model_ptr)) + if rc != 0 or not model_ptr.value: + raise RuntimeError( + f"frt_model_runtime_open_v1 failed for Pi0 (rc={rc})") + + self._model = FrtModelRuntimeV1.from_address(model_ptr.value) + # ABI gate: refuse to drive a struct laid out for a different ABI + # version (mirrors the check in runtime/bindings/runtime_pybind.cpp). + if self._model.abi_version != _FRT_MODEL_RUNTIME_ABI_VERSION: + raise RuntimeError( + f"frt_model_runtime_v1 abi_version={self._model.abi_version}, " + f"expected {_FRT_MODEL_RUNTIME_ABI_VERSION}; the provider " + f".so was built against a different FlashRT runtime ABI.") + if self._model.struct_size < ctypes.sizeof(FrtModelRuntimeV1): + raise RuntimeError( + f"frt_model_runtime_v1 struct_size={self._model.struct_size} " + f"< ctypes sizeof {ctypes.sizeof(FrtModelRuntimeV1)}; the " + f"provider .so is older than the Python frontend expects.") + self._model_ptr = model_ptr # keep the uintptr for sanity + + # -- VLAModel.predict contract ------------------------------------------- + + def set_prompt(self, prompt_text): + # Note: this signature does NOT accept `state` — VLAModel.predict + # detects that and routes state through observation["state"] instead. + if isinstance(prompt_text, bytes): + self._prompt = prompt_text + else: + self._prompt = (prompt_text or "").encode("utf-8") + + def infer(self, observation, debug=False): + if self._model is None: + raise RuntimeError("Pi0JetsonPiFrontend is closed") + _ = debug # accepted for VLAModel.predict signature parity; unused + # Collect images: predict() passes obs with 'images' list + 'image'/ + # 'wrist_image' legacy keys. Prefer the explicit list. + if "images" in observation: + images = list(observation["images"]) + else: + images = [observation["image"]] + if "wrist_image" in observation: + images.append(observation["wrist_image"]) + if len(images) != self.num_views: + raise ValueError( + f"expected {self.num_views} images, got {len(images)}") + views = self._make_image_views(images) + + state = observation.get("state") + if state is None: + raise ValueError("observation['state'] is required for the Jetson-PI Pi0 frontend") + state = np.asarray(state, dtype=np.float32).reshape(-1) + if state.size == 0: + raise ValueError("observation['state'] must contain at least one value") + if not np.all(np.isfinite(state)): + raise ValueError("observation['state'] values must be finite") + state = np.ascontiguousarray(state) + + verbs = self._model.verbs + self_ = self._model.self + + rc = verbs.set_input(self_, PORT_IMAGES, ctypes.cast(views, ctypes.c_void_p), + ctypes.sizeof(FrtImageView) * len(images), -1) + self._check(rc, "set_input images") + rc = verbs.set_input(self_, PORT_PROMPT, self._prompt, + len(self._prompt), -1) + self._check(rc, "set_input prompt") + rc = verbs.set_input(self_, PORT_STATE, state.ctypes.data, + state.nbytes, -1) + self._check(rc, "set_input state") + + rc = verbs.step(self_) + self._check(rc, "step infer") + + capacity = self.action_steps * self.action_dim * _F32_BYTES + out = (ctypes.c_char * capacity)() + written = ctypes.c_uint64(0) + rc = verbs.get_output(self_, PORT_ACTIONS, out, capacity, + ctypes.byref(written), -1) + self._check(rc, "get_output actions") + need = capacity + if written.value != need: + raise RuntimeError( + f"get_output wrote {written.value} bytes, expected {need}") + actions = np.frombuffer(out, dtype=np.float32, + count=self.action_steps * self.action_dim + ).reshape(self.action_steps, self.action_dim).copy() + return {"actions": actions} + + def context(self, observation): + """Run the Pi0 context stage and retain provider-private encoded state.""" + if self._model is None: + raise RuntimeError("Pi0JetsonPiFrontend is closed") + images = list(observation["images"]) if "images" in observation else [ + observation["image"]] + ( + [observation["wrist_image"]] + if "wrist_image" in observation else []) + if len(images) != self.num_views: + raise ValueError( + f"expected {self.num_views} images, got {len(images)}") + state = observation.get("state") + if state is None: + raise ValueError( + "observation['state'] is required for the Jetson-PI Pi0 frontend") + state = np.asarray(state, dtype=np.float32).reshape(-1) + if state.size == 0: + raise ValueError("observation['state'] must contain at least one value") + if not np.all(np.isfinite(state)): + raise ValueError("observation['state'] values must be finite") + state = np.ascontiguousarray(state) + views = self._make_image_views(images) + verbs = self._model.verbs + self_ = self._model.self + self._check(verbs.set_input( + self_, PORT_IMAGES, ctypes.cast(views, ctypes.c_void_p), + ctypes.sizeof(FrtImageView) * len(images), -1), + "set_input images") + self._check(verbs.set_input( + self_, PORT_PROMPT, self._prompt, len(self._prompt), -1), + "set_input prompt") + self._check(verbs.set_input( + self_, PORT_STATE, state.ctypes.data, + state.nbytes, -1), "set_input state") + self._check(_run_generic_stage(self._model, "context"), + "generic stage context") + + def action(self): + """Consume the pending Pi0 context and return one action chunk.""" + if self._model is None: + raise RuntimeError("Pi0JetsonPiFrontend is closed") + verbs = self._model.verbs + self._check(_run_generic_stage(self._model, "action"), + "generic stage action") + capacity = self.action_steps * self.action_dim * _F32_BYTES + out = (ctypes.c_char * capacity)() + written = ctypes.c_uint64(0) + self._check(verbs.get_output( + self._model.self, PORT_ACTIONS, out, capacity, + ctypes.byref(written), -1), "get_output actions") + if written.value != capacity: + raise RuntimeError( + f"get_output wrote {written.value} bytes, expected {capacity}") + return np.frombuffer( + out, dtype=np.float32, + count=self.action_steps * self.action_dim).reshape( + self.action_steps, self.action_dim).copy() + + def close(self): + if getattr(self, "_model", None) is not None: + if self._model.release: + self._model.release(self._model.owner) + self._model = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + # -- helpers -------------------------------------------------------------- + + def _make_image_views(self, images): + views = (FrtImageView * len(images))() + for i, im in enumerate(images): + arr = np.ascontiguousarray(im, dtype=np.uint8) + if arr.ndim != 3 or arr.shape[2] != 3: + raise ValueError( + f"image {i} must be HxWx3 uint8, got shape {arr.shape}") + if arr.shape[0] != self.image_height or arr.shape[1] != self.image_width: + raise ValueError( + f"image {i} must be {self.image_height}x{self.image_width}, " + f"got {arr.shape[0]}x{arr.shape[1]}") + views[i].struct_size = ctypes.sizeof(FrtImageView) + views[i].pixel_format = FRT_RT_PIXEL_RGB8 + views[i].data = ctypes.c_void_p(arr.ctypes.data) + views[i].bytes = arr.nbytes + views[i].width = int(arr.shape[1]) + views[i].height = int(arr.shape[0]) + views[i].stride_bytes = int(arr.strides[0]) + views[i].reserved = 0 + views[i].timestamp_ns = 0 + # Keep the array alive for the duration of the infer call by + # stashing it on the views object (engine copies on set_input). + setattr(views, f"_keepalive_{i}", arr) + return views + + def _check(self, rc, what): + if rc != 0: + err = b"" + try: + err = self._model.verbs.last_error(self._model.self) or b"" + except Exception: + pass + raise RuntimeError( + f"{what} failed (rc={rc}): " + f"{(err.decode(errors='replace') if err else 'no error')}") diff --git a/flash_rt/tests/__init__.py b/flash_rt/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/flash_rt/tests/test_jetson_pi_llm_python.py b/flash_rt/tests/test_jetson_pi_llm_python.py new file mode 100644 index 00000000..06c99daf --- /dev/null +++ b/flash_rt/tests/test_jetson_pi_llm_python.py @@ -0,0 +1,96 @@ +"""End-to-end smoke test for the Jetson-PI generic GGUF LLM provider through +the Python ``flash_rt.load_model(framework="jetson_pi", config="llm")`` entry. + +Skips (returns early) when FLASHRT_LLM_MODEL is unset. + +Env: + FLASHRT_LLM_MODEL path to a GGUF LLM (e.g. qwen3-0.6b-q4_k_m.gguf) + FLASHRT_LLM_LIB (optional) path to libflashrt_cpp_llama_cpp_provider_c.so + FLASHRT_LLM_BACKEND (optional) backend for the Jetson-PI engine; default + "cpu" (byte-identical to the original test). Set to + "cuda" to run the real forward pass on the GPU (the + engine maps backend=="cuda" to full-layer GPU offload). + CUDA_VISIBLE_DEVICES selects the physical card. +""" + +import os +import sys + +import numpy as np + + +def _skip(msg): + print(f"SKIP - {msg}") + return 0 + + +def main(): + model_env = os.environ.get("FLASHRT_LLM_MODEL") + if not model_env or not os.path.exists(model_env): + return _skip("FLASHRT_LLM_MODEL not set or missing") + + import flash_rt + + # Default "cpu" keeps the original behavior; set FLASHRT_LLM_BACKEND=cuda + # to exercise the real GPU forward pass through the Jetson-PI engine. + backend = os.environ.get("FLASHRT_LLM_BACKEND", "cpu") or "cpu" + + fe = flash_rt.load_model( + model_env, + framework="jetson_pi", + config="llm", + backend=backend, + n_ctx=2048, + n_threads=0, + temp=0.0, # greedy for deterministic test output + top_k=0, + top_p=0.0, + seed=1, + max_tokens=16, + lib_path=os.environ.get("FLASHRT_LLM_LIB")) + + text = fe.generate("What is 2 plus 2? The answer is") + + failed = 0 + def check(cond, msg): + nonlocal failed + if cond: + print(f"ok : {msg}") + else: + print(f"FAIL: {msg}") + failed = 1 + + check(isinstance(text, str) and len(text) > 0, "generated text is non-empty str") + print(f" generated: {text!r}") + if isinstance(text, str) and text: + printable = any(0x20 <= ord(c) < 0x7f for c in text) + check(printable, "generated text contains printable chars") + + prompt = "What is 2 plus 2? The answer is" + fe.reset() + logits = fe.prefill(prompt) + check(logits.ndim == 1 and logits.size > 0, + "prefill returns a non-empty logits vector") + check(bool((logits == logits).all()), "prefill logits contain no NaN") + step = None + for _ in range(16): + step = fe.decode() + check(isinstance(step["token"], int), "decode returns token id") + check(isinstance(step["is_eog"], bool), "decode returns EOG flag") + if step["is_eog"]: + break + check(step is not None and step["text"] == text, + "host-driven decode text matches one-shot generate") + + token_logits = fe.prefill(tokens=[0]) + check(token_logits.shape == logits.shape and np.isfinite(token_logits).all(), + "optional int32 tokens input produces finite logits") + + del fe + + print("\n== JETSON_PI LLM PYTHON " + ("PASSED" if not failed else "FAILED") + " ==") + return failed + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/flash_rt/tests/test_jetson_pi_mllm_python.py b/flash_rt/tests/test_jetson_pi_mllm_python.py new file mode 100644 index 00000000..bf3461e3 --- /dev/null +++ b/flash_rt/tests/test_jetson_pi_mllm_python.py @@ -0,0 +1,101 @@ +"""End-to-end smoke test for the Jetson-PI multimodal LLM provider through +the Python ``flash_rt.load_model(framework="jetson_pi", config="mllm")`` entry. + +Skips (returns early) when FLASHRT_MLLM_MODEL / FLASHRT_MLLM_MMPROJ are unset. + +Env: + FLASHRT_MLLM_MODEL path to a VLM GGUF (e.g. Qwen2.5-VL-3B-Instruct-q4_0.gguf) + FLASHRT_MLLM_MMPROJ path to the VIT mmproj GGUF + FLASHRT_MLLM_LIB (optional) path to libflashrt_cpp_llama_cpp_provider_c.so + FLASHRT_MLLM_BACKEND (optional) backend for the Jetson-PI engine; default + "cpu" (byte-identical to the original test). Set to + "cuda" to run the real forward pass on the GPU (the + engine maps backend=="cuda" to full-layer GPU offload + for both the LLM and the VIT/mmproj encoder). + CUDA_VISIBLE_DEVICES selects the physical card. +""" + +import os +import sys + +import numpy as np + + +def _skip(msg): + print(f"SKIP - {msg}") + return 0 + + +def main(): + model_env = os.environ.get("FLASHRT_MLLM_MODEL") + mmproj_env = os.environ.get("FLASHRT_MLLM_MMPROJ") + if not model_env or not os.path.exists(model_env): + return _skip("FLASHRT_MLLM_MODEL not set or missing") + if not mmproj_env or not os.path.exists(mmproj_env): + return _skip("FLASHRT_MLLM_MMPROJ not set or missing") + + import flash_rt + + # Default "cpu" keeps the original behavior; set FLASHRT_MLLM_BACKEND=cuda + # to exercise the real GPU forward pass through the Jetson-PI engine. + backend = os.environ.get("FLASHRT_MLLM_BACKEND", "cpu") or "cpu" + + fe = flash_rt.load_model( + model_env, + framework="jetson_pi", + config="mllm", + mmproj_path=mmproj_env, + backend=backend, + n_ctx=2048, + n_threads=0, + temp=0.0, + top_k=0, + top_p=0.0, + seed=1, + max_tokens=16, + lib_path=os.environ.get("FLASHRT_MLLM_LIB")) + + red_image = np.zeros((224, 224, 3), dtype=np.uint8) + red_image[:, :, 0] = 255 + + text = fe.generate([red_image], "Describe this image in one sentence.") + + failed = 0 + def check(cond, msg): + nonlocal failed + if cond: + print(f"ok : {msg}") + else: + print(f"FAIL: {msg}") + failed = 1 + + check(isinstance(text, str) and len(text) > 0, "generated text is non-empty str") + print(f" generated: {text!r}") + if isinstance(text, str) and text: + printable = any(0x20 <= ord(c) < 0x7f for c in text) + check(printable, "generated text contains printable chars") + + prompt = "Describe this image in one sentence." + fe.reset() + logits = fe.prefill([red_image], prompt) + check(logits.ndim == 1 and logits.size > 0, + "prefill returns a non-empty logits vector") + check(bool(np.isfinite(logits).all()), "prefill logits are finite") + step = None + for _ in range(16): + step = fe.decode() + check(isinstance(step["token"], int), "decode returns token id") + check(isinstance(step["is_eog"], bool), "decode returns EOG flag") + if step["is_eog"]: + break + check(step is not None and step["text"] == text, + "host-driven decode text matches one-shot generate") + + del fe + + print("\n== JETSON_PI MLLM PYTHON " + ("PASSED" if not failed else "FAILED") + " ==") + return failed + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/flash_rt/tests/test_jetson_pi_pi0_python.py b/flash_rt/tests/test_jetson_pi_pi0_python.py new file mode 100644 index 00000000..8f39a158 --- /dev/null +++ b/flash_rt/tests/test_jetson_pi_pi0_python.py @@ -0,0 +1,112 @@ +"""End-to-end smoke test for the Jetson-PI Pi0 provider through the Python +``flash_rt.load_model(framework="jetson_pi", ...)`` entry. + +Skips (returns early) when the weights / fixture env vars are unset, so CI +without weights still passes. + +Env: + FLASHRT_PI0_MODEL path to Pi0 policy GGUF + FLASHRT_PI0_MMPROJ path to VIT mmproj GGUF + FLASHRT_PI0_FIXTURE_DIR dir with image.png, wrist_image.png, state.bin, prompt.txt; + state.bin contains policy proprioception values + (8 floats for PI0.5), not action_dim padding + FLASHRT_PI0_LIB (optional) path to libflashrt_cpp_llama_cpp_provider_c.so + FLASHRT_PI0_ACTION_STEPS required model action horizon (pi0_base is 50). + FLASHRT_PI0_ACTION_DIM required model action width (pi0_base is 32). + FLASHRT_PI0_BACKEND required backend for the Jetson-PI engine, e.g. + "cpu" or "cuda". CUDA_VISIBLE_DEVICES selects the + physical card for a CUDA run. + +Run from the repository root after installing the provider or adding its +install directory to the platform dynamic-loader search path. +""" + +import os +import sys + + +def _skip(msg): + print(f"SKIP - {msg}") + return 0 + + +def main(): + model_env = os.environ.get("FLASHRT_PI0_MODEL") + mmproj_env = os.environ.get("FLASHRT_PI0_MMPROJ") + fixture_env = os.environ.get("FLASHRT_PI0_FIXTURE_DIR") + backend = os.environ.get("FLASHRT_PI0_BACKEND") + action_steps_env = os.environ.get("FLASHRT_PI0_ACTION_STEPS") + action_dim_env = os.environ.get("FLASHRT_PI0_ACTION_DIM") + if not model_env or not os.path.exists(model_env): + return _skip("FLASHRT_PI0_MODEL not set or missing") + if not mmproj_env or not os.path.exists(mmproj_env): + return _skip("FLASHRT_PI0_MMPROJ not set or missing") + if not fixture_env or not os.path.isdir(fixture_env): + return _skip("FLASHRT_PI0_FIXTURE_DIR not set or missing") + if not backend: + return _skip("FLASHRT_PI0_BACKEND not set") + if not action_steps_env or not action_dim_env: + return _skip("FLASHRT_PI0_ACTION_STEPS/ACTION_DIM not set") + for name in ("image.png", "wrist_image.png", "state.bin", "prompt.txt"): + if not os.path.exists(os.path.join(fixture_env, name)): + return _skip(f"fixture {name} missing in {fixture_env}") + + import numpy as np + from PIL import Image + + action_steps = int(action_steps_env) + action_dim = int(action_dim_env) + + import flash_rt + model = flash_rt.load_model( + model_env, + framework="jetson_pi", + mmproj_path=mmproj_env, + backend=backend, + num_views=2, + action_steps=action_steps, + action_dim=action_dim, + lib_path=os.environ.get("FLASHRT_PI0_LIB")) + + image = np.asarray(Image.open(os.path.join(fixture_env, "image.png")).convert("RGB"), dtype=np.uint8) + wrist = np.asarray(Image.open(os.path.join(fixture_env, "wrist_image.png")).convert("RGB"), dtype=np.uint8) + with open(os.path.join(fixture_env, "state.bin"), "rb") as f: + state = np.frombuffer(f.read(), dtype=np.float32) + if state.size == 0: + print("FAIL: state.bin must contain at least one float") + return 1 + with open(os.path.join(fixture_env, "prompt.txt")) as f: + prompt = f.read().rstrip("\n") + + actions = model.predict([image, wrist], prompt=prompt, state=state) + + failed = 0 + def check(cond, msg): + nonlocal failed + if cond: + print(f"ok : {msg}") + else: + print(f"FAIL: {msg}") + failed = 1 + + check(actions.shape == (action_steps, action_dim), + f"actions shape == ({action_steps},{action_dim}), got {actions.shape}") + if not np.any(np.isnan(actions)) and not np.any(np.isinf(actions)): + check(True, "actions contain no NaN/Inf") + else: + check(False, "actions contain NaN/Inf") + check(bool(np.any(actions != 0)), "actions are not all zero") + + model._pipe.context({"images": [image, wrist], "state": state}) + split_actions = model._pipe.action() + check(np.array_equal(split_actions, actions), + "context/action stages are bit-identical to predict") + + del model + + print("\n== JETSON_PI PYTHON " + ("PASSED" if not failed else "FAILED") + " ==") + return failed + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_jetson_pi_pi0_state_contract.py b/tests/test_jetson_pi_pi0_state_contract.py new file mode 100644 index 00000000..e890d044 --- /dev/null +++ b/tests/test_jetson_pi_pi0_state_contract.py @@ -0,0 +1,83 @@ +import ctypes + +import numpy as np +import pytest + +from flash_rt.frontends.jetson_pi import pi0 + + +def test_pi0_frontend_rejects_more_than_three_views(): + with pytest.raises(ValueError, match=r"num_views must be in \[1, 3\]"): + pi0.Pi0JetsonPiFrontend( + "/models/pi0.gguf", + mmproj_path="/models/pi0-mmproj.gguf", + num_views=4, + action_steps=50, + action_dim=32, + ) + + +def test_pi0_frontend_passes_policy_state_without_action_padding(monkeypatch): + received_states = [] + + @pi0._SetInputFn + def set_input(_self, port, data, size, _stream): + if port == pi0.PORT_STATE: + count = size // ctypes.sizeof(ctypes.c_float) + values = (ctypes.c_float * count).from_address(data) + received_states.append(np.ctypeslib.as_array(values).copy()) + return 0 + + @pi0._GetOutputFn + def get_output(_self, port, out, capacity, written, _stream): + assert port == pi0.PORT_ACTIONS + assert capacity == 32 * ctypes.sizeof(ctypes.c_float) + ctypes.memset(out, 0, capacity) + written[0] = capacity + return 0 + + @pi0._StepFn + def step(_self): + return 0 + + @pi0._LastErrorFn + def last_error(_self): + return b"" + + frontend = object.__new__(pi0.Pi0JetsonPiFrontend) + frontend.num_views = 1 + frontend.image_height = 1 + frontend.image_width = 1 + frontend.action_steps = 1 + frontend.action_dim = 32 + frontend._prompt = b"pick up the block" + frontend._model = pi0.FrtModelRuntimeV1() + frontend._model.self = ctypes.c_void_p(1) + frontend._model.verbs = pi0._FrtVerbsV1( + ctypes.sizeof(pi0._FrtVerbsV1), 0, set_input, get_output, + pi0._PrepareFn(), step, last_error) + frontend._callback_keepalive = (set_input, get_output, step, last_error) + + state = np.arange(8, dtype=np.float32) + observation = { + "images": [np.zeros((1, 1, 3), dtype=np.uint8)], + "state": state, + } + + frontend.infer(observation) + monkeypatch.setattr(pi0, "_run_generic_stage", lambda _model, _name: 0) + frontend.context(observation) + + assert len(received_states) == 2 + assert np.array_equal(received_states[0], state) + assert np.array_equal(received_states[1], state) + + with pytest.raises(ValueError, match="at least one value"): + frontend.infer({**observation, "state": np.empty(0, dtype=np.float32)}) + with pytest.raises(ValueError, match="must be finite"): + frontend.infer({**observation, "state": np.array([np.nan], dtype=np.float32)}) + with pytest.raises(ValueError, match="must be finite"): + frontend.context({ + **observation, + "state": np.array([np.inf], dtype=np.float32), + })