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 1b097a85..a215d0d2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -22,6 +22,9 @@ include(CTest) option(FLASHRT_CPP_WITH_EXEC "Build/link the exec layer for native replay" ON) option(FLASHRT_CPP_WITH_CUDA_STAGING "Enable conservative CUDA H2D/D2H modality staging" ON) option(FLASHRT_CPP_WITH_CUDA_KERNELS "Enable CUDA modality kernels" ON) +option(FLASHRT_CPP_WITH_JETSON_PI "Build optional Jetson-PI llama.cpp/mtmd integration checks" OFF) +set(JETSON_PI_ROOT "" CACHE PATH "Jetson-PI llama.cpp fork SOURCE root (dev add_subdirectory path)") +set(JetsonPI_ROOT "" CACHE PATH "Jetson-PI INSTALL prefix (production find_package path)") if(FLASHRT_CPP_WITH_CUDA_STAGING) find_package(CUDAToolkit REQUIRED) endif() @@ -29,6 +32,78 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) endif() +if(FLASHRT_CPP_WITH_JETSON_PI AND NOT BUILD_TESTING) + message(STATUS + "FLASHRT_CPP_WITH_JETSON_PI is exercised through test targets; building anyway.") +endif() + +# §5.2 production deploy: two integration routes. +# (1) Dev: -DJETSON_PI_ROOT= -> add_subdirectory (in-tree build, +# the GPU-verified path). The provider links the in-tree jetson_pi_* targets. +# (2) Prod: -DJetsonPI_ROOT= (no JETSON_PI_ROOT) -> find_package +# an INSTALLED Jetson-PI via cmake/FindJetsonPI.cmake; the provider links the +# imported JetsonPI::jetson_pi_* targets. Used for `cmake --install` layouts. +# Exactly one of the two must be set. Configuring both is an error: selecting a +# route implicitly could hide a stale cache entry and link the wrong Jetson-PI. +# The provider's link expression below uses a variable so the two routes are +# symmetric and the dev path stays byte-identical to before. +set(_frt_jetson_pi_pi0 jetson_pi_pi0) +set(_frt_jetson_pi_llm jetson_pi_llm) +set(_frt_jetson_pi_mllm jetson_pi_mllm) +if(FLASHRT_CPP_WITH_JETSON_PI) + if(JETSON_PI_ROOT AND JetsonPI_ROOT) + message(FATAL_ERROR + "Set exactly one Jetson-PI integration root: " + "-DJETSON_PI_ROOT= for the dev add_subdirectory route, or " + "-DJetsonPI_ROOT= for the production find_package route. " + "Both are currently set.") + endif() + if(NOT JETSON_PI_ROOT AND NOT JetsonPI_ROOT) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_JETSON_PI requires exactly one explicit integration " + "root: -DJETSON_PI_ROOT= for dev, or " + "-DJetsonPI_ROOT= for production.") + endif() + if(JETSON_PI_ROOT AND EXISTS "${JETSON_PI_ROOT}/CMakeLists.txt") + # (1) dev: add_subdirectory the source tree + # FORCE: as a subdirectory embedding, FlashRT must own these Jetson-PI build + # options regardless of stale cache or conflicting -D flags from the caller, + # so the "mtmd without the full tools tree" isolation holds. + set(LLAMA_BUILD_COMMON ON CACHE BOOL "llama: build common utils library" FORCE) + set(LLAMA_CURL OFF CACHE BOOL "llama: use libcurl to download model from an URL" FORCE) + set(LLAMA_HTTPLIB OFF CACHE BOOL "llama: disable httplib (FlashRT integration needs no HTTP)" FORCE) + set(LLAMA_BUILD_TESTS OFF CACHE BOOL "llama: build tests" FORCE) + set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "llama: build tools" FORCE) + set(JETSON_PI_BUILD_MTMD ON CACHE BOOL "Jetson-PI: build the mtmd library without the full tools tree" FORCE) + set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "llama: build examples" FORCE) + set(LLAMA_BUILD_SERVER OFF CACHE BOOL "llama: build server example" FORCE) + set(LLAMA_TOOLS_INSTALL OFF CACHE BOOL "llama: install tools" FORCE) + add_subdirectory("${JETSON_PI_ROOT}" + "${CMAKE_CURRENT_BINARY_DIR}/jetson-pi" + EXCLUDE_FROM_ALL) + message(STATUS "Jetson-PI: dev add_subdirectory path (${JETSON_PI_ROOT})") + elseif(JETSON_PI_ROOT) + # JETSON_PI_ROOT was set but does not point at a Jetson-PI source tree. + # Fail loudly rather than silently fall through to the prod find_package + # path (which would then error with a confusing "Jetson-PI not found"). + message(FATAL_ERROR + "JETSON_PI_ROOT is set to '${JETSON_PI_ROOT}' but it is not a Jetson-PI " + "source tree (no CMakeLists.txt found there). Fix the path, or unset " + "JETSON_PI_ROOT and use -DJetsonPI_ROOT= for the " + "production find_package route.") + else() + # (2) prod: find an installed Jetson-PI. Scope the module-path addition to + # this find so FindJetsonPI.cmake does not leak into unrelated find_package + # calls later in the same configure. + list(PREPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../cmake) + find_package(JetsonPI REQUIRED) + list(POP_FRONT CMAKE_MODULE_PATH) + set(_frt_jetson_pi_pi0 JetsonPI::jetson_pi_pi0) + set(_frt_jetson_pi_llm JetsonPI::jetson_pi_llm) + set(_frt_jetson_pi_mllm JetsonPI::jetson_pi_mllm) + message(STATUS "Jetson-PI: production find_package path (${JetsonPI_ROOT})") + endif() +endif() if(FLASHRT_CPP_WITH_EXEC AND NOT TARGET flashrt_exec) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../exec ${CMAKE_CURRENT_BINARY_DIR}/exec) @@ -95,18 +170,84 @@ target_link_libraries(flashrt_cpp_pi05_c target_include_directories(flashrt_cpp_pi05_c PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include) +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_link_libraries(flashrt_cpp_llama_cpp_provider + PUBLIC flashrt_runtime) +target_include_directories(flashrt_cpp_llama_cpp_provider + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/providers/llama_cpp/include) +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_link_libraries(flashrt_cpp_llama_cpp_provider + PRIVATE ${_frt_jetson_pi_pi0} ${_frt_jetson_pi_llm} ${_frt_jetson_pi_mllm}) + target_compile_definitions(flashrt_cpp_llama_cpp_provider + PRIVATE FLASHRT_CPP_WITH_JETSON_PI=1) + + # SHARED variant for Python ctypes consumption (Phase 2). Same sources as the + # STATIC provider minus the test-only link check; exports + # frt_llama_cpp_default_engine_factory + frt_llama_cpp_pi0_runtime_open_with_engine_factory + # for dlopen. Linux default visibility exposes the extern "C" symbols. + # + # Jetson-PI is linked PUBLIC here (vs PRIVATE on the STATIC provider above) + # intentionally: the SHARED _c.so is consumed by ctypes dlopen, so its NEEDED + # entries must name libjetson_pi_*.so / libmtmd / libllama / libggml so the + # dynamic loader resolves them at dlopen time. The STATIC provider is linked + # into test executables that name those libs directly, so PRIVATE suffices. + 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) + target_link_libraries(flashrt_cpp_llama_cpp_provider_c + PUBLIC flashrt_runtime ${_frt_jetson_pi_pi0} ${_frt_jetson_pi_llm} ${_frt_jetson_pi_mllm}) + 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) + set_target_properties(flashrt_cpp_llama_cpp_provider_c PROPERTIES + POSITION_INDEPENDENT_CODE ON) + + # §5.2 production install: lay the SHARED provider lib so a deployed FlashRT + # has the ctypes dlopen target on the install prefix. Dev add_subdirectory + # builds are unaffected (install() is a no-op unless cmake --install runs). + include(GNUInstallDirs) + if(JetsonPI_ROOT) + set_target_properties(flashrt_cpp_llama_cpp_provider_c PROPERTIES + INSTALL_RPATH "$ORIGIN;${JetsonPI_LIBRARY_DIRS}") + else() + set_target_properties(flashrt_cpp_llama_cpp_provider_c PROPERTIES + INSTALL_RPATH "$ORIGIN") + endif() + install(TARGETS flashrt_cpp_llama_cpp_provider_c flashrt_runtime + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) +endif() + if(BUILD_TESTING) add_executable(test_cpp_modalities tests/test_modalities.cpp) target_link_libraries(test_cpp_modalities PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) add_test(NAME cpp_modalities COMMAND test_cpp_modalities) - add_executable(test_pi05_runtime tests/test_pi05_runtime.cpp) - target_link_libraries(test_pi05_runtime - PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) - add_test(NAME pi05_runtime COMMAND test_pi05_runtime) + # test_pi05_runtime / test_device_staging transitively reference + # frt_graph_replay (via flashrt_cpp_pi05 -> runtime.cpp), which lives in + # flashrt_exec. They cannot link with FLASHRT_CPP_WITH_EXEC=OFF, so gate + # them on EXEC — this also keeps the Jetson-PI CPU-only provider build + # (FLASHRT_CPP_WITH_EXEC=OFF, no CUDA) clean under BUILD_TESTING=ON. + if(FLASHRT_CPP_WITH_EXEC) + add_executable(test_pi05_runtime tests/test_pi05_runtime.cpp) + target_link_libraries(test_pi05_runtime + PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities flashrt_exec) + add_test(NAME pi05_runtime COMMAND test_pi05_runtime) + endif() - if(FLASHRT_CPP_WITH_CUDA_STAGING) + if(FLASHRT_CPP_WITH_EXEC AND FLASHRT_CPP_WITH_CUDA_STAGING) add_executable(test_device_staging tests/test_device_staging.cpp) target_link_libraries(test_device_staging PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities CUDA::cudart) @@ -122,4 +263,72 @@ if(BUILD_TESTING) PRIVATE flashrt_cpp_pi05_c flashrt_exec flashrt_runtime CUDA::cudart) add_test(NAME pi05_model_runtime COMMAND test_pi05_model_runtime) endif() + + 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(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_link_libraries(test_llama_cpp_jetson_pi_engine + PRIVATE flashrt_cpp_llama_cpp_provider ${_frt_jetson_pi_pi0}) + target_include_directories(test_llama_cpp_jetson_pi_engine + PRIVATE ${JETSON_PI_ROOT}/vendor) # stb_image.h (dev tree) + add_test(NAME llama_cpp_jetson_pi_engine + COMMAND test_llama_cpp_jetson_pi_engine) + 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 ${_frt_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 ${_frt_jetson_pi_mllm}) + add_test(NAME llama_cpp_jetson_pi_mllm + COMMAND test_llama_cpp_jetson_pi_mllm) + + # Numerical-parity tests (jetsonpi迁移.txt §14): FlashRT provider path vs + # a DIRECT jetson_pi_* call, same inputs/backend. Proves the port/stage/ + # verb plumbing does not perturb outputs. Pi0 uses action tolerance; + # LLM/MLLM use greedy sampling and exact generated-text equality. + if(JETSON_PI_ROOT AND EXISTS "${JETSON_PI_ROOT}/vendor") + add_executable(test_llama_cpp_jetson_pi_parity + tests/test_llama_cpp_jetson_pi_parity.cpp) + target_link_libraries(test_llama_cpp_jetson_pi_parity + PRIVATE flashrt_cpp_llama_cpp_provider ${_frt_jetson_pi_pi0}) + target_include_directories(test_llama_cpp_jetson_pi_parity + PRIVATE ${JETSON_PI_ROOT}/vendor) # stb_image.h (dev tree) + add_test(NAME llama_cpp_jetson_pi_parity + COMMAND test_llama_cpp_jetson_pi_parity) + endif() + + 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 ${_frt_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 ${_frt_jetson_pi_mllm}) + add_test(NAME llama_cpp_jetson_pi_mllm_parity + COMMAND test_llama_cpp_jetson_pi_mllm_parity) + endif() endif() 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..d4396534 --- /dev/null +++ b/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/c_api.h @@ -0,0 +1,240 @@ +#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; + 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)) + +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_v2** 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_v2** 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_v2** 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_v2** 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_v2** 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_v2** 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..67d38d57 --- /dev/null +++ b/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/jetson_pi_engine.h @@ -0,0 +1,51 @@ +#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 v2 runtime, which releases it on drop. +// +// Engine error codes (returned by set_input/run_infer/get_output, surfaced +// through verbs_v2): 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); + +// Phase 6 memory-domain token verbs for the Pi0 actions OUT port. The paired +// frt_memory_token handle is frt_llama_cpp_engine_v1::self (the provider's +// engine pointer, round-tripped opaquely). location_kind is HOST_VISIBLE: +// the backing store is the engine's host-resident actions_buf, read live at +// copy_to_host call time (NOT a mint-time snapshot). Returns a borrowed +// pointer to a process-global verb table; valid for the process lifetime. +// destroy is null (the engine owns the buffer; release frees it after the +// holder's token-destroy loop). +const frt_memory_token_verbs* +frt_jetson_pi_actions_token_verbs(void); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* FLASHRT_PROVIDERS_LLAMA_CPP_JETSON_PI_ENGINE_H */ 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..25f37515 --- /dev/null +++ b/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp @@ -0,0 +1,1287 @@ +// 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 +#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, size == action_dim + std::vector actions_buf; // output action chunk, size == action_steps*action_dim + bool images_set = false; + bool prompt_set = false; + bool state_set = false; + uint32_t actions_host_maps = 0; + std::unordered_map actions_host_map_counts; + + 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(); + if (e->actions_host_maps != 0) { + e->set_error("cannot replace Pi0 input while actions host view is mapped"); + return -6; + } + 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; + } + const uint64_t expect_bytes = + static_cast(e->action_dim) * sizeof(float); + if (bytes != expect_bytes) { + e->set_error("state bytes != action_dim*sizeof(float)"); + return -1; + } + e->state.assign(static_cast(data), + static_cast(data) + e->action_dim); + 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->actions_host_maps != 0) { + e->set_error("cannot run Pi0 infer while actions host view is mapped"); + return -6; + } + 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->actions_host_maps != 0) { + e->set_error("cannot run Pi0 stage while actions host view is mapped"); + return -6; + } + 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; +} + +// --------------------------------------------------------------------------- +// Phase 6 memory-domain token: the Pi0 actions OUT port carries an opaque +// token (handle == this Engine*) so a consumer can read the action chunk +// through frt_memory_token_verbs instead of get_output. The verbs read the +// LIVE actions_buf at call time — NOT a pointer captured at mint time — +// because std::vector reallocates on clear()/assign() (set_input clears it +// for staleness; run_infer refills it). The bounds/return codes mirror +// engine_get_output: -7 not-ready (buf empty), -5 too-small (partial read). +// Single-threaded, same contract as engine_get_output (no locking on the +// engine vtable). location_kind is HOST_VISIBLE: actions_buf is host memory, +// readable via a plain memcpy, no device sync. destroy is null — the engine +// owns actions_buf and frees it in release(), which fires AFTER the holder's +// token-destroy loop (runtime_export.cpp). +// --------------------------------------------------------------------------- +int pi0_token_copy_to_host(frt_memory_token token, void * dst, + uint64_t dst_off, uint64_t src_off, + uint64_t bytes) { + Engine * e = reinterpret_cast(token); + if (!e) return -1; + if (bytes == 0) return 0; + if (!dst) return -1; + const size_t have_bytes = e->actions_buf.size() * sizeof(float); + if (have_bytes == 0) return -7; // not ready (set_input cleared it) + if (src_off > have_bytes || bytes > have_bytes - src_off) return -5; // buffer too small + std::memcpy(static_cast(dst) + dst_off, + reinterpret_cast(e->actions_buf.data()) + src_off, + bytes); + return 0; +} + +int pi0_token_copy_from_host(frt_memory_token /*token*/, const void * /*src*/, + uint64_t /*src_off*/, uint64_t /*dst_off*/, + uint64_t /*bytes*/) { + return -3; // unsupported: actions is an OUT port — writing into the model's output is not meaningful +} + +int pi0_token_sync(frt_memory_token /*token*/) { + return 0; // HOST_VISIBLE: host can read without a device sync +} + +int pi0_token_map_host(frt_memory_token token, uint64_t offset, + uint64_t bytes, uint32_t access, void ** out_ptr) { + Engine * e = reinterpret_cast(token); + if (!e || !out_ptr || access != FRT_RT_HOST_MAP_READ || bytes == 0) return -1; + *out_ptr = nullptr; + const uint64_t have_bytes = e->actions_buf.size() * sizeof(float); + if (have_bytes == 0) return -7; + if (offset > have_bytes || bytes > have_bytes - offset) return -5; + *out_ptr = reinterpret_cast(e->actions_buf.data()) + offset; + e->actions_host_maps += 1; + e->actions_host_map_counts[*out_ptr] += 1; + return 0; +} + +int pi0_token_unmap_host(frt_memory_token token, void * ptr) { + Engine * e = reinterpret_cast(token); + if (!e || !ptr || e->actions_host_maps == 0) return -1; + auto it = e->actions_host_map_counts.find(ptr); + if (it == e->actions_host_map_counts.end()) return -1; + if (--it->second == 0) e->actions_host_map_counts.erase(it); + e->actions_host_maps -= 1; + return 0; +} + +const frt_memory_token_verbs kPi0ActionsTokenVerbs = { + /*struct_size=*/sizeof(frt_memory_token_verbs), + /*reserved=*/0, + /*copy_to_host=*/pi0_token_copy_to_host, + /*copy_from_host=*/pi0_token_copy_from_host, + /*sync=*/pi0_token_sync, + /*destroy=*/nullptr, // engine owns actions_buf; freed in release() after the destroy loop + /*map_host=*/pi0_token_map_host, + /*unmap_host=*/pi0_token_unmap_host, +}; + +} // namespace + +const frt_memory_token_verbs * frt_jetson_pi_actions_token_verbs(void) { + return &kPi0ActionsTokenVerbs; +} + +// --------------------------------------------------------------------------- +// 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->image_height || + !config->image_width || !config->image_channels || + !config->action_steps || !config->action_dim) { + set_create_error("create_pi0 config has empty/zero fields"); + return -1; + } + 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; + 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..8ce7c4ac --- /dev/null +++ b/cpp/providers/llama_cpp/src/jetson_pi_link_check.cpp @@ -0,0 +1,50 @@ +#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 Phase 1 Pi0 whole-graph entry point + * (CLAUDE.md: mtmd_helper_eval_chunks_pi0 is the priority provider entry). + * The default-param symbols above do not by themselves exercise this TU's + * dependence on the Pi0 infer path. The `volatile` store defeats DCE so + * the address must be materialized as an undefined symbol resolved from + * libmtmd at link time. */ + 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_v2* 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..d455a245 --- /dev/null +++ b/cpp/providers/llama_cpp/src/llm_runtime.cpp @@ -0,0 +1,505 @@ +// frt_llama_cpp_llm_runtime — v2 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; +}; + +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_stage(void* self, uint32_t stage, int stream) { + (void)stream; + 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_stage(self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER, -1); +} + +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_v2** 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 >= sizeof(frt_llama_cpp_engine_v1) && + owner->engine.run_stage; + if (owner->engine.retain) owner->engine.retain(owner->engine.self); + + frt_runtime_builder b = frt_runtime_builder_create_provider_owned(); + 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); + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "infer", 0, nullptr, 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_callback_stage_v2( + b, "reset", 0, nullptr, 0); + const uint32_t prefill_after[1] = { + FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET}; + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "prefill", 0, prefill_after, 1); + const uint32_t decode_after[1] = { + FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL}; + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "decode", 0, decode_after, 1); + } + 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()); + if (rc != 0) { + frt_runtime_builder_discard(b); + destroy_owner(owner); + return -1; + } + + frt_model_runtime_verbs_v2 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; + verbs.run_stage = run_stage; + + frt_model_runtime_v2* model = frt_runtime_builder_finish_model_v2( + b, &verbs, owner, owner, nullptr, destroy_owner); + if (!model) { + 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_v2** 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_v2* 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..5fa955cb --- /dev/null +++ b/cpp/providers/llama_cpp/src/mllm_runtime.cpp @@ -0,0 +1,512 @@ +// frt_llama_cpp_mllm_runtime — v2 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; +}; + +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_stage(void* self, uint32_t stage, int stream) { + (void)stream; + 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_stage(self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER, -1); +} + +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_v2** 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 >= sizeof(*engine) && + owner->engine.run_stage; + if (owner->engine.retain) owner->engine.retain(owner->engine.self); + + frt_runtime_builder b = frt_runtime_builder_create_provider_owned(); + 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); + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "infer", 0, nullptr, 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_callback_stage_v2( + b, "reset", 0, nullptr, 0); + const uint32_t prefill_after[1] = { + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_RESET}; + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "prefill", 0, prefill_after, 1); + const uint32_t decode_after[1] = { + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL}; + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "decode", 0, decode_after, 1); + } + 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()); + if (rc != 0) { + frt_runtime_builder_discard(b); + destroy_owner(owner); + return -1; + } + + frt_model_runtime_verbs_v2 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; + verbs.run_stage = run_stage; + + frt_model_runtime_v2* model = frt_runtime_builder_finish_model_v2( + b, &verbs, owner, owner, nullptr, destroy_owner); + if (!model) { + 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_v2** 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_v2* 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..0c12a32c --- /dev/null +++ b/cpp/providers/llama_cpp/src/pi0_runtime.cpp @@ -0,0 +1,458 @@ +#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; +}; + +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_stage(void* self, uint32_t stage, int stream) { + (void)stream; + 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_stage(self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, -1); +} + +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_v2** 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->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 >= sizeof(*engine) && + owner->engine.run_stage; + 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); + // Pi0 state shares the model action_dim: llama_set_pi0_state requires + // n_values == hparams.action_dim, and the caller zero-pads real + // proprioception into that width. Exposing it on the state port keeps + // the host-visible shape honest. + owner->state_shape[0] = static_cast(config->action_dim); + owner->action_shape[0] = static_cast(config->action_steps); + owner->action_shape[1] = static_cast(config->action_dim); + + frt_runtime_builder b = frt_runtime_builder_create_provider_owned(); + 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); + // Phase 6: the actions OUT port carries a memory-domain token (handle = + // engine.self, the provider Engine*) so a consumer can read the action + // chunk through frt_memory_token_verbs as well as get_output. HOST_VISIBLE: + // actions_buf is host memory, read live at copy_to_host call time. The + // port desc this produces is byte-identical to add_port (STAGED, no raw + // buffer), so the model identity/fingerprint and get_output are unchanged. + // Only minted when a real Jetson-PI engine is linked; without + // FLASHRT_CPP_WITH_JETSON_PI there is no backing store to advertise (the + // fake-engine unit test path uses add_port and reads via get_output). +#if defined(FLASHRT_CPP_WITH_JETSON_PI) + const uint64_t action_window_bytes = + static_cast(config->action_steps) * + static_cast(config->action_dim) * sizeof(float); + rc |= frt_runtime_builder_add_port_token( + b, "actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, FRT_RT_LAYOUT_FLAT, + FRT_RT_PORT_OUT, 0, owner->action_shape, 2, 0, + reinterpret_cast(owner->engine.self), + frt_jetson_pi_actions_token_verbs(), + /*offset=*/0, action_window_bytes, FRT_RT_LOCATION_HOST_VISIBLE); +#else + 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); +#endif + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "infer", 0, nullptr, 0); + if (owner->staged_context_action) { + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "context", 0, nullptr, 0); + const uint32_t action_after[1] = { + FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT}; + rc |= frt_runtime_builder_add_callback_stage_v2( + b, "action", 0, action_after, 1); + } + 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); + if (rc != 0) { + frt_runtime_builder_discard(b); + destroy_owner(owner); + return -1; + } + + frt_model_runtime_verbs_v2 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; + verbs.run_stage = run_stage; + + frt_model_runtime_v2* model = frt_runtime_builder_finish_model_v2( + b, &verbs, owner, owner, nullptr, destroy_owner); + if (!model) { + 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_v2** 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) { + 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) { + flashrt::providers::llama_cpp::set_runtime_open_error( + "factory returned an invalid Pi0 engine"); + return -1; + } + frt_model_runtime_v2* 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/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/test_llama_cpp_jetson_pi_engine.cpp b/cpp/tests/test_llama_cpp_jetson_pi_engine.cpp new file mode 100644 index 00000000..d5e0c1c7 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_engine.cpp @@ -0,0 +1,362 @@ +// 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 (optional) override; default 50 (LIBERO base). +// FLASHRT_PI0_ACTION_DIM (optional) override; default 32. + +#include "flashrt/providers/llama_cpp/c_api.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_v2 * 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 (default 50x32 for LIBERO base; pi0_base is 10x32). + const char * steps_env = std::getenv("FLASHRT_PI0_ACTION_STEPS"); + const char * dim_env = std::getenv("FLASHRT_PI0_ACTION_DIM"); + // 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 = steps_env ? std::atol(steps_env) : 50; + long action_dim = dim_env ? std::atol(dim_env) : 32; + 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 ---- + // Real model is 10x32 (or whatever env says); 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_v2 * 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_v2 * 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_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, + &views[0], sizeof(views), -1) == 0, + "set_input images"); + CHECK(model->verbs_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "set_input prompt"); + CHECK(model->verbs_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), + state_bytes.size(), -1) == 0, + "set_input state"); + + rc = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, -1); + if (rc != 0) { + std::printf("FAIL: run_stage infer (rc=%d): %s\n", rc, + model->verbs_v2.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_v2.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"); + } + + // ---- sub-test E: the actions OUT port carries a Phase 6 memory-domain + // token. Assert the parallel-array invariant, that only the actions + // port mints a real token (HOST_VISIBLE), and that copy_to_host reads + // the LIVE provider buffer (byte-identical to get_output; -7 after + // set_input clears it). This is the real token consumer that closes + // the Phase 6 gap (the contract was previously paper-only). ---- + CHECK(model->n_port_tokens == model->n_ports, + "port_tokens parallel array aligns with ports"); + CHECK(model->n_port_tokens == 4, + "pi0 exposes 4 ports (images/prompt/state/actions)"); + CHECK(model->port_tokens[FRT_LLAMA_CPP_PI0_PORT_IMAGES].handle == nullptr && + model->port_tokens[FRT_LLAMA_CPP_PI0_PORT_PROMPT].handle == nullptr && + model->port_tokens[FRT_LLAMA_CPP_PI0_PORT_STATE].handle == nullptr, + "non-actions ports carry no token (null handle)"); + const auto & tk = model->port_tokens[FRT_LLAMA_CPP_PI0_PORT_ACTIONS]; + CHECK(tk.handle != nullptr, "actions port carries a real token handle"); + CHECK(tk.location_kind == FRT_RT_LOCATION_HOST_VISIBLE, + "actions token advertises HOST_VISIBLE"); + CHECK(tk.verbs != nullptr && tk.verbs->copy_to_host != nullptr && + tk.verbs->copy_from_host != nullptr && + tk.verbs->sync != nullptr, + "actions token verbs table is complete"); + if (tk.handle && tk.verbs && tk.verbs->copy_to_host) { + const uint64_t action_bytes = + static_cast(action_steps) * action_dim * sizeof(float); + std::vector via_token(actions.size()); + int trc = tk.verbs->copy_to_host(tk.handle, via_token.data(), + 0, 0, action_bytes); + CHECK(trc == 0, "copy_to_host reads the actions chunk via token"); + CHECK(tk.verbs->copy_to_host(tk.handle, nullptr, 0, 0, + sizeof(float)) == -1, + "copy_to_host rejects a null destination"); + if (trc == 0) { + CHECK(std::memcmp(via_token.data(), actions.data(), + static_cast(action_bytes)) == 0, + "token copy_to_host is byte-identical to get_output"); + } + float scratch = 0.0f; + CHECK(tk.verbs->copy_from_host(tk.handle, &scratch, 0, 0, + sizeof(float)) == -3, + "copy_from_host returns -3 unsupported (OUT port)"); + CHECK(tk.verbs->sync(tk.handle) == 0, + "sync is a no-op (HOST_VISIBLE)"); + } + + // ---- 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_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, + &views[0], sizeof(views), -1) == 0, + "tick2 set_input images"); + CHECK(model->verbs_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "tick2 set_input prompt"); + CHECK(model->verbs_v2.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_v2.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"); + } + // The token must read the LIVE buffer: after set_input cleared + // actions_buf, copy_to_host must fail too (-7). This proves the token + // is not a mint-time snapshot — it reflects the provider's current + // state. + { + const auto & tk2 = model->port_tokens[FRT_LLAMA_CPP_PI0_PORT_ACTIONS]; + const uint64_t action_bytes = + static_cast(action_steps) * action_dim * sizeof(float); + int trc2 = tk2.verbs->copy_to_host(tk2.handle, actions.data(), + 0, 0, action_bytes); + CHECK(trc2 != 0, + "token copy_to_host after set_input (before run_infer) fails"); + } + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, -1) == 0, + "tick2 run_stage infer"); + std::vector actions2( + static_cast(action_steps) * action_dim); + written = 0; + rc = model->verbs_v2.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..e6236bc2 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_llm.cpp @@ -0,0 +1,402 @@ +// 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 "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_v2 * 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_v2 * 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"); + std::printf(" runtime stages (%llu):", + static_cast(model->n_stages_v2)); + for (uint64_t i = 0; i < model->n_stages_v2; ++i) { + std::printf(" %s", model->stages_v2[i].name); + } + std::printf("\n"); + CHECK(model->n_stages == 0 && model->n_stages_v2 == 4, + "LLM runtime exposes infer/reset/prefill/decode callback stages"); + 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_v2.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_v2.set_input(model->self, + FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1); + CHECK(rc == 0, "set_input prompt"); + + rc = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER, -1); + if (rc != 0) { + std::printf("FAIL: run_stage infer (rc=%d): %s\n", rc, + model->verbs_v2.last_error(model->self)); + g_fail = 1; + } else { + std::printf("ok : run_stage infer\n"); + } + + uint64_t written = 0; + rc = model->verbs_v2.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_v2.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_v2.set_input( + model->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + nullptr, 0, -1) != 0, + "invalid prompt replacement is rejected"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER, -1) != 0, + "failed prompt replacement invalidates the previous prompt"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0, + "restore valid prompt after failed replacement"); + + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET, -1) == 0, + "run_stage reset"); + const auto prefill_start = std::chrono::steady_clock::now(); + rc = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1); + 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_v2.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_v2.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 = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1); + 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_v2.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_v2.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_v2.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_v2.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_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER, -1); + CHECK(rc == 0, "one-shot infer after staged decode"); + uint64_t stale_written = 0; + rc = model->verbs_v2.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_v2* 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_v2.set_input( + peer->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + peer_prompt, std::strlen(peer_prompt), -1) == 0, + "peer set_input distinct prompt"); + CHECK(peer->verbs_v2.run_stage( + peer->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET, -1) == 0 && + peer->verbs_v2.run_stage( + peer->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1) == 0, + "prepare peer standalone baseline"); + std::vector peer_baseline_tokens; + std::vector peer_baseline_eog; + for (uint32_t i = 0; i < 16; ++i) { + CHECK(peer->verbs_v2.run_stage( + peer->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1) == 0, + "decode peer standalone baseline"); + int32_t token = 0; + int32_t is_eog = 0; + uint64_t scalar_bytes = 0; + CHECK(peer->verbs_v2.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, &token, + sizeof(token), &scalar_bytes, -1) == 0 && + scalar_bytes == sizeof(token) && + peer->verbs_v2.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_v2.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_v2.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(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET, -1) == 0 && + peer->verbs_v2.run_stage( + peer->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET, -1) == 0, + "reset two independent sessions"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1) == 0 && + peer->verbs_v2.run_stage( + peer->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1) == 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(model->verbs_v2.run_stage( + model->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1) == 0 && + model->verbs_v2.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &token, sizeof(token), &scalar_bytes, -1) == 0 && + model->verbs_v2.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(peer->verbs_v2.run_stage( + peer->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1) == 0 && + peer->verbs_v2.get_output( + peer->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &token, sizeof(token), &scalar_bytes, -1) == 0 && + peer->verbs_v2.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_v2.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, + nullptr, 0, &model_written, -1) == -5 && + model_written > 0 && + peer->verbs_v2.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_v2.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_TEXT, model_text.data(), + model_text.size(), &model_written, -1) == 0 && + peer->verbs_v2.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_v2 * 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_v2.set_input( + budget_model->self, FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0 && + budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1) == 0 && + budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1) == 0, + "max_tokens=1 session permits exactly one decode"); + int32_t first_is_eog = 1; + uint64_t scalar_written = 0; + CHECK(budget_model->verbs_v2.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(budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1) != 0 && + std::strstr(budget_model->verbs_v2.last_error( + budget_model->self), + "max_tokens"), + "staged decode rejects calls beyond max_tokens"); + int32_t stale_scalar = 0; + CHECK(budget_model->verbs_v2.get_output( + budget_model->self, FRT_LLAMA_CPP_LLM_PORT_NEXT_TOKEN, + &stale_scalar, sizeof(stale_scalar), &scalar_written, -1) == -7 && + budget_model->verbs_v2.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(budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1) == 0 && + budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, -1) == 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..96250a0e --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_llm_parity.cpp @@ -0,0 +1,247 @@ +// Numerical-parity test: FlashRT's Jetson-PI LLM provider (via +// frt_model_runtime_v2) 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 "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_v2 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_v2 * 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_v2.set_input(model->self, + FRT_LLAMA_CPP_LLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0, + "FlashRT set_input prompt"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER, -1) == 0, + "FlashRT run_stage infer"); + uint64_t written = 0; + rc = model->verbs_v2.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_v2.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(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_RESET, -1) == 0 && + model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_PREFILL, -1) == 0, + "FlashRT staged prefill"); + uint64_t logits_bytes = 0; + CHECK(model->verbs_v2.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_v2.get_output( + model->self, FRT_LLAMA_CPP_LLM_PORT_LOGITS, + logits_flashrt.data(), logits_bytes, &logits_bytes, + -1) == 0 && + model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_LLM_STAGE_INDEX_DECODE, + -1) == 0, + "FlashRT read logits and decode first token"); + uint64_t token_bytes = 0; + CHECK(model->verbs_v2.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..bd961cef --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_mllm.cpp @@ -0,0 +1,312 @@ +// 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 "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_v2 * 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_v2 * 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"); + CHECK(model->n_stages == 0 && model->n_stages_v2 == 4, + "MLLM runtime exposes infer/reset/prefill/decode callback stages"); + CHECK(model->n_ports == 6, + "MLLM runtime exposes images/prompt/text/token/logits/eog ports"); + CHECK(model->verbs_v2.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_v2.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_v2.set_input(model->self, + FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1); + CHECK(rc == 0, "set_input prompt"); + + rc = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER, -1); + if (rc != 0) { + std::printf("FAIL: run_stage infer (rc=%d): %s\n", rc, + model->verbs_v2.last_error(model->self)); + g_fail = 1; + } else { + std::printf("ok : run_stage infer\n"); + } + + uint64_t written = 0; + rc = model->verbs_v2.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_v2.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_v2.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_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER, -1) != 0, + "failed image replacement invalidates previous images"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &view, sizeof(view), -1) == 0, + "restore valid images after failed replacement"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + nullptr, 0, -1) != 0, + "invalid prompt replacement is rejected"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER, -1) != 0, + "failed prompt replacement invalidates previous prompt"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0, + "restore valid prompt after failed replacement"); + + rc = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_RESET, -1); + CHECK(rc == 0, "run_stage reset"); + const auto prefill_start = std::chrono::steady_clock::now(); + rc = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL, -1); + 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_v2.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_v2.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 = model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE, -1); + 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_v2.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_v2.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_v2.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_v2.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_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER, -1); + CHECK(rc == 0, "one-shot infer after staged decode"); + uint64_t stale_written = 0; + rc = model->verbs_v2.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_v2 * 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_v2.set_input( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &view, sizeof(view), -1) == 0 && + budget_model->verbs_v2.set_input( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0 && + budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL, -1) == 0 && + budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE, -1) == 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_v2.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(budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE, -1) != 0 && + std::strstr(budget_model->verbs_v2.last_error( + budget_model->self), + "max_tokens"), + "staged decode rejects calls beyond max_tokens"); + int32_t stale_scalar = 0; + CHECK(budget_model->verbs_v2.get_output( + budget_model->self, FRT_LLAMA_CPP_MLLM_PORT_NEXT_TOKEN, + &stale_scalar, sizeof(stale_scalar), &scalar_written, -1) == -7 && + budget_model->verbs_v2.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(budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_PREFILL, -1) == 0 && + budget_model->verbs_v2.run_stage( + budget_model->self, + FRT_LLAMA_CPP_MLLM_STAGE_INDEX_DECODE, -1) == 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..767c7ec7 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_mllm_parity.cpp @@ -0,0 +1,129 @@ +// Numerical parity: FlashRT MLLM provider versus direct jetson_pi_mllm. + +#include "flashrt/providers/llama_cpp/c_api.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_v2* 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_v2.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_IMAGES, + &image, sizeof(image), -1) == 0 && + model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_MLLM_PORT_PROMPT, + prompt, std::strlen(prompt), -1) == 0 && + model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_MLLM_STAGE_INDEX_INFER, + -1) == 0, + "run FlashRT MLLM image+text infer"); + uint64_t written = 0; + rc = model->verbs_v2.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_v2.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..0adf7821 --- /dev/null +++ b/cpp/tests/test_llama_cpp_jetson_pi_parity.cpp @@ -0,0 +1,440 @@ +// Numerical-parity test: FlashRT's Jetson-PI Pi0 provider (via +// frt_model_runtime_v2) 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, marker injection + zero-pad-to-action_dim +// happen INSIDE jetson_pi_pi0_infer for both, backend->n_gpu_layers/use_gpu +// mapping is identical, n_threads=0 on both -> 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 (optional) override; default 50 (LIBERO base). +// pi0_base is 10 — set FLASHRT_PI0_ACTION_STEPS=10 +// or the open fails with an action-shape mismatch. +// FLASHRT_PI0_ACTION_DIM (optional) override; default 32. +// FLASHRT_PI0_BACKEND (optional) "cpu" (default) 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 "flashrt/providers/llama_cpp/jetson_pi_engine.h" +#include "jetson_pi_pi0.h" + +#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"); + 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; + } + + const char * steps_env = std::getenv("FLASHRT_PI0_ACTION_STEPS"); + const char * dim_env = std::getenv("FLASHRT_PI0_ACTION_DIM"); + 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 = steps_env ? std::atol(steps_env) : 50; + long action_dim = dim_env ? std::atol(dim_env) : 32; + 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); + 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) { + 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_v2 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_v2 * 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"); + CHECK(model->n_stages_v2 == 3, + "FlashRT Pi0 runtime exposes infer/context/action stages"); + + 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_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_IMAGES, + &views[0], sizeof(views), -1) == 0, + "FlashRT set_input images"); + CHECK(model->verbs_v2.set_input(model->self, + FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "FlashRT set_input prompt"); + CHECK(model->verbs_v2.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_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, -1) == 0, + "FlashRT run_stage infer"); + + uint64_t written = 0; + rc = model->verbs_v2.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"); + const frt_memory_token_desc & actions_token = + model->port_tokens[FRT_LLAMA_CPP_PI0_PORT_ACTIONS]; + void * mapped_actions = nullptr; + CHECK(actions_token.verbs->struct_size >= + sizeof(frt_memory_token_verbs) && + actions_token.verbs->map_host( + actions_token.handle, actions_token.offset, + actions_token.bytes, FRT_RT_HOST_MAP_READ, + &mapped_actions) == 0 && mapped_actions != nullptr && + std::memcmp(mapped_actions, actions_flashrt.data(), n_bytes) == 0, + "Pi0 actions token maps a byte-identical zero-copy host view"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == -6, + "Pi0 input mutation is rejected while host view is mapped"); + CHECK(actions_token.verbs->unmap_host( + actions_token.handle, mapped_actions) == 0, + "Pi0 zero-copy host view unmaps explicitly"); + CHECK(actions_token.verbs->unmap_host( + actions_token.handle, mapped_actions) != 0, + "Pi0 zero-copy host view rejects duplicate unmap"); + + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT, -1) == 0, + "FlashRT run_stage context"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_PROMPT, + prompt.data(), prompt.size(), -1) == 0, + "new Pi0 input discards pending context"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION, -1) == -7, + "action after replacement input is not ready"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT, -1) == 0, + "FlashRT rebuilds context after replacement input"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_STATE, + state_bytes.data(), state_bytes.size() - sizeof(float), -1) != 0, + "invalid replacement input is rejected"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT, -1) != 0, + "failed state replacement invalidates the previous state"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION, -1) != 0, + "failed replacement input blocks action execution"); + CHECK(model->verbs_v2.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_v2.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_IMAGES, + bad_views, sizeof(bad_views), -1) != 0, + "undersized Pi0 image stride is rejected"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT, -1) != 0, + "failed image replacement invalidates previous images"); + CHECK(model->verbs_v2.set_input( + model->self, FRT_LLAMA_CPP_PI0_PORT_IMAGES, + views, sizeof(views), -1) == 0, + "restore valid Pi0 images after failed replacement"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_CONTEXT, -1) == 0, + "FlashRT rebuilds context after failed replacement"); + CHECK(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION, -1) == 0, + "FlashRT run_stage action"); + written = 0; + rc = model->verbs_v2.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(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_ACTION, -1) == -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_provider.cpp b/cpp/tests/test_llama_cpp_provider.cpp new file mode 100644 index 00000000..f6603083 --- /dev/null +++ b/cpp/tests/test_llama_cpp_provider.cpp @@ -0,0 +1,561 @@ +#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 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 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; +} + +} // namespace + +int main() { + frt_model_runtime_v2* 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_v2* 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_v2* 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 && old_llm->n_stages_v2 == 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_v2* 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 && old_mllm->n_stages_v2 == 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; + + CHECK(frt_llama_cpp_pi0_runtime_create_with_engine( + &cfg, &old_engine_api, &old_pi0) == 0 && old_pi0 && + old_pi0->n_stages_v2 == 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_v2* 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_v2* 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"); + CHECK(model->abi_version == FRT_MODEL_RUNTIME_ABI_VERSION_V2 && + model->n_ports == 4 && model->n_stages == 0 && + model->n_stages_v2 == 1, + "runtime exposes v2 provider-owned shape"); + 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] == 2, + "state port schema"); + 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(model->stages_v2[0].name, "infer") == 0 && + model->stages_v2[0].kind == FRT_RT_STAGE_CALLBACK && + model->stages_v2[0].callback == 0, + "infer callback 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, "stage_v2:0:infer:"), + "identity carries provider, model family, and callback stage"); + + CHECK(model->verbs_v2.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(model->verbs_v2.run_stage( + model->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, -1) == 0 && + engine.infer == 1, + "run_stage delegates infer to engine"); + float out[4] = {}; + uint64_t written = 0; + CHECK(model->verbs_v2.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_v2.prepare(model->self, 0, 0) == -3 && + std::strstr(model->verbs_v2.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"); + + 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_v2* 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_v2.get_output( + null_error_model->self, FRT_LLAMA_CPP_PI0_PORT_PROMPT, + out, sizeof(out), &written, -1) == -1 && + std::strstr(null_error_model->verbs_v2.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_v2* 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_v2.run_stage( + opened->self, FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER, -1) == 0 && + factory_engine.infer == 1, + "opened runtime delegates infer to factory engine"); + 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_v2* 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_v2* 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_v2* 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_v2* 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_v2* 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"); + 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, + "open rejects factory engines missing hot-path hooks without release"); + 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, + "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/exec_contract.md b/docs/exec_contract.md index f9cf5b8a..b3bf4398 100644 --- a/docs/exec_contract.md +++ b/docs/exec_contract.md @@ -397,6 +397,34 @@ The Rust shell, OpenAI, and scheduler all live in serving / the upper layer, off - Therefore the common layer = replay-time execution of Buffer/Graph/Plan + a C ABI, **and not one bit more.** +### 7a. Two memory-domain contracts (Phase 6) + +FlashRT has two coexisting, non-unified memory-domain contracts. They are NOT +merged under one backend abstraction — each is honest about what it is: + +- **(A) CUDA-export path (same-backend, CUDA↔CUDA).** The `frt_buffer` + + `frt_buffer_dptr`/`frt_buffer_wrap` + `frt_graph_bind` + port SWAP window + + `FRT_RT_ROLE_INPUT|OUTPUT` role bitmask described throughout this doc. Two + graphs sharing one buffer on matching ports = zero-copy hand-off. This is + CUDA-specific by construction (device pointer, `cudaStream_t` `native_handle`) + and stays confined to the `exec/backend/cuda/` TU. **Unchanged by Phase 6.** + +- **(B) Provider-owned token path (cross-provider, FlashRT↔GGML).** A small + `frt_memory_token` (opaque) + `frt_memory_token_verbs` + (`copy_to_host`/`copy_from_host`/`sync`/`destroy`) + `frt_rt_location_kind` + (`HOST_VISIBLE`/`DEVICE_LOCAL`), attached to a provider-owned port via + `frt_runtime_builder_add_port_token`. FlashRT **never dereferences** the + token — zero-copy is an *advertised* capability (`location_kind=HOST_VISIBLE`), + not an assumption. `destroy` fires once at runtime release. This is the only + buffer form permitted on provider-owned ports; the builder still rejects a raw + `frt_buffer` there. See `docs/phase6_backend_vtable_eval.md`. + +The `exec/backend/backend.h` free-function surface is NOT promoted to a runtime +vtable in Phase 6 — link-time backend selection already suffices, and GGML is +integrated as a model provider (not a bare exec backend), so a full backend +vtable would have no consumer. `frt::be::import` is declared (commented) as a +future cross-domain buffer-registration seam but not implemented for non-CUDA. + --- ## 8. v1 acceptance checklist diff --git a/docs/jetson_pi_usage.md b/docs/jetson_pi_usage.md new file mode 100644 index 00000000..4664c4c8 --- /dev/null +++ b/docs/jetson_pi_usage.md @@ -0,0 +1,622 @@ +# Jetson-PI providers (Python) + +Jetson-PI is the inference engine introduced in +["Jetson-PI: Towards Onboard Real-Time Robot Control via Foresight-Aligned +Asynchronous Inference"](https://arxiv.org/abs/2607.12659). It targets +onboard, real-time robot control with PI0 and PI0.5 vision-language-action +models by running the policy through llama.cpp/GGML and exposing optimized +foreground inference paths across multiple backends, including desktop NVIDIA +GPUs, Jetson Orin/Thor-class devices, CPU-only builds, and NPU-oriented edge +integrations. The Jetson-PI codebase is available at +[PKU-SEC-Lab/Jetson-PI-Edge, branch `Jetson-PI-flashrt`](https://github.com/PKU-SEC-Lab/Jetson-PI-Edge/tree/Jetson-PI-flashrt). + +This FlashRT provider integrates the Jetson-PI runtime so applications can +load the same Jetson-PI GGUF model artifacts through FlashRT's Python API +instead of running the Jetson-PI HTTP foreground server. + +`flash_rt.load_model(..., framework="jetson_pi")` drives Jetson-PI +llama.cpp/GGML providers through the FlashRT `frt_model_runtime_v2` C ABI +via ctypes. No torch/jax, no GPU arch detection. + +The completed correctness/backend/performance gate record is maintained in +`docs/jetson_pi_validation_matrix.md`. + +Three providers, selected by `config=`: +- **`config="pi0"`** (default for VLA) — Pi0/Pi0.5 whole-graph infer via + `jetson_pi_pi0`. The narrow API detects the model family from the GGUF pair + and rejects a conflicting model/mmproj pair. Returns a `VLAModel` + (`.predict(images, prompt, state)`). +- **`config="llm"`** — generic GGUF text completion via `jetson_pi_llm`. + Returns an `LlmJetsonPiFrontend` directly (`.generate(prompt)`); not a VLA. +- **`config="mllm"`** — multimodal LLM (vision+text) via `jetson_pi_mllm`. + Returns an `MllmJetsonPiFrontend` directly (`.generate(images, prompt)`); + not a VLA. + +## Build + +Jetson-PI is a public repository: +https://github.com/PKU-SEC-Lab/Jetson-PI-Edge, branch `Jetson-PI-flashrt`. + +```bash +git clone --branch Jetson-PI-flashrt https://github.com/PKU-SEC-Lab/Jetson-PI-Edge.git Jetson-PI +``` + +Two integration routes (§5.2): a **dev** `add_subdirectory` path and a +**production** `find_package` path. Pick one. + +### Dev: `add_subdirectory` (builds Jetson-PI inline from its source tree) + +```bash +cmake -S FlashRT/cpp -B FlashRT/cpp/build-jetson-pi \ + -DCMAKE_C_COMPILER=.../x86_64-conda-linux-gnu-cc \ + -DCMAKE_CXX_COMPILER=.../x86_64-conda-linux-gnu-c++ \ + -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON \ + -DFLASHRT_CPP_WITH_JETSON_PI=ON \ + -DJETSON_PI_ROOT=/path/to/Jetson-PI +cmake --build FlashRT/cpp/build-jetson-pi -j32 --target flashrt_cpp_llama_cpp_provider_c +``` + +Produces `FlashRT/cpp/build-jetson-pi/libflashrt_cpp_llama_cpp_provider_c.so`. + +### Production: `find_package(JetsonPI)` (link a pre-installed Jetson-PI) + +Build and install Jetson-PI as a standalone tree (NOT via FlashRT's +`add_subdirectory`). With `LLAMA_BUILD_COMMON=ON` + `JETSON_PI_BUILD_MTMD=ON` +the narrow `jetson_pi_*` libs and `mtmd`/`llama`/`ggml` deps build. The default +direct-link layout installs `lib64/libjetson_pi_{pi0,llm,mllm}.so`, +`lib64/lib{llama,mtmd,ggml*}.so`, and +`include/jetson_pi_{pi0,llm,mllm}.h`. The installed MTMD public-header closure +also contains `mtmd-helper.h` and its `pi-model.h` dependency, so an +out-of-tree static or shared consumer does not depend on source-tree headers: + +```bash +cmake -S Jetson-PI -B Jetson-PI/build-prod \ + -DCMAKE_C_COMPILER=.../x86_64-conda-linux-gnu-cc \ + -DCMAKE_CXX_COMPILER=.../x86_64-conda-linux-gnu-c++ \ + -DCMAKE_CUDA_COMPILER=.../nvcc -DCMAKE_CUDA_ARCHITECTURES=89 \ + -DCMAKE_BUILD_TYPE=Release \ + -DGGML_CUDA=ON -DLLAMA_BUILD_COMMON=ON -DJETSON_PI_BUILD_MTMD=ON \ + -DLLAMA_CURL=OFF -DLLAMA_HTTPLIB=OFF \ + -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_TOOLS=OFF \ + -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=OFF \ + -DLLAMA_TOOLS_INSTALL=OFF \ + -DCMAKE_INSTALL_PREFIX=/path/to/jetson-pi-install-prefix +cmake --build Jetson-PI/build-prod -j16 \ + --target jetson_pi_pi0 jetson_pi_llm jetson_pi_mllm +cmake --install Jetson-PI/build-prod +``` + +> `LLAMA_TOOLS_INSTALL=OFF` is load-bearing: with `JETSON_PI_BUILD_MTMD=ON`, +> `tools/mtmd/CMakeLists.txt` unconditionally `add_executable(llama-mtmd-cli)` +> but only builds it when the full tools tree is built; the CLI's +> `install(TARGETS)` rule still runs unless `LLAMA_TOOLS_INSTALL=OFF`, aborting +> `cmake --install` with `cannot find llama-mtmd-cli` — and that abort happens +> *before* `libllama.so` is laid down, so `find_package(JetsonPI)` then fails on +> the missing `llama` lib. + +Then configure FlashRT against the installed prefix (no `JETSON_PI_ROOT`): + +```bash +cmake -S FlashRT/cpp -B FlashRT/cpp/build-jetson-pi-prod \ + -DCMAKE_C_COMPILER=.../x86_64-conda-linux-gnu-cc \ + -DCMAKE_CXX_COMPILER=.../x86_64-conda-linux-gnu-c++ \ + -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON \ + -DFLASHRT_CPP_WITH_JETSON_PI=ON \ + -DJetsonPI_ROOT=/path/to/jetson-pi-install-prefix +cmake --build FlashRT/cpp/build-jetson-pi-prod -j32 --target flashrt_cpp_llama_cpp_provider_c +``` + +`FlashRT/cmake/FindJetsonPI.cmake` resolves the installed libs/headers into +imported targets `JetsonPI::jetson_pi_pi0` / `::jetson_pi_llm` / `::jetson_pi_mllm` +(+ `::mtmd` / `::llama` / `::ggml`; direct-link packages also expose imported +per-backend targets). For a direct-link package, `ldd` on the built +`libflashrt_cpp_llama_cpp_provider_c.so` must resolve to the installed prefix's +`libjetson_pi_pi0.so.0` / `libllama.so.0` / `libggml-cuda.so.0`, not to any +in-tree copy. `cmake --install` of FlashRT lays down the SHARED provider +`_c.so` and `libflashrt_runtime.so` to `${CMAKE_INSTALL_LIBDIR}`. The installed +provider has an install RPATH containing `$ORIGIN` plus the exact library +directories resolved under `JetsonPI_ROOT`, so it can be loaded without an +ambient `LD_LIBRARY_PATH`. Set `-DJETSON_PI_ROOT=` for dev OR +`-DJetsonPI_ROOT=` for prod — exactly one is required. Configuration +fails explicitly when neither or both are set, preventing an implicit search +from selecting a stale system installation. + +When the installed prefix contains `libggml-cuda`, FlashRT's production +configure requires the matching CUDA Toolkit CMake package so the imported +`ggml-cuda` target carries its driver/runtime/cuBLAS link interface. This is a +build-host requirement; the installed provider itself was verified with only +its recorded RPATH and the normal NVIDIA runtime loader environment. + +#### Dynamic GGML backend package + +`GGML_BACKEND_DL=ON` keeps backend modules out of the core/provider ELF link +graph so one installed Jetson-PI package can select CPU or CUDA at runtime. A +production package must use an explicit absolute backend directory; relying on +the process executable directory or current working directory is rejected by +FlashRT's production finder: + +```bash +prefix=/absolute/path/to/jetson-pi-dynamic +cmake -S Jetson-PI -B Jetson-PI/build-prod-dynamic \ + -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=ON \ + -DGGML_BACKEND_DL=ON -DGGML_NATIVE=OFF \ + -DGGML_CUDA=ON -DGGML_CUDA_GRAPHS=ON \ + -DGGML_BACKEND_DIR="${prefix}/bin" \ + -DLLAMA_BUILD_COMMON=ON -DJETSON_PI_BUILD_MTMD=ON \ + -DLLAMA_CURL=OFF -DLLAMA_HTTPLIB=OFF \ + -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_TOOLS=OFF \ + -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=OFF \ + -DLLAMA_TOOLS_INSTALL=OFF -DCMAKE_INSTALL_PREFIX="${prefix}" +cmake --build Jetson-PI/build-prod-dynamic -j16 +cmake --install Jetson-PI/build-prod-dynamic +``` + +The resulting core/narrow libraries live in `lib64/`; runtime modules live in +`bin/libggml-{cpu,cuda}.so`. `FindJetsonPI.cmake` reads the same prefix's +`ggml-config.cmake`, requires its absolute `GGML_BACKEND_DIR` to remain inside +`JetsonPI_ROOT`, and verifies every module named by `GGML_AVAILABLE_BACKENDS`. +The modules are not added to `INTERFACE_LINK_LIBRARIES`. Consequently +`readelf -d` on `libggml.so`, `libjetson_pi_llm.so`, and FlashRT's provider has +no `NEEDED` entry for `libggml-cpu.so` or `libggml-cuda.so`; runtime logs show +`load_backend: loaded ... from /bin/...` instead. The installed FlashRT +`_c.so` was exercised from outside that directory through both C++ and Python: +CPU reports `offloaded 0/29`, CUDA on GPU6 reports `offloaded 29/29`, and both +pass staged decode, KV-isolation, and token-budget gates without a backend +fallback. The pure-CPU run sets `CUDA_VISIBLE_DEVICES=`; leaving a CUDA device +visible makes llama.cpp create a CUDA context even when all model weights remain +on CPU, so that configuration is not used as pure-CPU evidence. + +The narrow LLM/MLLM one-shot APIs support a no-inference size-query. The +reported capacity is `max_tokens` times the largest serialized token piece in +the loaded vocabulary, so callers do not assume a fixed UTF-8 byte count per +token. FlashRT stores the actual generated text and its Python frontends query +that completed output size before allocating the readback buffer. + +> **The Pi0 parity + engine tests are dev-path only.** `test_llama_cpp_jetson_pi_parity` +> and `test_llama_cpp_jetson_pi_engine` `#include "stb/stb_image.h"`, which lives +> only in Jetson-PI's `vendor/` source tree; the prod path has no `JETSON_PI_ROOT`, +> so `cpp/CMakeLists.txt` does not define those two targets on the prod path. +> The link, LLM, MLLM, and LLM-parity tests do not use stb and remain available +> with prod `BUILD_TESTING=ON`. A prod-path GPU smoke via the Python frontend +> (`FLASHRT_PI0_LIB=`) exercises the installed artifact +> end-to-end instead — verified on GPU6/RTX 4090 on 2026-07-10 (`offloaded +> 37/37`, actions `(10,32)`, no NaN/Inf, nonzero; installed `_c.so` resolves +> `libflashrt_runtime.so` through `$ORIGIN` and Jetson-PI/mtmd/llama/GGML through +> its exact install-prefix RPATH, without `LD_LIBRARY_PATH`). + +### Pinned Jetson-PI commit (§15.1) + +The Jetson-PI fork must be used as a locked version with its own GGML — do not +mix another llama.cpp/GGML version into the same provider (ABI/symbol/layout +conflicts). + +Jetson-PI is a public repository: https://github.com/PKU-SEC-Lab/Jetson-PI-Edge, +branch `Jetson-PI-flashrt`. The provider is GPU-verified against the tip of +that branch (commit `5de3f9e210086f7bc04c5a434990bd28e7ed2240` as of +2026-07-15). Tracked by https://github.com/flashrt-project/FlashRT/issues/143. + +The original migration was developed against a private fork (branch +`flashrt-migration-merge`, commit +`9c8e8b30e629a2475958c7b8a0bfb06f6f05295d`, based on upstream `origin/merge` +commit `436fdb2aceaf564e152be6e0779180f56a279074`); that history has since +been opened as the public `Jetson-PI-flashrt` branch above, so the private +commit SHAs are no longer reachable and are kept here only as a provenance +record. + +When upgrading the fork, re-run the parity tests +(`test_llama_cpp_jetson_pi_parity` / `test_llama_cpp_jetson_pi_llm_parity`) to +confirm FlashRT's provider path still matches the direct `jetson_pi_*` path +byte-for-byte. + +### Symbol visibility (§15.2) + +Each narrow C API lib (`libjetson_pi_pi0` / `libjetson_pi_llm` / +`libjetson_pi_mllm`) exports only its `jetson_pi_*` C symbols — no `ggml_*` / +`llama_*` / `mtmd_*` C API symbols are defined in the narrow libs (they are +imported via the PUBLIC link to `mtmd`/`llama`/`ggml`). Verified by +`for lib in libjetson_pi_{pi0,llm,mllm}.so; do nm -D --defined-only "$lib" | +grep -E ' (ggml_|llama_|mtmd_)[A-Za-z]'; done` → empty. (A +`std::vector` stdlib template weak symbol appears in +the pi0/mllm libs; that is a benign COMDAT stdlib instantiation referencing the +`mtmd_bitmap` type, not the mtmd C API.) + + +## Usage + +```python +import flash_rt + +model = flash_rt.load_model( + "/path/to/Pi0_Base-2.8B-F16.gguf", + framework="jetson_pi", + mmproj_path="/path/to/mmproj-model-f16.gguf", + backend="cpu", # or "cuda" + num_views=2, + action_steps=10, # pi0_base metadata; see CLAUDE.md weights note + # (pi0_libero_base fails this fork's check_tensor_dims) + action_dim=32, + lib_path=None, # auto-discover, or set FLASHRT_PI0_LIB +) + +actions = model.predict( + [image_rgb_224x224, wrist_rgb_224x224], # list of HxWx3 uint8 numpy + prompt="put the mug on the plate", + state=robot_state_floats, # 1-D float32, ≤ action_dim +) +# actions: np.ndarray shape (action_steps, action_dim), float32, raw (no unnorm) + +# Optional fine-grained execution. Context owns prompt/image encoding and +# provider-private encoded KV; action consumes that pending context once. +model._pipe.set_prompt("put the mug on the plate") +model._pipe.context({"images": [image, wrist_image], "state": robot_state}) +actions = model._pipe.action() + +# Read-only zero-copy host SWAP view over the latest action chunk. Versioned +# DLPack consumers retain the provider's read-only access contract. +actions_view = model._pipe.action_view() +numpy_actions = np.from_dlpack(actions_view) +actions_view.close() # DLPack consumer keeps its own map alive +del numpy_actions # required before replacing inputs or running a stage +``` + +Pi0.5 uses the same FlashRT face. Point `checkpoint` and `mmproj_path` at the +converted Pi0.5 pair; no process-global `PI_MODEL` override is required: + +```python +pi05 = flash_rt.load_model( + "/path/to/Pi05_Libero-2.8B-F16.gguf", + framework="jetson_pi", + config="pi0", + mmproj_path="/path/to/pi05/vit/mmproj-model-f16.gguf", + backend="cuda", + num_views=2, + action_steps=10, + action_dim=32, +) +actions = pi05.predict(images, prompt=prompt, state=state) +``` + +Jetson-PI detects Pi0 versus Pi0.5 from strong GGUF metadata/tensor features, +pins that choice to the handle for model load, VIT preprocessing, graph build, +and decode, and rejects an unknown or inconsistent pair during `open`. +The Pi0.5 converter also recognizes adaptive output-norm tensors in the raw +safetensors key namespace and omits an obsolete base output norm when both are +present. + +The Python smoke test verifies NumPy's versioned DLPack path is zero-copy, +read-only, and keeps the mapping live. PyTorch 2.6 requests the legacy DLPack +protocol, which cannot signal read-only storage; that export is intentionally +rejected instead of exposing a writable tensor over provider-owned actions. +No implicit copy fallback is installed. + +The Pi0 runtime exposes callback stages `infer`, `context`, and `action`, with +an explicit `context -> action` dependency. `infer` remains the compatibility +whole-tick face and is implemented by the same narrow context/action API. + +`LD_LIBRARY_PATH` must include the conda lib dir (for `libgomp.so`, needed by +`libggml-cpu.so`) and the build dir (for `libjetson_pi_pi0.so` / +`libmtmd.so` / `libllama.so` / `libggml*.so`): + +```bash +export LD_LIBRARY_PATH=.../miniconda3/lib:FlashRT/cpp/build-jetson-pi +``` + +## Run the smoke test + +```bash +python FlashRT/cpp/tests/fixtures/prepare_pi0_fixture.py \ + --npz --prompt --out-dir /tmp/pi0_fixture --action-dim 32 + +FLASHRT_PI0_MODEL=.../Pi0_Base-2.8B-F16.gguf \ +FLASHRT_PI0_MMPROJ=.../mmproj-model-f16.gguf \ +FLASHRT_PI0_FIXTURE_DIR=/tmp/pi0_fixture \ +FLASHRT_PI0_LIB=FlashRT/cpp/build-jetson-pi/libflashrt_cpp_llama_cpp_provider_c.so \ +FLASHRT_PI0_ACTION_STEPS=10 FLASHRT_PI0_ACTION_DIM=32 \ +python -m flash_rt.tests.test_jetson_pi_pi0_python +``` + +## Native Pi0 CLI baseline + +Jetson-PI's interactive `llama-mtmd-cli` now emits the complete policy output +after a successful Pi0 request: + +```text +Pi0 action shape: [10, 32] +Pi0 action: [0.491152287,...] +``` + +The action is row-major F32 and uses nine significant digits, which preserve a +round trip back to F32. `test_llama_cpp_jetson_pi_parity` accepts +`FLASHRT_PI0_CLI_ACTION_LOG=` and then requires the parsed CLI action to be +byte-identical to the FlashRT whole-infer result. The exact GPU6 command, +including the explicit test template needed because Pi0 Base has no embedded +chat template, is recorded in `docs/jetson_pi_validation_matrix.md`. + +## Limitations +- **Pi0 raw action chunk.** `predict` returns the model's `action_steps × + action_dim` output without unnormalization or LIBERO 7-D slicing. The caller + is responsible for post-processing (use `meta/stats.json` to unnormalize). +- **LLM raw prompt.** `generate(prompt)` takes a raw prompt; the caller must + apply the chat template (e.g. `llama_chat_apply_template`) before calling. + `generate()` returns one blob and starts an independent session; + `reset()`/`prefill()`/repeatable `decode()` expose the token boundary and keep + provider-private KV state for host-driven generation. +- **CPU/CUDA/SYCL clean-exit and Vulkan inference paths verified.** `backend="cuda"` is verified + end-to-end on an RTX 4090 (sm_89) for all provider faces — Pi0 and Pi0.5 + (`offloaded 37/37 layers` plus VIT), LLM (`29/29`, qwen3-0.6b-q4_k_m), and MLLM + (`29/29`, Qwen3-VL-2B-Instruct-Q4_K_M; the VIT/mmproj encoder offloads too). + For Pi0 and LLM the generated output is coherent; for MLLM the GPU execution + path (load → VIT encode → forward → sample) is exercised, but the smoke test + feeds a raw prompt so the output is template-token noise rather than a + caption — output coherence requires a caller-applied chat template (see the + MLLM note below). `backend="vulkan"` runs the intended model and VIT paths on + the RTX 4090 for Pi0, Pi0.5, LLM, and MLLM. Both policy variants reach + numerical parity with `max_diff = 0`, LLM retains exact direct-path text + parity, MLLM retains exact text parity, and all four final provider tests + reach their success markers. Each process then exits 139 in + the local NVIDIA 550.54.14 ICD during process-exit teardown, so none is a + clean-exit Vulkan validation. Per §6 of + the migration plan, compiling a backend does not guarantee every model op is + supported on it — each model×backend combo is verified separately. The CUDA + build needs its own build dir with `-DGGML_CUDA=ON` (it defaults OFF), plus + the same Jetson-PI flags: + + ```bash + cmake -S FlashRT/cpp -B FlashRT/cpp/build-jetson-pi-cuda \ + -DCMAKE_C_COMPILER=.../x86_64-conda-linux-gnu-cc \ + -DCMAKE_CXX_COMPILER=.../x86_64-conda-linux-gnu-c++ \ + -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \ + -DCMAKE_CUDA_ARCHITECTURES=89 \ + -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON \ + -DFLASHRT_CPP_WITH_JETSON_PI=ON \ + -DJETSON_PI_ROOT=/path/to/Jetson-PI \ + -DGGML_CUDA=ON -DGGML_CUDA_FA=ON + cmake --build FlashRT/cpp/build-jetson-pi-cuda -j32 \ + --target flashrt_cpp_llama_cpp_provider_c + ``` + + Then point `FLASHRT_PI0_LIB` at `build-jetson-pi-cuda/libflashrt_cpp_llama_cpp_provider_c.so`, + set `FLASHRT_PI0_BACKEND=cuda`, and ensure `LD_LIBRARY_PATH` includes the + CUDA build's `bin/` (for `libggml-cuda.so`) and `runtime/`. The smoke test + then offloads all model layers to the GPU (37/37 for pi0_base). The same + `_c.so` and env-override recipe apply to the LLM (`FLASHRT_LLM_*`) and MLLM + (`FLASHRT_MLLM_*`) smoke tests — see their sections below. + + Note: the Jetson-PI engine accepts `backend` `"cpu"` / `"cuda"` / `"vulkan"` + / `"opencl"` / `"sycl"` by exact string match. Any other value (including + typos) is **rejected** by `open` (returns `INVALID` + a null handle), not a + silent CPU fallback. `"cpu"` sets `n_gpu_layers=0`. A non-CPU value resolves + the exact GGML registry (`CUDA`, `Vulkan`, `OpenCL`, or `SYCL`), requires at + least one non-CPU registered device, and selects registry device 0 as the + llama model's offload device and MTMD/CLIP's primary accelerator. A build + missing the requested backend or a host with no accelerator therefore fails + during open instead of silently treating the request as CPU or another + compiled backend. GGML's normal heterogeneous scheduler remains intact: CPU + is still present for host work and for individual operations an accelerator + backend does not implement. This is per-operation scheduling inside the + explicitly selected provider, not backend-request fallback. Environment variables + such as `CUDA_VISIBLE_DEVICES` and `GGML_VK_VISIBLE_DEVICES` determine which + physical accelerator appears as that registry's device 0. Treat the + `offloaded N/N layers` and `CLIP using ` log lines as confirmation + that the intended backend was exercised. + + **Vulkan backend** (Pi0/Pi0.5/LLM/MLLM inference verified; all four have the NVIDIA + process-exit teardown failure below). Build in a separate + dir with `-DGGML_VULKAN=ON -DGGML_CUDA=OFF` (Vulkan is SPIR-V; no CUDA arch + needed) and conda compilers; the `vulkan-shaders-gen` build-time tool needs a + C++17 host compiler, so export `CC`/`CXX` to the conda compilers for the + build step: + + ```bash + cmake -S FlashRT/cpp -B FlashRT/cpp/build-jetson-pi-vulkan \ + -DCMAKE_C_COMPILER=.../x86_64-conda-linux-gnu-cc \ + -DCMAKE_CXX_COMPILER=.../x86_64-conda-linux-gnu-c++ \ + -DCMAKE_PREFIX_PATH=.../envs/migrate \ + -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON \ + -DFLASHRT_CPP_WITH_JETSON_PI=ON \ + -DJETSON_PI_ROOT=/path/to/Jetson-PI \ + -DGGML_VULKAN=ON -DGGML_CUDA=OFF + CC=.../x86_64-conda-linux-gnu-cc CXX=.../x86_64-conda-linux-gnu-c++ \ + cmake --build FlashRT/cpp/build-jetson-pi-vulkan -j32 \ + --target flashrt_cpp_llama_cpp_provider_c + ``` + + `CMAKE_PREFIX_PATH` must point at a prefix providing the Vulkan loader + (`libvulkan.so`), `glslc` (conda-forge `vulkan-loader` + `shaderc`), and the + Vulkan **headers**. Header caveat: conda-forge `vulkan-headers` (1.3.231) + lacks the `vk_video/` av1 codec headers and `vulkan_hpp_macros.hpp` that + `vulkan_core.h` unconditionally `#include`s, so a build against just the conda + headers fails with `fatal error: vk_video/vulkan_video_codec_av1std.h`. Use + the current Khronos `Vulkan-Headers` (≥1.3.295, which bundles `vk_video/` + + the split hpp files) — `git clone --depth 1 https://github.com/KhronosGroup/Vulkan-Headers.git` + and point `CMAKE_PREFIX_PATH` at its `include` parent (or stage its `vulkan/` + + `vk_video/` ahead of the conda env's). Do NOT rely on overwriting the conda + env's `include/vulkan` in place — a later `conda install/update vulkan-headers` + clobbers it; a separate prefix is robust. The Vulkan device is selected by + `GGML_VK_VISIBLE_DEVICES=` (emulates `CUDA_VISIBLE_DEVICES`). Then set + `FLASHRT_PI0_BACKEND=vulkan` / `FLASHRT_LLM_BACKEND=vulkan`, point + `FLASHRT_PI0_LIB`/`FLASHRT_LLM_LIB` at the vulkan build's `_c.so`, and ensure + `LD_LIBRARY_PATH` includes the vulkan build's `bin/` (for `libggml-vulkan.so`). + Verified: Pi0 `pi0_base` and the converted Pi0.5 LIBERO model each offload + 37/37 layers + VIT to Vulkan0, produce finite `(10,32)` actions, and retain + whole/split/direct parity `max_diff = 0`; + LLM `qwen3-0.6b-q4_k_m` offloads 29/29 layers to Vulkan0, greedy text + byte-identical to the direct `jetson_pi_llm` call. Qwen3-VL MLLM also + completes one-shot and staged generation with all 29 text layers and VIT on + Vulkan0. On final rebuilt GPU6 runs, all four tests print their success marker + and subsequently segfault inside the NVIDIA 550.54.14 Vulkan ICD's process-exit + destructors with exit code 139. Earlier GDB capture places the failure after + the test success marker, outside provider inference and explicit object + destruction; attempts to force Vulkan teardown made the driver race more + reproducible and were not retained. Treat these as inference-success and + teardown-fail results, not clean-exit Vulkan validation. + + **OpenCL backend.** A dedicated `-DGGML_OPENCL=ON` build succeeds with the + Khronos headers and NVIDIA ICD. The runtime enumerates platform `NVIDIA CUDA` + and device `NVIDIA GeForce RTX 4090 (OpenCL 3.0 CUDA)`, but this fork's + `ggml-opencl` implementation accepts only Adreno/Qualcomm and Intel GPU + families. It prints `Unsupported GPU`, drops the NVIDIA device, and explicit + `backend="opencl"` then fails with + `requested GGML backend has no registered device: opencl`. This is an honest + backend capability gate, not a missing download and not a model execution + result; no CPU fallback is installed. + + On this host the ICD loader must be pointed at the system vendor files with + `OCL_ICD_VENDORS=/etc/OpenCL/vendors`; without that environment setting the + conda loader does not enumerate the installed NVIDIA ICD. Enumeration still + ends at the explicit unsupported-family gate described above. + + **SYCL backend.** A separate `jetsonpi-sycl-cuda` conda environment uses + DPC++ 2025.3.0, GCC 13.4, CUDA 12.4, and oneMath cuBLAS to build + `-DGGML_SYCL=ON -DGGML_SYCL_TARGET=NVIDIA -DGGML_SYCL_DNN=OFF`. The matching + Unified Runtime loader and CUDA adapter are selected explicitly at runtime: + + ```bash + export LD_PRELOAD=/path/to/ur-cuda-2025.3.0/lib64/libur_loader.so.0 + export UR_ADAPTERS_FORCE_LOAD=/path/to/ur-cuda-2025.3.0/lib64/libur_adapter_cuda.so + export ONEAPI_DEVICE_SELECTOR=cuda:6 + export FLASHRT_LLM_BACKEND=sycl # likewise PI0/MLLM + ``` + + GPU6/RTX 4090 end-to-end results: LLM offloads 29/29 layers; MLLM offloads + 29/29 language layers and CLIP/VIT to SYCL0; Pi0 and Pi0.5 each offload + 37/37 model layers and VIT. LLM staged decode, MLLM staged/image parity, and + both policy variants' whole-versus-context/action parity all pass with + `max_abs_diff=0`. A focused SYCL + backend-op test passes nearest and F32 bilinear upscale (half-pixel and + align-corners), transpose, and up/downsample cases, 11/11 total. Bicubic is + still explicitly unsupported and is not claimed. The installed-prefix build + links `libsycl`, `libimf`, `libsvml`, `libirng`, `libintlc`, oneMath cuBLAS, + and the CUDA driver explicitly; `FindJetsonPI.cmake` validates the equivalent + direct-link dependencies instead of creating a partially linked target. +- **No calibration.** The frontends have no `calibrate`/`calibrated`; the + Jetson-PI providers do not need FlashRT-style FP8 calibration. +- **`state` is a separate port** for Pi0, not encoded into the prompt (unlike + Pi0.5). `VLAModel.predict` detects that `set_prompt` does not accept `state` + and routes it through `observation["state"]` automatically. +- **Deployment identity includes the checkpoint contents.** The standard JSON + open path streams each complete model/mmproj file through SHA-256, appends + the complete 64-digit hexadecimal digest to the canonical runtime identity, + and fails the open if any checkpoint cannot be opened or read. There is no + path-only fallback. The fingerprint also includes backend, port schema, callback-stage + DAG, and execution layout, so weights, quantization, backend, or stage-schema + changes alter the deployment fingerprint. The config structs carry these + digests in append-only tails, so existing explicit-engine callers using the + old prefix remain ABI-compatible. Full-file hashing adds deterministic model + open work in exchange for satisfying deployment identity exactly; it is not + part of the inference hot path. + +## Generic GGUF LLM (Phase 3) + +```python +import flash_rt + +fe = flash_rt.load_model( + "/path/to/qwen3-0.6b-q4_k_m.gguf", + framework="jetson_pi", + config="llm", + backend="cpu", + n_ctx=2048, n_threads=0, + temp=0.8, top_k=40, top_p=0.9, seed=1, max_tokens=512, + lib_path=None, # auto-discover, or set FLASHRT_LLM_LIB +) + +text = fe.generate("What is 2 plus 2? The answer is") +# text: str, the generated completion (no chat template applied by the engine) + +# Host-driven KV-cache session. prefill accepts exactly one of prompt/tokens. +logits = fe.prefill("What is 2 plus 2? The answer is") +step = fe.decode() # {"token": int, "is_eog": bool, "text": accumulated str} +logits = fe.get_logits() +fe.reset() +logits = fe.prefill(tokens=[151644, 198]) # caller-owned int32 token IDs +``` + +The returned object is an `LlmJetsonPiFrontend` (not a `VLAModel` — LLMs are +not VLA). `fe.infer({"prompt": ...})` returns `{"text": ...}` for callers that +want a dict-shaped interface. + +The provider runtime exposes callback stages `infer`, `reset`, `prefill`, and +repeatable `decode`, with STAGED `prompt`, optional `tokens`, `next_token`, +`logits`, `is_eog`, and accumulated `text` ports. KV state and sampling remain +provider-private. The host interrupts at a token boundary by stopping decode +calls. Each staged session enforces `max_tokens`; the narrow API starts a fresh +decode budget on `prefill`, while the FlashRT DAG convention uses +`reset`+`prefill` to start a new session. + +### Run the LLM smoke test + +```bash +FLASHRT_LLM_MODEL=.../qwen3-0.6b-q4_k_m.gguf \ +FLASHRT_LLM_LIB=FlashRT/cpp/build-jetson-pi/libflashrt_cpp_llama_cpp_provider_c.so \ +LD_LIBRARY_PATH=.../miniconda3/lib:FlashRT/cpp/build-jetson-pi \ +python -m flash_rt.tests.test_jetson_pi_llm_python +``` + +`FLASHRT_LLM_BACKEND=cuda` switches the test to the GPU forward pass; point +`FLASHRT_LLM_LIB` at the `build-jetson-pi-cuda/libflashrt_cpp_llama_cpp_provider_c.so` +and add the cuda build's `bin/` to `LD_LIBRARY_PATH` (same recipe as the Pi0 +CUDA section above). Verified on an RTX 4090 (sm_89): `offloaded 29/29 layers +to GPU` for qwen3-0.6b-q4_k_m. The same staged/one-shot smoke also passes on +CPU and CUDA for TinyLlama 1.1B (`llama`, 23 layers) and Gemma 2 2B (`gemma2`, +27 layers), covering three distinct GGUF architectures rather than a single +Qwen-specific path. + +## Multimodal LLM (Phase 4) + +```python +import flash_rt + +fe = flash_rt.load_model( + "/path/to/Qwen2.5-VL-3B-Instruct-q4_0.gguf", + framework="jetson_pi", + config="mllm", + mmproj_path="/path/to/Qwen2.5-VL-3B-Instruct-mmproj-f16.gguf", + backend="cpu", + n_ctx=2048, n_threads=0, + temp=0.0, top_k=0, top_p=0.0, seed=1, max_tokens=64, + lib_path=None, # auto-discover, or set FLASHRT_MLLM_LIB +) + +logits = fe.prefill([image], "Describe this image.") +step = fe.decode() +fe.reset() + +text = fe.generate( + [image_rgb_224x224], # list of HxWx3 uint8 numpy (may be empty for text-only) + "Describe this image in one sentence.", +) +# text: str, the generated completion +``` + +The returned object is an `MllmJetsonPiFrontend` (not a `VLAModel` — MLLMs +output text, not actions). `fe.infer({"images": [...], "prompt": ...})` returns +`{"text": ...}` for callers that want a dict-shaped interface. + +The caller is responsible for applying the chat template; the engine only does +raw prompt + media markers → text. Each `generate` call clears KV (independent +completion, no multi-turn). Staged decode uses the same token-boundary +interruption and `max_tokens` budget contract as the text LLM face. + +### Run the MLLM smoke test + +```bash +FLASHRT_MLLM_MODEL=.../Qwen2.5-VL-3B-Instruct-q4_0.gguf \ +FLASHRT_MLLM_MMPROJ=.../Qwen2.5-VL-3B-Instruct-mmproj-f16.gguf \ +FLASHRT_MLLM_LIB=FlashRT/cpp/build-jetson-pi/libflashrt_cpp_llama_cpp_provider_c.so \ +LD_LIBRARY_PATH=.../miniconda3/lib:FlashRT/cpp/build-jetson-pi \ +python -m flash_rt.tests.test_jetson_pi_mllm_python +``` + +`FLASHRT_MLLM_BACKEND=cuda` switches the test to the GPU forward pass (LLM +layers + VIT/mmproj encoder both offloaded); point `FLASHRT_MLLM_LIB` at the +`build-jetson-pi-cuda/libflashrt_cpp_llama_cpp_provider_c.so` (same CUDA recipe +as the Pi0 section above). Verified on an RTX 4090 (sm_89): Qwen3-VL 2B +offloads 29/29 language layers and the CLIP/VIT encoder to CUDA0. + +Note: the engine injects one media marker per supplied image *after* the prompt +text (see `jetson_pi_mllm.cpp`); the smoke test feeds a raw prompt, so the +generated text is a connectivity check (non-empty, printable), not a +task-accurate caption. Callers wanting a well-formed answer must apply the +model's chat template (with the image placeholder positioned by the template) +themselves. diff --git a/docs/jetson_pi_validation_matrix.md b/docs/jetson_pi_validation_matrix.md new file mode 100644 index 00000000..66c15488 --- /dev/null +++ b/docs/jetson_pi_validation_matrix.md @@ -0,0 +1,296 @@ +# Jetson-PI Provider Validation Matrix + +This document is the reviewable record for `jetsonpi迁移.txt` §14. It +separates correctness gates from diagnostic performance measurements. Timings +are single-run wall-clock diagnostics from the named logs, not statistically +stable benchmarks and not cross-backend speed claims. + +## Validated Revisions + +- FlashRT validation baseline: + branch `jetson-pi-link-check` (fork `heiheiha798/FlashRT`), tracked by + https://github.com/flashrt-project/FlashRT/issues/143. The provider + implementation was complete at `0080e3daa383b077955d3063f2ab755de1996dca`; + later commits through this baseline corrected and finalized the validation + record. (The pre-rebase history recorded baseline `91a5cc3`; commit SHAs + shift on rebase, so the branch + issue are the stable reference.) +- Jetson-PI revision: + public repository https://github.com/PKU-SEC-Lab/Jetson-PI-Edge, branch + `Jetson-PI-flashrt`, tip `5de3f9e210086f7bc04c5a434990bd28e7ed2240` + (2026-07-15). Originally developed against a private fork (branch + `flashrt-migration-merge`, commit + `9c8e8b30e629a2475958c7b8a0bfb06f6f05295d`, based on + `origin/merge@436fdb2aceaf564e152be6e0779180f56a279074`; the prior migration + is preserved at `origin/flashrt-migration-master-baseline@68dd395b3f89dbd031ae564e335780f702fbd1e7`). + That history has since been opened as the public `Jetson-PI-flashrt` branch, + so the private SHAs are kept here only as provenance. +- CUDA/Vulkan/SYCL device used for final accelerator runs: physical GPU 6, + NVIDIA GeForce RTX 4090. CUDA used `CUDA_VISIBLE_DEVICES=6`; Vulkan used + `GGML_VK_VISIBLE_DEVICES=6`; SYCL used `ONEAPI_DEVICE_SELECTOR=cuda:6` with + the matching Unified Runtime CUDA adapter. + +## Migration Requirement Coverage + +| Source requirement | Delivered result | Review gate | +|---|---|---| +| Phase 0: baseline and contracts | Locked Jetson-PI/GGML revision and build identity; full-file model/mmproj SHA-256; explicit port shape/dtype/ownership; native Pi0 CLI input, action output, and timing baseline. | CLI and FlashRT use the same two RGB fixtures, 32-F32 state, effective prompt, CUDA backend, and fixed-seed policy; all 320 action elements are bit-identical. | +| Phase 1: in-process Pi0 whole graph | `frt_model_runtime_v2` callback `infer` stage with staged images/prompt/state/actions; provider owns Jetson-PI/MTMD/GGML state and exposes no GGML type in FlashRT ABI. The migrated narrow API detects and pins Pi0 versus Pi0.5 per handle. | Pi0 and Pi0.5 CPU/CUDA/SYCL clean-exit model runs, Vulkan inference-success/teardown-fail runs, and direct narrow-C-API parity. | +| Phase 2: Python and C API | One C++ implementation is reached through the FlashRT C ABI, C++ factory, and ctypes `flash_rt.load_model(framework="jetson_pi")`; lifecycle, errors, stale output, and multiple instances are tested. | Native C++ and Python smoke suites, including installed-provider loading. | +| Phase 3: generic text LLM | Prompt/optional tokens, next-token, EOG, full logits, accumulated text, `reset -> prefill -> decode`, private KV/sampler state, host interruption, and hard token budget. | Qwen3 0.6B, TinyLlama 1.1B, and Gemma 2 2B on CPU/CUDA; Qwen3 on clean-exit SYCL and inference-success/teardown-fail Vulkan; direct parity and interleaved-session isolation. | +| Phase 4: multimodal LLM | Images + prompt, provider-private MTMD/VIT embeddings, one-shot and repeatable prefill/decode, next-token/EOG/logits/text outputs. | Qwen3-VL 2B CPU/CUDA/SYCL exact direct parity; Vulkan inference reaches success before the recorded NVIDIA ICD teardown fault. | +| Phase 5: finer Pi0 stages | Stable narrow `context`/`action` C API and FlashRT `context -> action` DAG; pending context is single-use; encoded cross-KV and strictly guarded decoder graph reuse remain provider-private. | Whole/split/direct parity is exact; stale/replaced/failed input invalidation is tested; graph profiling confirms one rebuild plus nine reuses. | +| Phase 6: backend and zero-copy evolution | Explicit CPU/CUDA/Vulkan/OpenCL/SYCL selection; opaque memory-domain token; host map/unmap; read-only versioned DLPack; dynamic GGML backend package. | Model-by-backend matrix below; OpenCL is an explicit capability gate; pointer identity, mapping lifetime, module loading, and hard-failure paths are tested. | +| §5.2/§15 production delivery | Dev `add_subdirectory` and installed `find_package(JetsonPI)` routes, exact-prefix RPATH, dynamic backend module validation, locked fork, and narrow-library symbol audit. | Standalone install, out-of-tree FlashRT consumer, `readelf`/`nm` audits, C++ CPU/CUDA, and installed `_c.so` Python CUDA runs. | +| §10.2 and §12 recommendations | A full generic `exec/backend` vtable and the illustrative directory reshuffle are not required for the provider contract. The implemented minimal memory-domain contract preserves the existing CUDA Graph path without moving verified code for style alone. | Design decision and re-entry criteria are recorded in `phase6_backend_vtable_eval.md`; no functional requirement depends on the suggested file layout. | +| §14 correctness/performance matrix | Shape/value/finite checks, repeatability, context/KV isolation, port-schema identity, checkpoint/backend/stage fingerprinting, latency breakdowns, token/s, and sampled accelerator memory. | Matrix, invariants, logs, and reproduction commands below. | + +## Correctness Matrix + +| Model face | Model | Backend | Result | Evidence | +|---|---|---|---|---| +| Pi0 | Pi0 Base F16 + VIT mmproj | CPU | PASS | FlashRT versus direct `jetson_pi_pi0` action parity, `max_abs_diff=0`; shape `(10,32)`; finite actions. | +| Pi0 | Pi0 Base F16 + VIT mmproj | CUDA | PASS | 37/37 model layers and VIT on CUDA; whole infer repeatability; no context leak; FlashRT/direct parity `max_abs_diff=0`; whole versus context/action and native CLI versus FlashRT are bit-identical. | +| Pi0 | Pi0 Base F16 + VIT mmproj | Vulkan | INFERENCE/PARITY PASS; TEARDOWN FAIL | 37/37 model layers plus VIT on Vulkan; FlashRT/direct parity `max_abs_diff=0` reaches `== PI0 PARITY PASSED ==`, then process exit is 139. | +| Pi0 | Pi0 Base F16 + VIT mmproj | SYCL-on-CUDA | PASS | 37/37 model layers plus VIT on SYCL0; whole/direct and whole/context-action parity `max_abs_diff=0`; repeated inference and host-SWAP/DLPack gates pass. | +| Pi0.5 | Pi05 LIBERO F16 (371 tensors) + VIT mmproj | CPU | PASS | Automatic Pi0.5 detection selects the Pi0.5 text/VIT adapters; finite `(10,32)` actions; whole/split/direct `max_abs_diff=0`; host cross-KV is refreshed on every denoise step. | +| Pi0.5 | Pi05 LIBERO F16 (371 tensors) + VIT mmproj | CUDA | PASS | 37/37 model layers plus VIT on CUDA; all 371 model tensors load; whole/split/direct `max_abs_diff=0`; installed Python host-SWAP and read-only DLPack gates pass. | +| Pi0.5 | Pi05 LIBERO F16 (371 tensors) + VIT mmproj | Vulkan | INFERENCE/PARITY PASS; TEARDOWN FAIL | 37/37 model layers plus VIT on Vulkan; whole/split/direct `max_abs_diff=0` reaches `== PI0 PARITY PASSED ==`, then process exit is 139. | +| Pi0.5 | Pi05 LIBERO F16 (371 tensors) + VIT mmproj | SYCL-on-CUDA | PASS | 37/37 model layers plus VIT on SYCL0; whole/split/direct `max_abs_diff=0`; finite actions and clean exit. | +| Text LLM | Qwen3 0.6B Q4_K_M | CPU | PASS | FlashRT/direct greedy text, first token, and complete prefill logits parity; logits `max_abs_diff=0`. | +| Text LLM | Qwen3 0.6B Q4_K_M | CUDA | PASS | 29/29 layers offloaded; FlashRT/direct text, first token, and complete prefill logits parity with `max_abs_diff=0`; two distinct interleaved sessions reproduce standalone token/EOG sequences. | +| Text LLM | Qwen3 0.6B Q4_K_M | Vulkan | INFERENCE/PARITY PASS; TEARDOWN FAIL | 29/29 layers offloaded; FlashRT/direct greedy text exact parity is retained, and the final rebuilt provider test reaches `== JETSON_PI LLM PASSED ==`, then process exit is 139. | +| Text LLM | Qwen3 0.6B Q4_K_M | SYCL-on-CUDA | PASS | 29/29 layers on SYCL0; staged decode, interleaved session isolation, logits, interruption, and token-budget gates pass with clean exit. | +| Text LLM | TinyLlama 1.1B | CPU/CUDA | PASS | 23/23 CPU and 23/23 CUDA model matrix runs complete. | +| Text LLM | Gemma 2 2B | CPU/CUDA | PASS | 27/27 CPU and 27/27 CUDA model matrix runs complete. | +| MLLM | Qwen3-VL 2B Q4_K_M + F16 mmproj | CPU | PASS | FlashRT and direct `jetson_pi_mllm` produce exact image+text output for the same generated RGB fixture and greedy sampling. | +| MLLM | Qwen3-VL 2B Q4_K_M + F16 mmproj | CUDA | PASS | 29/29 language layers plus CLIP/VIT on CUDA; FlashRT/direct image+text output is exact; one-shot equals staged prefill/decode; finite logits. | +| MLLM | Qwen3-VL 2B Q4_K_M + F16 mmproj | Vulkan | INFERENCE PASS; TEARDOWN FAIL | 29/29 language layers plus CLIP/VIT on Vulkan; one-shot/staged, finite-logit, and budget gates reach `== JETSON_PI MLLM PASSED ==`, then process exit is 139. | +| MLLM | Qwen3-VL 2B Q4_K_M + F16 mmproj | SYCL-on-CUDA | PASS | 29/29 language layers plus CLIP/VIT on SYCL0; FlashRT/direct output exact; one-shot/staged decode, finite logits, and token budget pass. | +| LLM provider | Qwen3 0.6B | OpenCL | BACKEND CAPABILITY GATE | NVIDIA OpenCL 3.0 platform/device is enumerated, but this fork accepts only Adreno/Qualcomm and Intel families, drops the RTX 4090, and explicit open fails with zero registered devices. | +| Production package | Qwen3 0.6B plus installed Python Pi0/Pi0.5/MLLM | direct-link and `GGML_BACKEND_DL` CPU/CUDA | PASS | Direct-link dependencies resolve to the locked install prefix and CUDA reports 29/29; dynamic-package core/provider ELF has no backend-module `DT_NEEDED`, CPU reports 0/29, CUDA reports 29/29, and C++ plus installed Python gates pass. The installed MTMD header closure includes `pi-model.h`, and out-of-tree static/shared consumers build without source-tree headers. | + +## Known Driver-Lifecycle Result + +- With the final rebuilt binaries, Pi0, Pi0.5, Qwen3 0.6B, and Qwen3-VL 2B all reach + their test success markers after running the intended Vulkan model/VIT paths. + Each process then segfaults during NVIDIA 550.54.14 driver teardown and exits + 139. No Vulkan row is counted as a clean PASS, and the provider does not + install a fallback or speculative teardown workaround. + +## §14 Invariants + +- Consecutive Pi0 requests invalidate stale actions and reproduce the same + fixed-seed output for identical inputs; no previous context leaks. +- CPU Pi0/Pi0.5 graph reuse refreshes host-resident encoded cross-KV on every + denoise step. The GPU-resident path retains its explicit ready guard; the + host path no longer reuses stale graph-input storage and produces finite, + exact-parity actions. +- A failed replacement of any known Pi0/LLM/MLLM input port clears that port's + ready state before returning. A later stage therefore hard-fails instead of + reusing the previous request's value. RGB conversion rejects negative or + undersized row strides and verifies the complete last-row byte span. +- Pi0 actions are checked for shape, nonzero values, and NaN/Inf absence. +- Pi0 host SWAP is consumed zero-copy through NumPy's versioned, read-only + DLPack path; pointer identity and the live-consumer stage guard are tested. + PyTorch 2.6 legacy DLPack export is rejected because it cannot preserve the + read-only provider mapping, and no implicit copy fallback is used. +- Pi0 token `copy_to_host` rejects a null destination for nonzero copies. +- LLM/MLLM one-shot size-query uses `max_tokens` times the largest serialized + token piece in the loaded vocabulary, not a fixed bytes-per-token guess. + FlashRT and its Python frontends then query and read the actual completed + text byte count from the staged output port. +- The narrow LLM API rejects prompt lengths that cannot be represented by the + llama tokenizer's `int32_t` length. Pi0 and both MLLM entry points reject + zero or overflowing RGB dimensions before bitmap construction; staged MLLM + `prefill` enforces the same boundary as one-shot inference. +- LLM/MLLM prefill exposes finite vocabulary logits; decode is host-repeatable. +- Host interruption is explicit: callers stop generation by not issuing another + decode stage. Each staged session enforces its configured `max_tokens`; one + extra decode hard-fails. The narrow API restores the budget on `prefill`, + while FlashRT conventionally starts sessions with reset+prefill. +- Same-session KV is reused across decode steps; distinct LLM sessions are + tested with different prompts and compared token-by-token with standalone + baselines. +- Backend selection is an explicit string and unavailable/unknown backends + hard-fail. A provider unit test verifies that changing backend changes the + fingerprint while preserving the complete port schema. +- OpenCL validation sets `OCL_ICD_VENDORS=/etc/OpenCL/vendors` so the NVIDIA + OpenCL 3.0 ICD is actually enumerated. `ggml-opencl` then rejects the NVIDIA + family, registers zero devices, and the requested OpenCL provider hard-fails. +- A dynamic GGML package validates every backend module declared by its own + `ggml-config.cmake`; runtime modules are not linked into the provider. Missing, + relative, or cross-prefix backend directories fail production configuration. +- Pure-CPU production validation hides CUDA devices with + `CUDA_VISIBLE_DEVICES=`. Otherwise llama.cpp creates contexts for visible CUDA + devices even when model weights report `offloaded 0/29`; that is not a pure + CPU execution claim and is not used as evidence here. +- Standard JSON open computes full-file SHA-256 for model and mmproj files. + A same-size edit after byte 64 KiB changes the fingerprint. Missing or + unreadable checkpoints fail before factory creation and expose a thread-local + open error; there is no path-only fallback. +- Canonical runtime identity also fingerprints quantization-bearing checkpoint + bytes, backend, port schema, callback-stage DAG, and executable layout. + +## Diagnostic Performance Snapshot + +These values come from existing single-run logs and show that every requested +measurement point is observable. They are not acceptance thresholds. + +| Model/backend | Diagnostic values | +|---|---| +| Pi0 CPU | VIT about 6.4-6.6 s per view, encode about 2.17 s, ten-step action decode about 39.38 s. | +| Pi0 CUDA | Warm VIT about 7.6-9.8 ms per view, encode about 102-114 ms, ten-step action decode about 1.10-1.52 s. Graph reuse triggers one rebuild plus nine reuses; build+alloc is below 1% of step time, so compute remains dominant. | +| Pi0 native CLI CUDA | Exact-parity fixture: two image slices 19/8 ms, VIT diagnostic 19.09 ms, encode 481.17 ms, ten-step action decode 117.08 ms, context total 969.96 ms, and process wall time 4.51 s including model/mmproj load and built-in warmup. | +| Pi0 SYCL-on-CUDA | Warm VIT 25.82/27.83 ms for two views, encode 189.22 ms, and ten-step action decode 215.52 ms; the first cold tick was 209.74/38.59 ms, 497.99 ms, and 496.37 ms respectively. | +| Qwen3 0.6B CPU | Dynamic-package staged prefill 4365.797 ms; 16-token repeatable decode 0.25 token/s (63784.955 ms). | +| Qwen3 0.6B CUDA | Staged prefill 5.909 ms; repeatable decode 433.98 token/s (16 tokens in 36.868 ms). The independently installed dynamic package measured 6.076 ms and 429.75 token/s. | +| Qwen3 0.6B SYCL-on-CUDA | Staged prefill 12.802 ms; repeatable decode 149.02 token/s (16 tokens in 107.368 ms). | +| TinyLlama 1.1B CPU/CUDA | 12-token batch diagnostics about 2.66-2.79 s on CPU and 4.1-30.6 ms on CUDA. | +| Gemma 2 2B CPU/CUDA | 12-13-token batch diagnostics about 3.92-7.79 s on CPU and 6.5-35.2 ms on CUDA. | +| Qwen3-VL 2B CPU | Warm VIT 315 ms, staged multimodal prefill 13248.101 ms, and repeatable decode 0.25 token/s (16 tokens in 64684.120 ms). | +| Qwen3-VL 2B CUDA | Warm VIT 5 ms, staged multimodal prefill 22.881 ms, and repeatable decode 295.27 token/s (16 tokens in 54.188 ms). | +| Qwen3-VL 2B SYCL-on-CUDA | Warm VIT 12-16 ms, staged multimodal prefill 61.547 ms, and repeatable decode 152.04 token/s (16 tokens in 105.233 ms). | +| Qwen3-VL 2B Vulkan | 58-59-token multimodal prompt batches about 39-98 ms before the known process-exit driver teardown failure. | + +At 50 ms `nvidia-smi` sampling on physical GPU6, the complete Qwen3 LLM CUDA +test peaked at 2914 MiB and the SYCL-on-CUDA test at 2844 MiB from a 0 MiB +baseline; both returned to 0 MiB after clean exit. These are process-level +observations for the full multi-session test, not allocator-internal counters. + +Final-run logs used by this matrix include: + +- `/tmp/merge-pi0-parity-cpu-fixed.log` +- `/tmp/merge-pi0-parity-cuda-fixed.log` +- `/tmp/merge-pi0-parity-sycl-gpu6.log` +- `/tmp/merge-pi0-parity-vulkan-gpu6.log` +- `/tmp/merge-pi05-parity-cpu.log` +- `/tmp/merge-pi05-parity-cuda-fixed.log` +- `/tmp/merge-pi05-parity-sycl-gpu6.log` +- `/tmp/merge-pi05-parity-vulkan-gpu6.log` +- `/tmp/merge-llm-parity-{cpu,sycl-gpu6,vulkan-gpu6}.log` +- `/tmp/merge-llm-parity-cuda-fixed.log` +- `/tmp/merge-mllm-parity-{cpu,sycl-gpu6,vulkan-gpu6}.log` +- `/tmp/merge-mllm-parity-cuda-fixed.log` +- `/tmp/merge-llm-opencl-nvidia-gate.log` +- `/tmp/merge-sycl-upscale-backend-ops-gpu6.log` +- `/tmp/merge-prod-direct-llm-cuda-gpu6.log` +- `/tmp/merge-prod-dynamic-llm-{cpu,cuda-gpu6}.log` +- `/tmp/merge-installed-python-{llm,pi0,pi05,mllm}-cuda-gpu6.log` +- `/tmp/merge-reviewfix-pi05-parity-cuda-gpu6.log` +- `/tmp/merge-reviewfix-mllm-parity-cuda-gpu6.log` + +The following older logs remain the historical pre-`origin/merge` evidence: + +- `/tmp/llm-token-logit-parity-cpu-final.log` +- `/tmp/llm-token-logit-parity-cuda-final-gpu6.log` +- `/tmp/mllm-parity-cpu-final.log` +- `/tmp/mllm-parity-cuda-final-gpu6.log` +- `/tmp/full-digest-pi0-gpu6.log` +- `/tmp/final-mllm-vulkan-gpu6.log` +- `/tmp/final-audit-pi0-vulkan-gpu6.log` +- `/tmp/final-audit-llm-vulkan-gpu6.log` +- `/tmp/final-audit-mllm-vulkan-gpu6.log` +- `/tmp/final-opencl-hard-gate.log` +- `/tmp/final-sycl-gate3.log` +- `/tmp/decode-budget-llm-gpu6.log` +- `/tmp/decode-budget-mllm-gpu6.log` +- `/tmp/llm-cuda-perf-gpu6.log` +- `/tmp/llm-cuda-memory-gpu6.csv` +- `/tmp/llm-sycl-perf-gpu6.log` +- `/tmp/llm-sycl-memory-gpu6.csv` +- `/tmp/pi0-engine-sycl-allweights-gpu6.log` +- `/tmp/pi0-parity-sycl-gpu6.log` +- `/tmp/mllm-cpu-perf-final.log` +- `/tmp/mllm-cuda-perf-gpu6.log` +- `/tmp/mllm-sycl-perf-gpu6.log` +- `/tmp/final-sycl-upscale-backend-ops-gpu6.log` +- `/tmp/llm-opencl-nvidia-gpu6.log` +- `/tmp/final-backend-dl-llm-pure-cpu.log` +- `/tmp/final-backend-dl-llm-cuda-gpu6.log` +- `/tmp/final-backend-dl-python-llm-cuda-gpu6.log` +- `/tmp/final-prod-direct-python-llm-cuda-gpu6.log` +- `/tmp/final-cuda-provider-ctest-gpu6.log` +- `/tmp/final-sycl-provider-ctest-gpu6.log` +- `/tmp/final-python-{pi0,llm,mllm}-cuda-gpu6.log` +- `/tmp/final-python-pi0-sycl-gpu6.log` +- `/tmp/final-review-fix-cuda-provider-ctest-gpu6.log` +- `/tmp/review-fix-sycl-provider-ctest-gpu6.log` +- `/tmp/review-fix-python-{llm,mllm}-cuda-gpu6.log` +- `/tmp/review-fix-backend-dl-llm-pure-cpu.log` +- `/tmp/review-fix-backend-dl-llm-cuda-gpu6.log` +- `/tmp/review-fix-backend-dl-python-llm-cuda-gpu6.log` +- `/tmp/review-fix-prod-direct-python-llm-cuda-gpu6.log` +- `/tmp/review-fix-nonmodel-ctest.log` +- `/tmp/pi0-cli-baseline-gpu6.log` +- `/tmp/pi0-cli-parity-gpu6.log` +- `/tmp/pi0-cli-parity-gate-gpu6.log` +- `/tmp/pi0-flashrt-cli-parity-gpu6.log` + +## Reproduction + +Check for an idle device first; prefer GPU 6 when several devices are idle. + +```bash +nvidia-smi --query-gpu=index,memory.used,memory.total,utilization.gpu \ + --format=csv,noheader + +export CUDA_VISIBLE_DEVICES=6 +export LD_LIBRARY_PATH="$PWD/cpp/build-jetson-pi-cuda/bin:$PWD/cpp/build-jetson-pi-cuda/runtime:$LD_LIBRARY_PATH" + +FLASHRT_LLM_MODEL=/data/pretrained_models/qwen3-0.6b-q4_k_m.gguf \ +FLASHRT_LLM_BACKEND=cuda \ +cpp/build-jetson-pi-cuda/test_llama_cpp_jetson_pi_llm + +FLASHRT_PI0_MODEL=/data/pretrained_models/pi0_model/pi0_base/Pi0_Base-2.8B-F16.gguf \ +FLASHRT_PI0_MMPROJ=/data/pretrained_models/pi0_model/pi0_base/vit/mmproj-model-f16.gguf \ +FLASHRT_PI0_FIXTURE_DIR=/tmp/pi0_fixture \ +FLASHRT_PI0_ACTION_STEPS=10 FLASHRT_PI0_ACTION_DIM=32 \ +FLASHRT_PI0_BACKEND=cuda \ +cpp/build-jetson-pi-cuda/test_llama_cpp_jetson_pi_engine + +FLASHRT_MLLM_MODEL=/data/pretrained_models/Qwen3-VL-2B-Instruct-GGUF/Qwen3-VL-2B-Instruct-Q4_K_M.gguf \ +FLASHRT_MLLM_MMPROJ=/data/pretrained_models/Qwen3-VL-2B-Instruct-GGUF/mmproj-F16.gguf \ +FLASHRT_MLLM_BACKEND=cuda \ +cpp/build-jetson-pi-cuda/test_llama_cpp_jetson_pi_mllm +``` + +### Native Pi0 CLI parity + +The Pi0 GGUF has no embedded chat template, so the CLI requires an explicit +template rather than choosing one implicitly. For the numerical gate, the +test-only Jinja expression rotates the two 11-byte media markers inserted by +interactive mode from the front to the end. The resulting token input is +exactly the narrow API's `prompt + marker + marker` convention. + +From the migration workspace root: + +```bash +state="$(od -An -t f4 -v /tmp/pi0_fixture/state.bin | \ + tr -s ' ' '\n' | sed '/^$/d' | paste -sd, -)" +prompt="$(tr -d '\r\n' < /tmp/pi0_fixture/prompt.txt)" + +printf '/image %s\n/image %s\n/state %s\n%s\n/exit\n' \ + /tmp/pi0_fixture/image.png /tmp/pi0_fixture/wrist_image.png \ + "$state" "$prompt" | \ +CUDA_VISIBLE_DEVICES=6 Jetson-PI/build-backend-dl-cuda/bin/llama-mtmd-cli \ + -m /data/pretrained_models/pi0_model/pi0_base/Pi0_Base-2.8B-F16.gguf \ + --mmproj /data/pretrained_models/pi0_model/pi0_base/vit/mmproj-model-f16.gguf \ + -ngl 999 --jinja \ + --chat-template "{{ messages[0]['content'][22:] }}{{ messages[0]['content'][:22] }}" \ + -v > /tmp/pi0-cli-parity-gpu6.log 2>&1 + +CUDA_VISIBLE_DEVICES=6 \ +FLASHRT_PI0_MODEL=/data/pretrained_models/pi0_model/pi0_base/Pi0_Base-2.8B-F16.gguf \ +FLASHRT_PI0_MMPROJ=/data/pretrained_models/pi0_model/pi0_base/vit/mmproj-model-f16.gguf \ +FLASHRT_PI0_FIXTURE_DIR=/tmp/pi0_fixture \ +FLASHRT_PI0_ACTION_STEPS=10 FLASHRT_PI0_ACTION_DIM=32 \ +FLASHRT_PI0_BACKEND=cuda \ +FLASHRT_PI0_CLI_ACTION_LOG=/tmp/pi0-cli-parity-gpu6.log \ +FlashRT/cpp/build-jetson-pi-cuda/test_llama_cpp_jetson_pi_parity +``` + +The final GPU6 gate reports 320 CLI elements and +`CLI max abs diff = 0`, then verifies byte identity with `memcmp`. + +The Vulkan, OpenCL, and SYCL build commands and runtime dependency notes are +recorded in `docs/jetson_pi_usage.md`. diff --git a/docs/model_runtime_api.md b/docs/model_runtime_api.md index 7730bfc4..1d5d3047 100644 --- a/docs/model_runtime_api.md +++ b/docs/model_runtime_api.md @@ -36,11 +36,17 @@ port, never advertise-and-refuse. `bytes` (null buffer = staged-only). Strings/arrays are owned by the runtime object and stay valid while a reference is held. -`frt_runtime_stage_desc` — one schedulable stage: `graph` (index into the +`frt_runtime_stage_desc` — v1 graph-backed stage: `graph` (index into the export's graphs) plus `after[n_after]` (earlier stage indices). Declared array order is the sequential order `step` uses. Stage streams are not a separate field: replay stream comes from the referenced graph descriptor. +`frt_runtime_stage_desc_v2` — v2 backend-neutral stage: `name`, `kind` +(`GRAPH` or `CALLBACK`), `graph` (for graph-backed stages), `callback` (a +provider-private executable id), and `after[n_after]`. `CALLBACK` stages are +executed through `frt_model_runtime_verbs_v2::run_stage`; GGML/llama.cpp +objects remain provider-private and are not exposed in the ABI. + Graph stream placement and the stage DAG are deployment identity. A cut from `full` to `context_action`, a stream move, or a dependency change changes the fingerprint; this is intentional state-safety, not a cache miss. @@ -58,6 +64,12 @@ frt_model_runtime_v1 { } ``` +`frt_model_runtime_v2` keeps the v1 prefix concepts (`exp`, `ports`, +v1-compatible verbs, lifetime) and appends `stages_v2 / n_stages_v2` plus +`verbs_v2`. Its legacy `stages / n_stages` view is empty when callback stages +are present, so graph-only consumers do not accidentally treat a provider-owned +stage as a CUDA graph. + **Verbs** (`frt_model_runtime_verbs`; every entry is always callable — absent producer verbs are filled with unsupported stubs returning `-3`): @@ -69,6 +81,10 @@ producer verbs are filled with unsupported stubs returning `-3`): | `step(self)` | HOT (sugar) | fire all stages in declared order; scheduling hosts fire stages themselves | | `last_error(self)` | — | message for the most recent failure | +v2 adds `run_stage(self, stage, stream)`, which executes one v2 stage by index. +It is required for `CALLBACK` stages. A white-box host may still schedule +`GRAPH` stages through `exp->graphs`. + Status codes follow the pi05 C face: `0` ok, `-1` invalid, `-2` not found, `-3` unsupported, `-4` shape mismatch, `-5` insufficient storage, `-6` backend. @@ -94,6 +110,34 @@ frt_model_runtime_v1* m = frt_runtime_builder_finish_model( b, &verbs, verbs_self, owner, retain_owner, release_owner); ``` +**Provider-owned v2** — the provider owns execution and exposes staged ports +plus callback stages, without pretending to have FlashRT exec graphs: + +```c +frt_runtime_builder b = frt_runtime_builder_create_provider_owned(); +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, shape, rank, 0, + NULL, 0, 0); +frt_runtime_builder_add_callback_stage_v2(b, "infer", 0, NULL, 0); +frt_model_runtime_v2* m = frt_runtime_builder_finish_model_v2( + b, &verbs_v2, verbs_self, owner, retain_owner, release_owner); +``` + +Provider-owned v2 ports are `STAGED` by default. The builder rejects a raw +SWAP window (`frt_buffer` + offset + bytes) on this path — that would imply +same-backend device sharing with no contract. **Phase 6** adds the ONE +permitted buffer form on provider-owned ports: a memory-domain **token** +(`frt_memory_token` + `frt_memory_token_verbs` + `frt_rt_location_kind`), +attached via `frt_runtime_builder_add_port_token`. The token is an opaque +provider-minted handle; FlashRT never dereferences it — it only calls the +provider's `copy_to_host` / `copy_from_host` / `sync` / `destroy` verbs. +Zero-copy is an ADVERTISED capability (`location_kind = HOST_VISIBLE`), not an +assumption FlashRT makes. `destroy` fires exactly once when the runtime's +refcount hits zero. See `docs/phase6_backend_vtable_eval.md` for the full +contract and the vtable-vs-contract decision. The CUDA-export SWAP path +(`frt_buffer` on a non-provider-owned runtime) is unchanged. + Identity covers each port's schema **and its bound window** (buffer index into the declared buffers array, offset, bytes) plus the stage DAG; only `cadence_hint_hz` stays out. A port-schema or window change therefore changes @@ -103,6 +147,7 @@ the fingerprint, and stored state is refused. Canonical record formats: port:::::::::::: graph:: stage::: +stage_v2:::::: ``` **Adapter** — wrap an existing export with ports/verbs; identity inherited, @@ -137,6 +182,12 @@ identity is unchanged. (`flash_rt/models/pi05/runtime_export.py`, via `flash_rt.runtime.export.build_model_runtime`) and the native Pi0.5 verb overlay `frt_pi05_model_runtime_create_over` (`cpp/models/pi05/`). +The initial llama.cpp/GGML Pi0 provider boundary lives at +`cpp/providers/llama_cpp/`: it exposes a `frt_model_runtime_v2` wrapper with +`images`, `prompt`, `state`, `actions` staged ports and one provider-owned +`infer` callback stage. The actual Jetson-PI/GGML engine remains behind an +explicit injected engine interface; until that engine is wired, the wrapper does +not fabricate inference or expose GGML types. ## Producer layout: model contract vs hardware pipeline diff --git a/docs/phase5_pi0_stages_eval.md b/docs/phase5_pi0_stages_eval.md new file mode 100644 index 00000000..2e09be55 --- /dev/null +++ b/docs/phase5_pi0_stages_eval.md @@ -0,0 +1,319 @@ +# Phase 5 Evaluation — Finer Pi0 Stages (context / encoder / diffusion / action) + +> Implementation update (2026-07-10): the earlier no-go below is superseded at +> the stable boundary required by the complete migration. Jetson-PI now exposes +> separate `prepare_context` and `run_action` narrow C API calls, and FlashRT +> exposes callback stages `context -> action` while preserving `infer` as the +> compatibility whole-tick face. The context stage owns image preprocessing, +> VIT, language/image encode, and provider-private encoded cross-KV; action owns +> the complete ten-step diffusion loop and consumes the pending context once. +> Whole infer versus split execution is bit-identical (`max_abs_diff=0`) on +> CUDA and SYCL. This does not claim intermediate diffusion actions or an +> intra-tick scheduling speedup; the historical analysis remains authoritative +> for any split finer than the stable context/action boundary. + +> Historical evaluation dated 2026-07-07. The implementation update above is +> authoritative; the remaining text records the original performance gate. + +## 1. Executive summary + recommendation + +**Recommendation: NO-GO. Defer Phase 5.** + +Four independent lenses (Pi0 internal compute structure, FlashRT provider surface, FlashRT ABI expressiveness, scheduling-benefit) converge on the same conclusion: splitting Pi0 into finer FlashRT-visible stages yields **zero latency win under current conditions** and adds ABI surface, new ports, new sync points, and new error paths for no return. The dominant optimizations live *inside* Jetson-PI's diffusion phase (Phase C), not at a new FlashRT stage boundary. + +The single most important structural fact: Pi0 is **flow-matching diffusion with atomic chunk emission**, not autoregressive generation. The K=10 denoise steps refine a *latent* `[50, action_dim]` in lockstep; `cross.action_ready` is set `true` only after the full loop completes (`Jetson-PI/src/llama-context.cpp:1389`), and `llama_get_pi0_action` refuses early reads (`:798`). Intermediate latents are not valid actions (each velocity contribution is scaled by `1/n_inference_steps`, `Jetson-PI/src/models/pi0_ae.cpp:109`). **Time-to-first-action == time-to-full-chunk.** There is no streaming win to capture, which removes the most attractive reason to split. + +The second structural fact: the encoder (Phase B) already runs **once per tick, outside** the K-loop (`llama-context.cpp:975` encode, then `:1241` K decode calls). The encoder→diffusion dependency is already serialized once and non-redundant. Splitting B and C into two FlashRT stages cannot overlap them within a tick because diffusion's first step consumes encoder cross-KV (`pi0_ae.cpp:13,68-69`). The only theoretical scheduling win is *cross-tick* pipelining (tick N+1's VIT+encoder behind tick N's K diffusion), and that is a provider-internal streaming change on separate GGML streams with double-buffered encoder KV — **not a FlashRT stage decomposition**. + +The Phase 5 gate in CLAUDE.md requires all four of: scheduling benefit, buffer contract, sync semantics, profiling evidence. On current evidence **none of the first three pass**, and the fourth (profiling) is to be backfilled by the main process using the plan in this document. The profiling plan's purpose is to *confirm* the no-go (diffusion dominates, graph-reuse is the real lever) and to set the numeric bar any future revisit must clear — not to search for a justification to split. + +A no-go here is the correct, defensible outcome: it prevents adding ABI surface and stage-boundary overhead for zero latency gain on the actual Jetson deployment, and it correctly routes the real optimization work (graph reuse across the 10 denoise steps; eliminate the per-step D2H recurrence; eliminate the cross-KV host round-trip) to where it belongs — inside Jetson-PI, behind the existing single `infer` stage. + +--- + +## 2. Pi0 internal phase map (one policy tick) + +Entry: `jetson_pi_pi0_infer` (`Jetson-PI/src/jetson_pi_pi0.cpp:205-344`). The whole compute is delegated to `mtmd_helper_eval_chunks_pi0` (`Jetson-PI/tools/mtmd/mtmd-helper.cpp:446-794`); only the final action readout is a separate C-API call. Five phases: + +| Phase | What | Where | Cost shape | +|---|---|---|---| +| A — VIT image encoder | Per-view CLIP-VIT pass; output `ctx->image_embd_v` `[n_img_tokens=196, n_mmproj_embd=1152]` F32 per view | `mtmd-helper.cpp:575` → `mtmd_encode` → `clip_image_batch_encode` (`tools/mtmd/mtmd.cpp:780,809,835`); graph `clip_ctx::build_pi0` (`tools/mtmd/clip.cpp:626-653`) | 1 pass per view (LIBERO: 2 views). Already timed ("Vit took"). | +| B — Language+image context encoder | Single non-causal prefill over all image+text tokens; first half of 27 layers (il=0..12). NO KV cache (`build_attn_inp_no_cache_pi0`). | `llama_encode` at `mtmd-helper.cpp:766` → `llama_context::encode` (`llama-context.cpp:911`) → `process_ubatch(LLM_GRAPH_TYPE_ENCODER)` (`:975`); graph `llm_build_pi0` (`src/models/pi0.cpp:4-151`), dispatched `llama-model.cpp:7180-7182`. | 1× per tick. Side outputs: `cross.encoded_kv_data[i]` (per-layer cross-KV) + `cross.action` (Gaussian noise latent `[50,32]`) + `cross.state`. | +| C — Diffusion flow-matching denoise loop | K=10 Euler steps; each builds+computes the AE decoder graph (second half of layers, il=13..26) over `action_steps+1=51` query tokens with cross-attention to encoder KV. | Loop `for (i=0; iget_action()` is copied device→host via `ggml_backend_tensor_get_async` (`llama-context.cpp:1254-1260`), then `cross.action[j] -= action_data[action_dim+j]` (`:1263-1264`). The `1/n_inference_steps` scale is baked into the graph at `pi0_ae.cpp:109`. +- **State carried between steps: ONLY `cross.action`** (updated in place). `cross.encoded_kv_data` and `cross.state` are CONSTANT across all 10 steps. No KV cache, no recurrent state. +- **The decoder graph is REBUILT + reallocated on every one of the 10 steps.** `can_reuse` returns `false` for `llm_graph_input_state` (`llama-graph.cpp:109`), `llm_graph_input_action` (`:137`), `llm_graph_input_sinusoidal_embedding` (`:192`), and `llm_graph_input_attn_no_cache_ae` inherits the base default `return false` (`llama-graph.h:100-105`, no override). So `process_ubatch` falls into the rebuild branch (`llama-context.cpp:864-887`) 10× per tick. This is the single most actionable inefficiency, and it is a Jetson-PI fix — not a FlashRT stage split. + +**Phase-boundary buffers (the contract surface a split would have to specify):** + +| Boundary | Buffer | Shape | Lives in | file:line | +|---|---|---|---|---| +| A→B (per view) | `ctx->image_embd_v` | `[196, 1152]` F32 | mtmd ctx host | `mtmd.cpp:817`, `mtmd-helper.cpp:589` | +| B→C (encoder→decoder) | `cross.encoded_kv_data[i]`, i=0..n_layer-1 (per-layer K and V) | each `[n_embd_head=256, n_head_kv=16, n_enc]` F32, n_enc=n_tokens | host (D2H via `ggml_backend_tensor_get_async`) | `llama-graph.h:75`, `llama-context.cpp:1066-1071`, `llama-graph.cpp:441-457` | +| B→C | `cross.action` (noise latent x_T) | `[50, 32]` F32 | host | `llama-context.cpp:1077` | +| B→C | `cross.state` (proprioception) | `[32]` F32 (padded) | host | `llama-context.cpp:1079-1081`, `llama-graph.cpp:102-107` | +| C internal (per step) | `res->action` (velocity v_theta) | `[32*(50+1)]` F32 | GPU tensor | `pi0_ae.cpp:113`, read `llama-context.cpp:1251` | +| C→D / D→E | `cross.action` (final x_0) | `[50, 32]` F32 | host | `llama-context.cpp:1264,1389,804` | +| cross-step constant | `cross.encoded_kv_data`, `cross.state` | as above | host | unchanged across 10 steps | + +The B→C cross-KV is the heaviest contract: n_layer·2 tensors, currently materialized host-side via an explicit D2H with a `// TODO: tmp` hack comment (`llama-graph.h:56-60`). This device→host→device round-trip (`llama-context.cpp:1069` D2H, `llama-graph.cpp:455` H2D feed-back per denoise step) is itself a serialization cost and a known buffer-contract sore point. + +**Streaming potential: effectively NONE.** (1) The chunk is produced all-at-once at the END (`action_ready` only after the loop, `llama-context.cpp:1389`; early reads rejected `:798`). (2) Even intermediate `cross.action` after each Euler step is the latent at time t_i, NOT a usable action — it must complete all 10 steps to converge to x_0. The per-step velocity is a gradient, not an action. There is no "early action step available before later ones" semantics; the 50 action timesteps emit together as one `[50,32]` block. The only streaming-shaped opportunity is running fewer denoise steps (K'<10) for a lower-quality chunk — that is a real per-tick latency reduction, but it trades quality for latency and is orthogonal to stage splitting: the whole-graph path already exposes it via `pi0.num_inference_steps` in the GGUF, no stage boundary needed. Stage splitting does not unlock it. + +**Cross-tick context reuse: zero.** The Pi0 encoder is stateless/no-cache (`build_attn_inp_no_cache_pi0`), KV is cleared every tick (`jetson_pi_pi0.cpp:234`), and the camera images change every tick in a robot deployment — so there is no encoder output to reuse across ticks even when the prompt is unchanged. A `prefill_context` stage split therefore cannot amortize across ticks. + +--- + +## 3. Current whole-graph path summary (FlashRT provider surface) + +The provider exposes **four STAGED ports and exactly ONE callback stage**: + +- Ports (`FlashRT/cpp/providers/llama_cpp/src/pi0_runtime.cpp:142-157`): `images` IN (IMAGE/U8/NHWC `[n_views,H,W,C]`), `prompt` IN (TEXT/U8/FLAT `{-1}`), `state` IN (STATE/F32/FLAT `[action_dim]`), `actions` OUT (ACTION/F32/FLAT `[action_steps,action_dim]`). All STAGED, none with a SWAP device window (buffer=nullptr, offset=0, bytes=0). +- Stage: `frt_runtime_builder_add_callback_stage_v2(b, "infer", 0, nullptr, 0)` (`pi0_runtime.cpp:158-159`). Stage index enum frozen at `FRT_LLAMA_CPP_PI0_STAGE_INDEX_INFER=0` (`c_api.h:25-27`). +- `run_stage` (`pi0_runtime.cpp:58-73`) accepts ONLY `INFER` (`:62`) and forwards to `owner->engine.run_infer` (`:66`); any other index returns "unknown stage" (`:63`). `step` (`:75-77`) is sugar calling `run_stage(INFER,-1)`. +- Engine vtable `frt_llama_cpp_engine_v1` (`c_api.h:96-115`) has only `set_input` / `run_infer` / `get_output` — a SINGLE monolithic infer verb. `run_infer` (`jetson_pi_engine.cpp:229-255`) gates on `images_set/prompt_set/state_set` (`:233`), allocates a local `actions` vector (`:239-240`), calls **one opaque blocking C call** `jetson_pi_pi0_infer(...)` (`:242-247`), and copies the result into `e->actions_buf` (`:253`). The `stream` parameter is dropped at every engine verb (`jetson_pi_engine.cpp:151,265`). +- `get_output` (`jetson_pi_engine.cpp:264-291`) serves only ACTIONS, requires whole-chunk capacity (`:277-284`), gates on `actions_buf.size()==need_elems` (`:285`, else `-7 ACTION_NOT_READY`), and memcpys the whole chunk (`:289`). +- Invalidation is blunt: `e->actions_buf.clear()` on **every** `set_input` regardless of port (`jetson_pi_engine.cpp:157`). + +**The only FlashRT-visible seam is the single `jetson_pi_pi0_infer` call at `jetson_pi_engine.cpp:242-247`.** Everything inside that call (VIT, context prefill, KV fill, diffusion loop, action decode) is opaque to FlashRT. A split into `[encode_images, prefill_context, diffusion_loop, decode_actions]` therefore requires **Jetson-PI to expose separate sub-phase C entry points**; FlashRT cannot fabricate the boundary from its side. Profiling the sub-phases from the FlashRT side is impossible for the same reason — the only FlashRT-side instrumentation point is bracketing `jetson_pi_pi0_infer`. + +--- + +## 4. FlashRT ABI capability assessment + +**Can the ABI DECLARE finer stages today? YES — append-only, no breakage.** `frt_runtime_stage_desc_v2` (`FlashRT/runtime/include/flashrt/model_runtime.h:172-179`) carries `{name, kind, graph, callback, n_after, after[]}`; `frt_runtime_builder_add_callback_stage_v2` (`:323-327`) appends one stage per call; `run_stage(self, stage, stream)` (`:232`) fires one stage by index. A mixed graph+callback 2-stage DAG is proven by `test_model_runtime.cpp:394-409`. So 4 callback stages with `after` edges are declarable today. The single-stage Pi0 provider is a provider choice, not an ABI ceiling. + +**Can the ABI host OBSERVABLE finer stages with inter-stage buffer contracts today? NO.** Four concrete gaps: + +1. **No intermediate port direction.** `frt_rt_port_direction` is a strict 2-value enum, NOT a bitmask: `FRT_RT_PORT_IN=0 / FRT_RT_PORT_OUT=1` (`model_runtime.h:101`). A port cannot be stage-N-out AND stage-(N+1)-in. STAGED ports only flow caller→engine (IN via `set_input`) and engine→caller (OUT via `get_output`). Every inter-stage tensor (image embeddings, encoder KV, diffusion latent, action logits) must therefore be a PRIVATE provider buffer; FlashRT cannot mediate, fingerprint, or sync it. + +2. **Provider-owned path forbids buffers and zero-copy windows.** The builder rejects all streams/graphs/buffers/regions on provider-owned runtimes (`FlashRT/runtime/src/model_runtime.cpp:203-208`) and rejects any SWAP window / offset / bytes on ports (`:209-217`). The export-level buffer `role` IS a bitmask allowing INPUT|OUTPUT (`runtime.h:45-49`, comment "an in-place diffusion noise/action buffer") and `frt_graph_bind` enables zero-copy hand-off between CUDA graphs sharing one buffer — but that mechanism belongs to the **CUDA-graph export path only**, not provider-owned v2. The docs state this explicitly: "Provider-owned v2 ports are STAGED only in this first contract. The builder rejects raw SWAP windows on this path until FlashRT has an explicit memory-domain contract for cross-provider buffers" (`docs/model_runtime_api.md:127-129`). That contract is the **Phase 6 deferral**. + +3. **No partial / per-step / streaming output.** `get_output(self, port, out, capacity, written, stream)` (`model_runtime.h:201-203,222-224`) is a full-snapshot read of ONE port by index. No step-k/K parameter, no per-step ready flag, no streaming token. If diffusion were split into K denoise sub-stages, each would be an opaque `run_stage` call with NO observable intermediate output until the final decode. The Pi0 engine confirms the all-or-nothing shape: `set_input` clears `actions_buf` (`jetson_pi_engine.cpp:157`), `get_output` returns the whole chunk only after `run_infer` completes (`:264-291`). + +4. **Engine vtable is monolithic.** `frt_llama_cpp_engine_v1` (`c_api.h:96-115`) exposes only `run_infer`. `jetson_pi_pi0_infer` (`jetson_pi_engine.cpp:242-247`) is one opaque call. Backing 4 stages needs either a generic `run_stage(stage_idx)` engine verb or per-sub-stage verbs, AND Jetson-PI must expose the sub-phase C entry points. Both are provider-internal changes. + +**Implication.** Finer stages can be prototyped TODAY only as **provider-internal sub-steps hidden behind the existing single `infer` stage** (the Phase 1 whole-graph shape stays the externally declared surface). Promoting them to first-class FlashRT-observable stages — where FlashRT mediates/fingerprints/syncs the buffers between encode→prefill→diffusion→decode and the host can read partial action progress — requires Phase 6-style work first (memory-domain/SWAP extension to provider-owned ports, or new append-only ABI concepts: intermediate port direction, per-step `get_output`, stage-ready flag). The current ABI alone cannot host observable finer Pi0 stages without new abstractions. + +--- + +## 5. Scheduling-benefit analysis (with honest counterargument) + +### 5.1 Action-chunk streaming — NO win +Pi0 is flow-matching diffusion, not autoregressive. The K steps refine a latent in lockstep; the chunk is emitted atomically after the loop. Time-to-first-action == time-to-full-chunk. **Theoretical latency reduction from action streaming: 0.** (See phase map §2.) + +### 5.2 Prefill/diffusion overlap — within-tick blocked; cross-tick bounded +Within one tick: VIT → encoder (13 layers, non-causal, once) → K-step diffusion (14 layers, K=10). The encoder already runs **once, outside** the K-loop (`llama-context.cpp:975` then `:1241`). Diffusion's first step consumes encoder cross-KV, so encoder→diffusion is a strict prerequisite — but only once per tick, not once per K step. The architectural seam already exists; there is no *redundant* work to remove by splitting B and C into two FlashRT stages. + +Cross-tick pipelining (tick N+1's VIT+encoder overlapping tick N's K diffusion) is the **only genuine scheduling opportunity**. On a single Jetson GPU it requires: (a) encoder and decoder on independent streams with no shared backend-scheduler mutex; (b) double-buffered encoder KV so tick N's diffusion reads buffer A while tick N+1's encoder writes buffer B; (c) **elimination of the per-step D2H+host-recurrence serialization barrier** (`llama-context.cpp:1254` D2H of `res->action`, then `:1263-1264` host Euler update) so the diffusion loop is GPU-resident. None of (a)(b)(c) exist today. Even with all three, the win is bounded by `min(encoder_time, K×step_time)` and lives entirely inside the provider — it is not a FlashRT stage decomposition. The correct exposure even then is ONE callback stage wrapping a provider-internal pipeline. + +### 5.3 Batching — NO dimension on a single Jetson +`cparams.n_seq_max=1` (`jetson_pi_pi0.cpp:135`), KV cleared every tick (`jetson_pi_pi0.cpp:234`), explicitly stateless single-tick. Batch=1 is the deployment norm on one Jetson controlling one robot. There is no second concurrent tick to batch. Finer stages do not unlock batching on the target deployment. + +(VIT *cross-view* batching — combining LIBERO's 2 camera views into one VIT pass — is a potential latency win, but it is an mtmd/clip-internal change to the per-view loop at `mtmd-helper.cpp:575`; a FlashRT `encode_images` stage split would still call the same per-view mtmd path and gain nothing. It is orthogonal to staging and is a Jetson-PI-internal optimization.) + +### 5.4 Critical path — diffusion dominates; only a future tick's encoder is hideable +Tick critical path: VIT → encoder (1×) → K× decoder (10×). The critical phase is the K-step diffusion (10× decoder passes + 10× per-step D2H+host recurrence syncs). The only non-critical, hideable phase is a *future* tick's VIT+encoder — i.e. cross-tick pipelining (§5.2), not intra-tick stage splitting. Within a single tick, encoder→diffusion is already serial and non-redundant, so splitting it adds a boundary with zero overlap benefit. + +### 5.5 Honest counterargument (strongest case AGAINST splitting) +1. **Sequential data dependency in diffusion**: step i+1's input `cross.action` is step i's output (`llama-graph.cpp:133` uploads `cross->action.data()`; `llama-context.cpp:1264` updates it). The K steps are a strict serial chain. No intra-tick parallelism across K. +2. **No incremental action emission**: `action_ready` set once after all K steps (`llama-context.cpp:1389`); intermediate latents are not valid actions. Zero time-to-first-action win. +3. **Encoder already amortized once per tick**: not inside the K-loop. Splitting encoder/diffusion into two FlashRT stages cannot overlap them within a tick — diffusion depends on encoder KV. +4. **The whole-graph path already amortizes graph build in principle** (`llm_graph_result::can_reuse`, `llama-graph.cpp:813`), BUT `llm_graph_input_action::can_reuse` returns `false` (`llama-graph.cpp:137-139`), forcing a rebuild every K step today. **This is a Jetson-PI-side inefficiency to fix, NOT a reason to add FlashRT stage boundaries.** Adding stages does not fix it; fixing `can_reuse` (or pinning the action input as a `set_input` rather than a rebuild trigger) does. +5. **Stage-boundary overhead + ABI complexity for no win**: each new FlashRT stage = new ports, new staged-buffer contracts, new sync points, new error paths, all constrained by CLAUDE.md's no-GGML-types-in-public-ABI rule. The single `infer` callback stage is the minimal correct surface. Splitting adds surface area proportional to stage count with zero latency return unless cross-tick pipelining is ALSO built — and that pipelining lives inside the provider, not in FlashRT's stage graph. +6. **The per-step D2H+host recurrence** (`llama-context.cpp:1254,1263`) is the real serialization bug, independent of staging. Fix it (move recurrence on-device, keep `cross.action` resident) before considering any stage split. +7. **Multi-callback middle-ground considered and rejected**: one could declare multiple callback stages sharing provider-internal state via the engine `self` pointer (no observable inter-stage port needed, no new ABI) — e.g. `encode_images` → `prefill` → `diffusion` → `decode` as separate `run_stage` calls on one engine. This IS expressible today. But it creates no overlap (the intra-tick data dependencies above still force strict serial execution), and the only thing it buys is host-side insertion points between phases (telemetry, early abort). Those insertions carry no latency win and add stage-boundary dispatch overhead. It is therefore cosmetic decomposition under current conditions, not a scheduling win. + +**Splitting would NOT help when**: batch=1 (the Jetson norm), single-robot single-tick latency is the objective, cross-tick pipelining is not built, and the per-step D2H sync remains. These are exactly the current conditions. Under them, finer stages are cosmetic decomposition. + +--- + +## 6. The four gates assessed individually + +CLAUDE.md Phase 5: "拆分前必须明确调度收益、buffer contract、同步语义和 profiling 证据." + +### Gate 1 — Scheduling benefit: **FAIL** +No intra-tick benefit. K diffusion steps are a strict serial chain. Encoder already runs once outside the loop. The only theoretical win (cross-tick encoder/diffusion pipelining) requires provider-internal work (separate streams, double-buffered KV, on-device recurrence) that is not a FlashRT stage split. Action streaming is a hard NO (atomic chunk emission). Batching is a NO on a single Jetson. + +### Gate 2 — Buffer contract: **FAIL (gated on Jetson-PI + Phase 6)** +The B→C boundary passes n_layer cross-KV tensors currently forced through a device→host→device round-trip with an explicit TODO hack (`llama-graph.h:56-60`). Any split placing B and C in different memory domains must specify ownership/lifetime of these tensors; today the host copy is the de-facto contract. A|B and C|D/E are trivial host float buffers and do not block. But the harder problem is on the FlashRT side: the provider-owned ABI path forbids inter-stage buffers and SWAP windows (`model_runtime.cpp:203-217`), so observable inter-stage contracts cannot be expressed without Phase 6 memory-domain work. + +### Gate 3 — Sync semantics: **FAIL** +The per-step `ggml_backend_tensor_get_async` D2H of the action tensor (`llama-context.cpp:1254`) followed by the host-side Euler update (`:1263-1264`) is itself a serialization barrier that must be eliminated before ANY overlap — intra- or cross-tick — can materialize. This is a Jetson-PI fix, independent of staging. On the FlashRT side, the only readiness signal is `actions_buf.size()==need_elems` (`jetson_pi_engine.cpp:285`); invalidation is the blunt `clear()` on any `set_input` (`:157`). A multi-stage split would need per-stage readiness flags and stage-aware invalidation, neither of which exists, and each new sync point can erase the scheduling benefit if the stage is short. + +### Gate 4 — Profiling evidence: **NOT YET COLLECTED — backfill plan in §8** +The code has printf instrumentation (Vit/Encode/Decode at `mtmd-helper.cpp:579,781-783`) but the per-step `step_decode_ms` computed at `llama-context.cpp:1270` is **NOT printed** (verified: only `printf("inference_time: %d over\n",i)` at `:1247` and a separator at `:1248`). No encoder-vs-diffusion ratio is reported. The profiling-evidence requirement is unmet; the plan in §8 is what the main process executes to backfill it. **The structural analysis is strong enough to recommend no-go now; profiling is to confirm and to set the numeric bar for any future revisit, not to search for a justification.** + +--- + +## 7. What would change the answer (conditions for revisiting) + +A real, non-cosmetic scheduling benefit exists **only if ALL of these hold**: + +- **(a) On-device recurrence**: the per-step D2H action copy + host Euler update (`llama-context.cpp:1254,1263`) is eliminated so `cross.action` stays GPU-resident and the diffusion loop has no per-step host sync. (Jetson-PI fix.) +- **(b) Graph reuse across the 10 denoise steps**: `llm_graph_input_state/action/sinusoidal_embedding/attn_no_cache_ae::can_reuse` are made to return `true` under proper guards (`llama-graph.cpp:109,137,192`; `llama-graph.h:100-105`), so the decoder graph is built once and replayed 9×. (Jetson-PI fix. This is the highest-value, lowest-risk optimization and should be measured FIRST, before any stage work.) +- **(c) Cross-tick double-buffered encoder KV**: tick N+1's VIT+encoder runs concurrently with tick N's K diffusion on separate streams. (Provider-internal pipeline.) +- **(d) Profiling shows encoder_time is a non-trivial hideable fraction of K×step_time** (i.e. the encoder is actually hideable rather than negligible). (Backfilled by §8.) +- **(e) Phase 6 memory-domain contract exists** so provider-owned ports can carry SWAP windows, OR new append-only ABI concepts (intermediate port direction, per-step `get_output`, stage-ready flag) are added — otherwise finer stages can only ever be provider-internal sub-steps hidden behind one `infer` stage, which is the current shape. + +Even when (a)–(e) all hold, the correct implementation is a **provider-internal pipeline exposed to FlashRT as ONE callback stage** — not a FlashRT stage decomposition. The revisit question is therefore not "split into FlashRT stages" but "build the provider-internal cross-tick pipeline and re-measure." + +**Recommended next actions (in priority order, all inside Jetson-PI, none a FlashRT stage split):** +1. Enable graph reuse across the 10 denoise steps (fix `can_reuse`). +2. Eliminate the cross-KV host round-trip (`llama-context.cpp:1069` D2H / `llama-graph.cpp:455` H2D feed-back) or at minimum measure it. +3. Move the per-step Euler recurrence on-device. +4. Collect the per-phase/per-step profile (§8) and re-evaluate this gate only if, after 1–3, the encoder (B) or a denoise step (C-inner) still dominates AND a cross-stage pipeline could hide it. + +--- + +## 8. Profiling plan (for the main process to execute and backfill as profiling evidence) + +The main process should add high-resolution `std::chrono` timestamp pairs (or `ggml_time_us()`) at the sites below, run ≥5 LIBERO ticks on the target Jetson with the Pi0-2.8B-F16 GGUF, and report per-tick and per-step wall times. All sites are in the Jetson-PI repo unless noted; the FlashRT-side sites confirm the stage boundary is negligible today. + +### 8.1 Per-phase wall time (the critical-path attribution) + +| Site | file:line | What to timestamp | +|---|---|---| +| Phase A start / end | `Jetson-PI/tools/mtmd/mtmd-helper.cpp:574` (`t_vit_start`) / `:576` (`t_vit_end`) | Per-view VIT. Already prints "Vit took" at `:579` — aggregate per tick (sum over views). | +| Phase B start / end | `Jetson-PI/tools/mtmd/mtmd-helper.cpp:765` (`t_encode_start`) / `:771` (`t_encode_end`) | Encoder (llm_build_pi0). Already prints "Encode took" at `:781`. | +| Phase C (whole loop) start / end | `Jetson-PI/tools/mtmd/mtmd-helper.cpp:774` (`t_decode_start`) / `:778` (`t_decode_end`) | All 10 denoise steps. Already prints "Decode took" at `:782`. | +| Phase C end (authoritative) | `Jetson-PI/src/llama-context.cpp:1388-1389` (`action_ready=true`) | Pair with loop start at `:1241` for total diffusion time independent of the mtmd wrapper. | +| Phase E | `Jetson-PI/src/jetson_pi_pi0.cpp:328` (call) / `:342` (return) | Confirm readout is negligible (pure memcpy at `llama-context.cpp:804`). | + +### 8.2 Per-denoise-step wall time (THE critical measurement) + +| Site | file:line | What to timestamp | +|---|---|---| +| Step start | `Jetson-PI/src/llama-context.cpp:1243` (`step_decode_start`) | BEGIN of denoise step i. | +| Step end | `Jetson-PI/src/llama-context.cpp:1268` (`step_decode_end`) | END of denoise step i. | +| Step duration | `Jetson-PI/src/llama-context.cpp:1270` (`step_decode_ms`) | **Currently computed but NOT printed.** ADD a `printf("step %d decode: %.3f ms\n", i, step_decode_ms);` here. This is the single most important missing print. | + +**Measurement that confirms the critical-path hypothesis:** collect 10 `step_decode_ms` samples per tick across ≥5 ticks. Expectation: roughly constant across steps. If `sum(step_decode_ms) ≈ Decode took`, the diffusion loop is confirmed as the critical path and the per-step cost is attributable. + +### 8.3 Graph-build vs graph-compute split (the reuse-fix justification) + +Inside `process_ubatch` (`Jetson-PI/src/llama-context.cpp`), the rebuild branch is `:864-887`. To separate graph-build from graph-compute: + +| Site | file:line | What to timestamp | +|---|---|---| +| After `model.build_graph` | `Jetson-PI/src/llama-context.cpp:872` (after `gf = model.build_graph(gparams)`) | Graph construction time. | +| After `ggml_backend_sched_alloc_graph` | `Jetson-PI/src/llama-context.cpp:877` (after alloc) | Sched-alloc time. | +| After `graph_compute` | `Jetson-PI/src/llama-context.cpp:898` (after `graph_compute`) | Pure compute time. | +| Reuse branch taken? | `Jetson-PI/src/llama-context.cpp:860` (`can_reuse` check) | Log whether the reuse or rebuild branch fired per step. Expectation today: rebuild 10× (because `can_reuse` returns false for state/action/sinusoidal/attn_no_cache_ae). | + +**Measurement that confirms the reuse-fix lever:** if `(build_graph + sched_alloc)` at `:872-877` is a non-trivial fraction of per-step wall, the graph-reuse fix (making the four `can_reuse` paths return `true` under guards) is the real win and should be tried before any stage work. This refutes stage-splitting as the lever and routes the optimization to Jetson-PI. + +### 8.4 The per-step D2H serialization barrier (the cross-tick-overlap blocker) + +| Site | file:line | What to timestamp | +|---|---|---| +| D2H of `res->action` (velocity) | `Jetson-PI/src/llama-context.cpp:1254-1260` (`ggml_backend_tensor_get_async`) | Per-step device→host copy of `[action_dim*(action_steps+1)]` = `[32*51]` F32. Measure host-side wall (the `_get_async` name is misleading — it is consumed immediately by the Euler update, so it is effectively synchronous). | +| Host Euler update | `Jetson-PI/src/llama-context.cpp:1263-1264` (`cross.action[j] -= ...`) | Host recurrence. Usually cheap but verify it is not the bottleneck between steps. | +| H2D feed-back of `cross.action` next step | `Jetson-PI/src/llama-graph.cpp:129-134` (`llm_graph_input_action::set_input`, `ggml_backend_tensor_set`) | Per-step host→device re-upload of the latent `[50,32]`. | + +**Measurement that confirms the sync barrier:** if `D2H(:1254) + H2D(:129) + host Euler(:1263)` per step is a measurable fraction of `step_decode_ms`, on-device recurrence is justified and is a prerequisite for any cross-tick overlap. + +### 8.5 The cross-KV host round-trip (the B→C buffer-contract cost) + +| Site | file:line | What to timestamp | +|---|---|---| +| Encoder GPU sync before KV extract | `Jetson-PI/src/llama-context.cpp:1057` (`synchronize()` inside encode, Pi0 branch) | Boundary between encoder GPU compute and host-side KV extraction. | +| Per-layer cross-KV D2H | `Jetson-PI/src/llama-context.cpp:1066-1071` (loop; `ggml_backend_tensor_get_async` at `:1069`) | Aggregate over n_layer=27 layers. Log total bytes (`ggml_nbytes(t_encoded_kv[i])` summed) and total ms. This is the B→C contract materialization cost. | +| H2D feed-back per denoise step | `Jetson-PI/src/llama-graph.cpp:455` (`ggml_backend_tensor_set` in `llm_graph_input_cross_kv_pi0::set_input`) | Per-step host→device feed of cross-KV. Note: cross-KV is constant across the 10 steps, so this is repeated work that graph reuse could eliminate. | + +**Measurement that confirms the round-trip cost:** aggregate D2H bytes at `:1069` × 27 layers + aggregate H2D at `:455` × 10 steps. If non-trivial, eliminating the round-trip (keep cross-KV device-resident) is a Jetson-PI fix and a prerequisite for any zero-copy stage design. + +### 8.6 FlashRT-side stage-boundary overhead (confirm splitting only adds overhead) + +| Site | file:line | What to timestamp | +|---|---|---| +| Stage dispatch | `FlashRT/cpp/providers/llama_cpp/src/pi0_runtime.cpp:66` (`owner->engine.run_infer(...)`) | Whole-stage wall. | +| Monolithic infer | `FlashRT/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp:242-247` (`jetson_pi_pi0_infer`) | The block a split would break. | +| Output readback | `FlashRT/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp:264-290` (memcpy at `:289`) | Provider→host output copy. | +| Input staging (image swizzle) | `FlashRT/cpp/providers/llama_cpp/src/jetson_pi_engine.cpp:170-193` (`view_to_rgb` loop) | Host-side pixel swizzle a finer split would isolate into `encode_images` staging. | + +**Measurement that confirms the stage boundary is negligible today:** the gap between `pi0_runtime.cpp:66` dispatch and `jetson_pi_engine.cpp:242` infer should be sub-millisecond, i.e. splitting adds boundary overhead with no offsetting overlap. + +### 8.7 Decision-relevant aggregations to report + +1. `encoder_time / (K × step_time)` ratio — if `encoder_time << K×step_time`, cross-tick pipelining has bounded value and the only theoretical stage-split win is small. If `encoder_time` is a meaningful fraction, revisit condition (d) is plausible. +2. `graph_build_time / step_time` per step — if non-trivial, the reuse fix (§7 action 2) is the lever, not a stage split. +3. `cross_kv_D2H + cross_kv_H2D_per_step` — sizes the B→C round-trip; if non-trivial, on-device cross-KV is a prerequisite. +4. `per_step_D2H_action + host_Euler + H2D_action` per step — sizes the sync barrier; if non-trivial, on-device recurrence is a prerequisite. +5. `sum(step_decode_ms) vs Decode took` — confirms the loop is the critical path and steps are equal-cost. + +**Hypothesis to confirm or refute:** "Under current conditions, the diffusion K-loop dominates the tick, the dominant intra-loop inefficiency is graph rebuild + per-step D2H (both fixable inside Jetson-PI without a FlashRT stage split), and the encoder is not hideable in a useful intra-tick way." If the profile confirms this, the no-go stands. If it surprisingly shows the encoder is a large hideable fraction AND on-device recurrence + double-buffered KV are in place, only then revisit — and even then as a provider-internal pipeline behind one `infer` stage. + + +--- + +## 9. Phase 5 decision (recorded 2026-07-07) + +**Decision: DEFER Phase 5 (no-go under current conditions).** + +The four-gate test fails on scheduling benefit, buffer contract, and sync +semantics; profiling evidence is not yet collected on the Jetson target (the +structural case is strong enough to decide now). Phase 5 is therefore closed +as an evaluation/decision deliverable — no FlashRT stage split is implemented. + +### Routing of the real optimization work + +The latency wins the analysis identified all live **inside Jetson-PI**, behind +the existing single `infer` callback stage — they are NOT FlashRT stage +decompositions. Priority order (Jetson-PI-internal, future work, not part of +this migration phase): + +1. **Graph reuse across the 10 denoise steps** — make + `llm_graph_input_state`/`_action`/`_sinusoidal_embedding`/ + `_attn_no_cache_ae::can_reuse` return `true` under guards + (`Jetson-PI/src/llama-graph.cpp:109,137,192`; `llama-graph.h:100-105`) so + the decoder graph builds once and replays 9x. Highest value, lowest risk. +2. **On-device recurrence** — eliminate the per-step D2H of `res->action` + + host Euler update (`llama-context.cpp:1254,1263-1264`) so `cross.action` + stays GPU-resident. Prerequisite for any cross-tick overlap. +3. **Cross-KV round-trip** — eliminate or measure the device->host->device + cross-KV hop (`llama-context.cpp:1069` D2H / `llama-graph.cpp:455` H2D per + step; the `// TODO: tmp` hack at `llama-graph.h:56-60`). +4. **Cross-tick pipelining (only after 1-3)** — double-buffered encoder KV so + tick N+1 VIT+encoder overlaps tick N diffusion on separate streams. Even + then, expose to FlashRT as ONE callback stage. + +### Conditions for revisiting Phase 5 + +Revisit only if ALL hold (from section 7): (a) on-device recurrence, (b) graph +reuse across denoise steps, (c) cross-tick double-buffered encoder KV, (d) +profiling shows encoder_time is a non-trivial hideable fraction of +K*step_time, (e) Phase 6 memory-domain contract exists (or new append-only +ABI: intermediate port direction, per-step `get_output`, stage-ready flag). +Even then the correct implementation is a provider-internal pipeline behind one +`infer` stage — not a FlashRT stage decomposition. + +### Baseline sanity check (CPU, this host, 2026-07-07) + +A baseline Pi0 run on this host (CPU, pi0_base F16 + pi0_base mmproj, LIBERO +fixture, action_steps=10/action_dim=32) confirmed the provider runs end-to-end +and surfaced the coarse phase timings the code already prints. The two-view VIT +(Phase A) measured ~6.8s and ~6.7s per view on CPU — dominated by the lack of +GPU and by per-step graph rebuild, consistent with the structural finding that +the real levers are Jetson-PI-internal (graph reuse, on-device recurrence), not +a FlashRT stage split. The full per-step profile (section 8) is to be collected +on the Jetson target where the deployment runs; on this CPU host the absolute +numbers are not deployment-representative, so they are not recorded here as the +profiling-evidence gate. + +### FlashRT ABI gaps that would block finer stages (Phase 6 territory) + +- No intermediate (bidirectional) port direction: frt_rt_port_direction is a strict 2-value enum (FRT_RT_PORT_IN=0 / FRT_RT_PORT_OUT=1), not a bitmask (FlashRT/runtime/include/flashrt/model_runtime.h:101). A port cannot be stage-N-out AND stage-(N+1)-in, so inter-stage buffers (image embeddings, encoder KV, diffusion latent) cannot be expressed as observable ports. +- No partial / per-step / streaming output: get_output(self, port, out, capacity, written, stream) is a full-snapshot read of one port by index (FlashRT/runtime/include/flashrt/model_runtime.h:201-203/222-224). No step-k/K parameter, no per-step ready flag, no streaming token. Per-denoise-step or partial action-chunk progress cannot be surfaced. +- Provider-owned path forbids inter-stage buffers and zero-copy windows: the builder rejects all streams/graphs/buffers/regions on provider-owned runtimes (FlashRT/runtime/src/model_runtime.cpp:203-208) and rejects any SWAP window / offset / bytes on ports (model_runtime.cpp:209-217). Hot intermediate tensors (diffusion noise, encoder KV) cannot get a fast lane. The docs explicitly defer this to Phase 6 (docs/model_runtime_api.md:127-129). +- No stage-ready / stage-aware invalidation signal: the only readiness signal is the whole-actions size check (jetson_pi_engine.cpp:285); invalidation is the blunt actions_buf.clear() on any set_input (jetson_pi_engine.cpp:157). A multi-stage split needs per-stage readiness flags and stage-aware invalidation, neither present. +- Engine vtable is monolithic: frt_llama_cpp_engine_v1 exposes only set_input / run_infer / get_output (FlashRT/cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/c_api.h:96-115); run_infer is a single opaque verb. Backing N stages needs either a generic run_stage(stage_idx) engine verb or per-sub-stage verbs, and Jetson-PI must expose sub-phase C entry points (jetson_pi_pi0_encode_images / _prefill_context / _diffusion_step / _decode_actions) to replace the single jetson_pi_pi0_infer (jetson_pi_engine.cpp:242-247). +- No shape-bucket prepare for sub-stages: provider returns -3 unsupported_prepare (pi0_runtime.cpp:19-23). If a split introduced variable sub-stage shapes (e.g. per-step), there is no capture-on-miss path. diff --git a/docs/phase6_backend_vtable_eval.md b/docs/phase6_backend_vtable_eval.md new file mode 100644 index 00000000..19235246 --- /dev/null +++ b/docs/phase6_backend_vtable_eval.md @@ -0,0 +1,271 @@ +# Phase 6 Evaluation — Backend-neutral exec/backend vtable + memory-domain contract + +## 1. Executive summary + recommendation + +**Recommendation: `go-minimal-contract` — build a small, append-only memory-domain contract; do NOT build a full backend-neutral exec/backend vtable in this phase.** + +Four independent lenses (exec/ coupling depth, provider/zero-copy path needs, mllm reference patterns, ABI/memory-domain constraints) converge on the same conclusion. The decisive findings, all verified against the source: + +1. **exec/ already has a backend-neutral vtable SHAPE.** `exec/backend/backend.h` is a flat namespace `frt::be::*` of 13 free functions (malloc/free, stream/event, memcpy_dtod_async, capture_begin/end, graph_launch, graph_exec_destroy). The core ABI in `exec/src/*.cpp` links ONLY these — never a CUDA symbol, zero `virtual`. ALL CUDA Graph API calls are confined to ONE file, `exec/backend/cuda/cuda_backend.cpp` (107 lines). Adding a second exec backend is one new TU + a CMake selection line — `exec/src/` and `exec/include/` change ZERO files. A "full vtable" rewrite of exec/ would be solving a problem that the existing free-function seam already solves at link time. + +2. **The provider-owned callback-stage bypass ALREADY EXISTS and is in production use.** `FRT_RT_STAGE_CALLBACK` (`model_runtime.h:109-112`), `frt_runtime_stage_desc_v2` (`:172-179`), the `run_stage` verb (`:232`), and `frt_runtime_builder_create_provider_owned` (`:317`) let a provider-owned runtime host a non-CUDA-Graph stage with NO `frt_ctx`/`frt_graph`/`frt_buffer`. The Jetson-PI Pi0 provider (`cpp/providers/llama_cpp/src/pi0_runtime.cpp`) does NOT include `flashrt/exec.h` (verified by grep: zero includes across `cpp/providers/`) and never touches exec/. CUDA-Graph vs provider-owned coexistence is already achieved by leaving the v1 `stages` array empty when callback stages are present (`model_runtime.cpp:232-235`). CLAUDE.md's "provider-owned stage coexistence with CUDA Graph" constraint is already met — Phase 6 does not need to unify the two paths under one backend abstraction. + +3. **GGML is self-contained; FlashRT does not execute GGML ops.** A full exec/backend vtable (alloc/free/stream/event/executable create/replay) for the GGML path would have NO consumer — FlashRT would be mediating execution it does not perform. This directly contradicts CLAUDE.md's "Jetson-PI as model provider, not bare exec/ hardware backend" and "mllm is reference architecture, not transplant material." + +4. **The ONLY genuinely missing piece for the stated Phase 6 goal** ("FlashRT 与 GGML 的 zero-copy 或 shared-buffer 路径") is a memory-domain CONTRACT: an opaque FlashRT-owned token + ownership/lifetime/location/copy/sync rules, attachable to a provider-owned port in place of today's forbidden SWAP window. The provider-owned builder today rejects all buffers/SWAP on provider-owned ports (`model_runtime.cpp:203-217`), and the docs explicitly defer this to Phase 6 (`docs/model_runtime_api.md:127-129`: "The builder rejects raw SWAP windows on this path until FlashRT has an explicit memory-domain contract for cross-provider buffers"). This is Phase 5's Gap 2 — the single gap that is unambiguously Phase 6's job. + +CLAUDE.md Phase 6 text reads "设计更完整的 backend-neutral exec/backend vtable. **可能**包含... 在 contract 明确**后**再考虑 FlashRT 与 GGML 的 zero-copy 或 shared-buffer 路径." The hedge "可能" (possibly/may) and the "after the contract is clear" ordering are load-bearing: the contract is the prerequisite, and it is the deliverable that unblocks the stated zero-copy goal. The minimal honest deliverable that fulfills Phase 6's intent, unblocks Phase 5's future revisit, preserves Phase 1-4, and leaks no GGML types is the **memory-domain contract** — not a full vtable. The exec/ backend seam is already backend-neutral in shape and CUDA-isolated to one TU; documenting its split (shared memory/stream/event primitives vs CUDA-only graph-capture) is in-scope as clarification, but rewriting it into a registered function-pointer vtable is over-design with no consumer. + +## 2. exec/ current-state summary + CUDA Graph coupling (lens 1) + +exec/ defines exactly five abstractions, all opaque C-ABI handles backed by plain C++ structs in `exec/src/internal.h`: + +- **`frt_ctx`** (`internal.h:57-75`) — stream/event pool + arena owner. Owns all buffers; frees them at `frt_ctx_destroy`. +- **`frt_buffer`** (`internal.h:14-20`) — `{ctx, name, dptr, bytes, owned}`. Pure device-pointer wrapper. `frt_buffer_dptr` is the ONLY way the runtime layer touches device memory. +- **`frt_graph`** (`internal.h:32-43`) — ShapeKey→graph-exec variant table with LRU. The executable is a raw `void* exec` (`internal.h:27-30`), NOT a virtual interface. +- **`frt_plan`** (`internal.h:45-55`) — dumb DAG of `(graph, key, stream_id)` nodes. Graph-only by construction (`plan.cpp:75` unconditionally calls `frt_graph_replay`). +- **`frt_event`** (`internal.h:22-25`) — cross-stream sync, wraps `void* handle`. + +There is NO separate allocator abstraction — allocation is two free functions `frt::be::malloc/free`. There is NO stage abstraction in exec/ itself; the exec-level replay primitive is `frt_graph_replay(graph, key, stream_id)` (`graph.cpp:105-113`). "Stages" live one layer up in the model-runtime ABI. + +**CUDA Graph coupling is confined to ONE TU.** All `cudaStreamBeginCapture`/`cudaStreamEndCapture`/`cudaGraphInstantiate`/`cudaGraphLaunch`/`cudaGraphExecDestroy` calls are in `exec/backend/cuda/cuda_backend.cpp` (`:78-103`). The graph object is a raw `cudaGraphExec_t` cast to `void*` — NOT wrapped in a FlashRT type. The core ABI (`src/*.cpp`) only ever holds it as `void* exec` and passes it to `frt::be::*`. This is the key design property: the core never sees CUDA types, only the backend TU does. + +**Where stage replay assumes CUDA Graph** — exactly two sites, BOTH in the model-runtime layer (NOT in exec/): `cpp/models/pi05/src/model_runtime.cpp:231` (`step()` → `frt_graph_replay` for every v1 stage) and `cpp/models/pi05/src/runtime.cpp:203` (`default_replay` → `frt_graph_replay`). These are used ONLY by the pi05 CUDA-Graph provider. The llama_cpp provider supplies its own `step`/`run_stage` verbs and never enters that loop. The CLAUDE.md warning "stage semantics are CUDA-Graph-coupled" is therefore already half-obsolete: the v2 callback-stage seam decoupled it. + +**Blast radius of a "vtable over exec/":** the seam already exists as free functions. Backend selection is LINK-TIME (one TU per build; `exec/CMakeLists.txt:24-34` hardcodes `cuda/cuda_backend.cpp`). To add a second exec backend you change ZERO files in `exec/src/` or `exec/include/`; you add one TU and one CMake line. The ONLY structural change the current design does NOT already accommodate is RUNTIME multi-backend selection in one process (linking both a CUDA and a CPU exec backend) — that would require promoting `frt::be::*` from free functions to a registered function-pointer table. No Phase 6 requirement demands this. Promoting it now is over-design. + +## 3. Provider path needs + zero-copy export path analysis (lens 2) + +### Jetson-PI provider TODAY (pure host-staged callback execution) + +The provider never touches exec/. Verified: `cpp/providers/llama_cpp/src/pi0_runtime.cpp` includes only `flashrt/providers/llama_cpp/c_api.h` and `flashrt/model_runtime.h`. Every input is memcpy'd into provider-owned host buffers; every output is memcpy'd out. + +- **Inputs (host memcpy IN):** images swizzled into host `std::vector rgb_scratch`; prompt into `std::string`; state into `std::vector`. +- **Infer:** allocates a local `std::vector actions`, calls one opaque `jetson_pi_pi0_infer(...)` that writes into `actions.data()`, then assigns into `e->actions_buf`. GGML owns all device memory internally; FlashRT never sees a device pointer. +- **Output (host memcpy OUT):** `std::memcpy(out, e->actions_buf.data(), need_bytes)`. +- The `stream` parameter is dropped at every engine verb. All four ports are STAGED with null buffer / 0 offset / 0 bytes (`pi0_runtime.cpp:142-157`). + +**What the provider needs from a backend abstraction: nothing.** Data volumes are batch=1 robot control (~300KB images, 32B state, 280B actions, 6.4KB diffusion noise). Host memcpy is negligible against multi-second (CPU) / hundreds-of-ms (Jetson GPU) diffusion compute. GGML's own internal per-step D2H recurrence (`llama-context.cpp:1254`) is a bigger serialization cost than any FlashRT-boundary copy would ever be. Zero-copy at the FlashRT boundary would not fix the in-provider serialization and is premature. + +### CUDA-graph export zero-copy path (the de-facto same-backend memory-domain contract) + +This path ALREADY works and is used end-to-end by pi05: + +- **Buffer role is a bitmask** `FRT_RT_ROLE_INPUT|OUTPUT|STATE|SCRATCH` (`runtime.h:45-49`), with the comment "a buffer can be both input and output (e.g. an in-place diffusion noise/action buffer)." +- **Port SWAP window:** `frt_runtime_port_desc {buffer, offset, bytes}` (`model_runtime.h:154-156`). A non-null buffer turns the port into a raw device-window the host writes/reads directly. `frt_buffer_dptr(buffer)` (`exec.h:102`) returns the "stable device pointer"; `frt_buffer_wrap(ctx, name, void* dptr, bytes)` (`exec.h:101`) wraps an externally-owned device pointer (e.g. a torch tensor's `data_ptr`). +- **`frt_graph_bind`:** "Two graphs sharing one buffer on matching ports = zero-copy hand-off (this is the entire multi-subgraph / multi-model wiring mechanism)" (`exec.h:146`). +- **pi05 uses it:** `tensor_from_port` extracts `frt_buffer_dptr(p.buffer)` and tags `MemoryPlace::kDevice` (`model_runtime.cpp:74-96`), feeds it as `image_input_override`/`action_output_override` (`:414-415`); the `noise` port is `FRT_RT_PORT_SWAP` (`:335`). + +**Is it abstract or CUDA-specific?** The HANDLE (`frt_buffer`) is opaque, but the SEMANTICS are CUDA-specific: `frt_buffer_dptr` returns a "stable device pointer," `frt_buffer_copy` is "device-to-device copy on a stream," streams/events are CUDA-shaped (`native_handle` is "raw backend stream, e.g. cudaStream_t" at `runtime.h:72-75`). It is NOT a multi-backend memory-domain abstraction; it is a CUDA-graph wiring mechanism that happens to use opaque handles. This is the canonical same-backend (CUDA) memory-domain contract and should be documented as such — it does not need to be generalized. + +### Decision: do NOT build a device-buffer + sync abstraction for the provider-owned path + +The CUDA export path already has one (CUDA-specific by construction); GGML must not be forced through it (would violate the no-GGML-types rule and the no-zero-copy-without-contract rule). The minimal, conservative Phase 6 deliverable is a memory-domain CONTRACT: a port annotation that lets a provider-owned port OPTIONALLY carry a buffer (opaque token) with honest ownership/lifetime/location/copy/sync, in place of today's forbidden SWAP window. + +## 4. mllm reference patterns worth borrowing (lens 3) + +> NOTE: the `mllm/` reference repo was deleted from the workspace on 2026-07-07 after this distillation. The file:line citations below (`X2XOp.cpp:13-15` etc.) are historical references to that repo, retained as provenance; reclone `git@github.com:UbiquitousLearning/mllm.git` if you need to re-open them. The distilled patterns themselves are self-contained in this section. + +mllm is reference architecture, NOT transplant material (CLAUDE.md). Worth borrowing for Phase 6: + +1. **Device-tagged storage with tag-driven alloc/free routing** (`Storage.device_` + `MemoryManager::alloc` lookup). This is the minimal viable "memory domain" primitive: a buffer carries a domain tag, and a registry maps tag→allocator. FlashRT's analog: a `frt_memory_token` carries a `location_kind` tag; copy/sync dispatch by that tag. +2. **The three-collaborator split** (Backend / Allocator / Dispatcher as distinct objects). The analog for FlashRT: the `frt_memory_token` (memory) is distinct from the `executable` (create/replay) is distinct from the `provider` (model load/infer). mllm got this boundary right; Phase 6 should keep the token's verb set narrow (memory only), not let it grow into an execution vtable. +3. **Explicit `initXxxBackend()` registration into a singleton** rather than macro self-registration. For FlashRT with ~2 providers, manual registration is simpler and debuggable. Append-only and predictable. (NOT needed in Phase 6 — there is no second exec backend to register — but the pattern informs any future registry.) +4. **Copy-as-explicit-routed-op / honesty about cross-device transfer.** mllm's CPU X2X literally warns "should be implemented in device backends" (`X2XOp.cpp:13-15`). The right instinct: don't pretend zero-copy; route copies through the token's `copy_to_host`/`copy_from_host` verbs with explicit location_kind. + +NOT worth borrowing (mismatch with FlashRT's low-latency CUDA-Graph focus): + +- mllm's per-OpType op-factory table (60+ factories) — FlashRT providers execute whole graphs, not per-op. +- mllm's per-backend `Dispatcher` + `static_thread_pool` + stdexec senders — adds latency to a synchronous CUDA-Graph-replay hot path. +- mllm's `dlopen` `PluginOpPackageDescriptor` plugin system — FlashRT has two compile-time providers, not a plugin ecosystem. +- mllm's `BuddyMemPool` / large-tensor-threshold memory manager — FlashRT should use CUDA's native graph memory pools, not a framework buddy allocator. +- Any stream/event/executable-create/replay generality beyond what the CUDA-Graph provider concretely needs — mllm proves these are NOT naturally backend-neutral (each backend keeps them private). + +Crucially: **mllm has NO stream/event abstraction, NO executable create/replay, NO host-callback, and NO memory-domain/import-handle concept in its backend-neutral interface.** Each backend manages streams/events privately; cross-device copy is a regular op, not a vtable method; there is no CUDA-Graph capture/replay analogue. This is direct evidence that "backend-neutral vtable for stream/event/executable" is NOT a naturally extractable abstraction — Phase 6 should not attempt it. + +## 5. ABI/memory-domain constraints + the "vtable vs contract" decision (lens 4) + +### ABI versioning story (already established, append-only) + +FlashRT's ABI uses a two-field POD gate on every hand-off struct: `abi_version` + `struct_size`. Consumers gate with `>=` on struct_size (tail-extension probe): + +- `frt_model_runtime_wrap` requires `exp->abi_version == FRT_RUNTIME_ABI_VERSION && exp->struct_size >= sizeof(frt_runtime_export_v1)` (`model_runtime.cpp:308-309`). +- `valid_model_runtime` requires `m->abi_version == FRT_MODEL_RUNTIME_ABI_VERSION && m->struct_size >= sizeof(frt_model_runtime_v1)` (`model_runtime.cpp:384-385`). +- `copy_verbs` reads v1 fields only if `verbs->struct_size >= sizeof(frt_model_runtime_verbs)` (`model_runtime.cpp:40`); `copy_verbs_v2` reads the v2 superset only if `verbs->struct_size >= sizeof(frt_model_runtime_verbs_v2)` (`:59`). Null entries are stub-filled with -3 unsupported (`:47-51, 67-72`). + +This is the established pattern for adding vtable methods without breaking providers: a future `verbs_v3` with appended function pointers is readable by a v3-aware builder via the same `struct_size >=` probe; a v2 producer stays valid because its smaller `struct_size` satisfies the v2 reader. Phase 1-4 evolved strictly append-only — v2 structs alongside v1, never mutating v1, enums "ABI-frozen after v1 (append-only)" (`runtime.h:40`, `model_runtime.h:62`). Phase 6 MUST follow the same discipline. + +### Phase 5 ABI gaps — triage against the memory-domain contract + +The Phase 5 eval lists five gaps. Triaging each: + +- **Gap 2 — provider-owned path forbids buffers and SWAP windows** (`model_runtime.cpp:203-217`; docs defer at `model_runtime_api.md:127-129`). THE Phase 6 core. Without lifting this, provider-owned ports cannot carry any buffer, so there is nothing for a memory-domain contract to annotate. **REQUIRED for Phase 6.** +- **Gap 1 — no intermediate/bidirectional port direction** (`model_runtime.h:101`, strict 2-value enum). Phase 5-revisit enabler for OBSERVABLE finer stages, NOT a memory-domain requirement. SEPARABLE. Phase 5's job. +- **Gap 3 — no partial/per-step/streaming output** (`model_runtime.h:201-203, 222-224`). Phase 5-revisit enabler. SEPARABLE. +- **Gap 4 — no stage-ready/stage-aware invalidation.** Phase 5-revisit enabler. SEPARABLE. +- **Gap 5 — monolithic engine vtable** (`c_api.h:96-115`, provider-internal, NOT FlashRT public ABI). The FlashRT public side ALREADY supports multi-stage via `run_stage(stage, stream)` and `add_callback_stage_v2` (a mixed graph+callback DAG is proven declarable). Gap 5 is a provider-internal split needing Jetson-PI sub-phase C entry points. SEPARABLE. Phase 5's job. + +**Phase 6 minimal set = Gap 2 only.** Gaps 1, 3, 4, 5 are Phase 5-revisit enablers for OBSERVABLE finer stages; they do not block an honest memory-domain contract and should not be bundled into Phase 6 (scope creep with no Phase 6 consumer). + +### The vtable-vs-memory-domain-contract decision + +A **"backend vtable"** = alloc/free/stream/event/executable create/replay generalized to non-CUDA backends — the big abstraction CLAUDE.md Phase 6 lists as a *possibility*. This is the `exec.h` surface generalized. Building it would mean FlashRT gains a backend-registry/dispatcher layer akin to mllm's — which CLAUDE.md explicitly frames as "reference architecture, not transplant material." + +A **"memory-domain contract"** = a small PORT ANNOTATION: an opaque FlashRT-owned token + ownership/lifetime/location/copy/sync rules, attached to a provider-owned port in place of today's forbidden SWAP window. + +**Phase 6 should build the contract, NOT the vtable.** Three reasons: + +1. **GGML is self-contained.** FlashRT does not run GGML ops; a FlashRT backend vtable would have no GGML ops to dispatch. The provider path's only FlashRT-visible need is letting a port OPTIONALLY carry a buffer for honest zero-copy or host readback — a port annotation, not an execution layer. +2. **Coexistence is already solved.** The v2 callback stage + provider-owned builder already lets CUDA-Graph and provider-owned runtimes coexist WITHOUT a shared backend abstraction. Phase 6 does not need to unify them. +3. **The only unblocking need is Gap 2.** Phase 5 explicitly defers the provider-owned SWAP rejection to Phase 6. Lifting it requires a contract that says "this provider-owned port may carry a buffer, here is how it is owned/located/copied/synced" — the memory-domain contract. Nothing in the Phase 5 gap list or the CLAUDE.md Phase 6 description requires FlashRT to gain alloc/free/stream/event/executable verbs for non-CUDA backends in this migration. + +The temptation to build the full backend abstraction should be resisted: it adds surface area (new vtable, new registry, new dispatcher, new error paths) constrained by the no-GGML-types-in-public-ABI rule, with no consumer in this migration, and contradicts CLAUDE.md's "Jetson-PI as model provider, not bare exec/ hardware backend" and "mllm is reference architecture, not transplant material." + +## 6. The memory-domain contract (ownership/lifetime/location/copy/sync; opaque-token vs device-pointer) + +The contract must pin five facts (CLAUDE.md: "不要在没有 memory-domain contract 前假装可以零拷贝共享 buffer"): + +1. **Ownership** — who allocates and who frees. The provider mints the token; FlashRT never allocates provider-owned backing store. The provider supplies a `destroy` verb; FlashRT calls it when the port/export refcount hits zero. +2. **Lifetime** — the token is valid only while the holder retains a reference (the export pattern at `runtime.h:138-145`). A provider-owned buffer's lifetime ties to the port/export retain-release, not to an implicit provider-internal pointer. +3. **Location** — host vs device, expressed as an OPAQUE FlashRT-defined `frt_rt_location_kind` enum (`HOST_VISIBLE`, `DEVICE_LOCAL`, ...), NOT a `ggml_backend_buffer_t` or `cudaMallocAsync` pointer. CLAUDE.md forbids leaking `ggml_tensor`/`ggml_cgraph`/`ggml_backend_t`/`ggml_backend_sched_t` into the public ABI. +4. **Copy semantics** — verbs `copy_to_host(dst, dst_off, src_off, bytes)` and `copy_from_host(src, src_off, dst_off, bytes)`. The existing CUDA `frt_buffer_copy` (`exec.h:110-111`) is the model; the token needs an equivalent implemented by the provider against its own backend. +5. **Sync point** — a `sync` verb (or "host-visible, no sync needed" declared via `location_kind`). Zero-copy becomes an ADVERTISED capability (`location_kind = HOST_VISIBLE`), NOT an assumption FlashRT makes. + +**Opaque-token vs device-pointer: OPAQUE TOKEN.** Exposing a raw device pointer would (a) couple FlashRT to the provider's device-memory manager, (b) risk leaking GGML/CUDA specifics, (c) let a consumer dereference memory it does not own — violating the "no pretend zero-copy" rule. The existing `frt_buffer_dptr` accessor stays INTERNAL to the CUDA exec layer. For provider-owned ports, FlashRT sees only: token (opaque), bytes, location_kind, copy_to_host, copy_from_host, sync, destroy. The provider vouches for the buffer's contents and lifetime through these verbs; FlashRT never dereferences a provider-owned pointer. + +## 7. How this layers over exec/ without rewriting it + +The memory-domain contract lives at the **model-runtime / port layer**, NOT in exec/. Concretely: + +- **exec/ is UNTOUCHED.** No changes to `exec/src/*.cpp` (`context.cpp`, `buffer.cpp`, `graph.cpp`, `plan.cpp`) or `exec/include/flashrt/exec.h`. The CUDA-graph export path (`frt_buffer` + `frt_graph_bind` + SWAP window + INPUT|OUTPUT role bitmask) remains the canonical same-backend (CUDA) memory-domain contract, documented as such. +- **backend.h gets a documentation split** (NOT a rewrite): clarify that the 13 free functions decompose into (a) shared memory/stream/event primitives (malloc/free, stream_*, event_*, memcpy_dtod_async, memset_async — every backend implements) and (b) CUDA-only graph-capture primitives (capture_begin/end, graph_exec_destroy, graph_launch — optional; a backend returning `false` from `capture_begin` simply cannot host graph-capture stages and must use the callback-stage path). GGML has no `cudaGraph_t` analogue; that subset is correctly CUDA-only and stays isolated in `cuda_backend.cpp`. +- **The contract is added at the model-runtime layer:** a new opaque `frt_memory_token` handle + a small `frt_memory_token_verbs` struct (copy_to_host, copy_from_host, sync, destroy) + a `frt_rt_location_kind` enum, plus an append-only extension to the port descriptor so a provider-owned port MAY carry a token in place of the today-forbidden `frt_buffer/offset/bytes` triple. +- **The provider-owned builder rejection (`model_runtime.cpp:209-217`) is relaxed for the token form ONLY.** A provider-owned port may carry a non-null `frt_memory_token` + offset/bytes + location_kind; it STILL rejects a raw `frt_buffer` (which would imply cross-backend device sharing that has no contract). The token is the ONLY allowed buffer form on provider-owned ports. + +## 8. ABI evolution (append-only, versioning, what must not break) + +- **New types are additive:** `frt_memory_token` (opaque typedef), `frt_memory_token_verbs` (new struct with its own `struct_size`), `frt_rt_location_kind` (new enum, values appended after existing enums — "ABI-frozen after v1 (append-only)"). +- **Port descriptor extension is tail-only:** add a `frt_runtime_port_desc_v2` that EMBEDS the v1 layout and APPENDS `{frt_memory_token token; uint64_t offset; uint64_t bytes; uint32_t location_kind; uint32_t reserved}`. A v1 producer's smaller `struct_size` continues to satisfy the v1 reader; a v2-aware builder reads the trailing fields only when `struct_size >= sizeof(frt_runtime_port_desc_v2)` — the identical pattern `copy_verbs`/`copy_verbs_v2` already use. Old consumers and old providers keep working unchanged. +- **Builder API is appended, not mutated:** add `frt_runtime_builder_add_port_v2` (or an overload) accepting the token form; the existing `frt_runtime_builder_add_port` is untouched. The v2 `finish_model_v2` path's provider-owned rejection (`model_runtime.cpp:203-217`) is RELAXED in-place to permit the token form while still rejecting raw `frt_buffer` — this is a behavioral relaxation on an existing v2 path, not a struct-layout change, so v2 binaries are not broken (a v2 producer that supplies no token sees identical behavior). +- **Token verbs follow the stub-fill discipline:** absent verbs are replaced by unsupported stubs returning -3 (`model_runtime.cpp:47-51, 67-72`), so every entry is always callable. +- **What MUST NOT break:** (a) every Phase 1-4 provider binary that fills `frt_model_runtime_verbs`/`verbs_v2` with v1/v2-sized structs; (b) `frt_model_runtime_wrap`'s `abi_version`+`struct_size` gate (`model_runtime.cpp:308-309`); (c) `valid_model_runtime`'s gate (`:384-385`); (d) the CUDA-graph export path's `frt_buffer`/`frt_graph_bind`/SWAP mechanism (byte-identical). The memory-domain contract is a NEW, parallel surface; it does not alter the CUDA export path's contract. + +## 9. In-scope (minimal) vs out-of-scope (deferred) + +See the structured `in_scope` / `out_of_scope` fields. The boundary is: Phase 6 delivers the memory-domain CONTRACT (token + verbs + port extension + builder relaxation + backend.h doc split + smoke test + documentation). It does NOT deliver a full exec/backend vtable, a runtime-selectable backend registry, cross-backend zero-copy implementation, or any of the Phase 5-revisit ABI gaps (1, 3, 4, 5). + +## 10. Implementation plan (ordered) + +See the structured `implementation_plan` field. The order is: (1) document the two existing contracts first (no code risk, clarifies the boundary); (2) add the token types + location_kind enum (purely additive headers); (3) add the port descriptor v2 + builder entry (append-only); (4) relax the provider-owned rejection for the token form only (behavioral, gated); (5) add the token verbs stub-fill + lifetime wiring to retain/release; (6) backend.h doc split + optional `import` declaration (no exec/src changes); (7) smoke test (provider-owned runtime with a host-visible token, round-trip copy, verify v1 consumers unchanged). + +## 11. Risks + mitigations + +See the structured `risks` field. Headline risks: (a) scope creep into a backend vtable — mitigated by explicit out-of-scope list and CLAUDE.md constraints; (b) token verb set too narrow for future needs — mitigated by the struct_size append-only probe (verbs_v2/v3); (c) provider mis-advertises location_kind — mitigated by FlashRT never dereferencing the token (only copy verbs); (d) lifetime mismatch — mitigated by tying destroy to port/export retain-release; (e) Phase 5 revisit needs MORE than the token — acknowledged, Phase 6 satisfies condition (e) only, not (a)-(d). + +## 12. Phase 5 coupling — does this satisfy revisit condition (e)? + +**Yes, this satisfies Phase 5 revisit condition (e) — and ONLY (e).** Phase 5's condition (e) is an OR: "Phase 6 memory-domain contract exists (or new append-only ABI: intermediate port direction, per-step `get_output`, stage-ready flag)." The memory-domain contract is the first disjunct. By delivering the token + port extension + builder relaxation, Phase 6 removes the FlashRT-side ABI blocker (`model_runtime.cpp:209-217`, the Gap 2 deferral) that today prevents provider-owned ports from carrying any buffer. + +**Important honesty caveat:** satisfying (e) does NOT by itself make finer Pi0 stages worth doing. Conditions (a)-(d) are Jetson-PI-INTERNAL: (a) on-device recurrence, (b) graph reuse across the 10 denoise steps, (c) cross-tick double-buffered encoder KV, (d) profiling showing the encoder is a non-trivial hideable fraction. None of those are FlashRT's concern, and none are unblocked by Phase 6. The Phase 5 eval's conclusion stands: even when (a)-(e) ALL hold, the correct implementation is a **provider-internal pipeline exposed to FlashRT as ONE callback stage** — not a FlashRT stage decomposition. Phase 6's memory-domain contract is the necessary-but-not-sufficient FlashRT-side enabler; it lets a FUTURE provider-internal pipeline (should one materialize) optionally expose honest host-visible or device-local buffers through provider-owned ports WITHOUT forcing the provider to fake a CUDA graph or leak GGML types. That is the smallest honest contribution Phase 6 can make to the Phase 5 revisit, and it is exactly the contribution CLAUDE.md Phase 6 asks for ("在 contract 明确后再考虑 zero-copy"). + +--- + +## 13. Phase 6 decision (recorded 2026-07-07) + +**Decision: GO-MINIMAL-CONTRACT.** Build the small, append-only memory-domain +contract; do NOT build a full backend-neutral exec/backend vtable in this phase. + +The four lenses converge: exec/ already has a backend-neutral *shape* (13 free +functions, CUDA isolated to one 107-line TU, zero `virtual`); the provider-owned +callback-stage bypass already exists and is in production use by the llama_cpp +provider (which never includes exec.h); GGML is self-contained so a full +backend vtable would have no consumer; the ONLY genuinely missing piece for the +stated Phase 6 zero-copy goal is the memory-domain CONTRACT (Gap 2, explicitly +deferred to Phase 6 at docs/model_runtime_api.md:127-129). Building a full +backend vtable now would add surface area with no consumer and contradict +CLAUDE.md ("Jetson-PI as model provider, not bare exec/ hardware backend"; +"mllm is reference architecture, not transplant material"). + +### The memory-domain contract (opaque token, NOT raw device pointer) +- Ownership: provider mints `frt_memory_token` + supplies `destroy`; FlashRT + never allocates provider-owned backing store; `destroy` called once when + port/export refcount hits zero. +- Lifetime: token valid only while holder retains a reference (export pattern + runtime.h:138-145). +- Location: opaque `frt_rt_location_kind` enum {HOST_VISIBLE, DEVICE_LOCAL} -- + never a ggml_backend_buffer_t / cudaMallocAsync pointer. +- Copy: `copy_to_host` / `copy_from_host` verbs (model: exec.h frt_buffer_copy). +- Sync: `sync` verb, or "host-visible, no sync" declared via location_kind. +- Zero-copy becomes an ADVERTISED capability (location_kind=HOST_VISIBLE), not + an assumption FlashRT makes. FlashRT NEVER dereferences the token. + +### Two coexisting contracts +- (A) CUDA-export path: opaque `frt_buffer` + `frt_buffer_dptr`/`wrap` + + `frt_graph_bind` + port SWAP window + INPUT|OUTPUT role bitmask. CUDA-specific + semantics, canonical same-backend hand-off. UNTOUCHED by Phase 6. +- (B) Provider-owned token path: the new opaque-token contract above, the only + buffer form permitted on provider-owned ports. + +### ABI evolution (strictly append-only; Phase 1-4 binaries unchanged) +- New enum `frt_rt_location_kind` (values appended; ABI-frozen-after-v1 respected). +- Port descriptor extension via the established `struct_size >=` tail probe + (identical to copy_verbs / copy_verbs_v2 at model_runtime.cpp:40,59). +- Token verbs struct carries its own `struct_size`; absent verbs stub-filled -3. +- Provider-owned builder rejection (model_runtime.cpp:209-217) relaxed for the + token form ONLY; raw `frt_buffer` rejection stays. +- What MUST NOT break: abi_version+struct_size gates (model_runtime.cpp:308-309, + 384-385); every Phase 1-4 provider binary; the CUDA-export path byte-identical. + +### Phase 5 coupling +Satisfies Phase 5 revisit condition (e) -- and ONLY (e). Conditions (a)-(d) are +Jetson-PI-internal (on-device recurrence, graph reuse, double-buffered KV, +profiling). The token is the necessary-but-not-sufficient FlashRT-side enabler; +Phase 5 ABI gaps 1/3/4/5 (intermediate port direction, per-step get_output, +stage-ready flag, engine run_stage-by-index) stay OUT of Phase 6 scope. + +### In-scope deliverables (minimal) +1. Document the two existing contracts (no code risk). +2. Add token types + `frt_rt_location_kind` enum (additive headers). +3. Port descriptor v2 + builder entry (append-only). +4. Relax provider-owned rejection for token form only (gated behavioral). +5. Token verbs stub-fill + lifetime wiring to retain/release. +6. backend.h documentation split (shared primitives vs CUDA-only graph-capture); + declare `frt::be::import` extension point (not implemented for non-CUDA). +7. Smoke test: provider-owned runtime with a host-visible token, round-trip copy, + v1-consumer unchanged, raw frt_buffer still rejected. + +### Explicitly out of scope (over-design / Phase 5-revisit territory) +- Full exec/backend vtable (alloc/free/stream/event/executable) for non-CUDA. +- Runtime-selectable backend registry (link-time selection already suffices). +- Cross-backend zero-copy implementation (batch=1 robot control is premature). +- Phase 5 ABI gaps 1/3/4/5. +- Any exec/src/ struct refactor; mllm op-factory / dispatcher / plugin patterns. + +### Implementation revision note (recorded 2026-07-07, post-review) + +The §8/§13 text names `frt_runtime_port_desc_v2` + a per-port `struct_size` +tail-probe as the token carrier. The shipped implementation chose a **simpler, +parallel-array carrier** instead, recorded here so history does not contradict +the code: + +- Tokens ride a NEW trailing `port_tokens` / `n_port_tokens` array on + `frt_model_runtime_v2`, index-parallel with `ports` + (`n_port_tokens == n_ports` always: a port added via `add_port` gets a + null-handle entry; a port added via `frt_runtime_builder_add_port_token` + gets the minted token). This matches how v2 already carries parallel + `ports`/`stages`/`stages_v2` arrays and keeps `frt_runtime_port_desc` + byte-identical. The `frt_runtime_port_desc_v2` tail-probe form is NOT used + (it was a dead declaration in the first draft; removed on review). +- `destroy` fires once when the **holder refcount** hits zero — not a separate + per-port/per-export counter. On the provider-owned v2 path the holder IS the + token's lifetime owner (there is no v2 wrap/override path, so the holder + refcount is the single relevant count). The §6/§13 "port/export refcount" + wording was wider than the implementation; this note narrows it to "holder + refcount" as actually shipped. +- Append-only is honored by `struct_size >=` / `<` probes in the C ABI + (`frt_model_runtime_wrap`) and the pybind readers (`as_model_v2` uses `<`). + This revision also fixed a pre-existing strict-`!=` pybind gate that the v2 + growth would have made reject newer producers. + +The decision (GO-MINIMAL-CONTRACT) and the vtable-vs-contract reasoning are +unchanged; only the token *carrier* mechanism differs from the §8 sketch. diff --git a/docs/phase7_llm_repeatable_decode_eval.md b/docs/phase7_llm_repeatable_decode_eval.md new file mode 100644 index 00000000..dbe20be0 --- /dev/null +++ b/docs/phase7_llm_repeatable_decode_eval.md @@ -0,0 +1,94 @@ +# Phase 7 Evaluation — LLM Repeatable Decode (prefill / decode-step / logits) + +> Implementation update (2026-07-10): the earlier no-go below is superseded by +> the complete-migration requirement. Jetson-PI LLM and MLLM narrow APIs now +> expose `reset`, `prefill`, repeatable `decode_step`, and `get_logits`; one-shot +> generate is implemented through the same staged path. FlashRT exposes +> `reset -> prefill -> decode`, optional token input, next-token/EOG/logits/text +> outputs, and Python `reset()`/`prefill()`/`decode()` methods. Tests cover exact +> one-shot/staged parity, finite/full logits parity, host interruption by +> stopping at token boundaries, distinct interleaved KV sessions, and strict +> `max_tokens` budgets on CPU, CUDA, Vulkan, and SYCL. + +> Historical evaluation dated 2026-07-07. The implementation update above is +> authoritative; the remaining text records why this work was initially +> deferred before the complete-migration scope was requested. + +## 1. Executive summary + recommendation + +**Recommendation: NO-GO. Defer TODO-2 under current conditions.** + +Four lenses (scheduling benefit, buffer contract, sync semantics, ABI cost) do **not** converge on GO. The split is *implementable* (lens 3 passes cleanly — unlike Phase 5, where Pi0's flow-matching has no usable intermediate state, an LLM's per-token decode IS the natural boundary), but it is not *justified under current callers*: no FlashRT caller multiplexes LLM sessions or interleaves LLM decode with another provider, so the "host interrupts/schedules at token boundaries" benefit (the entire point of §8.1's repeatable decode) is theoretical. Meanwhile the cost is real: a Jetson-PI narrow-API extension (`prefill`/`decode_step`/`reset`), an append-only FlashRT engine-vtable ABI extension (`run_infer` → `run_stage`), new ports/stages, and a refactor of the one-shot `generate` path that TODO-1 (commit `0cee8c7`) just proved byte-identical against the native API. + +The single most important structural fact: the Jetson-PI narrow C API `jetson_pi_llm` exposes only `open/generate/close/open_error` (`Jetson-PI/include/jetson_pi_llm.h`). `generate` hardcodes `llama_memory_clear(..., true)` (`Jetson-PI/src/jetson_pi_llm.cpp:200`) and `llama_sampler_reset` (`:201`) every call — KV and sampler state are wiped per completion. There is **no** `prefill`/`decode_step`/`reset`/`get_logits` symbol. FlashRT cannot fabricate the prefill/decode boundary from its side; the only seam is `generate`, and everything inside it is opaque. This is the same gate Phase 5 hit: a stage split requires Jetson-PI to expose separate sub-phase entry points. + +A DEFER here is the correct, defensible outcome — it prevents adding ABI surface, new ports, and a 2-repo refactor for a benefit no caller exercises, exactly as Phase 5's DEFER prevented Pi0 stage-splitting for zero latency gain. It routes the work to **when a real caller needs it** (revisit conditions a–d, §6). + +## 2. Gate finding (verified against source) + +**Jetson-PI narrow API surface** (`Jetson-PI/include/jetson_pi_llm.h`, 98 lines): `jetson_pi_llm_open` (h:55), `close` (h:59), `last_error` (h:64), `open_error` (h:69), `generate` (h:89). Header doc (h:11-13, h:88) is explicit: "Each generate call clears KV state and runs an independent first-turn completion (no multi-turn state)." Config (`jetson_pi_llm_config`, h:28-42) carries `temp/top_k/top_p/seed/max_tokens` — no flags for logits exposure, sampler rebuild, or step mode. + +**`generate` internals** (`Jetson-PI/src/jetson_pi_llm.cpp:174-277`): +- KV wipe: `llama_memory_clear(llama_get_memory(e->lctx), true)` at **:200**. +- Sampler wipe: `llama_sampler_reset(e->smpl)` at **:201**. +- Prefill: `llama_batch_get_one(prompt_tokens)` + `llama_decode` at **:220-225** (whole-prompt single batch). +- Decode loop (**:231-266**): `llama_sampler_sample` (:232) → EOG check via `llama_vocab_is_eog` (:233, break internally — host never sees it) → `llama_token_to_piece` detokenize (:242-256) → `llama_sampler_accept(new_token_id)` (:258, feeds token back into sampler state) → `llama_batch_get_one(&new_token_id,1)` + `llama_decode` (:261-265, single-token step). +- **Logits are never exposed**: zero `llama_get_logits`/`llama_get_logits_ith` calls in `jetson_pi_llm.cpp`. The host gets only the final detokenized text blob. + +The llama.cpp primitives needed for a split are **all already linked**; all except `llama_get_logits_ith` are already used piecewise inside `generate` (`llama_decode` :222/:262, `llama_sampler_sample` :232, `llama_sampler_accept` :258, `llama_memory_clear` :200, `llama_vocab_is_eog` :233). `llama_get_logits_ith` is available (declared in `llama-context.h`) but currently uncalled — the work is *exposing* these as host-driven entry points, not inventing capability. + +## 3. Current FlashRT LLM provider surface + +- **C ABI** (`cpp/providers/llama_cpp/include/flashrt/providers/llama_cpp/c_api.h`): LLM ports `FRT_LLAMA_CPP_LLM_PORT_PROMPT=0`, `PORT_TEXT=1` (h:29-32) — **no** `tokens` in, `next_token` out, or `logits` out. Stage `FRT_LLAMA_CPP_LLM_STAGE_INDEX_INFER=0` only (h:34-36) — **no** prefill/decode split. Engine vtable `frt_llama_cpp_engine_v1` (h:96-115) has `run_infer(void*)` only — **no stage parameter**. +- **Runtime** (`cpp/providers/llama_cpp/src/llm_runtime.cpp`): two STAGED ports (prompt in, text out, :131-138), one `infer` callback stage (:139-140), `run_stage` (:64-79) collapses every stage to `engine.run_infer`. +- **Engine** (`cpp/providers/llama_cpp/src/jetson_pi_engine.cpp`, LLM section :346-462): `LlmEngine` (:354-367) holds `jetson_pi_llm*`, prompt stash, `text_buf`. `llm_engine_run_infer` (:401-426) calls `jetson_pi_llm_generate` **once** (:415-417). Each `run_infer` is a full independent completion. +- **Python** (`flash_rt/frontends/jetson_pi/llm.py`): `generate(prompt)` (:103-127) = `set_input(PROMPT)` → `run_stage(INFER)` → `get_output(TEXT)`. Single whole-prompt completion. + +§8.1 gap: §8.1 specifies `prompt`/`tokens`(opt in)/`next_token`(out)/`logits`(opt out) ports + `prefill once → decode (repeatable)` stages. Current has `prompt`+`text`+single `infer`. Every one of the five §8.1 items (tokens-in, next_token, logits, prefill/decode split, repeatable decode) is absent. + +## 4. Four-lens evaluation + +### Lens 1 — Scheduling benefit + +§8.1's stated benefit (`jetsonpi迁移.txt:302`): "host 在 token 边界重复调用 decode，从而保留中断和调度能力". For a **single-stream completion** (one prompt → one text blob), there is no scheduling overlap to win: prefill and decode are strictly sequential within one session, and FlashRT has no second concurrent stream to pipeline against. The win only materializes when FlashRT actually schedules multiple LLM sessions or interleaves LLM decode with another provider's work — which no **jetson_pi** caller does. The one-shot `generate` already returns the full blob; a host wanting interrupt/resume would need FlashRT-level session multiplexing first. **Lens 1: FAILS under current callers.** + +> Scope note: FlashRT's **torch** LLM frontends (`flash_rt/frontends/torch/qwen3_rtx.py` `set_prompt`/`reset_state`/`forward_own_decode_nvfp4`/`decode_step_with_graph` + n-gram speculative decode; `qwen36_rtx.py` MTP speculative decode `generate_own_speculative_K*` + committed-stream prefill/decode) already implement host-driven per-token decode — but on a different provider (`framework="torch"`, NVFP4 weights, not GGUF), constructed directly (not via `load_model(..., framework="jetson_pi")`). §8.1/TODO-2 is explicitly the **jetson_pi/llama.cpp** provider spec. The torch precedent shows the pattern is implementable and useful where there is demand; it does not create a jetson_pi caller need. A GGUF-LLM caller cannot use the torch path. This scoping is legitimate, not a dodge — but it must be stated, since the torch frontends are the obvious counterexample to a blanket "FlashRT has no per-step LLM decode" claim. + +### Lens 2 — Buffer contract + +A `next_token` (I32 out) port is trivial (4 bytes). `logits` (F32 out, vocab-sized) is the invasive piece: qwen3-0.6b vocab ~150k → ~600KB/step, with ownership/lifetime questions (per-step snapshot vs ring). The Phase 6 memory-domain token contract (commit `95d7175`) could carry a `logits` token, but no consumer needs it today. **Lens 2: contract defined (STAGED ports + token contract exist), but no consumer pulls logits.** + +### Lens 3 — Sync semantics (the lens that PASSES, unlike Phase 5) + +The boundary is clean: a `decode_step` would sample one token, accept it into the sampler, decode it, return `(token_id, eog_flag)`. KV and sampler state persist across steps (the whole point). EOG is observable (`llama_vocab_is_eog`). The semantics are well-defined and llama.cpp supports them directly. **This is where LLM differs from Pi0**: Pi0's flow-matching has no usable intermediate state (TTFA==TTC, Phase 5 §1); an LLM's per-token decode IS the natural boundary. **Lens 3: PASSES — TODO-2 is implementable, not incoherent.** + +### Lens 4 — ABI cost + +The engine vtable `frt_llama_cpp_engine_v1` (c_api.h:96-115) has only `run_infer(void*)`. To expose `prefill` + `decode` as two observable stages, the vtable needs an append-only `run_stage(void*, uint32_t stage)` verb (struct_size-gated, consistent with §10.1's append-only guidance, `jetsonpi迁移.txt:392-395`). Plus new port enums (`PORT_NEXT_TOKEN`, `PORT_LOGITS`) and stage enums (`STAGE_PREFILL`, `STAGE_DECODE`). The runtime wrapper `llm_runtime.cpp:64` already has a `run_stage` dispatcher that collapses every stage to `engine.run_infer` — so the vtable extension is the one FlashRT-ABI change. Append-only and safe, but it's a vtable shape change on a structure TODO-1's parity test and the fake-engine test both exercise. **Lens 4: real cost — one vtable ABI extension + 2-repo refactor of a path TODO-1 just stabilized.** + +## 5. Cost / blast radius (if implemented — option iii MVP) + +- **Jetson-PI** (`jetson_pi_llm.h/.cpp`): add `jetson_pi_llm_prefill`/`decode_step`/`reset` (+ optional `get_logits`); refactor `generate` into `prefill` + loop(`decode_step`) + `reset` (one code path). New Jetson-PI tests. +- **FlashRT C ABI** (`c_api.h`): `PORT_NEXT_TOKEN`/`PORT_LOGITS` enums + `STAGE_PREFILL`/`STAGE_DECODE` enums + engine-vtable `run_stage` verb (append-only, struct_size-gated). +- **FlashRT runtime** (`llm_runtime.cpp`): declare new ports + 2 stages; wire `run_stage` to the new engine verb. +- **FlashRT engine** (`jetson_pi_engine.cpp` LlmEngine): rework verbs for prefill/decode/`get_output(next_token, logits)`; stop forcing KV reset between calls. +- **Python** (`llm.py`): `prefill()`/`decode_step()`/`reset()`; `generate()` becomes a wrapper. +- **Tests**: streaming smoke + §14 KV-isolation gate (consecutive sessions must not leak — the §14 principle `jetsonpi迁移.txt:599` is stated for Pi0 context; the same isolation requirement applies to LLM KV across sessions). Existing `test_llama_cpp_jetson_pi_llm.cpp` (Phase 3, `083f6f8`) + `test_llama_cpp_jetson_pi_llm_parity.cpp` (TODO-1, `0cee8c7`) drive the old `STAGE_INFER`→`PORT_TEXT` path — must keep working (the `text` convenience face stays; if implemented, `generate()` becomes a wrapper over `prefill`+loop(`decode_step`)+`reset`). + +Three repos touched; one vtable ABI extension; one refactor of a path TODO-1 just proved byte-identical. Non-trivial but bounded. + +## 6. Recommendation + revisit conditions + +**DEFER.** The four lenses do not converge on GO: lens 1 (scheduling benefit) fails under current callers; lens 2 (buffer contract) is half-met (no logits consumer); lens 4 (ABI cost) is real. Only lens 3 (sync semantics) passes — which makes TODO-2 *implementable* (not incoherent like a naive Pi0 split would be) but not *justified*. + +**Revisit (GO when ANY holds):** +- **a.** A real FlashRT caller needs host-driven token looping — e.g. an LLM serving path with early-stop / speculative decode / multi-session KV reuse. Lens 1 → GO. +- **b.** FlashRT scheduler gains multi-session LLM multiplexing — prefill/decode becomes the natural scheduling unit. Lens 1 → GO. *(Caveat: as of 2026-07-07 no jetson_pi multi-session scheduler exists in FlashRT — grep for scheduler/multiplex/session in `flash_rt/`/`runtime/`/`cpp/` returns only diffusion samplers and unrelated comments; the only LLM serving path `examples/qwen3_openai_server.py` uses the torch frontend. So (b) is contingent on a jetson_pi serving path that does not exist today, and is near-redundant with (a) — if (a) fires, (b) is moot. It is a legitimate future condition, not an escape hatch.)* +- **c.** A caller needs logits (external sampler, classifier head, constrained decoding). Lens 2 → GO; the `logits` port + a `get_logits` Jetson-PI symbol become justified. +- **d.** LLM streaming UX becomes a requirement (token-by-token output to a downstream consumer that can't wait for the full blob). A product requirement, not a scheduling one, but a legitimate GO trigger. + +Until one of a–d, the current one-shot `generate` (verified byte-parity by TODO-1, commit `0cee8c7`) is the honest, sufficient LLM face. + +## 7. Note on the Phase 5 precedent + +This mirrors Phase 5 (commit `8a4cb9d`, `docs/phase5_pi0_stages_eval.md`): a phase resolved by **evaluation/decision**, not implementation, with explicit revisit conditions. The Phase 5 gate required all four of (scheduling benefit, buffer contract, sync semantics, profiling evidence); here lens 3 passes where Phase 5's sync-semantics lens failed, but lens 1 fails where Phase 5's also failed — so the DEFER conclusion holds on the same "no current caller benefits" ground. The difference is that TODO-2's revisit is gated on a *caller need* (a–d), whereas Phase 5's is gated on a *profiling signal + Jetson-PI internal optimization* — because LLM decode is genuinely splittable (lens 3) in a way Pi0 diffusion is not. diff --git a/docs/phase8_zerocopy_dlpack_eval.md b/docs/phase8_zerocopy_dlpack_eval.md new file mode 100644 index 00000000..f98fad5a --- /dev/null +++ b/docs/phase8_zerocopy_dlpack_eval.md @@ -0,0 +1,72 @@ +# Phase 8 — Zero-copy / DLPack / host SWAP for provider-owned ports + +> Implementation update (2026-07-10): the earlier no-go evaluation below is +> superseded by the complete-migration requirement. `frt_memory_token_verbs` +> now has append-only `map_host`/`unmap_host` verbs. Pi0 actions exposes a +> read-only HOST_VISIBLE SWAP window; the Python frontend returns a zero-copy +> NumPy view whose native `__dlpack__` exports the same storage. New input or +> stage execution is rejected while the view is mapped, preventing vector +> reallocation and stale pointers. The original copy_to_host path remains. + +> Historical evaluation dated 2026-07-08. The implementation update above is +> authoritative; the remaining text records why this work was previously +> deferred before the complete-migration scope was requested. + +## 1. Executive summary + recommendation + +**Recommendation: NO-GO. Defer TODO-5 under current conditions.** + +The Phase 6 memory-domain contract (commit `5d25672`/`b9919be`) plus its real token consumer (commit `95d7175`, Pi0 actions OUT port, `HOST_VISIBLE`) already cover the honest zero-copy-adjacent case: a host can read the actions buffer without a device sync, via `copy_to_host`. No FlashRT caller reads the actions buffer via DLPack, via a raw SWAP window, or as a device-resident tensor today. TODO-5's own prerequisite in CLAUDE.md — "明确真实 caller 需求（谁要零拷贝读 actions？当前 consumer 都是 copy_to_host）" — answers itself: there is no such caller. + +Adding DLPack export or a host SWAP window now would be to add a **second, half-tested buffer form for a nonexistent consumer**, which directly violates the contract's founding principle ("不要在没有 memory-domain contract 前假装可以零拷贝共享 buffer" — and once you HAVE the contract, don't add a parallel half-finished one either). The four lenses do not converge on GO: only sync semantics (DLPack is a well-defined capsule) passes; scheduling benefit, buffer-contract-need, and ABI cost all fail under current callers. This is the same gate Phase 5 (Pi0 stage split) and Phase 7 (LLM repeatable decode) hit: *implementable ≠ justified*. + +A DEFER here routes the work to **when a real caller needs it** (revisit conditions a–c, §6) and keeps the provider surface honest: one buffer form (`HOST_VISIBLE` token + `copy_to_host`), exercised and parity-verified, not two. + +## 2. Gate finding (verified against source) + +**The only token consumer today** is the Pi0 actions OUT port (commit `95d7175`, `cpp/providers/llama_cpp/src/jetson_pi_engine.cpp`): +- `pi0_token_copy_to_host` (:307) — memcpy from the live `actions_buf` into a caller host buffer. +- `pi0_token_copy_from_host` (:321) — returns `-3` unsupported (OUT port). +- `pi0_token_sync` (:327) — `return 0` (HOST_VISIBLE no-op). +- `destroy` = nullptr. + +`location_kind = FRT_RT_LOCATION_HOST_VISIBLE` (`model_runtime.h:120`: "copy_to_host is a plain read; no sync"). The backing store is a host `std::vector`. **Zero callers** use a DLPack capsule, a raw `frt_buffer` SWAP window, or a `DEVICE_LOCAL` token against this port — verified by `grep -rnE 'dlpack|DLPack|swap_window|cap_swap'` across `cpp/providers` + `flash_rt/frontends/jetson_pi` returning nothing relevant (the only `from_dlpack` hits are in `csrc/attention/flash_attn_4_src/flashrt_fa4/cute/` — CUTLASS/CuTe internals for FA4 kernels, unrelated to the Jetson-PI provider). + +**Provider-owned ports are forced STAGED** (`model_runtime.cpp:184` `d.buffer = nullptr` in `frt_runtime_builder_add_port_token`): there is no raw `frt_buffer` SWAP window on a token port by construction. The `HOST_VISIBLE` token already advertises the honest capability ("host can read without a device sync") that every actual consumer uses. + +## 3. Current surface + +- **Phase 6 contract**: opaque `frt_memory_token` + `frt_memory_token_verbs` (`copy_to_host`/`copy_from_host`/`sync`/`destroy`) + `frt_rt_location_kind` (`HOST_VISIBLE`/`DEVICE_LOCAL`), parallel `port_tokens` array on `frt_model_runtime_v2` (`model_runtime.h:183-229`, `:378`). +- **Real consumer**: Pi0 actions port mints a `HOST_VISIBLE` token (commit `95d7175`); the C++ e2e test sub-test E reads it via `copy_to_host` and asserts byte-identity vs `get_output`. +- **No DLPack path** exists for provider-owned ports; no `__dlpack_managed`/`from_dlpack` bridge in the Python frontend (`flash_rt/frontends/jetson_pi/pi0.py` returns a NumPy array copied out of `get_output`). + +§11 gap: §11 says "后续如果需要 CPU SWAP 或零拷贝，FlashRT buffer descriptor 需要增加 memory domain" and "在没有 memory-domain 契约前，不应将普通 CPU 指针伪装成 FlashRT device buffer". The memory-domain contract now EXISTS (Phase 6); the question is whether to **use** it for zero-copy yet. The answer is no, because no consumer asks. + +## 4. Four-lens evaluation + +### Lens 1 — Scheduling / performance benefit +None under current callers. The `copy_to_host` of a `[10, 32]` F32 action chunk (1.28 KB) is negligible vs the ~1.1 s Pi0 tick; a zero-copy read would save a 1.28 KB memcpy — unmeasurable. A device-resident read would matter only if a downstream consumer runs on the same GPU and currently pays a D2H it could avoid — but no such consumer exists. **Lens 1: FAILS.** + +### Lens 2 — Buffer contract +Half-met. The contract is in place and COULD carry a DLPack capsule (a new verb `to_dlpack` minting a managed tensor, or a `DEVICE_LOCAL` token for a CUDA-resident buffer). But adding it now creates a second buffer form alongside the exercised `HOST_VISIBLE`+`copy_to_host` path — a half-tested parallel surface for a nonexistent consumer. The contract's own principle ("don't pretend zero-copy; route copies through explicit verbs with explicit location_kind") argues against adding a SWAP window until a consumer needs it. **Lens 2: half-met → don't add a half-finished SWAP.** + +### Lens 3 — Sync semantics (the lens that PASSES) +DLPack is a well-defined, language-agnostic capsule (`DLManagedTensor` with `deleter`); `HOST_VISIBLE` already means "no device sync needed"; `DEVICE_LOCAL` would require a `sync` before host read (the contract already models this). The semantics are clean and implementable. **Lens 3: PASSES** — TODO-5 is *implementable*, not *incoherent*. (Same position as Phase 7's lens 3.) + +### Lens 4 — ABI cost +Real. A DLPack export on a provider-owned port is new ABI surface: a `to_dlpack`/`from_dlpack` verb pair (or a new `location_kind` + capsule-manage verbs), dtype/shape/strides negotiation, and a managed-lifetime contract (who calls `deleter`, how it interlocks with the holder refcount). For zero consumers, this is uncompensated ABI surface that future append-only extensions must stay compatible with. **Lens 4: real cost, no consumer to justify it.** + +## 5. Recommendation + revisit conditions + +**DEFER.** The four lenses do not converge on GO: lens 1 (benefit) fails, lens 2 (contract) is half-met and the contract's own principle argues against a premature SWAP, lens 4 (ABI) is real. Only lens 3 (sync semantics) passes — making TODO-5 *implementable* but not *justified*, exactly the Phase 5 / Phase 7 pattern. + +**Revisit (GO when ANY of):** +- **a.** A real caller needs NumPy/PyTorch **zero-copy of the actions buffer** — e.g. a downstream policy controller running at kHz that cannot afford the (currently 1.28 KB) `copy_to_host` memcpy, or a learning loop that materializes the action as a trainable tensor without copy. Lens 1 → GO. +- **b.** A caller needs a **device-resident actions read** — a second provider (or a FlashRT CUDA-Graph stage) consuming Pi0's actions as its input on the same GPU, avoiding a D2H+H2D round-trip. Then a `DEVICE_LOCAL` token (with a real `sync` verb) is justified; lens 2 → GO. +- **c.** **Cross-provider tensor hand-off via DLPack** becomes a FlashRT requirement — two providers exchanging a tensor through a DLPack capsule rather than a host staging copy. Lens 2/4 → GO (the ABI cost is paid by a real interop need). + +Until one of a–c, the current `HOST_VISIBLE` token + `copy_to_host` (byte-parity-verified by sub-test E and TODO-1 parity) is the honest, sufficient zero-copy-adjacent face. + +## 6. Note on the Phase 5 / Phase 7 precedent + +This mirrors Phase 5 (commit `8a4cb9d`) and Phase 7/TODO-2 (commit `8dfb52d`): a phase resolved by **evaluation/decision**, not implementation, with explicit revisit conditions. The shared gate is "no current caller benefits / no current consumer needs the surface". The difference: Phase 5's sync-semantics lens FAILED (Pi0 flow-matching has no usable intermediate state); Phase 7's and Phase 8's lens 3 PASS (LLM decode / DLPack capsules are clean boundaries) — so TODO-5 is *implementable* where Phase 5 was *incoherent*, but the DEFER conclusion holds on the same "no caller need" ground. The revisit is gated on a **caller need** (a–c), like Phase 7, not on a profiling signal (Phase 5's was). diff --git a/exec/backend/backend.h b/exec/backend/backend.h index b9d7bd1b..0cca4363 100644 --- a/exec/backend/backend.h +++ b/exec/backend/backend.h @@ -7,6 +7,32 @@ * * All functions return null / false on failure; the core layer turns that into * the public frt_status codes. No exceptions cross this boundary. + * + * Phase 6 clarification — the surface splits into two groups: + * + * (a) SHARED primitives (every backend implements): malloc/free, stream_*, + * event_*, memcpy_dtod_async, memset_async. A provider-owned runtime + * (e.g. Jetson-PI/GGML) does NOT consume these — GGML owns its device + * memory and streams internally — so FlashRT never calls this group for + * the provider-owned path. It is the CUDA-export path's vocabulary only. + * + * (b) CUDA-GRAPH-CAPTURE primitives (optional): capture_begin/end, + * graph_exec_destroy, graph_launch. A backend whose capture_begin + * returns false cannot host graph-capture stages; it must use the + * provider-owned callback-stage path (frt_runtime_stage_desc_v2 with + * kind=CALLBACK) instead. GGML has no cudaGraph_t analogue, so a GGML- + * facing exec backend would implement group (a) only and never (b) — + * but no such backend is needed in this migration because GGML is + * integrated as a model provider (not a bare exec backend). See + * docs/phase6_backend_vtable_eval.md. + * + * Backend selection is LINK-TIME: one TU per build (exec/CMakeLists.txt links + * cuda/cuda_backend.cpp). Runtime multi-backend selection in one process would + * require promoting this namespace to a registered function-pointer table; that + * is out of scope for Phase 6 (no consumer — see the eval). The import() seam + * below is declared so a future cross-domain buffer registration path has a + * named extension point, but it is NOT implemented for non-CUDA backends in + * this phase. */ #ifndef FLASHRT_EXEC_BACKEND_H #define FLASHRT_EXEC_BACKEND_H @@ -16,6 +42,8 @@ namespace frt { namespace be { +/* ── (a) SHARED primitives — every backend implements ─────────────────── */ + /* --- device memory --- */ void* malloc(size_t bytes); /* device alloc; null on failure */ void free(void* dptr); @@ -35,15 +63,26 @@ bool stream_wait_event(void* stream, void* event); bool memcpy_dtod_async(void* dst, const void* src, size_t bytes, void* stream); bool memset_async(void* dptr, int value, size_t bytes, void* stream); -/* --- CUDA-graph-style capture/replay --- +/* ── (b) CUDA-GRAPH-CAPTURE primitives — optional ─────────────────────── * capture_begin/end wrap stream-level capture (RELAXED mode); capture_end * instantiates and returns an executable graph handle, freeing the transient - * non-executable graph. graph_exec_destroy frees the executable. */ + * non-executable graph. graph_exec_destroy frees the executable. A backend + * that cannot capture returns false from capture_begin; the caller then uses + * the provider-owned callback-stage path instead (no graph object involved). */ bool capture_begin(void* stream); void* capture_end(void* stream); /* returns graph-exec handle; null = fail */ void graph_exec_destroy(void* graph_exec); bool graph_launch(void* graph_exec, void* stream); +/* ── future extension point (DECLARED ONLY, Phase 6) ──────────────────── + * Register an externally-owned memory region (e.g. a provider's device buffer) + * so a CUDA-export graph could bind it as a SWAP window without a copy. NOT + * implemented for non-CUDA backends in this phase — cross-backend zero-copy + * needs a memory-domain contract that FlashRT only just gained at the + * provider-owned port layer (frt_memory_token). This declaration exists so a + * later phase can add the implementation without retroactive header churn. */ +/* void* import(void* external_handle, size_t bytes, uint32_t location_kind); */ + } // namespace be } // namespace frt diff --git a/flash_rt/api.py b/flash_rt/api.py index 1ebe3875..7cfc486a 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: @@ -432,15 +445,68 @@ 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": + # The Jetson-PI provider serves Pi0 (VLA) and generic GGUF LLMs. config + # picks the path. Neither uses torch/jax or GPU arch detection. + if config not in ("pi0", "pi05", "llm", "mllm"): + logger.warning( + "framework='jetson_pi' serves Pi0 / LLM / MLLM; config=%r ignored.", + config) + elif config not in ("pi05", "groot", "groot_n17", "pi0", "pi0fast", "motus", "wan22_ti2v_5b", "cosmos3_video", "nexn2"): raise ValueError( f"Unknown config: {config}. " f"Supported: pi05, groot, groot_n17, pi0, pi0fast, motus, " f"wan22_ti2v_5b, cosmos3_video, 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") + + # ── Jetson-PI (llama.cpp/GGML) provider — Phase 2 ── + # Drives the Jetson-PI Pi0 provider through the frt_model_runtime_v2 C ABI + # via ctypes. No torch/jax, no GPU arch detection. The Pi0 action chunk + # shape is passed explicitly by the caller (e.g. 10x32 for pi0_base). + if framework == "jetson_pi": + # The Jetson-PI provider serves Pi0 (VLA) and generic GGUF LLMs. config + # picks the path. Neither uses torch/jax or GPU arch detection. + 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..315c8f8c --- /dev/null +++ b/flash_rt/frontends/jetson_pi/llm.py @@ -0,0 +1,292 @@ +"""Jetson-PI generic GGUF LLM frontend — drives the Jetson-PI llama.cpp LLM +provider through the FlashRT ``frt_model_runtime_v2`` C ABI via ctypes. + +This is the Phase 3 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 + FrtModelRuntimeV2, + _FactoryV1, + _FRT_MODEL_RUNTIME_ABI_VERSION_V2, + _find_lib, +) + +# 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_llama_cpp_default_engine_factory.argtypes = [] + self._lib.frt_llama_cpp_default_engine_factory.restype = ctypes.c_void_p + try: + self._lib.frt_llama_cpp_runtime_open_error.argtypes = [] + self._lib.frt_llama_cpp_runtime_open_error.restype = ctypes.c_char_p + except AttributeError as exc: + raise RuntimeError( + "provider .so is older than the Python frontend expects: " + "missing frt_llama_cpp_runtime_open_error") from exc + self._lib.frt_llama_cpp_llm_runtime_open_with_engine_factory.argtypes = [ + ctypes.c_char_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ] + self._lib.frt_llama_cpp_llm_runtime_open_with_engine_factory.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") + + factory_ptr = self._lib.frt_llama_cpp_default_engine_factory() + if not factory_ptr: + raise RuntimeError( + "frt_llama_cpp_default_engine_factory returned NULL " + "(FlashRT built without FLASHRT_CPP_WITH_JETSON_PI?)") + factory = _FactoryV1.from_address(factory_ptr) + if factory.struct_size < ctypes.sizeof(_FactoryV1): + raise RuntimeError( + f"frt_llama_cpp_engine_factory_v1 struct_size=" + f"{factory.struct_size} < ctypes sizeof " + f"{ctypes.sizeof(_FactoryV1)}; provider .so is older than the " + f"Python frontend expects.") + + model_ptr = ctypes.c_void_p(0) + rc = self._lib.frt_llama_cpp_llm_runtime_open_with_engine_factory( + config_json, factory_ptr, ctypes.byref(model_ptr)) + if rc != 0 or not model_ptr.value: + err = self._lib.frt_llama_cpp_runtime_open_error() or b"" + if not err: + err = factory.last_error(factory.self) or b"" + raise RuntimeError( + f"frt_llama_cpp_llm_runtime_open_with_engine_factory failed " + f"(rc={rc}): {(err.decode(errors='replace') if err else 'no error')}") + + self._model = FrtModelRuntimeV2.from_address(model_ptr.value) + if self._model.abi_version != _FRT_MODEL_RUNTIME_ABI_VERSION_V2: + raise RuntimeError( + f"frt_model_runtime_v2 abi_version={self._model.abi_version}, " + f"expected {_FRT_MODEL_RUNTIME_ABI_VERSION_V2}.") + if self._model.struct_size < ctypes.sizeof(FrtModelRuntimeV2): + raise RuntimeError( + f"frt_model_runtime_v2 struct_size={self._model.struct_size} " + f"< ctypes sizeof {ctypes.sizeof(FrtModelRuntimeV2)}.") + 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) + v2 = self._model.verbs_v2 + self_ = self._model.self + + rc = v2.set_input(self_, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + rc = v2.run_stage(self_, STAGE_INFER, -1) + self._check(rc, "run_stage infer") + + written = ctypes.c_uint64(0) + rc = v2.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 = v2.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 = self._model.verbs_v2.run_stage( + self._model.self, STAGE_RESET, -1) + self._check(rc, "run_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") + v2 = self._model.verbs_v2 + 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 = v2.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 = v2.set_input(self._model.self, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + rc = v2.run_stage(self._model.self, STAGE_PREFILL, -1) + self._check(rc, "run_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") + v2 = self._model.verbs_v2 + rc = v2.run_stage(self._model.self, STAGE_DECODE, -1) + self._check(rc, "run_stage decode") + + next_token = ctypes.c_int32() + written = ctypes.c_uint64(0) + rc = v2.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 = v2.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 = v2.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 = v2.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") + v2 = self._model.verbs_v2 + required = ctypes.c_uint64(0) + rc = v2.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 = v2.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_v2.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..e5fa70ba --- /dev/null +++ b/flash_rt/frontends/jetson_pi/mllm.py @@ -0,0 +1,320 @@ +"""Jetson-PI multimodal LLM frontend — drives the Jetson-PI llama.cpp MLLM +provider through the FlashRT ``frt_model_runtime_v2`` C ABI via ctypes. + +This is the Phase 4 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, + FrtModelRuntimeV2, + FRT_RT_PIXEL_RGB8, + _FactoryV1, + _FRT_MODEL_RUNTIME_ABI_VERSION_V2, + _find_lib, +) + +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_llama_cpp_default_engine_factory.argtypes = [] + self._lib.frt_llama_cpp_default_engine_factory.restype = ctypes.c_void_p + try: + self._lib.frt_llama_cpp_runtime_open_error.argtypes = [] + self._lib.frt_llama_cpp_runtime_open_error.restype = ctypes.c_char_p + except AttributeError as exc: + raise RuntimeError( + "provider .so is older than the Python frontend expects: " + "missing frt_llama_cpp_runtime_open_error") from exc + self._lib.frt_llama_cpp_mllm_runtime_open_with_engine_factory.argtypes = [ + ctypes.c_char_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ] + self._lib.frt_llama_cpp_mllm_runtime_open_with_engine_factory.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") + + factory_ptr = self._lib.frt_llama_cpp_default_engine_factory() + if not factory_ptr: + raise RuntimeError( + "frt_llama_cpp_default_engine_factory returned NULL " + "(FlashRT built without FLASHRT_CPP_WITH_JETSON_PI?)") + factory = _FactoryV1.from_address(factory_ptr) + if factory.struct_size < ctypes.sizeof(_FactoryV1): + raise RuntimeError( + f"frt_llama_cpp_engine_factory_v1 struct_size=" + f"{factory.struct_size} < ctypes sizeof " + f"{ctypes.sizeof(_FactoryV1)}; provider .so is older than the " + f"Python frontend expects.") + + model_ptr = ctypes.c_void_p(0) + rc = self._lib.frt_llama_cpp_mllm_runtime_open_with_engine_factory( + config_json, factory_ptr, ctypes.byref(model_ptr)) + if rc != 0 or not model_ptr.value: + err = self._lib.frt_llama_cpp_runtime_open_error() or b"" + if not err: + err = factory.last_error(factory.self) or b"" + raise RuntimeError( + f"frt_llama_cpp_mllm_runtime_open_with_engine_factory failed " + f"(rc={rc}): {(err.decode(errors='replace') if err else 'no error')}") + + self._model = FrtModelRuntimeV2.from_address(model_ptr.value) + if self._model.abi_version != _FRT_MODEL_RUNTIME_ABI_VERSION_V2: + raise RuntimeError( + f"frt_model_runtime_v2 abi_version={self._model.abi_version}, " + f"expected {_FRT_MODEL_RUNTIME_ABI_VERSION_V2}.") + if self._model.struct_size < ctypes.sizeof(FrtModelRuntimeV2): + raise RuntimeError( + f"frt_model_runtime_v2 struct_size={self._model.struct_size} " + f"< ctypes sizeof {ctypes.sizeof(FrtModelRuntimeV2)}.") + 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") + v2 = self._model.verbs_v2 + self_ = self._model.self + + views = self._make_image_views(images) + rc = v2.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 = v2.set_input(self_, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + + rc = v2.run_stage(self_, STAGE_INFER, -1) + self._check(rc, "run_stage infer") + + written = ctypes.c_uint64(0) + rc = v2.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 = v2.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 = self._model.verbs_v2.run_stage( + self._model.self, STAGE_RESET, -1) + self._check(rc, "run_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") + v2 = self._model.verbs_v2 + views = self._make_image_views(images) + rc = v2.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 = v2.set_input(self._model.self, PORT_PROMPT, prompt_bytes, + len(prompt_bytes), -1) + self._check(rc, "set_input prompt") + rc = v2.run_stage(self._model.self, STAGE_PREFILL, -1) + self._check(rc, "run_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") + v2 = self._model.verbs_v2 + rc = v2.run_stage(self._model.self, STAGE_DECODE, -1) + self._check(rc, "run_stage decode") + next_token = ctypes.c_int32() + written = ctypes.c_uint64(0) + rc = v2.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 = v2.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 = v2.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 = v2.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") + v2 = self._model.verbs_v2 + required = ctypes.c_uint64(0) + rc = v2.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 = v2.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_v2.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..65ae14f9 --- /dev/null +++ b/flash_rt/frontends/jetson_pi/pi0.py @@ -0,0 +1,586 @@ +"""Jetson-PI Pi0 frontend — drives the Jetson-PI llama.cpp/GGML Pi0 provider +through the FlashRT ``frt_model_runtime_v2`` C ABI via ctypes. + +This frontend is the Phase 2 Python entry for the Jetson-PI provider. 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_llama_cpp_default_engine_factory`` + ``open_with_engine_factory``, and +drives one whole-graph 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 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_v2 verbs + struct mirrors (only the fields we touch) ---- +# +# We only need: abi_version, struct_size (sanity), self, verbs_v2 (for +# set_input/get_output/run_stage/last_error), owner, release. The full struct +# has more fields (exp, ports, stages, verbs v1, retain, stages_v2) but ctypes +# only requires correct field ORDER up to the last one we read; trailing fields +# can be omitted as long as we never touch them. To stay robust against layout +# drift we mirror the full v2 layout. + +_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) +_RunStageFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_int) +_RetainReleaseFn = ctypes.CFUNCTYPE(None, ctypes.c_void_p) +_TokenCopyToHostFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64) +_TokenCopyFromHostFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64) +_TokenSyncFn = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p) +_TokenDestroyFn = ctypes.CFUNCTYPE(None, ctypes.c_void_p) +_TokenMapHostFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint64, + ctypes.c_uint32, ctypes.POINTER(ctypes.c_void_p)) +_TokenUnmapHostFn = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p) + + +class _FrtMemoryTokenVerbs(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("reserved", ctypes.c_uint32), + ("copy_to_host", _TokenCopyToHostFn), + ("copy_from_host", _TokenCopyFromHostFn), + ("sync", _TokenSyncFn), + ("destroy", _TokenDestroyFn), + ("map_host", _TokenMapHostFn), + ("unmap_host", _TokenUnmapHostFn), + ] + + +class _FrtMemoryTokenDesc(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("handle", ctypes.c_void_p), + ("verbs", ctypes.POINTER(_FrtMemoryTokenVerbs)), + ("offset", ctypes.c_uint64), + ("bytes", ctypes.c_uint64), + ("location_kind", ctypes.c_uint32), + ("reserved", ctypes.c_uint32), + ] + + +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 _FrtVerbsV2(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), + ("run_stage", _RunStageFn), + ] + + +class FrtModelRuntimeV2(ctypes.Structure): + _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), + ("stages_v2", ctypes.c_void_p), + ("n_stages_v2", ctypes.c_uint64), + ("verbs_v2", _FrtVerbsV2), + ("port_tokens", ctypes.POINTER(_FrtMemoryTokenDesc)), + ("n_port_tokens", ctypes.c_uint64), + ] + + +class _MappedActions(np.ndarray): + def __new__(cls, frontend, token, pointer): + count = frontend.action_steps * frontend.action_dim + buffer = (ctypes.c_float * count).from_address(pointer) + view = np.ctypeslib.as_array(buffer).reshape( + frontend.action_steps, frontend.action_dim).view(cls) + view._frontend = frontend + view._token = token + view._pointer = pointer + view._closed = False + view._owns_mapping = True + view._owner = frontend._model.owner + view._release = frontend._model.release + frontend._model.retain(view._owner) + view.setflags(write=False) + return view + + def __array_finalize__(self, source): + if source is None: + return + self._frontend = getattr(source, "_frontend", None) + self._token = getattr(source, "_token", None) + self._pointer = getattr(source, "_pointer", None) + self._closed = True + self._owns_mapping = False + self._owner = None + self._release = None + + def close(self): + if self._closed or not self._owns_mapping: + return + rc = self._token.verbs.contents.unmap_host( + self._token.handle, self._pointer) + if rc != 0: + raise RuntimeError(f"unmap_host actions failed (rc={rc})") + self._closed = True + self._release(self._owner) + + def __dlpack__(self, *, stream=None, max_version=None, + dl_device=None, copy=None): + if self._closed: + raise RuntimeError("cannot export a closed actions host view") + dlpack_view = self._frontend._map_actions() + return np.ndarray.__dlpack__( + dlpack_view, stream=stream, max_version=max_version, + dl_device=dl_device, copy=copy) + + def __del__(self): + try: + self.close() + except Exception: + pass + + +# 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_v2 abi_version (model_runtime.h: FRT_MODEL_RUNTIME_ABI_VERSION_V2) +_FRT_MODEL_RUNTIME_ABI_VERSION_V2 = 2 +_F32_BYTES = np.dtype(np.float32).itemsize + + +def _find_lib(lib_path, env_var="FLASHRT_PI0_LIB"): + """Resolve the SHARED provider .so path. + + Priority: explicit ``lib_path`` kwarg > ``env_var`` env > build-dir + convention. An explicit ``lib_path`` that does not exist is a hard error + (no silent fallback to env/build-dirs) so callers get deterministic loads. + """ + 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 + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.dirname(os.path.dirname(os.path.dirname(here))) + candidates = [ + os.path.join(repo, "cpp", "build-jetson-pi", + "libflashrt_cpp_llama_cpp_provider_c.so"), + os.path.join(repo, "cpp", "build-container", + "libflashrt_cpp_llama_cpp_provider_c.so"), + os.path.join(repo, "cpp", "build", + "libflashrt_cpp_llama_cpp_provider_c.so"), + ] + for c in candidates: + if os.path.exists(c): + return c + raise RuntimeError( + "libflashrt_cpp_llama_cpp_provider_c.so not found. Build it with " + "-DFLASHRT_CPP_WITH_JETSON_PI=ON, or pass lib_path=, or set " + f"{env_var}.") + + +# frt_llama_cpp_engine_factory_v1 layout: +# uint32 struct_size; uint32 reserved; void* self; +# int (*create_pi0)(void*, const cfg*, engine*); +# int (*create_llm)(void*, const cfg*, engine*); // Phase 3, between pi0 and last_error +# const char* (*last_error)(void*); +# self is nullptr for the default factory; last_error ignores it (returns the +# thread-local create-error sink). Defined at module scope so it is not +# re-created per instance. Field order MUST match c_api.h exactly — ctypes +# reads by offset, so a missing field shifts every later field. +class _FactoryV1(ctypes.Structure): + _fields_ = [ + ("struct_size", ctypes.c_uint32), + ("reserved", ctypes.c_uint32), + ("self", ctypes.c_void_p), + ("create_pi0", ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p)), + ("create_llm", ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p)), + ("create_mllm", ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p)), + ("last_error", ctypes.CFUNCTYPE(ctypes.c_char_p, ctypes.c_void_p)), + ] + + +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 (e.g. 10x32 " + "for pi0_base, 50x32 for pi0_libero_base)") + self.num_views = int(num_views) + 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) + + # C ABI signatures. + self._lib.frt_llama_cpp_default_engine_factory.argtypes = [] + self._lib.frt_llama_cpp_default_engine_factory.restype = ctypes.c_void_p + try: + self._lib.frt_llama_cpp_runtime_open_error.argtypes = [] + self._lib.frt_llama_cpp_runtime_open_error.restype = ctypes.c_char_p + except AttributeError as exc: + raise RuntimeError( + "provider .so is older than the Python frontend expects: " + "missing frt_llama_cpp_runtime_open_error") from exc + + # frt_llama_cpp_engine_factory_v1 { struct_size, reserved, self, + # create_pi0(self, config*, engine*) -> int, last_error(self) -> char* } + # We only call create_pi0 and last_error through the returned pointer. + self._lib.frt_llama_cpp_pi0_runtime_open_with_engine_factory.argtypes = [ + ctypes.c_char_p, # config_json + ctypes.c_void_p, # factory* + ctypes.POINTER(ctypes.c_void_p), # out model** + ] + self._lib.frt_llama_cpp_pi0_runtime_open_with_engine_factory.restype = ( + ctypes.c_int) + + config = { + "model_family": "pi0", + "model_path": str(checkpoint), + "mmproj_path": str(mmproj_path), + "backend": backend, + "n_views": int(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") + + factory_ptr = self._lib.frt_llama_cpp_default_engine_factory() + if not factory_ptr: + raise RuntimeError( + "frt_llama_cpp_default_engine_factory returned NULL " + "(FlashRT built without FLASHRT_CPP_WITH_JETSON_PI?)") + + factory = _FactoryV1.from_address(factory_ptr) + if factory.struct_size < ctypes.sizeof(_FactoryV1): + raise RuntimeError( + f"frt_llama_cpp_engine_factory_v1 struct_size=" + f"{factory.struct_size} < ctypes sizeof " + f"{ctypes.sizeof(_FactoryV1)}; provider .so is older than the " + f"Python frontend expects.") + + model_ptr = ctypes.c_void_p(0) + rc = self._lib.frt_llama_cpp_pi0_runtime_open_with_engine_factory( + config_json, factory_ptr, ctypes.byref(model_ptr)) + if rc != 0 or not model_ptr.value: + err = self._lib.frt_llama_cpp_runtime_open_error() or b"" + if not err: + err = factory.last_error(factory.self) or b"" + raise RuntimeError( + f"frt_llama_cpp_pi0_runtime_open_with_engine_factory failed " + f"(rc={rc}): {(err.decode(errors='replace') if err else 'no error')}") + + self._model = FrtModelRuntimeV2.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_V2: + raise RuntimeError( + f"frt_model_runtime_v2 abi_version={self._model.abi_version}, " + f"expected {_FRT_MODEL_RUNTIME_ABI_VERSION_V2}; the provider " + f".so was built against a different FlashRT runtime ABI.") + if self._model.struct_size < ctypes.sizeof(FrtModelRuntimeV2): + raise RuntimeError( + f"frt_model_runtime_v2 struct_size={self._model.struct_size} " + f"< ctypes sizeof {ctypes.sizeof(FrtModelRuntimeV2)}; 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 > self.action_dim: + raise ValueError( + f"state has {state.size} values, more than action_dim={self.action_dim}") + state_padded = np.zeros(self.action_dim, dtype=np.float32) + state_padded[:state.size] = state + + v2 = self._model.verbs_v2 + self_ = self._model.self + + rc = v2.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 = v2.set_input(self_, PORT_PROMPT, self._prompt, + len(self._prompt), -1) + self._check(rc, "set_input prompt") + rc = v2.set_input(self_, PORT_STATE, state_padded.ctypes.data, + state_padded.nbytes, -1) + self._check(rc, "set_input state") + + rc = v2.run_stage(self_, STAGE_INFER, -1) + self._check(rc, "run_stage infer") + + capacity = self.action_steps * self.action_dim * _F32_BYTES + out = (ctypes.c_char * capacity)() + written = ctypes.c_uint64(0) + rc = v2.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 > self.action_dim: + raise ValueError( + f"state has {state.size} values, more than action_dim={self.action_dim}") + state_padded = np.zeros(self.action_dim, dtype=np.float32) + state_padded[:state.size] = state + views = self._make_image_views(images) + v2 = self._model.verbs_v2 + self_ = self._model.self + self._check(v2.set_input( + self_, PORT_IMAGES, ctypes.cast(views, ctypes.c_void_p), + ctypes.sizeof(FrtImageView) * len(images), -1), + "set_input images") + self._check(v2.set_input( + self_, PORT_PROMPT, self._prompt, len(self._prompt), -1), + "set_input prompt") + self._check(v2.set_input( + self_, PORT_STATE, state_padded.ctypes.data, + state_padded.nbytes, -1), "set_input state") + self._check(v2.run_stage(self_, STAGE_CONTEXT, -1), + "run_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") + v2 = self._model.verbs_v2 + self._check(v2.run_stage(self._model.self, STAGE_ACTION, -1), + "run_stage action") + capacity = self.action_steps * self.action_dim * _F32_BYTES + out = (ctypes.c_char * capacity)() + written = ctypes.c_uint64(0) + self._check(v2.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 action_view(self): + """Map the latest actions as a read-only zero-copy NumPy/DLPack view. + + The returned ndarray implements ``__dlpack__``. Close it before + setting new inputs or running another stage; those operations reject + while a host view is mapped so the backing pointer cannot go stale. + """ + return self._map_actions() + + def _map_actions(self): + if self._model is None: + raise RuntimeError("Pi0JetsonPiFrontend is closed") + if self._model.n_port_tokens != self._model.n_ports: + raise RuntimeError("provider runtime has invalid port token layout") + token = self._model.port_tokens[PORT_ACTIONS] + if not token.handle or not token.verbs: + raise RuntimeError("actions port does not expose a memory token") + verbs = token.verbs.contents + if verbs.struct_size < ctypes.sizeof(_FrtMemoryTokenVerbs): + raise RuntimeError("actions token does not support host mapping") + pointer = ctypes.c_void_p() + rc = verbs.map_host( + token.handle, token.offset, token.bytes, 1, ctypes.byref(pointer)) + if rc != 0 or not pointer.value: + raise RuntimeError(f"map_host actions failed (rc={rc})") + return _MappedActions(self, token, pointer.value) + + 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_v2.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/runtime/export.py b/flash_rt/runtime/export.py index 05437dc0..164b2e6e 100644 --- a/flash_rt/runtime/export.py +++ b/flash_rt/runtime/export.py @@ -83,6 +83,10 @@ def _import_native(): "swap": int(_c.PORT_SWAP), "staged": int(_c.PORT_STAGED), "setup": int(_c.PORT_SETUP), } +STAGE_KIND = { + "graph": int(_c.STAGE_GRAPH), + "callback": int(_c.STAGE_CALLBACK), +} def _enum(table: Mapping[str, int], value: int | str) -> int: @@ -166,6 +170,15 @@ class StageSpec: after: Sequence[int] = () +@dataclass +class CallbackStageSpec: + """One provider-owned v2 callback stage.""" + + name: str + callback: int + after: Sequence[int] = () + + @dataclass class RuntimeExport: """A finished export. ``ptr`` is the ``frt_runtime_export_v1*`` to hand to a @@ -215,6 +228,29 @@ def release(self) -> None: self.ptr = 0 +@dataclass +class ModelRuntimeV2: + """A backend-neutral model runtime with explicit v2 stages.""" + + ptr: int + export_ptr: int + fingerprint: int + identity: str + manifest: str | None + _anchor: Any = field(repr=False, default=None) + + def ports(self) -> list: + return list(_c.model_v2_ports(self.ptr)) + + def stages(self) -> list: + return list(_c.model_stages_v2(self.ptr)) + + def release(self) -> None: + if self.ptr: + _c.model_v2_release(self.ptr) + self.ptr = 0 + + class _Anchor: """Keeps the exec-binding wrappers (Ctx/Graph/Buffer) and the producer's owner object alive for the lifetime of the export. This is the object the @@ -307,6 +343,77 @@ def build_model_runtime( ) +def build_provider_model_runtime_v2( + *, + ports: Sequence[PortSpec] = (), + stages: Sequence[CallbackStageSpec], + identity: Mapping[str, str], + manifest_extra: Mapping[str, Any] | None = None, + owner: Any = None, + set_input=None, + get_output=None, + prepare=None, + step=None, + run_stage=None, +) -> ModelRuntimeV2: + """Assemble a provider-owned model runtime. + + This path declares staged ports and callback stages only; it does not + create or require a FlashRT exec context, buffers, streams, or graphs. + """ + if run_stage is None: + raise ValueError("provider-owned v2 runtimes require run_stage") + if not stages: + raise ValueError("provider-owned v2 runtimes require at least one stage") + b = _c.Builder.provider_owned() + for p in ports: + if _enum(UPDATE, p.update) != UPDATE["staged"]: + raise ValueError("provider-owned ports must use staged updates") + if p.buffer is not None: + raise ValueError("provider-owned ports must not expose SWAP buffers") + if p.offset or p.nbytes is not None: + raise ValueError("provider-owned ports must not declare raw windows") + b.add_port(p.name, _enum(MODALITY, p.modality), _enum(DTYPE, p.dtype), + _enum(LAYOUT, p.layout), _enum(DIRECTION, p.direction), + _enum(UPDATE, p.update), int(bool(p.required)), + list(p.shape), p.cadence_hz, 0, 0, 0) + for st in stages: + b.add_callback_stage_v2(st.name, st.callback, list(st.after)) + for k, v in identity.items(): + b.add_identity(str(k), str(v)) + + manifest = { + "ports": [ + {"name": p.name, "modality": _enum(MODALITY, p.modality), + "dtype": _enum(DTYPE, p.dtype), "layout": _enum(LAYOUT, p.layout), + "direction": _enum(DIRECTION, p.direction), + "update": _enum(UPDATE, p.update), "required": bool(p.required), + "shape": list(p.shape), "cadence_hz": p.cadence_hz} + for p in ports], + "stages_v2": [ + {"name": st.name, "kind": STAGE_KIND["callback"], + "callback": st.callback, "after": list(st.after)} + for st in stages], + } + if manifest_extra: + manifest.update(dict(manifest_extra)) + manifest_json = json.dumps(manifest, sort_keys=True) + b.set_manifest(manifest_json) + + anchor = _Anchor([owner, (set_input, get_output, prepare, step, run_stage)]) + ptr = b.finish_model_v2(anchor, set_input=set_input, get_output=get_output, + prepare=prepare, step=step, run_stage=run_stage) + export_ptr = int(_c.model_v2_export_ptr(ptr)) + return ModelRuntimeV2( + ptr=ptr, + export_ptr=export_ptr, + fingerprint=int(_c.export_fingerprint(export_ptr)), + identity=_c.export_identity(export_ptr), + manifest=manifest_json, + _anchor=anchor, + ) + + def _assemble(ctx, *, streams, graphs, buffers, regions, ports, stages, identity, manifest_extra, owner): if not streams: @@ -379,11 +486,11 @@ def _assemble(ctx, *, streams, graphs, buffers, regions, ports, stages, __all__ = [ - "RuntimeExport", "ModelRuntime", + "RuntimeExport", "ModelRuntime", "ModelRuntimeV2", "StreamSpec", "GraphSpec", "BufferSpec", "RegionSpec", "PortSpec", - "StageSpec", - "build_export", "build_model_runtime", + "StageSpec", "CallbackStageSpec", + "build_export", "build_model_runtime", "build_provider_model_runtime_v2", "ROLE_INPUT", "ROLE_OUTPUT", "ROLE_STATE", "ROLE_SCRATCH", "REGION_SNAPSHOT", "REGION_RESTORE", "REGION_DEFAULT", - "MODALITY", "DTYPE", "LAYOUT", "DIRECTION", "UPDATE", + "MODALITY", "DTYPE", "LAYOUT", "DIRECTION", "UPDATE", "STAGE_KIND", ] 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..2afff526 --- /dev/null +++ b/flash_rt/tests/test_jetson_pi_pi0_python.py @@ -0,0 +1,151 @@ +"""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 + FLASHRT_PI0_LIB (optional) path to libflashrt_cpp_llama_cpp_provider_c.so + FLASHRT_PI0_ACTION_STEPS (optional) override; default 10 (pi0_base). + FLASHRT_PI0_ACTION_DIM (optional) override; default 32. + FLASHRT_PI0_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 n_gpu_layers=9999 + and use_gpu for the mmproj). CUDA_VISIBLE_DEVICES + selects the physical card. + +Run from the repo root: + FLASHRT_PI0_MODEL=... FLASHRT_PI0_MMPROJ=... FLASHRT_PI0_FIXTURE_DIR=/tmp/pi0_fixture \ + LD_LIBRARY_PATH=.../miniconda3/lib:FlashRT/cpp/build-jetson-pi \ + python flash_rt/tests/test_jetson_pi_pi0_python.py +""" + +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") + 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") + 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(os.environ.get("FLASHRT_PI0_ACTION_STEPS", "10")) + action_dim = int(os.environ.get("FLASHRT_PI0_ACTION_DIM", "32")) + # Default "cpu" keeps the original behavior; set FLASHRT_PI0_BACKEND=cuda to + # exercise the real GPU forward pass through the Jetson-PI engine. + backend = os.environ.get("FLASHRT_PI0_BACKEND", "cpu") or "cpu" + + 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 != action_dim: + print(f"FAIL: state.bin has {state.size} floats, expected {action_dim}") + 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") + + actions_view = model._pipe.action_view() + check(not actions_view.flags.writeable, + "actions host SWAP view is read-only") + check(np.shares_memory(actions_view, np.asarray(actions_view)), + "actions host SWAP view is zero-copy") + check(np.array_equal(actions_view, actions), + "actions host SWAP view matches copied output") + dlpack_actions = np.from_dlpack(actions_view) + check(not dlpack_actions.flags.writeable and + dlpack_actions.ctypes.data == actions_view.ctypes.data and + np.array_equal(dlpack_actions, actions), + "NumPy consumes actions DLPack zero-copy and read-only") + actions_view.close() + try: + model._pipe.set_prompt(prompt) + model._pipe.context({"images": [image, wrist], "state": state}) + check(False, "live DLPack consumer retains its host mapping") + except RuntimeError as exc: + check("mapped" in str(exc), + "live DLPack consumer retains its host mapping") + del dlpack_actions + + try: + import torch + except ImportError: + print("SKIP - PyTorch DLPack legacy-protocol rejection check") + torch_legacy_checked = False + else: + torch_legacy_checked = True + torch_view = model._pipe.action_view() + try: + torch.from_dlpack(torch_view) + check(False, "PyTorch legacy DLPack cannot erase read-only access") + except BufferError as exc: + check("readonly" in str(exc), + "PyTorch legacy DLPack cannot erase read-only access") + torch_view.close() + + model._pipe.context({"images": [image, wrist], "state": state}) + if torch_legacy_checked: + check(True, "failed legacy DLPack export releases its host mapping") + 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/runtime/bindings/runtime_pybind.cpp b/runtime/bindings/runtime_pybind.cpp index d73fca3a..41600069 100644 --- a/runtime/bindings/runtime_pybind.cpp +++ b/runtime/bindings/runtime_pybind.cpp @@ -37,7 +37,7 @@ void release_py_owner(void* owner) { * the GIL (a native consumer calls these fn pointers from its own threads) * and translate exceptions into negative status + last_error. */ struct PyVerbs { - py::object set_input, get_output, prepare, step; + py::object set_input, get_output, prepare, step, run_stage; std::string last_error; py::object owner; /* the producer object the export anchors */ }; @@ -119,6 +119,21 @@ int verb_step(void* self) { } } +int verb_run_stage(void* self, uint32_t stage, int stream) { + auto* v = static_cast(self); + py::gil_scoped_acquire gil; + if (!v->run_stage || v->run_stage.is_none()) { + v->last_error = "run_stage is not provided by this producer"; + return -3; + } + try { + return py::cast(v->run_stage(stage, stream)); + } catch (const std::exception& e) { + v->last_error = e.what(); + return -1; + } +} + const char* verb_last_error(void* self) { return static_cast(self)->last_error.c_str(); } @@ -134,6 +149,13 @@ struct PyRtBuilder { b = frt_runtime_builder_create(reinterpret_cast(ctx_raw)); if (!b) throw std::runtime_error("frt_runtime_builder_create failed (null ctx?)"); } + explicit PyRtBuilder(bool provider_owned) { + if (!provider_owned) { + throw std::runtime_error("internal error: provider-owned flag is required"); + } + b = frt_runtime_builder_create_provider_owned(); + if (!b) throw std::runtime_error("frt_runtime_builder_create_provider_owned failed"); + } ~PyRtBuilder() { /* A never-finished builder leaks its holder by design tradeoff: the * builder is consumed by finish(); reaching here without finish() is @@ -175,8 +197,42 @@ struct PyRtBuilder { frt_model_runtime_v1* mr = frt_runtime_builder_finish_model( b, &verbs, pv, pv, /*retain_owner=*/nullptr, &release_py_verbs); + if (!mr) { + release_py_verbs(pv); + throw std::runtime_error("finish_model failed"); + } + b = nullptr; + return reinterpret_cast(mr); + } + + std::uintptr_t finish_model_v2(py::object owner, py::object set_input, + py::object get_output, py::object prepare, + py::object step, py::object run_stage) { + need(); + auto* pv = new PyVerbs(); + pv->set_input = std::move(set_input); + pv->get_output = std::move(get_output); + pv->prepare = std::move(prepare); + pv->step = std::move(step); + pv->run_stage = std::move(run_stage); + pv->owner = std::move(owner); + + frt_model_runtime_verbs_v2 verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.set_input = &verb_set_input; + verbs.get_output = &verb_get_output; + verbs.prepare = &verb_prepare; + verbs.step = &verb_step; + verbs.last_error = &verb_last_error; + verbs.run_stage = &verb_run_stage; + + frt_model_runtime_v2* mr = frt_runtime_builder_finish_model_v2( + b, &verbs, pv, pv, /*retain_owner=*/nullptr, &release_py_verbs); + if (!mr) { + release_py_verbs(pv); + throw std::runtime_error("finish_model_v2 failed"); + } b = nullptr; - if (!mr) { release_py_verbs(pv); throw std::runtime_error("finish_model failed"); } return reinterpret_cast(mr); } }; @@ -184,7 +240,7 @@ struct PyRtBuilder { frt_runtime_export_v1* as_export(std::uintptr_t p) { auto* e = reinterpret_cast(p); if (!e || e->abi_version != FRT_RUNTIME_ABI_VERSION || - e->struct_size != sizeof(frt_runtime_export_v1)) + e->struct_size < sizeof(frt_runtime_export_v1)) throw std::runtime_error("not a valid frt_runtime_export_v1 pointer"); return e; } @@ -192,11 +248,26 @@ frt_runtime_export_v1* as_export(std::uintptr_t p) { frt_model_runtime_v1* as_model(std::uintptr_t p) { auto* m = reinterpret_cast(p); if (!m || m->abi_version != FRT_MODEL_RUNTIME_ABI_VERSION || - m->struct_size != sizeof(frt_model_runtime_v1)) + m->struct_size < sizeof(frt_model_runtime_v1)) throw std::runtime_error("not a valid frt_model_runtime_v1 pointer"); return m; } +frt_model_runtime_v2* as_model_v2(std::uintptr_t p) { + auto* m = reinterpret_cast(p); + /* Append-only: accept any producer whose struct_size is AT LEAST the v2 + * size we were compiled against. A producer built with a newer header + * (larger struct_size, e.g. with Phase 6 port_tokens) must still be + * drivable; a producer built with this header matches exactly. The old + * strict `!=` gate would reject a newer producer and break the + * append-only contract — see frt_model_runtime_wrap (model_runtime.cpp) + * which already uses `>=`. */ + if (!m || m->abi_version != FRT_MODEL_RUNTIME_ABI_VERSION_V2 || + m->struct_size < sizeof(frt_model_runtime_v2)) + throw std::runtime_error("not a valid frt_model_runtime_v2 pointer"); + return m; +} + } // namespace PYBIND11_MODULE(_flashrt_runtime, m) { @@ -211,6 +282,7 @@ PYBIND11_MODULE(_flashrt_runtime, m) { m.attr("REGION_RESTORE") = (unsigned)FRT_RT_REGION_RESTORE; m.attr("MODEL_ABI_VERSION") = FRT_MODEL_RUNTIME_ABI_VERSION; + m.attr("MODEL_ABI_VERSION_V2") = FRT_MODEL_RUNTIME_ABI_VERSION_V2; m.attr("MOD_TENSOR") = (unsigned)FRT_RT_MOD_TENSOR; m.attr("MOD_IMAGE") = (unsigned)FRT_RT_MOD_IMAGE; m.attr("MOD_TEXT") = (unsigned)FRT_RT_MOD_TEXT; @@ -235,9 +307,14 @@ PYBIND11_MODULE(_flashrt_runtime, m) { m.attr("PORT_SWAP") = (unsigned)FRT_RT_PORT_SWAP; m.attr("PORT_STAGED") = (unsigned)FRT_RT_PORT_STAGED; m.attr("PORT_SETUP") = (unsigned)FRT_RT_PORT_SETUP; + m.attr("STAGE_GRAPH") = (unsigned)FRT_RT_STAGE_GRAPH; + m.attr("STAGE_CALLBACK") = (unsigned)FRT_RT_STAGE_CALLBACK; py::class_(m, "Builder") .def(py::init(), py::arg("ctx_raw")) + .def_static("provider_owned", []() { + return PyRtBuilder(true); + }) .def("add_stream", [](PyRtBuilder& s, const std::string& name, int stream_id, int priority, std::uintptr_t native_handle) { s.need(); @@ -309,6 +386,26 @@ PYBIND11_MODULE(_flashrt_runtime, m) { (uint32_t)after.size()), "add_stage"); }, py::arg("graph"), py::arg("after") = std::vector{}) + .def("add_graph_stage_v2", [](PyRtBuilder& s, const std::string& name, + unsigned graph, + const std::vector& after) { + s.need(); + check(frt_runtime_builder_add_graph_stage_v2( + s.b, name.c_str(), graph, after.data(), + (uint32_t)after.size()), + "add_graph_stage_v2"); + }, py::arg("name"), py::arg("graph"), + py::arg("after") = std::vector{}) + .def("add_callback_stage_v2", [](PyRtBuilder& s, const std::string& name, + unsigned callback, + const std::vector& after) { + s.need(); + check(frt_runtime_builder_add_callback_stage_v2( + s.b, name.c_str(), callback, after.data(), + (uint32_t)after.size()), + "add_callback_stage_v2"); + }, py::arg("name"), py::arg("callback"), + py::arg("after") = std::vector{}) .def("finish", &PyRtBuilder::finish, py::arg("owner"), "Consume the builder; returns the export pointer (uintptr). The export " "holds one reference; hand the pointer to a native consumer, which must " @@ -321,7 +418,15 @@ PYBIND11_MODULE(_flashrt_runtime, m) { "GIL, so a native consumer may call them from any thread. " "set_input(port, payload: bytes, stream) -> int; " "get_output(port, stream) -> bytes; prepare(graph, key) -> int; " - "step() -> int."); + "step() -> int.") + .def("finish_model_v2", &PyRtBuilder::finish_model_v2, + py::arg("owner"), + py::arg("set_input") = py::none(), py::arg("get_output") = py::none(), + py::arg("prepare") = py::none(), py::arg("step") = py::none(), + py::arg("run_stage") = py::none(), + "Consume the builder; returns the frt_model_runtime_v2 pointer " + "(uintptr). run_stage(stage, stream) executes provider-owned " + "callback stages."); /* Introspection over a raw export pointer (tests / mismatch tooling). */ m.def("export_fingerprint", [](std::uintptr_t p) { return as_export(p)->fingerprint; }); @@ -354,6 +459,9 @@ PYBIND11_MODULE(_flashrt_runtime, m) { m.def("model_export_ptr", [](std::uintptr_t p) { return reinterpret_cast(as_model(p)->exp); }); + m.def("model_v2_export_ptr", [](std::uintptr_t p) { + return reinterpret_cast(as_model_v2(p)->exp); + }); m.def("model_ports", [](std::uintptr_t p) { auto* mr = as_model(p); py::list out; @@ -384,12 +492,51 @@ PYBIND11_MODULE(_flashrt_runtime, m) { } return out; }); + m.def("model_stages_v2", [](std::uintptr_t p) { + auto* mr = as_model_v2(p); + py::list out; + for (std::uint64_t i = 0; i < mr->n_stages_v2; ++i) { + const frt_runtime_stage_desc_v2& d = mr->stages_v2[i]; + py::dict e; + e["name"] = std::string(d.name); + e["kind"] = d.kind; + e["graph"] = d.graph; + e["callback"] = d.callback; + e["after"] = std::vector(d.after, d.after + d.n_after); + out.append(e); + } + return out; + }); + m.def("model_v2_ports", [](std::uintptr_t p) { + auto* mr = as_model_v2(p); + py::list out; + for (std::uint64_t i = 0; i < mr->n_ports; ++i) { + const frt_runtime_port_desc& d = mr->ports[i]; + py::dict e; + e["name"] = std::string(d.name); + e["modality"] = d.modality; e["dtype"] = d.dtype; + e["layout"] = d.layout; e["direction"] = d.direction; + e["update"] = d.update; e["required"] = d.required; + e["shape"] = std::vector(d.shape, d.shape + d.rank); + e["cadence_hint_hz"] = d.cadence_hint_hz; + e["buffer"] = reinterpret_cast(d.buffer); + e["offset"] = d.offset; e["bytes"] = d.bytes; + out.append(e); + } + return out; + }); m.def("model_retain", [](std::uintptr_t p) { auto* mr = as_model(p); mr->retain(mr->owner); }); m.def("model_release", [](std::uintptr_t p) { auto* mr = as_model(p); py::gil_scoped_release nogil; /* release re-acquires internally */ mr->release(mr->owner); }); + m.def("model_v2_retain", [](std::uintptr_t p) { auto* mr = as_model_v2(p); mr->retain(mr->owner); }); + m.def("model_v2_release", [](std::uintptr_t p) { + auto* mr = as_model_v2(p); + py::gil_scoped_release nogil; + mr->release(mr->owner); + }); /* Drive the verbs THROUGH the C fn pointers (tests exercise the same * entry a native consumer uses). */ m.def("model_set_input", [](std::uintptr_t p, unsigned port, py::bytes data, @@ -418,8 +565,47 @@ PYBIND11_MODULE(_flashrt_runtime, m) { py::gil_scoped_release nogil; return mr->verbs.step(mr->self); }); + m.def("model_v2_step", [](std::uintptr_t p) { + auto* mr = as_model_v2(p); + py::gil_scoped_release nogil; + return mr->verbs_v2.step(mr->self); + }); + m.def("model_v2_set_input", [](std::uintptr_t p, unsigned port, + py::bytes data, int stream) { + auto* mr = as_model_v2(p); + std::string s = data; + py::gil_scoped_release nogil; + return mr->verbs_v2.set_input(mr->self, port, s.data(), s.size(), + stream); + }, py::arg("ptr"), py::arg("port"), py::arg("data"), + py::arg("stream") = -1); + m.def("model_v2_get_output", [](std::uintptr_t p, unsigned port, + std::uint64_t capacity, int stream) { + auto* mr = as_model_v2(p); + std::string buf(capacity, '\0'); + std::uint64_t written = 0; + int rc; + { + py::gil_scoped_release nogil; + rc = mr->verbs_v2.get_output(mr->self, port, buf.data(), capacity, + &written, stream); + } + return py::make_tuple(rc, py::bytes(buf.data(), + rc == 0 ? written : 0), written); + }, py::arg("ptr"), py::arg("port"), py::arg("capacity"), + py::arg("stream") = -1); + m.def("model_v2_run_stage", [](std::uintptr_t p, unsigned stage, + int stream) { + auto* mr = as_model_v2(p); + py::gil_scoped_release nogil; + return mr->verbs_v2.run_stage(mr->self, stage, stream); + }, py::arg("ptr"), py::arg("stage"), py::arg("stream") = -1); m.def("model_last_error", [](std::uintptr_t p) { auto* mr = as_model(p); return std::string(mr->verbs.last_error(mr->self)); }); + m.def("model_v2_last_error", [](std::uintptr_t p) { + auto* mr = as_model_v2(p); + return std::string(mr->verbs_v2.last_error(mr->self)); + }); } diff --git a/runtime/include/flashrt/model_runtime.h b/runtime/include/flashrt/model_runtime.h index 51dbe689..631898c1 100644 --- a/runtime/include/flashrt/model_runtime.h +++ b/runtime/include/flashrt/model_runtime.h @@ -56,6 +56,7 @@ extern "C" { #endif #define FRT_MODEL_RUNTIME_ABI_VERSION 1u +#define FRT_MODEL_RUNTIME_ABI_VERSION_V2 2u /* ------------------------------------------------------------------ */ /* Enums — values are ABI-frozen after v1 (append-only). */ @@ -105,6 +106,26 @@ enum frt_rt_port_update { FRT_RT_PORT_SETUP = 2 }; +enum frt_rt_stage_kind { + FRT_RT_STAGE_GRAPH = 0, + FRT_RT_STAGE_CALLBACK = 1 +}; + +/* Phase 6 — memory-domain contract. Where a provider-owned buffer lives, so + * FlashRT can reason about access (host-readable now, device-resident later) + * WITHOUT ever dereferencing the provider's memory or leaking a backend type. + * Values are ABI-frozen after v1 (append-only): new location kinds get higher + * integers. */ +enum frt_rt_location_kind { + FRT_RT_LOCATION_HOST_VISIBLE = 0, /* copy_to_host is a plain read; no sync */ + FRT_RT_LOCATION_DEVICE_LOCAL = 1 /* provider-owned device memory; sync first */ +}; + +enum frt_rt_host_map_access { + FRT_RT_HOST_MAP_READ = 1, + FRT_RT_HOST_MAP_WRITE = 2 +}; + /* ------------------------------------------------------------------ */ /* Payload types (STAGED lane). */ /* ------------------------------------------------------------------ */ @@ -149,6 +170,93 @@ typedef struct frt_runtime_port_desc { uint64_t offset, bytes; } frt_runtime_port_desc; +/* ------------------------------------------------------------------ */ +/* Phase 6 — memory-domain token (provider-owned buffer contract). */ +/* */ +/* A provider-owned runtime cannot use the CUDA-export SWAP window */ +/* (frt_buffer + offset + bytes) — that implies same-backend device */ +/* sharing with no contract. Instead, a provider-owned port MAY carry */ +/* an opaque memory token the provider mints, plus a small verb set */ +/* FlashRT calls to copy/sync/destroy. FlashRT NEVER dereferences the */ +/* token: zero-copy is an ADVERTISED capability (location_kind = */ +/* HOST_VISIBLE), not an assumption. This is the honest contract that */ +/* lets a future provider optionally expose host-visible or device- */ +/* local buffers through provider-owned ports WITHOUT faking a CUDA */ +/* graph or leaking GGML/CUDA types (see docs/phase6_backend_vtable_ */ +/* eval.md). */ +/* ------------------------------------------------------------------ */ +typedef struct frt_memory_token_s* frt_memory_token; + +/* Provider-supplied verbs against its own backing store. The original + * copy/sync prefix is mandatory and always callable. Append-only tail verbs + * such as map_host/unmap_host must be struct_size-probed before use. + * `destroy` is the only void entry; FlashRT calls it exactly + * once when the runtime's holder refcount hits zero (on the provider-owned + * v2 path the holder IS the token's lifetime owner — there is no separate + * per-port/per-export counter, and no v2 wrap/override path exists, so the + * holder refcount is the single relevant count). */ +typedef struct frt_memory_token_verbs { + uint32_t struct_size; /* = sizeof(frt_memory_token_verbs) */ + uint32_t reserved; + + /* Copy `bytes` from the token's backing store (at `src_off`) into a + * host buffer `dst` (at `dst_off`). Returns 0 on success. */ + int (*copy_to_host)(frt_memory_token token, void* dst, + uint64_t dst_off, uint64_t src_off, uint64_t bytes); + + /* Copy `bytes` from a host buffer `src` (at `src_off`) into the token's + * backing store (at `dst_off`). Returns 0 on success. */ + int (*copy_from_host)(frt_memory_token token, const void* src, + uint64_t src_off, uint64_t dst_off, uint64_t bytes); + + /* Block until the token's backing store is safe to read/write from the + * host (no-op for HOST_VISIBLE; required for DEVICE_LOCAL). Returns 0. */ + int (*sync)(frt_memory_token token); + + /* Release the provider's backing store. Called exactly once when the + * holder's refcount hits zero. May be null if the provider owns the + * store externally (then FlashRT never destroys it). */ + void (*destroy)(frt_memory_token token); + + /* Append-only host SWAP window. map_host returns a provider-owned host + * pointer for [offset, offset + bytes); the caller must keep the runtime + * alive and call unmap_host exactly once before any operation that can + * mutate/reallocate the backing store. Access is frt_rt_host_map_access. + * Providers compiled against the original prefix omit these fields. */ + int (*map_host)(frt_memory_token token, uint64_t offset, uint64_t bytes, + uint32_t access, void** out_ptr); + int (*unmap_host)(frt_memory_token token, void* ptr); +} frt_memory_token_verbs; + +#define FRT_MEMORY_TOKEN_VERBS_COPY_SYNC_SIZE \ + ((uint32_t)offsetof(frt_memory_token_verbs, map_host)) + +/* A token bound to one provider-owned port. The provider retains the + * backing store until `verbs.destroy` fires; FlashRT's only job is to fire + * it once at runtime release. `bytes` is the logical size of the window + * (not necessarily the whole backing store); `offset` is provider-relative. */ +typedef struct frt_memory_token_desc { + uint32_t struct_size; /* = sizeof(frt_memory_token_desc) */ + frt_memory_token handle; /* opaque provider handle */ + const frt_memory_token_verbs* verbs; /* borrowed for the port's life */ + uint64_t offset; /* provider-relative window offset */ + uint64_t bytes; /* logical window size */ + uint32_t location_kind; /* enum frt_rt_location_kind */ + uint32_t reserved; +} frt_memory_token_desc; + +/* Carrier: how a token reaches a provider-owned port. Rather than a second + * port-descriptor struct probed by struct_size, Phase 6 carries tokens as a + * PARALLEL array `port_tokens`/`n_port_tokens` on frt_model_runtime_v2, + * index-aligned with `ports` (port_tokens[i] corresponds to ports[i]; a port + * added via frt_runtime_builder_add_port gets a null-handle entry, a port + * added via frt_runtime_builder_add_port_token gets the minted token). This + * matches how v2 already carries parallel ports/stages/stages_v2 arrays and + * keeps the v1 frt_runtime_port_desc byte-identical. The eval's named + * frt_runtime_port_desc_v2 tail-probe form is NOT used — the parallel array + * is the chosen, simpler carrier (see docs/phase6_backend_vtable_eval.md §8 + * for the vtable-vs-contract decision; the carrier choice is recorded there). */ + /* One schedulable stage = one export graph + dependency edges. Declared * array order is the sequential firing order `step` uses; `after` lists * stage indices that must complete first (for hosts that overlap stages @@ -159,6 +267,19 @@ typedef struct frt_runtime_stage_desc { const uint32_t* after; /* stage indices */ } frt_runtime_stage_desc; +/* v2 stage descriptors keep graph-backed stages explicit while allowing a + * provider-owned callback executable. Callback indices are private to the + * provider and are executed through frt_model_runtime_verbs_v2::run_stage; + * no GGML/llama.cpp/backend type is exposed in the public ABI. */ +typedef struct frt_runtime_stage_desc_v2 { + const char* name; /* stable stage name: "infer", "decode"... */ + uint32_t kind; /* frt_rt_stage_kind */ + uint32_t graph; /* exp->graphs index, or UINT32_MAX */ + uint32_t callback; /* provider stage id, or UINT32_MAX */ + uint32_t n_after; + const uint32_t* after; /* stage indices */ +} frt_runtime_stage_desc_v2; + /* ------------------------------------------------------------------ */ /* Verbs — implemented by the producer, called by the host. */ /* set_input / get_output are HOT (contract above); prepare is WARM; */ @@ -194,6 +315,25 @@ typedef struct frt_model_runtime_verbs { const char* (*last_error)(void* self); } frt_model_runtime_verbs; +typedef struct frt_model_runtime_verbs_v2 { + uint32_t struct_size; /* = sizeof(frt_model_runtime_verbs_v2) */ + uint32_t reserved; + + int (*set_input)(void* self, uint32_t port, + const void* data, uint64_t bytes, int stream); + int (*get_output)(void* self, uint32_t port, + void* out, uint64_t capacity, uint64_t* written, + int stream); + int (*prepare)(void* self, uint32_t graph, frt_shape_key key); + int (*step)(void* self); + const char* (*last_error)(void* self); + + /* Execute one v2 stage by index. Required for CALLBACK stages; graph + * stages remain directly schedulable through exp->graphs by white-box + * hosts. `stream` = an exp stream id, or -1 for the provider default. */ + int (*run_stage)(void* self, uint32_t stage, int stream); +} frt_model_runtime_verbs_v2; + /* ------------------------------------------------------------------ */ /* The model runtime object. */ /* ------------------------------------------------------------------ */ @@ -219,12 +359,53 @@ typedef struct frt_model_runtime_v1 { void (*release)(void* owner); } frt_model_runtime_v1; +typedef struct frt_model_runtime_v2 { + uint32_t abi_version; /* = FRT_MODEL_RUNTIME_ABI_VERSION_V2 */ + uint32_t struct_size; /* = sizeof(frt_model_runtime_v2) */ + + const frt_runtime_export_v1* exp; + + const frt_runtime_port_desc* ports; uint64_t n_ports; + + /* Legacy graph-only view. Empty when v2 stages include provider-owned + * callbacks, so old graph schedulers do not misinterpret the DAG. */ + const frt_runtime_stage_desc* stages; uint64_t n_stages; + + void* self; + frt_model_runtime_verbs verbs; /* v1-compatible verb subset */ + + void* owner; + void (*retain)(void* owner); + void (*release)(void* owner); + + const frt_runtime_stage_desc_v2* stages_v2; + uint64_t n_stages_v2; + frt_model_runtime_verbs_v2 verbs_v2; + + /* Phase 6 (append-only): memory-domain tokens, a PARALLEL array to + * `ports` by index. n_port_tokens == n_ports ALWAYS: a port added via + * frt_runtime_builder_add_port gets a null-handle entry here, a port + * added via frt_runtime_builder_add_port_token gets the minted token. + * An entry with handle == nullptr means that port carries no token. + * These fields are trailing additions to the v2 struct, so an older + * producer (smaller struct_size) is still accepted by the >= / < probes + * used by frt_model_runtime_wrap and as_model_v2; such a consumer simply + * never reads the tail. FlashRT fires `token.verbs->destroy` exactly + * once when the holder refcount hits zero. */ + const frt_memory_token_desc* port_tokens; + uint64_t n_port_tokens; +} frt_model_runtime_v2; + /* Factory symbol convention for NATIVE model runtimes: a model-runtime .so * exports exactly this symbol. Returns a retained object (caller releases). */ #define FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL "frt_model_runtime_open_v1" typedef int (*frt_model_runtime_open_v1_fn)(const char* config_json, frt_model_runtime_v1** out); +#define FRT_MODEL_RUNTIME_OPEN_V2_SYMBOL "frt_model_runtime_open_v2" +typedef int (*frt_model_runtime_open_v2_fn)(const char* config_json, + frt_model_runtime_v2** out); + /* ------------------------------------------------------------------ */ /* Construction path 1 — INTEGRATED (preferred): the export builder */ /* assembles export + ports + stages in one identity. Port and stage */ @@ -245,6 +426,35 @@ int frt_runtime_builder_add_port(frt_runtime_builder, const char* name, int frt_runtime_builder_add_stage(frt_runtime_builder, uint32_t graph, const uint32_t* after, uint32_t n_after); +/* Add a provider-owned port carrying a Phase 6 memory-domain token in place + * of a CUDA-export SWAP window. `token` is borrowed only for this call (the + * builder copies the descriptor). On a provider-owned runtime this is the + * only permitted buffer form; the builder still rejects a raw frt_buffer. + * The token's `destroy` is fired exactly once when the runtime's refcount + * hits zero. */ +int frt_runtime_builder_add_port_token( + frt_runtime_builder, const char* name, + uint32_t modality, uint32_t dtype, uint32_t layout, uint32_t direction, + uint32_t required, const int64_t* shape, uint32_t rank, + uint32_t cadence_hint_hz, + frt_memory_token handle, const frt_memory_token_verbs* verbs, + uint64_t offset, uint64_t bytes, uint32_t location_kind); + +/* Create a builder for provider-owned model runtimes with no FlashRT exec + * context. It may only be finished through frt_runtime_builder_finish_model_v2. + */ +frt_runtime_builder frt_runtime_builder_create_provider_owned(void); + +int frt_runtime_builder_add_graph_stage_v2(frt_runtime_builder, + const char* name, uint32_t graph, + const uint32_t* after, + uint32_t n_after); +int frt_runtime_builder_add_callback_stage_v2(frt_runtime_builder, + const char* name, + uint32_t callback, + const uint32_t* after, + uint32_t n_after); + /* Like frt_runtime_builder_finish, but returns the model runtime whose * `exp` is the internally-built export (one object, one refcount). `verbs` * is copied; entries may be null (the runtime then reports them @@ -254,6 +464,11 @@ frt_model_runtime_v1* frt_runtime_builder_finish_model( const frt_model_runtime_verbs* verbs, void* verbs_self, void* owner, void (*retain_owner)(void*), void (*release_owner)(void*)); +frt_model_runtime_v2* frt_runtime_builder_finish_model_v2( + frt_runtime_builder, + const frt_model_runtime_verbs_v2* verbs, void* verbs_self, + void* owner, void (*retain_owner)(void*), void (*release_owner)(void*)); + /* ------------------------------------------------------------------ */ /* Construction path 2 — ADAPTER: wrap an EXISTING export with ports */ /* and verbs (e.g. a native C++ model runtime over a Python-built */ diff --git a/runtime/include/flashrt/runtime.h b/runtime/include/flashrt/runtime.h index ee3e81ff..2c435b7a 100644 --- a/runtime/include/flashrt/runtime.h +++ b/runtime/include/flashrt/runtime.h @@ -189,6 +189,7 @@ int frt_runtime_builder_add_region(frt_runtime_builder, const char* name, int frt_runtime_builder_add_identity(frt_runtime_builder, const char* key, const char* value); int frt_runtime_builder_set_manifest(frt_runtime_builder, const char* json); +void frt_runtime_builder_discard(frt_runtime_builder); /* Finish: builds the canonical identity string, computes the fingerprint, * flattens everything into one export object, and consumes the builder diff --git a/runtime/src/internal.h b/runtime/src/internal.h index cda71c8b..84d38a20 100644 --- a/runtime/src/internal.h +++ b/runtime/src/internal.h @@ -40,9 +40,17 @@ struct Holder { std::deque> after_arrays; std::vector ports; std::vector stages; + std::vector stages_v2; + + /* Phase 6 — memory-domain tokens attached to provider-owned ports. + * One entry per port that carries a token (parallel to `ports` by index, + * with token.handle == nullptr for ports without one). Copied verbatim + * from the builder; `verbs.destroy` is fired at release. */ + std::vector port_tokens; frt_runtime_export_v1 exp{}; frt_model_runtime_v1 model{}; + frt_model_runtime_v2 model_v2{}; }; extern "C" void frt_rt_holder_retain(void* owner); @@ -56,6 +64,7 @@ struct frt_runtime_builder_s { frt_ctx ctx = nullptr; frt_rt::Holder* h = nullptr; /* built up in place; adopted by finish */ std::string identity_pairs; + bool provider_owned = false; }; namespace frt_rt { diff --git a/runtime/src/model_runtime.cpp b/runtime/src/model_runtime.cpp index 9e82e706..5643718e 100644 --- a/runtime/src/model_runtime.cpp +++ b/runtime/src/model_runtime.cpp @@ -29,6 +29,7 @@ int stub_get_output(void*, uint32_t, void*, uint64_t, uint64_t*, int) { } int stub_prepare(void*, uint32_t, frt_shape_key) { return -3; } int stub_step(void*) { return -3; } +int stub_run_stage(void*, uint32_t, int) { return -3; } const char* stub_last_error(void*) { return "verb not provided by this producer"; } @@ -51,6 +52,55 @@ void copy_verbs(frt_model_runtime_v1* m, const frt_model_runtime_verbs* verbs, m->self = verbs_self; } +void copy_verbs_v2(frt_model_runtime_v2* m, + const frt_model_runtime_verbs_v2* verbs, + void* verbs_self) { + m->verbs_v2.struct_size = (uint32_t)sizeof(frt_model_runtime_verbs_v2); + if (verbs && verbs->struct_size >= sizeof(frt_model_runtime_verbs_v2)) { + m->verbs_v2.set_input = verbs->set_input; + m->verbs_v2.get_output = verbs->get_output; + m->verbs_v2.prepare = verbs->prepare; + m->verbs_v2.step = verbs->step; + m->verbs_v2.last_error = verbs->last_error; + m->verbs_v2.run_stage = verbs->run_stage; + } + if (!m->verbs_v2.set_input) m->verbs_v2.set_input = stub_set_input; + if (!m->verbs_v2.get_output) m->verbs_v2.get_output = stub_get_output; + if (!m->verbs_v2.prepare) m->verbs_v2.prepare = stub_prepare; + if (!m->verbs_v2.step) m->verbs_v2.step = stub_step; + if (!m->verbs_v2.last_error) m->verbs_v2.last_error = stub_last_error; + if (!m->verbs_v2.run_stage) m->verbs_v2.run_stage = stub_run_stage; + + m->verbs.struct_size = (uint32_t)sizeof(frt_model_runtime_verbs); + m->verbs.set_input = m->verbs_v2.set_input; + m->verbs.get_output = m->verbs_v2.get_output; + m->verbs.prepare = m->verbs_v2.prepare; + m->verbs.step = m->verbs_v2.step; + m->verbs.last_error = m->verbs_v2.last_error; + m->self = verbs_self; +} + +int add_stage_v2(frt_runtime_builder b, const char* name, uint32_t kind, + uint32_t graph, uint32_t callback, const uint32_t* after, + uint32_t n_after) { + if (!b || !name || !name[0] || (n_after && !after)) return -1; + Holder* h = b->h; + if (kind == FRT_RT_STAGE_GRAPH && graph >= h->graphs.size()) return -1; + if (kind != FRT_RT_STAGE_GRAPH && kind != FRT_RT_STAGE_CALLBACK) return -1; + for (uint32_t i = 0; i < n_after; ++i) + if (after[i] >= h->stages_v2.size()) return -1; + h->after_arrays.emplace_back(after, after + n_after); + frt_runtime_stage_desc_v2 d{}; + d.name = stored(h, name); + d.kind = kind; + d.graph = graph; + d.callback = callback; + d.after = h->after_arrays.back().data(); + d.n_after = n_after; + h->stages_v2.push_back(d); + return 0; +} + } // namespace extern "C" int frt_runtime_builder_add_port(frt_runtime_builder b, @@ -80,6 +130,71 @@ extern "C" int frt_runtime_builder_add_port(frt_runtime_builder b, d.offset = offset; d.bytes = bytes; h->ports.push_back(d); + + /* Phase 6: push a NULL token entry so `h->port_tokens` stays index-parallel + * with `h->ports` (the header invariant: port_tokens[i] aligns with + * ports[i], handle == nullptr for a port that carries no token). A port + * added through this entry never carries a token — only + * frt_runtime_builder_add_port_token mints a real one. The release-time + * destroy loop already no-ops on a null handle, so this is just alignment. + * Without this, a runtime that mixes add_port + add_port_token would have + * n_port_tokens != n_ports and port_tokens would be column-shifted. */ + frt_memory_token_desc null_tk{}; + null_tk.struct_size = (uint32_t)sizeof(frt_memory_token_desc); + h->port_tokens.push_back(null_tk); + return 0; +} + +/* Phase 6 — provider-owned port with a memory-domain token. The port is + * recorded with update=STAGED (no raw frt_buffer; the token is the only + * buffer form on provider-owned ports) and a parallel port_tokens entry + * whose handle is the provider-minted token. `valid_port_args` rejects bad + * name/direction/update/shape. We additionally require a non-null handle + + * verbs for a token port. */ +extern "C" int frt_runtime_builder_add_port_token( + frt_runtime_builder b, const char* name, + uint32_t modality, uint32_t dtype, uint32_t layout, uint32_t direction, + uint32_t required, const int64_t* shape, uint32_t rank, + uint32_t cadence_hint_hz, + frt_memory_token handle, const frt_memory_token_verbs* verbs, + uint64_t offset, uint64_t bytes, uint32_t location_kind) { + if (!b) return -1; + if (!valid_port_args(name, direction, FRT_RT_PORT_STAGED, shape, rank)) + return -1; + if (!handle || !verbs || + verbs->struct_size < FRT_MEMORY_TOKEN_VERBS_COPY_SYNC_SIZE || + !verbs->copy_to_host || !verbs->copy_from_host || !verbs->sync) { + return -1; /* a token port needs the full copy/sync verb set */ + } + if (location_kind > FRT_RT_LOCATION_DEVICE_LOCAL) return -1; + + Holder* h = b->h; + h->shape_arrays.emplace_back(shape, shape + rank); + frt_runtime_port_desc d{}; + d.name = stored(h, name); + d.modality = modality; + d.dtype = dtype; + d.layout = layout; + d.direction = direction; + d.update = FRT_RT_PORT_STAGED; /* token ports are STAGED; no raw window */ + d.required = required; + d.shape = h->shape_arrays.back().data(); + d.rank = rank; + d.cadence_hint_hz = cadence_hint_hz; + d.buffer = nullptr; /* never a raw frt_buffer on a token port */ + d.offset = 0; + d.bytes = 0; + h->ports.push_back(d); + + frt_memory_token_desc tk{}; + tk.struct_size = (uint32_t)sizeof(frt_memory_token_desc); + tk.handle = handle; + tk.verbs = verbs; + tk.offset = offset; + tk.bytes = bytes; + tk.location_kind = location_kind; + tk.reserved = 0; + h->port_tokens.push_back(tk); return 0; } @@ -101,12 +216,30 @@ extern "C" int frt_runtime_builder_add_stage(frt_runtime_builder b, return 0; } +extern "C" int frt_runtime_builder_add_graph_stage_v2( + frt_runtime_builder b, const char* name, uint32_t graph, + const uint32_t* after, uint32_t n_after) { + return add_stage_v2(b, name, FRT_RT_STAGE_GRAPH, graph, UINT32_MAX, after, + n_after); +} + +extern "C" int frt_runtime_builder_add_callback_stage_v2( + frt_runtime_builder b, const char* name, uint32_t callback, + const uint32_t* after, uint32_t n_after) { + return add_stage_v2(b, name, FRT_RT_STAGE_CALLBACK, UINT32_MAX, callback, + after, n_after); +} + extern "C" frt_model_runtime_v1* frt_runtime_builder_finish_model( frt_runtime_builder b, const frt_model_runtime_verbs* verbs, void* verbs_self, void* owner, void (*retain_owner)(void*), void (*release_owner)(void*)) { if (!b) return nullptr; + if (b->provider_owned) { + frt_runtime_builder_discard(b); + return nullptr; + } Holder* h = b->h; frt_rt::finish_export_into(h, b, owner, retain_owner, release_owner); @@ -125,6 +258,80 @@ extern "C" frt_model_runtime_v1* frt_runtime_builder_finish_model( return &h->model; } +extern "C" frt_model_runtime_v2* frt_runtime_builder_finish_model_v2( + frt_runtime_builder b, + const frt_model_runtime_verbs_v2* verbs, void* verbs_self, + void* owner, void (*retain_owner)(void*), + void (*release_owner)(void*)) { + if (!b) return nullptr; + Holder* h = b->h; + if (b->provider_owned && + (!h->streams.empty() || !h->graphs.empty() || !h->buffers.empty() || + !h->regions.empty() || h->stages_v2.empty())) { + frt_runtime_builder_discard(b); + return nullptr; + } + if (b->provider_owned) { + for (const auto& p : h->ports) { + if (p.update != FRT_RT_PORT_STAGED || p.buffer || p.offset || + p.bytes) { + frt_runtime_builder_discard(b); + return nullptr; + } + } + } + for (const auto& s : h->stages_v2) { + if (s.kind == FRT_RT_STAGE_GRAPH && s.graph >= h->graphs.size()) { + frt_runtime_builder_discard(b); + return nullptr; + } + if (s.kind == FRT_RT_STAGE_CALLBACK && + (!verbs || verbs->struct_size < sizeof(frt_model_runtime_verbs_v2) || + !verbs->run_stage)) { + frt_runtime_builder_discard(b); + return nullptr; + } + } + frt_rt::finish_export_into(h, b, owner, retain_owner, release_owner); + + bool has_callback_stage = false; + for (const auto& st : h->stages_v2) + if (st.kind == FRT_RT_STAGE_CALLBACK) has_callback_stage = true; + if (has_callback_stage) h->stages.clear(); + + const bool graph_only_v2 = h->stages.empty() && + !has_callback_stage && h->stages_v2.size() > 0; + if (graph_only_v2) { + for (const auto& st : h->stages_v2) { + if (st.kind != FRT_RT_STAGE_GRAPH) { + break; + } + frt_runtime_stage_desc d{}; + d.graph = st.graph; + d.n_after = st.n_after; + d.after = st.after; + h->stages.push_back(d); + } + if (h->stages.size() != h->stages_v2.size()) h->stages.clear(); + } + + frt_model_runtime_v2& m = h->model_v2; + m.abi_version = FRT_MODEL_RUNTIME_ABI_VERSION_V2; + m.struct_size = (uint32_t)sizeof(frt_model_runtime_v2); + m.exp = &h->exp; + m.ports = h->ports.data(); m.n_ports = h->ports.size(); + m.stages = h->stages.data(); m.n_stages = h->stages.size(); + m.stages_v2 = h->stages_v2.data(); m.n_stages_v2 = h->stages_v2.size(); + m.port_tokens = h->port_tokens.data(); m.n_port_tokens = h->port_tokens.size(); + copy_verbs_v2(&m, verbs, verbs_self); + m.owner = h; + m.retain = frt_rt::frt_rt_holder_retain; + m.release = frt_rt::frt_rt_holder_release; + + delete b; + return &h->model_v2; +} + /* ---- adapter path: wrap an existing export -------------------------------- */ namespace { diff --git a/runtime/src/runtime_export.cpp b/runtime/src/runtime_export.cpp index bf18cfa1..67cb2117 100644 --- a/runtime/src/runtime_export.cpp +++ b/runtime/src/runtime_export.cpp @@ -18,6 +18,15 @@ extern "C" void frt_rt_holder_retain(void* owner) { extern "C" void frt_rt_holder_release(void* owner) { Holder* h = static_cast(owner); if (h->refs.fetch_sub(1, std::memory_order_acq_rel) == 1) { + /* Phase 6: fire each port token's destroy exactly once before the + * holder (and its port descriptors) go away. A null handle or null + * destroy verb is a no-op. This is the ONLY place FlashRT touches a + * provider-owned token's lifetime. */ + for (const auto& tk : h->port_tokens) { + if (tk.handle && tk.verbs && tk.verbs->destroy) { + tk.verbs->destroy(tk.handle); + } + } if (h->user_release) h->user_release(h->user_owner); delete h; } @@ -91,6 +100,18 @@ void finish_export_into(Holder* h, frt_runtime_builder_s* b, } id += '\n'; } + for (size_t i = 0; i < h->stages_v2.size(); ++i) { + const auto& s = h->stages_v2[i]; + std::snprintf(line, sizeof(line), "stage_v2:%zu:%s:%u:%u:%u:", + i, s.name, s.kind, s.graph, s.callback); + id += line; + for (uint32_t d = 0; d < s.n_after; ++d) { + std::snprintf(line, sizeof(line), "%s%u", d ? "," : "", + s.after[d]); + id += line; + } + id += '\n'; + } h->identity = std::move(id); h->user_owner = owner; @@ -126,6 +147,13 @@ extern "C" frt_runtime_builder frt_runtime_builder_create(frt_ctx ctx) { return b; } +extern "C" frt_runtime_builder frt_runtime_builder_create_provider_owned(void) { + auto* b = new frt_runtime_builder_s(); + b->provider_owned = true; + b->h = new Holder(); + return b; +} + extern "C" int frt_runtime_builder_add_stream(frt_runtime_builder b, const char* name, int stream_id, int priority, void* native_handle) { @@ -204,6 +232,12 @@ extern "C" int frt_runtime_builder_set_manifest(frt_runtime_builder b, return 0; } +extern "C" void frt_runtime_builder_discard(frt_runtime_builder b) { + if (!b) return; + delete b->h; + delete b; +} + extern "C" uint64_t frt_runtime_fingerprint(const void* data, size_t len) { /* FNV-1a 64 — deterministic, dependency-free. An identity guard against * accidental mismatch, not an adversarial hash. */ @@ -222,7 +256,11 @@ extern "C" frt_runtime_export_v1* frt_runtime_builder_finish( if (!b) return nullptr; /* Ports/stages declare a MODEL runtime; a plain export cannot carry * them — use frt_runtime_builder_finish_model instead. */ - if (!b->h->ports.empty() || !b->h->stages.empty()) return nullptr; + if (!b->h->ports.empty() || !b->h->stages.empty() || + !b->h->stages_v2.empty() || b->provider_owned) { + frt_runtime_builder_discard(b); + return nullptr; + } Holder* h = b->h; frt_rt::finish_export_into(h, b, owner, retain_owner, release_owner); delete b; /* h lives on inside the export */ diff --git a/runtime/tests/test_model_runtime.cpp b/runtime/tests/test_model_runtime.cpp index 0c9b29a5..a7a82f26 100644 --- a/runtime/tests/test_model_runtime.cpp +++ b/runtime/tests/test_model_runtime.cpp @@ -11,6 +11,7 @@ #include #include #include +#include static int g_fail = 0; #define CHECK(cond, msg) do { \ @@ -24,7 +25,11 @@ static frt_graph FAKE_G0 = (frt_graph)0x20; static frt_graph FAKE_G1 = (frt_graph)0x21; static frt_buffer FAKE_B0 = (frt_buffer)0x30; -struct VerbLog { int set_input = 0, get_output = 0, prepare = 0, step = 0; }; +struct VerbLog { + int set_input = 0, get_output = 0, prepare = 0, step = 0, run_stage = 0; + uint32_t last_stage = UINT32_MAX; + int last_stream = -99; +}; static int v_set_input(void* s, uint32_t, const void*, uint64_t, int) { ((VerbLog*)s)->set_input++; return 0; } @@ -35,6 +40,13 @@ static int v_prepare(void* s, uint32_t, frt_shape_key) { ((VerbLog*)s)->prepare++; return 0; } static int v_step(void* s) { ((VerbLog*)s)->step++; return 0; } +static int v_run_stage(void* s, uint32_t stage, int stream) { + auto* log = (VerbLog*)s; + log->run_stage++; + log->last_stage = stage; + log->last_stream = stream; + return 0; +} static const char* v_last_error(void*) { return ""; } struct Owner { int retains = 0, releases = 0; }; @@ -87,7 +99,9 @@ int main() { add_ports_and_stages(b); CHECK(frt_runtime_builder_finish(b, nullptr, nullptr, nullptr) == nullptr, "plain finish refuses a builder that declared ports/stages"); - /* the builder survives that refusal — finish_model consumes it */ + + b = make_builder(); + add_ports_and_stages(b); frt_model_runtime_v1* m = frt_runtime_builder_finish_model( b, nullptr, nullptr, nullptr, nullptr, nullptr); CHECK(m != nullptr, "finish_model after refused finish"); @@ -295,6 +309,401 @@ int main() { "override release frees native owner and drops the base model"); } + /* --- v2 provider-owned callback runtime ----------------------------- */ + { + Owner provider_owner; + VerbLog provider_vlog; + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + + const int64_t text_shape[1] = {-1}; + const int64_t action_shape[2] = {50, 32}; + CHECK(frt_runtime_builder_add_port( + pb, "prompt", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, + 1, text_shape, 1, 0, nullptr, 0, 0) == 0, + "provider-owned v2 accepts staged text input"); + CHECK(frt_runtime_builder_add_port( + pb, "actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, + 0, action_shape, 2, 0, nullptr, 0, 0) == 0, + "provider-owned v2 accepts staged action output"); + CHECK(frt_runtime_builder_add_callback_stage_v2( + pb, "infer", 7, nullptr, 0) == 0, + "provider-owned v2 accepts callback stage"); + + frt_model_runtime_verbs_v2 verbs2{}; + verbs2.struct_size = sizeof(verbs2); + verbs2.set_input = v_set_input; + verbs2.get_output = v_get_output; + verbs2.prepare = v_prepare; + verbs2.step = v_step; + verbs2.last_error = v_last_error; + verbs2.run_stage = v_run_stage; + + frt_model_runtime_v2* mv2 = frt_runtime_builder_finish_model_v2( + pb, &verbs2, &provider_vlog, &provider_owner, owner_retain, + owner_release); + CHECK(mv2 != nullptr, "finish provider-owned model_runtime_v2"); + CHECK(mv2->abi_version == FRT_MODEL_RUNTIME_ABI_VERSION_V2 && + mv2->struct_size == sizeof(frt_model_runtime_v2), + "model runtime v2 ABI stamp"); + CHECK(mv2->exp && mv2->exp->ctx == nullptr && + mv2->exp->n_graphs == 0, + "provider-owned v2 export has no FlashRT exec graph"); + CHECK(mv2->n_stages == 0 && mv2->n_stages_v2 == 1, + "callback stage is visible only in the v2 stage view"); + CHECK(std::strcmp(mv2->stages_v2[0].name, "infer") == 0 && + mv2->stages_v2[0].kind == FRT_RT_STAGE_CALLBACK && + mv2->stages_v2[0].callback == 7, + "callback stage descriptor round-trips"); + std::string id2 = mv2->exp->identity; + CHECK(id2.find("stage_v2:0:infer:1:4294967295:7:") != + std::string::npos, + "identity carries v2 callback stage"); + + mv2->verbs_v2.run_stage(mv2->self, 0, -1); + mv2->verbs_v2.set_input(mv2->self, 0, nullptr, 0, -1); + CHECK(provider_vlog.run_stage == 1 && + provider_vlog.last_stage == 0 && + provider_vlog.last_stream == -1 && + provider_vlog.set_input == 1, + "v2 verbs dispatch through self"); + CHECK(provider_owner.retains == 1, "v2 finish retained owner once"); + mv2->release(mv2->owner); + CHECK(provider_owner.releases == 1, + "v2 final release frees owner exactly once"); + } + + { + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + CHECK(frt_runtime_builder_finish_model(pb, nullptr, nullptr, nullptr, + nullptr, nullptr) == nullptr, + "provider-owned builders cannot finish as v1 runtimes"); + } + + { + Owner mixed_owner; + VerbLog mixed_vlog; + frt_runtime_builder mb = make_builder(); + const int64_t scalar_shape[1] = {1}; + CHECK(frt_runtime_builder_add_port( + mb, "input", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, + 1, scalar_shape, 1, 0, nullptr, 0, 0) == 0, + "mixed v2 test port"); + const uint32_t after0[1] = {0}; + CHECK(frt_runtime_builder_add_graph_stage_v2( + mb, "prefill", 0, nullptr, 0) == 0, + "mixed v2 accepts graph stage"); + CHECK(frt_runtime_builder_add_callback_stage_v2( + mb, "decode", 3, after0, 1) == 0, + "mixed v2 accepts callback stage after graph stage"); + frt_model_runtime_verbs_v2 verbs2{}; + verbs2.struct_size = sizeof(verbs2); + verbs2.step = v_step; + verbs2.last_error = v_last_error; + verbs2.run_stage = v_run_stage; + frt_model_runtime_v2* mv2 = frt_runtime_builder_finish_model_v2( + mb, &verbs2, &mixed_vlog, &mixed_owner, owner_retain, + owner_release); + CHECK(mv2 && mv2->n_stages == 0 && mv2->n_stages_v2 == 2 && + mv2->stages_v2[1].after[0] == 0, + "mixed callback DAG has no legacy graph-only stage view"); + mv2->release(mv2->owner); + CHECK(mixed_owner.releases == 1, "mixed v2 release frees owner"); + } + + /* --- Phase 6: provider-owned memory-domain token --------------------- */ + /* A stub provider mints a host-backed token (plain malloc) and supplies + * the copy/sync/destroy verbs. The runtime must (a) accept the token + * port, (b) round-trip data through copy_to_host/copy_from_host, (c) + * expose it on the v2 struct, (d) fire destroy exactly once at release, + * and (e) STILL reject a raw frt_buffer on a provider-owned port. No + * Jetson-PI/GGML dependency — the token verbs operate on host memory. */ + { + struct HostToken { + std::vector store; + int destroys = 0; + }; + HostToken ht; + ht.store.assign(64, 0); + + auto copy_to_host = [](frt_memory_token t, void* dst, + uint64_t dst_off, uint64_t src_off, + uint64_t bytes) -> int { + auto* h = reinterpret_cast(t); + if (src_off + bytes > h->store.size()) return -1; + std::memcpy(static_cast(dst) + dst_off, + h->store.data() + src_off, bytes); + return 0; + }; + auto copy_from_host = [](frt_memory_token t, const void* src, + uint64_t src_off, uint64_t dst_off, + uint64_t bytes) -> int { + auto* h = reinterpret_cast(t); + if (dst_off + bytes > h->store.size()) return -1; + std::memcpy(h->store.data() + dst_off, + static_cast(src) + src_off, bytes); + return 0; + }; + auto sync = [](frt_memory_token) -> int { return 0; }; + auto destroy = [](frt_memory_token t) { + reinterpret_cast(t)->destroys += 1; + }; + auto map_host = [](frt_memory_token t, uint64_t offset, + uint64_t bytes, uint32_t access, + void** out_ptr) -> int { + auto* h = reinterpret_cast(t); + if (!out_ptr || access != FRT_RT_HOST_MAP_READ || + offset + bytes > h->store.size()) return -1; + *out_ptr = h->store.data() + offset; + return 0; + }; + auto unmap_host = [](frt_memory_token, void* ptr) -> int { + return ptr ? 0 : -1; + }; + + frt_memory_token_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.copy_to_host = copy_to_host; + verbs.copy_from_host = copy_from_host; + verbs.sync = sync; + verbs.destroy = destroy; + verbs.map_host = map_host; + verbs.unmap_host = unmap_host; + + Owner po; + VerbLog vl; + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + const int64_t out_shape[1] = {16}; + CHECK(frt_runtime_builder_add_port_token( + pb, "acts", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, 0, + out_shape, 1, 0, + (frt_memory_token)&ht, &verbs, + 0, 16 * sizeof(float), + FRT_RT_LOCATION_HOST_VISIBLE) == 0, + "provider-owned v2 accepts a memory-token port"); + CHECK(frt_runtime_builder_add_callback_stage_v2( + pb, "infer", 0, nullptr, 0) == 0, + "token port runtime needs a callback stage"); + frt_model_runtime_verbs_v2 v2{}; + v2.struct_size = sizeof(v2); + v2.run_stage = v_run_stage; + v2.last_error = v_last_error; + frt_model_runtime_v2* m = frt_runtime_builder_finish_model_v2( + pb, &v2, &vl, &po, owner_retain, owner_release); + CHECK(m != nullptr, "finish provider-owned token runtime"); + + /* (c) token surfaces on the v2 struct. */ + CHECK(m->n_port_tokens == 1 && + m->port_tokens[0].handle == (frt_memory_token)&ht && + m->port_tokens[0].location_kind == + FRT_RT_LOCATION_HOST_VISIBLE && + m->port_tokens[0].bytes == 16 * sizeof(float), + "v2 struct exposes the memory-token descriptor"); + + /* (b) round-trip through the provider verbs via the descriptor. */ + const float write_vals[4] = {1.0f, 2.0f, 3.0f, 4.0f}; + const auto& tk = m->port_tokens[0]; + CHECK(tk.verbs->copy_from_host(tk.handle, write_vals, 0, 0, + sizeof(write_vals)) == 0, + "copy_from_host writes into the token"); + float read_vals[4] = {0, 0, 0, 0}; + CHECK(tk.verbs->copy_to_host(tk.handle, read_vals, 0, 0, + sizeof(read_vals)) == 0 && + std::memcmp(read_vals, write_vals, sizeof(write_vals)) == 0, + "copy_to_host round-trips the token contents"); + CHECK(tk.verbs->sync(tk.handle) == 0, "token sync is a no-op for HOST"); + void* mapped = nullptr; + CHECK(tk.verbs->map_host(tk.handle, 0, sizeof(write_vals), + FRT_RT_HOST_MAP_READ, &mapped) == 0 && + mapped == ht.store.data() && + std::memcmp(mapped, write_vals, sizeof(write_vals)) == 0, + "HOST_VISIBLE token exposes a zero-copy host SWAP window"); + CHECK(tk.verbs->unmap_host(tk.handle, mapped) == 0, + "host SWAP window unmaps explicitly"); + + /* (a)/(e) the port itself is STAGED with no raw frt_buffer. */ + CHECK(m->n_ports == 1 && + m->ports[0].update == FRT_RT_PORT_STAGED && + m->ports[0].buffer == nullptr && + m->ports[0].offset == 0 && m->ports[0].bytes == 0, + "token port is STAGED with no raw frt_buffer window"); + + /* (d) destroy has NOT fired while retained. */ + CHECK(ht.destroys == 0, "token not destroyed while runtime live"); + m->release(m->owner); + CHECK(ht.destroys == 1, "token destroy fires exactly once at release"); + } + + /* Phase 6: provider-owned still rejects a raw frt_buffer (the token is + * the only permitted buffer form). add_port_token itself rejects a null + * handle / incomplete verbs. */ + { + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + const int64_t s[1] = {1}; + CHECK(frt_runtime_builder_add_port( + pb, "x", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, + 1, s, 1, 0, FAKE_B0, 0, 4) == 0, + "add_port records a raw-buffer SWAP port (rejected at finish)"); + frt_model_runtime_verbs_v2 v2{}; + v2.struct_size = sizeof(v2); + v2.run_stage = v_run_stage; + frt_model_runtime_v2* m = frt_runtime_builder_finish_model_v2( + pb, &v2, nullptr, nullptr, nullptr, nullptr); + CHECK(m == nullptr, + "provider-owned finish rejects a raw frt_buffer port"); + } + { + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + const int64_t s[1] = {1}; + frt_memory_token_verbs empty_verbs{}; + empty_verbs.struct_size = sizeof(empty_verbs); /* all null fns */ + CHECK(frt_runtime_builder_add_port_token( + pb, "x", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, 1, s, 1, 0, + (frt_memory_token)0x1, &empty_verbs, 0, 4, + FRT_RT_LOCATION_HOST_VISIBLE) != 0, + "add_port_token rejects a token with incomplete verbs"); + frt_runtime_builder_discard(pb); + } + { + uint8_t storage[4]{}; + frt_memory_token_verbs prefix_verbs{}; + prefix_verbs.struct_size = FRT_MEMORY_TOKEN_VERBS_COPY_SYNC_SIZE; + prefix_verbs.copy_to_host = [](frt_memory_token t, void* dst, + uint64_t dst_off, uint64_t src_off, + uint64_t bytes) -> int { + std::memcpy(static_cast(dst) + dst_off, + reinterpret_cast(t) + src_off, bytes); + return 0; + }; + prefix_verbs.copy_from_host = [](frt_memory_token t, const void* src, + uint64_t src_off, uint64_t dst_off, + uint64_t bytes) -> int { + std::memcpy(reinterpret_cast(t) + dst_off, + static_cast(src) + src_off, bytes); + return 0; + }; + prefix_verbs.sync = [](frt_memory_token) -> int { return 0; }; + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + const int64_t s[1] = {4}; + CHECK(frt_runtime_builder_add_port_token( + pb, "legacy", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_U8, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, 0, s, 1, 0, + reinterpret_cast(storage), &prefix_verbs, + 0, sizeof(storage), FRT_RT_LOCATION_HOST_VISIBLE) == 0, + "add_port_token accepts the original copy/sync verbs prefix"); + CHECK(frt_runtime_builder_add_callback_stage_v2( + pb, "infer", 0, nullptr, 0) == 0, + "legacy-prefix token runtime accepts callback stage"); + frt_model_runtime_verbs_v2 v2{}; + v2.struct_size = sizeof(v2); + v2.run_stage = v_run_stage; + v2.last_error = v_last_error; + frt_model_runtime_v2* m = frt_runtime_builder_finish_model_v2( + pb, &v2, nullptr, nullptr, nullptr, nullptr); + CHECK(m && m->port_tokens[0].verbs->struct_size == + FRT_MEMORY_TOKEN_VERBS_COPY_SYNC_SIZE, + "legacy-prefix token remains tail-probeable by consumers"); + if (m) m->release(m->owner); + } + + /* Phase 6: MIXED port build — one normal STAGED port via add_port + one + * token port via add_port_token on the SAME provider-owned builder. This + * is the scenario the original diff got wrong: n_port_tokens must equal + * n_ports, and port_tokens[i] must align with ports[i] (null handle for + * the non-token port, the minted handle for the token port). Before the + * fix, add_port pushed nothing into port_tokens, so n_port_tokens was 1 + * (not 2) and port_tokens[0] held the p1 token misaligned against p0. */ + { + struct HostToken { + std::vector store; + int destroys = 0; + }; + HostToken ht; + ht.store.assign(32, 0); + + auto copy_to_host = [](frt_memory_token t, void* dst, + uint64_t dst_off, uint64_t src_off, + uint64_t bytes) -> int { + auto* h = reinterpret_cast(t); + if (src_off + bytes > h->store.size()) return -1; + std::memcpy(static_cast(dst) + dst_off, + h->store.data() + src_off, bytes); + return 0; + }; + auto copy_from_host = [](frt_memory_token t, const void* src, + uint64_t src_off, uint64_t dst_off, + uint64_t bytes) -> int { + auto* h = reinterpret_cast(t); + if (dst_off + bytes > h->store.size()) return -1; + std::memcpy(h->store.data() + dst_off, + static_cast(src) + src_off, bytes); + return 0; + }; + auto sync = [](frt_memory_token) -> int { return 0; }; + auto destroy = [](frt_memory_token t) { + reinterpret_cast(t)->destroys += 1; + }; + frt_memory_token_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.copy_to_host = copy_to_host; + verbs.copy_from_host = copy_from_host; + verbs.sync = sync; + verbs.destroy = destroy; + + Owner po; + VerbLog vl; + frt_runtime_builder pb = frt_runtime_builder_create_provider_owned(); + + /* p0: normal STAGED input port via add_port (no token). */ + const int64_t in_shape[1] = {4}; + CHECK(frt_runtime_builder_add_port( + pb, "p0", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, + 0, in_shape, 1, 0, nullptr, 0, 0) == 0, + "mixed build: add_port records a non-token STAGED input"); + + /* p1: token output port via add_port_token. */ + const int64_t out_shape[1] = {8}; + CHECK(frt_runtime_builder_add_port_token( + pb, "p1", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, 0, + out_shape, 1, 0, + (frt_memory_token)&ht, &verbs, 0, 8 * sizeof(float), + FRT_RT_LOCATION_HOST_VISIBLE) == 0, + "mixed build: add_port_token records a token output"); + + CHECK(frt_runtime_builder_add_callback_stage_v2( + pb, "infer", 0, nullptr, 0) == 0, + "mixed build: callback stage accepted"); + frt_model_runtime_verbs_v2 v2{}; + v2.struct_size = sizeof(v2); + v2.run_stage = v_run_stage; + v2.last_error = v_last_error; + frt_model_runtime_v2* m = frt_runtime_builder_finish_model_v2( + pb, &v2, &vl, &po, owner_retain, owner_release); + CHECK(m != nullptr, "mixed build finishes provider-owned v2"); + + /* The exact invariant the original diff violated. */ + CHECK(m->n_ports == 2 && m->n_port_tokens == 2, + "mixed build: n_port_tokens == n_ports (index-parallel arrays)"); + CHECK(m->port_tokens[0].handle == nullptr, + "mixed build: non-token port p0 has a null token handle"); + CHECK(m->port_tokens[1].handle == (frt_memory_token)&ht && + m->port_tokens[1].location_kind == + FRT_RT_LOCATION_HOST_VISIBLE, + "mixed build: token port p1 aligns at index 1 with its handle"); + + /* destroy fires only for the real token, exactly once at release. */ + CHECK(ht.destroys == 0, "mixed build: token not destroyed while live"); + m->release(m->owner); + CHECK(ht.destroys == 1, "mixed build: token destroy fires once"); + } + std::printf(g_fail ? "\n== MODEL RUNTIME ABI FAILED ==\n" : "\n== MODEL RUNTIME ABI PASSED ==\n"); return g_fail; diff --git a/runtime/tests/test_model_runtime_py.py b/runtime/tests/test_model_runtime_py.py index 39dde2e7..0dccdf6f 100644 --- a/runtime/tests/test_model_runtime_py.py +++ b/runtime/tests/test_model_runtime_py.py @@ -17,6 +17,7 @@ import ctypes import gc +import json import weakref import _flashrt_exec as ex @@ -24,8 +25,9 @@ import flash_rt.runtime.export as export_mod from flash_rt.runtime.export import ( - BufferSpec, GraphSpec, PortSpec, RegionSpec, StageSpec, StreamSpec, - build_model_runtime, DTYPE, LAYOUT, MODALITY, UPDATE, + BufferSpec, CallbackStageSpec, GraphSpec, PortSpec, RegionSpec, + StageSpec, StreamSpec, build_model_runtime, build_provider_model_runtime_v2, + DTYPE, LAYOUT, MODALITY, STAGE_KIND, UPDATE, ) from flash_rt.subgraphs.stage_plan import ( StagePlan, @@ -272,6 +274,99 @@ def check_vjp_guided_port_lowering(setup): mr.release() +def check_provider_owned_v2(): + calls = {"set_input": [], "run_stage": [], "step": 0} + + def py_set_input(port, payload, stream): + calls["set_input"].append((port, bytes(payload), stream)) + return 0 + + def py_run_stage(stage, stream): + calls["run_stage"].append((stage, stream)) + return 0 + + def py_step(): + calls["step"] += 1 + return py_run_stage(0, -1) + + mr = build_provider_model_runtime_v2( + ports=[ + PortSpec("prompt", "text", "u8", "flat", "in", "staged", + required=True, shape=(-1,)), + PortSpec("state", "state", "f32", "flat", "in", "staged", + required=True, shape=(8,)), + PortSpec("actions", "action", "f32", "flat", "out", "staged", + shape=(50, 32)), + ], + stages=[CallbackStageSpec("infer", 0)], + identity={"model": "pi0", "provider": "llama_cpp_test"}, + manifest_extra={"provider": "llama_cpp", "model_family": "pi0"}, + set_input=py_set_input, + step=py_step, + run_stage=py_run_stage, + ) + try: + manifest = json.loads(mr.manifest) + ports = mr.ports() + stages = mr.stages() + check("provider-owned v2 exports staged ports", ( + len(ports) == 3 + and ports[0]["name"] == "prompt" + and ports[0]["modality"] == MODALITY["text"] + and ports[0]["update"] == UPDATE["staged"] + and ports[0]["buffer"] == 0)) + check("provider-owned v2 exposes callback infer stage", ( + stages == [{"name": "infer", "kind": STAGE_KIND["callback"], + "graph": 0xFFFFFFFF, "callback": 0, "after": []}])) + check("provider-owned v2 identity carries callback stage", ( + "stage_v2:0:infer:" in mr.identity + and "port:0:prompt:" in mr.identity)) + check("provider-owned v2 has no exec graph manifest", ( + "graphs" not in manifest + and manifest["provider"] == "llama_cpp")) + check("provider-owned v2 set_input reaches Python callable", ( + rt.model_v2_set_input(mr.ptr, 0, b"pick cube", -1) == 0 + and calls["set_input"] == [(0, b"pick cube", -1)])) + check("provider-owned v2 run_stage reaches Python callable", ( + rt.model_v2_run_stage(mr.ptr, 0, -1) == 0 + and calls["run_stage"][-1] == (0, -1))) + check("provider-owned v2 step can delegate to callback stage", ( + rt.model_v2_step(mr.ptr) == 0 and calls["step"] == 1)) + finally: + mr.release() + try: + build_provider_model_runtime_v2( + ports=[ + PortSpec("bad_swap", "tensor", "f32", "flat", "in", "swap", + shape=(1,)), + ], + stages=[CallbackStageSpec("infer", 0)], + identity={"model": "bad"}, + run_stage=py_run_stage, + ) + except ValueError as e: + rejected_swap = "staged updates" in str(e) + else: + rejected_swap = False + check("provider-owned v2 rejects non-staged ports", rejected_swap) + try: + build_provider_model_runtime_v2( + ports=[ + PortSpec("bad_window", "tensor", "f32", "flat", "in", + "staged", shape=(1,), nbytes=4), + ], + stages=[CallbackStageSpec("infer", 0)], + identity={"model": "bad"}, + run_stage=py_run_stage, + ) + except ValueError as e: + rejected_window = "raw windows" in str(e) + else: + rejected_window = False + check("provider-owned v2 rejects raw window declarations", + rejected_window) + + def main(): CHECKS.clear() setup = make_setup() @@ -333,6 +428,8 @@ def py_step(): print("== stage plan registry ==") check_stage_plan_registry() check_vjp_guided_port_lowering(setup) + print("== provider-owned v2 callback stage ==") + check_provider_owned_v2() print("== verbs through the C function pointers ==") rc = rt.model_set_input(mr.ptr, 1, b"\xAA\xBB", -1)