diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 00000000..6c5f513b --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,25 @@ +## Summary + + + +## Design boundaries + + + +## Compatibility + + + +## Validation + + + +- [ ] Focused tests cover success and rejection paths +- [ ] Affected CUDA-off/hardware configurations were checked or disclosed +- [ ] Numerical claims use a fixed, justified gate +- [ ] STAGED ports have real matching verbs +- [ ] Identity uses observed runtime facts and changes with contract changes +- [ ] Hot-path allocation/capture/rebind claims are measured at the right scope +- [ ] Documentation and migration notes are updated +- [ ] Diff contains no private paths, hosts, containers, credentials, or logs +- [ ] Shared kernel/CMake ownership and packaging were reviewed diff --git a/CMakeLists.txt b/CMakeLists.txt index 8152df92..183e6f61 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -767,8 +767,8 @@ if(ENABLE_NVFP4) endif() # ── Flash-Attention 2 vendored kernels (fp16 + bf16 fwd SM80) ── -# Built as an object library and linked into a SEPARATE pybind module -# flash_rt_fa2.so — same isolation pattern as flash_rt_fp4.so. The +# Built as an object library and linked into a Python-free raw library; +# flash_rt_fa2.so is a thin pybind adapter over that library. The # main flash_rt_kernels.so stays small (~3.6 MB) and its rebuild no # longer waits on FA2's heavy CUTLASS 3.x template codegen (~8 min). # @@ -831,9 +831,9 @@ if(ENABLE_FA2) csrc/attention/fa2_causal_inst/flash_fwd_split_hdim256_bf16_sm80_causal.cu ) endif() - if("bf16" IN_LIST FA2_DTYPES AND ("128" IN_LIST FA2_HDIMS OR "256" IN_LIST FA2_HDIMS)) - list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu) - endif() + # Always emit the causal C entry. Builds without a causal hdim retain a + # stable raw-library symbol that fails clearly instead of an unresolved ABI. + list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu) add_library(fa2_vendor_obj OBJECT ${FA2_SRCS}) set_target_properties(fa2_vendor_obj PROPERTIES @@ -954,6 +954,21 @@ if(ENABLE_FA2) math(EXPR _FA2_NKERN "${_FA2_NSRC} - 1") message(STATUS "FA2 vendor instantiations: hdim={${_FA2_H}} x dtype={${_FA2_D}} x {split,no-split} = ${_FA2_NKERN} kernel .cu files") + # Keep the raw C entries in a Python-free library. The native C++ runtime + # links this target directly; the pybind module below is only an adapter. + add_library(flashrt_fa2_raw SHARED) + target_sources(flashrt_fa2_raw PRIVATE $) + target_link_libraries(flashrt_fa2_raw PRIVATE CUDA::cudart) + if(NOT MSVC) + target_link_options(flashrt_fa2_raw PRIVATE + "LINKER:--no-undefined" -static-libstdc++ -static-libgcc) + endif() + set_target_properties(flashrt_fa2_raw PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN") + # Independent pybind module. Caller side: # from flash_rt import flash_rt_fa2 as fa2 # fa2.fwd_fp16(...) / fa2.fwd_bf16(...) @@ -961,12 +976,16 @@ if(ENABLE_FA2) set_target_properties(flash_rt_fa2 PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) - target_sources(flash_rt_fa2 PRIVATE $) target_include_directories(flash_rt_fa2 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/csrc ) - target_link_libraries(flash_rt_fa2 PRIVATE CUDA::cudart) - install(TARGETS flash_rt_fa2 DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) + target_link_libraries(flash_rt_fa2 PRIVATE flashrt_fa2_raw CUDA::cudart) + set_target_properties(flash_rt_fa2 PROPERTIES + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH ON) + install(TARGETS flash_rt_fa2 flashrt_fa2_raw + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) message(STATUS "FA2 pybind module: flash_rt_fa2 (separate .so)") endif() @@ -1440,8 +1459,8 @@ if(FLASHRT_ENABLE_MOTUS AND GPU_ARCH STREQUAL "120") ENABLE_TINYFP8_KERNELS=1) endif() -# (ENABLE_FA2 object-lib is linked into flash_rt_fa2.so only — see -# the dedicated pybind11_add_module(flash_rt_fa2 ...) above. The main +# (ENABLE_FA2 object-lib is linked into libflashrt_fa2_raw.so only; the +# dedicated pybind module and native C++ runtime both consume it. The main # flash_rt_kernels.so deliberately does NOT pull FA2 in, so rebuilds # of our hand-written kernels don't re-trigger the FA2 codegen tax.) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8acbe70b..328083a5 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -18,10 +18,15 @@ Before opening a PR: - New model integration: [`docs/adding_new_model.md`](docs/adding_new_model.md) - Kernel catalog: [`docs/kernel_catalog.md`](docs/kernel_catalog.md) - Calibration contract: [`docs/calibration.md`](docs/calibration.md) + - Native model producers: + [`docs/native_model_runtime_producer.md`](docs/native_model_runtime_producer.md) + - PR review standard: [`docs/pr_review_checklist.md`](docs/pr_review_checklist.md) 2. Build the extension modules locally. 3. Run the smallest test set that covers your change. -4. Include the exact GPU, CUDA, command lines, and latency/precision numbers - in the PR description when the change touches runtime behavior. +4. Include sanitized, reproducible build/test commands and the relevant public + hardware capability and latency/precision results when runtime behavior + changes. Never include private paths, host names, container names, tokens, + checkpoint locations, or internal dataset identifiers. ## Development Setup @@ -224,6 +229,13 @@ reviewers hold every PR to: [`docs/subgraph_stage_plans.md`](docs/subgraph_stage_plans.md). A structural cut is a re-ordering, not an approximation — split-vs-full replay must stay bit-exact (`cpp/tests/gate_pi05_model_runtime_export.py` is the gate). +- All three C construction paths mechanically reject a STAGED input without + `set_input` or a STAGED output without `get_output`. Do not bypass this with + a published declaration-only object. +- Derive hardware identity from the active runtime device. A requested build + target or configuration string is not proof of the executing architecture. +- Schema shared by multiple producers needs checked-in canonical records; + every producer compares independently against that golden face. ### Calibration And Precision @@ -354,6 +366,8 @@ Before requesting review: - Mention unsupported hardware or missing local fixtures explicitly. - Avoid committing generated build outputs, local checkpoints, logs, or `third_party/cutlass`. +- Search the diff for private absolute paths, user/host/container names, + credentials, internal URLs, and environment dumps before pushing. ## Reporting Hardware Results diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 1b097a85..c57f00e0 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -22,6 +22,10 @@ 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_SENTENCEPIECE "Enable native SentencePiece text tokenization" OFF) +option(FLASHRT_CPP_WITH_FA2 "Enable the Python-free vendored FA2 driver" OFF) +set(FLASHRT_CPP_FA2_LIBRARY "" CACHE FILEPATH + "Path to the Python-free libflashrt_fa2_raw shared library") if(FLASHRT_CPP_WITH_CUDA_STAGING) find_package(CUDAToolkit REQUIRED) endif() @@ -29,6 +33,50 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) endif() +if(FLASHRT_CPP_WITH_FA2) + if(NOT FLASHRT_CPP_WITH_CUDA_KERNELS) + message(FATAL_ERROR "FLASHRT_CPP_WITH_FA2 requires CUDA kernels") + endif() + if(NOT FLASHRT_CPP_FA2_LIBRARY) + unset(FLASHRT_CPP_FA2_LIBRARY CACHE) + find_library(FLASHRT_CPP_FA2_LIBRARY + NAMES flashrt_fa2_raw + PATHS ${CMAKE_CURRENT_SOURCE_DIR}/../flash_rt + NO_DEFAULT_PATH) + endif() + if(NOT FLASHRT_CPP_FA2_LIBRARY) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_FA2 requires libflashrt_fa2_raw; build the root " + "flashrt_fa2_raw target or set FLASHRT_CPP_FA2_LIBRARY") + endif() + add_library(flashrt_cpp_fa2_external SHARED IMPORTED GLOBAL) + set_target_properties(flashrt_cpp_fa2_external PROPERTIES + IMPORTED_LOCATION ${FLASHRT_CPP_FA2_LIBRARY}) +endif() +if(FLASHRT_CPP_WITH_SENTENCEPIECE) + find_path(FLASHRT_SENTENCEPIECE_INCLUDE_DIR sentencepiece_processor.h) + find_library(FLASHRT_SENTENCEPIECE_LIBRARY sentencepiece) + if(FLASHRT_SENTENCEPIECE_INCLUDE_DIR AND FLASHRT_SENTENCEPIECE_LIBRARY) + add_library(flashrt_sentencepiece_external INTERFACE) + target_include_directories(flashrt_sentencepiece_external + INTERFACE ${FLASHRT_SENTENCEPIECE_INCLUDE_DIR}) + target_link_libraries(flashrt_sentencepiece_external + INTERFACE ${FLASHRT_SENTENCEPIECE_LIBRARY}) + set(FLASHRT_CPP_SENTENCEPIECE_TARGET flashrt_sentencepiece_external) + else() + include(FetchContent) + set(SPM_ENABLE_SHARED OFF CACHE BOOL "" FORCE) + set(SPM_BUILD_TEST OFF CACHE BOOL "" FORCE) + set(SPM_BUILD_TESTS OFF CACHE BOOL "" FORCE) + FetchContent_Declare( + sentencepiece + GIT_REPOSITORY https://github.com/google/sentencepiece.git + GIT_TAG v0.2.1) + FetchContent_MakeAvailable(sentencepiece) + set(FLASHRT_SENTENCEPIECE_INCLUDE_DIR ${sentencepiece_SOURCE_DIR}/src) + set(FLASHRT_CPP_SENTENCEPIECE_TARGET sentencepiece-static) + endif() +endif() if(FLASHRT_CPP_WITH_EXEC AND NOT TARGET flashrt_exec) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../exec ${CMAKE_CURRENT_BINARY_DIR}/exec) @@ -50,9 +98,19 @@ endif() set(FLASHRT_CPP_MODALITY_SRCS modalities/src/types.cpp modalities/src/vision_cpu.cpp + modalities/src/text_cpu.cpp modalities/src/action_cpu.cpp) +if(FLASHRT_CPP_WITH_SENTENCEPIECE) + list(APPEND FLASHRT_CPP_MODALITY_SRCS + modalities/src/sentencepiece_tokenizer.cpp) +else() + list(APPEND FLASHRT_CPP_MODALITY_SRCS + modalities/src/tokenizer_unavailable.cpp) +endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) - list(APPEND FLASHRT_CPP_MODALITY_SRCS modalities/src/vision_cuda.cu) + list(APPEND FLASHRT_CPP_MODALITY_SRCS + modalities/src/vision_cuda.cu + modalities/src/text_cuda.cu) endif() add_library(flashrt_cpp_modalities STATIC ${FLASHRT_CPP_MODALITY_SRCS}) @@ -63,6 +121,14 @@ if(FLASHRT_CPP_WITH_CUDA_STAGING) PUBLIC FLASHRT_CPP_WITH_CUDA_STAGING=1) target_link_libraries(flashrt_cpp_modalities PUBLIC CUDA::cudart) endif() +if(FLASHRT_CPP_WITH_SENTENCEPIECE) + target_compile_definitions(flashrt_cpp_modalities + PUBLIC FLASHRT_CPP_HAS_SENTENCEPIECE=1) + target_include_directories(flashrt_cpp_modalities + PUBLIC ${FLASHRT_SENTENCEPIECE_INCLUDE_DIR}) + target_link_libraries(flashrt_cpp_modalities + PUBLIC ${FLASHRT_CPP_SENTENCEPIECE_TARGET}) +endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) target_compile_definitions(flashrt_cpp_modalities PRIVATE FLASHRT_CPP_WITH_CUDA_KERNELS=1) @@ -78,40 +144,293 @@ target_include_directories(flashrt_cpp_vla target_link_libraries(flashrt_cpp_vla INTERFACE flashrt_cpp_modalities flashrt_cpp) -add_library(flashrt_cpp_pi05 STATIC +add_library(flashrt_cpp_loader STATIC + loader/src/safetensors.cpp + loader/src/sha256.cpp) +set_property(SOURCE loader/src/sha256.cpp APPEND PROPERTY COMPILE_OPTIONS -O2) +target_include_directories(flashrt_cpp_loader + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/loader/include) + +set(FLASHRT_CPP_PI05_SRCS models/pi05/src/spec.cpp + models/pi05/src/native_weights.cpp + models/pi05/src/native_weight_ops.cpp + models/pi05/src/native_device_weights.cpp + models/pi05/src/native_weight_packer.cpp + models/pi05/src/native_weight_materializer.cpp + models/pi05/src/native_workspace.cpp + models/pi05/src/native_rtx_attention.cpp + models/pi05/src/prompt_format.cpp + models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp models/pi05/src/runtime.cpp) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + list(APPEND FLASHRT_CPP_PI05_SRCS + models/pi05/src/native_quantization.cu) +else() + list(APPEND FLASHRT_CPP_PI05_SRCS + models/pi05/src/native_quantization_unavailable.cpp) +endif() + +add_library(flashrt_cpp_pi05 STATIC ${FLASHRT_CPP_PI05_SRCS}) target_include_directories(flashrt_cpp_pi05 PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include) target_link_libraries(flashrt_cpp_pi05 - PUBLIC flashrt_cpp_modalities flashrt_cpp_vla) + PUBLIC flashrt_cpp_modalities flashrt_cpp_vla flashrt_cpp_loader) + +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + add_library(flashrt_cpp_pi05_kernels STATIC + models/pi05/src/native_bf16_forward.cpp + models/pi05/src/native_kernel_driver.cu + models/pi05/src/native_style_precompute.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/activation.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/elementwise.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/norm.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/patch_embed.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/rope.cu) + target_include_directories(flashrt_cpp_pi05_kernels + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels) + target_link_libraries(flashrt_cpp_pi05_kernels + PUBLIC flashrt_cpp_pi05 CUDA::cublas CUDA::cublasLt CUDA::cudart) + set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON) + target_compile_options(flashrt_cpp_pi05_kernels PRIVATE + $<$: + --expt-relaxed-constexpr -O3 + --ftz=true --prec-div=false --prec-sqrt=false>) + if(FLASHRT_CPP_WITH_FA2) + target_sources(flashrt_cpp_pi05_kernels PRIVATE + models/pi05/src/native_graph_owner.cpp + models/pi05/src/native_rtx_attention_driver.cu) + target_include_directories(flashrt_cpp_pi05_kernels PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc) + target_link_libraries(flashrt_cpp_pi05_kernels + PUBLIC flashrt_cpp_fa2_external) + target_compile_definitions(flashrt_cpp_pi05_kernels + PUBLIC FLASHRT_CPP_WITH_FA2=1) + endif() +endif() add_library(flashrt_cpp_pi05_c SHARED models/pi05/src/c_api.cpp - models/pi05/src/model_runtime.cpp) + models/pi05/src/model_runtime.cpp + models/pi05/src/native_model_runtime.cpp + models/pi05/src/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c PUBLIC flashrt_cpp_pi05 flashrt_runtime) target_include_directories(flashrt_cpp_pi05_c PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + target_link_libraries(flashrt_cpp_pi05_c + PRIVATE flashrt_cpp_pi05_kernels) +endif() if(BUILD_TESTING) + add_executable(test_safetensors_loader tests/test_safetensors_loader.cpp) + target_link_libraries(test_safetensors_loader PRIVATE flashrt_cpp_loader) + add_test(NAME safetensors_loader COMMAND test_safetensors_loader) + + add_executable(test_sha256 tests/test_sha256.cpp) + target_link_libraries(test_sha256 PRIVATE flashrt_cpp_loader) + add_test(NAME sha256 COMMAND test_sha256) + 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_text_modalities tests/test_text_modalities.cpp) + target_link_libraries(test_text_modalities PRIVATE flashrt_cpp_modalities) + add_test(NAME text_modalities COMMAND test_text_modalities) + + add_executable(test_text_tokenizer tests/test_text_tokenizer.cpp) + target_link_libraries(test_text_tokenizer PRIVATE flashrt_cpp_modalities) + add_test(NAME text_tokenizer COMMAND test_text_tokenizer) + 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) + add_executable(test_pi05_native_weight_ops + tests/test_pi05_native_weight_ops.cpp) + target_link_libraries(test_pi05_native_weight_ops PRIVATE flashrt_cpp_pi05) + add_test(NAME pi05_native_weight_ops COMMAND test_pi05_native_weight_ops) + + add_executable(pi05_native_weight_probe + tests/pi05_native_weight_probe.cpp) + target_link_libraries(pi05_native_weight_probe PRIVATE flashrt_cpp_pi05) + + if(FLASHRT_CPP_WITH_CUDA_KERNELS) + add_executable(test_pi05_native_quantization + tests/test_pi05_native_quantization.cpp) + target_link_libraries(test_pi05_native_quantization PRIVATE flashrt_cpp_pi05) + add_test(NAME pi05_native_quantization + COMMAND test_pi05_native_quantization) + + add_executable(pi05_native_quant_probe + tests/pi05_native_quant_probe.cpp) + target_link_libraries(pi05_native_quant_probe PRIVATE flashrt_cpp_pi05) + + add_executable(test_pi05_native_weight_packer + tests/test_pi05_native_weight_packer.cpp) + target_link_libraries(test_pi05_native_weight_packer + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_weight_packer + COMMAND test_pi05_native_weight_packer) + + add_executable(test_pi05_native_kernel_driver + tests/test_pi05_native_kernel_driver.cpp) + target_link_libraries(test_pi05_native_kernel_driver + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_kernel_driver + COMMAND test_pi05_native_kernel_driver) + + add_executable(test_pi05_native_forward_primitives + tests/test_pi05_native_forward_primitives.cpp) + target_link_libraries(test_pi05_native_forward_primitives + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_forward_primitives + COMMAND test_pi05_native_forward_primitives) + + add_executable(test_pi05_native_bf16_forward + tests/test_pi05_native_bf16_forward.cpp) + target_link_libraries(test_pi05_native_bf16_forward + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_bf16_forward + COMMAND test_pi05_native_bf16_forward) + + add_executable(pi05_native_encoder_qkv_probe + tests/pi05_native_encoder_qkv_probe.cpp) + target_link_libraries(pi05_native_encoder_qkv_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(test_pi05_native_style_precompute + tests/test_pi05_native_style_precompute.cpp) + target_link_libraries(test_pi05_native_style_precompute + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_style_precompute + COMMAND test_pi05_native_style_precompute) + + add_executable(pi05_native_style_probe + tests/pi05_native_style_probe.cpp) + target_link_libraries(pi05_native_style_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(test_pi05_native_workspace + tests/test_pi05_native_workspace.cpp) + target_link_libraries(test_pi05_native_workspace + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_workspace + COMMAND test_pi05_native_workspace) + + add_executable(test_pi05_native_rtx_attention + tests/test_pi05_native_rtx_attention.cpp) + target_link_libraries(test_pi05_native_rtx_attention + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_rtx_attention + COMMAND test_pi05_native_rtx_attention) + + if(FLASHRT_CPP_WITH_FA2) + add_executable(pi05_native_encoder_layer_probe + tests/pi05_native_encoder_layer_probe.cpp) + target_link_libraries(pi05_native_encoder_layer_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_encoder_probe + tests/pi05_native_encoder_probe.cpp) + target_link_libraries(pi05_native_encoder_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_vision_probe + tests/pi05_native_vision_probe.cpp) + target_link_libraries(pi05_native_vision_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_diffusion_probe + tests/pi05_native_diffusion_probe.cpp) + target_link_libraries(pi05_native_diffusion_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_graph_probe + tests/pi05_native_graph_probe.cpp) + target_link_libraries(pi05_native_graph_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_native_open_probe + tests/pi05_native_open_probe.cpp) + target_link_libraries(pi05_native_open_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + + add_executable(pi05_native_e2e_probe + tests/pi05_native_e2e_probe.cpp) + target_link_libraries(pi05_native_e2e_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + + add_executable(pi05_native_dlopen_probe + tests/pi05_native_dlopen_probe.cpp) + target_include_directories(pi05_native_dlopen_probe + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime/include + ${CMAKE_CURRENT_SOURCE_DIR}/../exec/include) + target_link_libraries(pi05_native_dlopen_probe + PRIVATE ${CMAKE_DL_LIBS}) + endif() + + add_executable(test_pi05_native_rtx_attention_driver + tests/test_pi05_native_rtx_attention_driver.cpp) + target_link_libraries(test_pi05_native_rtx_attention_driver + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_rtx_attention_driver + COMMAND test_pi05_native_rtx_attention_driver) + endif() + + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_tokenizer_corpus_probe + tests/pi05_tokenizer_corpus_probe.cpp) + target_link_libraries(pi05_tokenizer_corpus_probe + PRIVATE flashrt_cpp_pi05) + endif() + + add_executable(pi05_native_rope_probe + tests/pi05_native_rope_probe.cpp) + target_link_libraries(pi05_native_rope_probe + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + endif() + + add_executable(test_pi05_prompt_format tests/test_pi05_prompt_format.cpp) + target_link_libraries(test_pi05_prompt_format PRIVATE flashrt_cpp_pi05) + add_test(NAME pi05_prompt_format COMMAND test_pi05_prompt_format) + + add_executable(test_pi05_prompt_embed tests/test_pi05_prompt_embed.cpp) + target_link_libraries(test_pi05_prompt_embed PRIVATE flashrt_cpp_pi05) + add_test(NAME pi05_prompt_embed COMMAND test_pi05_prompt_embed) + if(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) add_test(NAME device_staging COMMAND test_device_staging) + add_executable(test_pi05_native_device_weights + tests/test_pi05_native_device_weights.cpp) + target_link_libraries(test_pi05_native_device_weights + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_device_weights + COMMAND test_pi05_native_device_weights) + + add_executable(test_pi05_native_weight_materializer + tests/test_pi05_native_weight_materializer.cpp) + target_link_libraries(test_pi05_native_weight_materializer + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_weight_materializer + COMMAND test_pi05_native_weight_materializer) + add_executable(test_pi05_c_api tests/test_pi05_c_api.cpp) target_link_libraries(test_pi05_c_api PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) @@ -121,5 +440,14 @@ if(BUILD_TESTING) target_link_libraries(test_pi05_model_runtime PRIVATE flashrt_cpp_pi05_c flashrt_exec flashrt_runtime CUDA::cudart) add_test(NAME pi05_model_runtime COMMAND test_pi05_model_runtime) + + add_executable(test_pi05_native_open tests/test_pi05_native_open.cpp) + target_link_libraries(test_pi05_native_open + PRIVATE flashrt_cpp_pi05_c flashrt_runtime) + if(FLASHRT_CPP_WITH_FA2 AND FLASHRT_CPP_WITH_SENTENCEPIECE) + target_compile_definitions(test_pi05_native_open + PRIVATE FLASHRT_CPP_PI05_NATIVE_OPEN_ENABLED=1) + endif() + add_test(NAME pi05_native_open COMMAND test_pi05_native_open) endif() endif() diff --git a/cpp/loader/include/flashrt/cpp/loader/safetensors.h b/cpp/loader/include/flashrt/cpp/loader/safetensors.h new file mode 100644 index 00000000..e0fa3c5e --- /dev/null +++ b/cpp/loader/include/flashrt/cpp/loader/safetensors.h @@ -0,0 +1,56 @@ +#pragma once + +#include +#include +#include +#include + +namespace flashrt { +namespace loader { + +struct SafetensorInfo { + std::string dtype; + std::vector shape; + std::uint64_t data_offset = 0; + std::uint64_t bytes = 0; +}; + +class SafetensorsFile { +public: + SafetensorsFile() = default; + ~SafetensorsFile(); + + SafetensorsFile(const SafetensorsFile&) = delete; + SafetensorsFile& operator=(const SafetensorsFile&) = delete; + SafetensorsFile(SafetensorsFile&& other) noexcept; + SafetensorsFile& operator=(SafetensorsFile&& other) noexcept; + + bool open(const std::string& path); + void close(); + + bool is_open() const { return mapping_ != nullptr; } + const std::string& path() const { return path_; } + const std::string& error() const { return error_; } + std::uint64_t file_bytes() const { return mapping_bytes_; } + std::uint64_t data_offset() const { return data_offset_; } + + const std::map& tensors() const { + return tensors_; + } + const SafetensorInfo* find(const std::string& name) const; + const void* data(const SafetensorInfo& tensor) const; + +private: + void move_from(SafetensorsFile&& other) noexcept; + + int fd_ = -1; + void* mapping_ = nullptr; + std::uint64_t mapping_bytes_ = 0; + std::uint64_t data_offset_ = 0; + std::string path_; + std::string error_; + std::map tensors_; +}; + +} // namespace loader +} // namespace flashrt diff --git a/cpp/loader/include/flashrt/cpp/loader/sha256.h b/cpp/loader/include/flashrt/cpp/loader/sha256.h new file mode 100644 index 00000000..635b94f3 --- /dev/null +++ b/cpp/loader/include/flashrt/cpp/loader/sha256.h @@ -0,0 +1,15 @@ +#ifndef FLASHRT_CPP_LOADER_SHA256_H +#define FLASHRT_CPP_LOADER_SHA256_H + +#include + +namespace flashrt { +namespace loader { + +bool sha256_file(const std::string& path, std::string* hex, + std::string* error = nullptr); + +} // namespace loader +} // namespace flashrt + +#endif // FLASHRT_CPP_LOADER_SHA256_H diff --git a/cpp/loader/src/safetensors.cpp b/cpp/loader/src/safetensors.cpp new file mode 100644 index 00000000..23eef574 --- /dev/null +++ b/cpp/loader/src/safetensors.cpp @@ -0,0 +1,456 @@ +#include "flashrt/cpp/loader/safetensors.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flashrt { +namespace loader { +namespace { + +constexpr std::uint64_t kMaxHeaderBytes = 128ull << 20; + +std::uint64_t dtype_bytes(const std::string& dtype) { + if (dtype == "F64" || dtype == "I64" || dtype == "U64") return 8; + if (dtype == "F32" || dtype == "I32" || dtype == "U32") return 4; + if (dtype == "F16" || dtype == "BF16" || dtype == "I16" || + dtype == "U16") { + return 2; + } + if (dtype == "I8" || dtype == "U8" || dtype == "BOOL" || + dtype == "F8_E4M3FN" || dtype == "F8_E5M2") { + return 1; + } + return 0; +} + +bool tensor_bytes(const SafetensorInfo& tensor, std::uint64_t* out) { + std::uint64_t bytes = dtype_bytes(tensor.dtype); + if (!bytes) return false; + for (std::uint64_t dim : tensor.shape) { + if (dim && bytes > std::numeric_limits::max() / dim) { + return false; + } + bytes *= dim; + } + if (out) *out = bytes; + return true; +} + +class HeaderParser { +public: + HeaderParser(const char* begin, const char* end) + : begin_(begin), cur_(begin), end_(end) {} + + bool parse(std::map* tensors) { + skip_ws(); + if (!consume('{')) return fail("safetensors header must be an object"); + skip_ws(); + if (consume('}')) return finish(tensors); + while (cur_ < end_) { + std::string key; + if (!parse_string(&key)) return false; + skip_ws(); + if (!consume(':')) return fail("expected ':' after tensor name"); + skip_ws(); + if (key == "__metadata__") { + if (!skip_value()) return false; + } else { + SafetensorInfo tensor; + if (!parse_tensor(&tensor)) return false; + if (!parsed_.emplace(std::move(key), std::move(tensor)).second) { + return fail("duplicate tensor name in safetensors header"); + } + } + skip_ws(); + if (consume('}')) return finish(tensors); + if (!consume(',')) return fail("expected ',' or '}' in header"); + skip_ws(); + } + return fail("unterminated safetensors header"); + } + + const std::string& error() const { return error_; } + +private: + bool finish(std::map* tensors) { + skip_ws(); + if (cur_ != end_) return fail("trailing data in safetensors header"); + if (parsed_.empty()) return fail("safetensors file contains no tensors"); + if (tensors) *tensors = std::move(parsed_); + return true; + } + + bool parse_tensor(SafetensorInfo* tensor) { + if (!consume('{')) return fail("tensor metadata must be an object"); + bool have_dtype = false; + bool have_shape = false; + bool have_offsets = false; + bool closed = false; + std::vector offsets; + skip_ws(); + if (consume('}')) return fail("tensor metadata is empty"); + while (cur_ < end_) { + std::string key; + if (!parse_string(&key)) return false; + skip_ws(); + if (!consume(':')) return fail("expected ':' in tensor metadata"); + skip_ws(); + if (key == "dtype") { + if (have_dtype || !parse_string(&tensor->dtype)) { + return fail("invalid tensor dtype"); + } + have_dtype = true; + } else if (key == "shape") { + if (have_shape || !parse_u64_array(&tensor->shape)) { + return fail("invalid tensor shape"); + } + have_shape = true; + } else if (key == "data_offsets") { + if (have_offsets || !parse_u64_array(&offsets)) { + return fail("invalid tensor data_offsets"); + } + have_offsets = true; + } else if (!skip_value()) { + return false; + } + skip_ws(); + if (consume('}')) { + closed = true; + break; + } + if (!consume(',')) return fail("expected ',' in tensor metadata"); + skip_ws(); + } + if (!closed) return fail("unterminated tensor metadata"); + if (!have_dtype || !have_shape || !have_offsets || offsets.size() != 2 || + offsets[1] < offsets[0]) { + return fail("incomplete tensor metadata"); + } + tensor->data_offset = offsets[0]; + tensor->bytes = offsets[1] - offsets[0]; + std::uint64_t expected = 0; + if (!tensor_bytes(*tensor, &expected)) { + return fail("unsupported tensor dtype or overflowing shape"); + } + if (expected != tensor->bytes) { + return fail("tensor byte range does not match dtype and shape"); + } + return true; + } + + bool parse_u64_array(std::vector* values) { + if (!consume('[')) return false; + std::vector parsed; + skip_ws(); + if (consume(']')) { + if (values) *values = std::move(parsed); + return true; + } + while (cur_ < end_) { + std::uint64_t value = 0; + if (!parse_u64(&value)) return false; + parsed.push_back(value); + skip_ws(); + if (consume(']')) { + if (values) *values = std::move(parsed); + return true; + } + if (!consume(',')) return false; + skip_ws(); + } + return false; + } + + bool parse_u64(std::uint64_t* out) { + if (cur_ >= end_ || !std::isdigit(static_cast(*cur_))) { + return false; + } + std::uint64_t value = 0; + while (cur_ < end_ && + std::isdigit(static_cast(*cur_))) { + const std::uint64_t digit = static_cast(*cur_ - '0'); + if (value > (std::numeric_limits::max() - digit) / + 10ull) { + return false; + } + value = value * 10ull + digit; + ++cur_; + } + if (out) *out = value; + return true; + } + + bool parse_string(std::string* out) { + if (!consume('"')) return fail("expected JSON string"); + std::string value; + while (cur_ < end_ && *cur_ != '"') { + unsigned char c = static_cast(*cur_++); + if (c < 0x20) return fail("control character in JSON string"); + if (c != '\\') { + value.push_back(static_cast(c)); + continue; + } + if (cur_ >= end_) return fail("unterminated JSON escape"); + switch (*cur_++) { + case '"': value.push_back('"'); break; + case '\\': value.push_back('\\'); break; + case '/': value.push_back('/'); break; + case 'b': value.push_back('\b'); break; + case 'f': value.push_back('\f'); break; + case 'n': value.push_back('\n'); break; + case 'r': value.push_back('\r'); break; + case 't': value.push_back('\t'); break; + default: return fail("unsupported JSON string escape"); + } + } + if (!consume('"')) return fail("unterminated JSON string"); + if (out) *out = std::move(value); + return true; + } + + bool skip_value() { + skip_ws(); + if (cur_ >= end_) return fail("missing JSON value"); + if (*cur_ == '"') return parse_string(nullptr); + if (*cur_ == '{') return skip_object(); + if (*cur_ == '[') return skip_array(); + const char* literals[] = {"true", "false", "null"}; + for (const char* literal : literals) { + const std::size_t n = std::strlen(literal); + if (static_cast(end_ - cur_) >= n && + std::strncmp(cur_, literal, n) == 0) { + cur_ += n; + return true; + } + } + return skip_number(); + } + + bool skip_object() { + if (!consume('{')) return false; + skip_ws(); + if (consume('}')) return true; + while (cur_ < end_) { + if (!parse_string(nullptr)) return false; + skip_ws(); + if (!consume(':')) return fail("expected ':' in JSON object"); + if (!skip_value()) return false; + skip_ws(); + if (consume('}')) return true; + if (!consume(',')) return fail("expected ',' in JSON object"); + skip_ws(); + } + return fail("unterminated JSON object"); + } + + bool skip_array() { + if (!consume('[')) return false; + skip_ws(); + if (consume(']')) return true; + while (cur_ < end_) { + if (!skip_value()) return false; + skip_ws(); + if (consume(']')) return true; + if (!consume(',')) return fail("expected ',' in JSON array"); + skip_ws(); + } + return fail("unterminated JSON array"); + } + + bool skip_number() { + const char* start = cur_; + if (cur_ < end_ && *cur_ == '-') ++cur_; + if (cur_ >= end_ || + !std::isdigit(static_cast(*cur_))) { + cur_ = start; + return fail("unsupported JSON value"); + } + if (*cur_ == '0') { + ++cur_; + } else { + while (cur_ < end_ && + std::isdigit(static_cast(*cur_))) ++cur_; + } + if (cur_ < end_ && *cur_ == '.') { + ++cur_; + const char* fractional = cur_; + while (cur_ < end_ && + std::isdigit(static_cast(*cur_))) ++cur_; + if (cur_ == fractional) return fail("invalid JSON number"); + } + if (cur_ < end_ && (*cur_ == 'e' || *cur_ == 'E')) { + ++cur_; + if (cur_ < end_ && (*cur_ == '+' || *cur_ == '-')) ++cur_; + const char* exponent = cur_; + while (cur_ < end_ && + std::isdigit(static_cast(*cur_))) ++cur_; + if (cur_ == exponent) return fail("invalid JSON number"); + } + return true; + } + + void skip_ws() { + while (cur_ < end_ && + std::isspace(static_cast(*cur_))) ++cur_; + } + + bool consume(char c) { + if (cur_ >= end_ || *cur_ != c) return false; + ++cur_; + return true; + } + + bool fail(const char* message) { + error_ = message; + error_ += " at header byte "; + error_ += std::to_string(static_cast(cur_ - begin_)); + return false; + } + + const char* begin_; + const char* cur_; + const char* end_; + std::string error_; + std::map parsed_; +}; + +} // namespace + +SafetensorsFile::~SafetensorsFile() { close(); } + +SafetensorsFile::SafetensorsFile(SafetensorsFile&& other) noexcept { + move_from(std::move(other)); +} + +SafetensorsFile& SafetensorsFile::operator=(SafetensorsFile&& other) noexcept { + if (this != &other) { + close(); + move_from(std::move(other)); + } + return *this; +} + +void SafetensorsFile::move_from(SafetensorsFile&& other) noexcept { + fd_ = other.fd_; + mapping_ = other.mapping_; + mapping_bytes_ = other.mapping_bytes_; + data_offset_ = other.data_offset_; + path_ = std::move(other.path_); + error_ = std::move(other.error_); + tensors_ = std::move(other.tensors_); + other.fd_ = -1; + other.mapping_ = nullptr; + other.mapping_bytes_ = 0; + other.data_offset_ = 0; +} + +bool SafetensorsFile::open(const std::string& path) { + close(); + fd_ = ::open(path.c_str(), O_RDONLY | O_CLOEXEC); + if (fd_ < 0) { + error_ = "unable to open safetensors file: "; + error_ += std::strerror(errno); + return false; + } + struct stat st {}; + if (::fstat(fd_, &st) != 0 || !S_ISREG(st.st_mode) || st.st_size < 9) { + error_ = "safetensors path is not a non-empty regular file"; + close(); + return false; + } + mapping_bytes_ = static_cast(st.st_size); + mapping_ = ::mmap(nullptr, static_cast(mapping_bytes_), + PROT_READ, MAP_PRIVATE, fd_, 0); + if (mapping_ == MAP_FAILED) { + mapping_ = nullptr; + error_ = "unable to mmap safetensors file: "; + error_ += std::strerror(errno); + close(); + return false; + } + + const auto* bytes = static_cast(mapping_); + std::uint64_t header_bytes = 0; + for (int i = 7; i >= 0; --i) { + header_bytes = (header_bytes << 8) | bytes[i]; + } + if (!header_bytes || header_bytes > kMaxHeaderBytes || + header_bytes > mapping_bytes_ - 8) { + error_ = "safetensors header length is invalid"; + close(); + return false; + } + data_offset_ = 8 + header_bytes; + const char* header = static_cast(mapping_) + 8; + HeaderParser parser(header, header + header_bytes); + if (!parser.parse(&tensors_)) { + error_ = parser.error(); + close(); + return false; + } + + std::vector> spans; + spans.reserve(tensors_.size()); + const std::uint64_t payload_bytes = mapping_bytes_ - data_offset_; + for (const auto& entry : tensors_) { + const SafetensorInfo& tensor = entry.second; + if (tensor.data_offset > payload_bytes || + tensor.bytes > payload_bytes - tensor.data_offset) { + error_ = "tensor byte range exceeds safetensors payload: "; + error_ += entry.first; + close(); + return false; + } + spans.emplace_back(tensor.data_offset, + tensor.data_offset + tensor.bytes); + } + std::sort(spans.begin(), spans.end()); + for (std::size_t i = 1; i < spans.size(); ++i) { + if (spans[i].first < spans[i - 1].second) { + error_ = "overlapping tensor byte ranges in safetensors payload"; + close(); + return false; + } + } + path_ = path; + error_.clear(); + return true; +} + +void SafetensorsFile::close() { + if (mapping_) { + ::munmap(mapping_, static_cast(mapping_bytes_)); + } + if (fd_ >= 0) ::close(fd_); + fd_ = -1; + mapping_ = nullptr; + mapping_bytes_ = 0; + data_offset_ = 0; + path_.clear(); + tensors_.clear(); +} + +const SafetensorInfo* SafetensorsFile::find(const std::string& name) const { + const auto it = tensors_.find(name); + return it == tensors_.end() ? nullptr : &it->second; +} + +const void* SafetensorsFile::data(const SafetensorInfo& tensor) const { + if (!mapping_ || tensor.data_offset > mapping_bytes_ - data_offset_ || + tensor.bytes > mapping_bytes_ - data_offset_ - tensor.data_offset) { + return nullptr; + } + return static_cast(mapping_) + data_offset_ + + tensor.data_offset; +} + +} // namespace loader +} // namespace flashrt diff --git a/cpp/loader/src/sha256.cpp b/cpp/loader/src/sha256.cpp new file mode 100644 index 00000000..780a4873 --- /dev/null +++ b/cpp/loader/src/sha256.cpp @@ -0,0 +1,175 @@ +#include "flashrt/cpp/loader/sha256.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace flashrt { +namespace loader { +namespace { + +constexpr std::uint32_t kRound[64] = { + 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, +}; + +std::uint32_t rotate(std::uint32_t value, unsigned bits) { + return (value >> bits) | (value << (32 - bits)); +} + +class Sha256 { +public: + void update(const unsigned char* data, std::size_t bytes) { + total_bytes_ += bytes; + while (bytes) { + const std::size_t count = + std::min(bytes, block_.size() - block_bytes_); + std::copy(data, data + count, block_.begin() + block_bytes_); + block_bytes_ += count; + data += count; + bytes -= count; + if (block_bytes_ == block_.size()) { + transform(block_.data()); + block_bytes_ = 0; + } + } + } + + std::array finish() { + const std::uint64_t bits = total_bytes_ * 8; + block_[block_bytes_++] = 0x80; + if (block_bytes_ > 56) { + std::fill(block_.begin() + block_bytes_, block_.end(), 0); + transform(block_.data()); + block_bytes_ = 0; + } + std::fill(block_.begin() + block_bytes_, block_.begin() + 56, 0); + for (int i = 0; i < 8; ++i) { + block_[63 - i] = static_cast(bits >> (8 * i)); + } + transform(block_.data()); + std::array output{}; + for (std::size_t i = 0; i < state_.size(); ++i) { + output[4 * i] = static_cast(state_[i] >> 24); + output[4 * i + 1] = static_cast(state_[i] >> 16); + output[4 * i + 2] = static_cast(state_[i] >> 8); + output[4 * i + 3] = static_cast(state_[i]); + } + return output; + } + +private: + void transform(const unsigned char* data) { + std::uint32_t words[64]{}; + for (int i = 0; i < 16; ++i) { + words[i] = (static_cast(data[4 * i]) << 24) | + (static_cast(data[4 * i + 1]) << 16) | + (static_cast(data[4 * i + 2]) << 8) | + static_cast(data[4 * i + 3]); + } + for (int i = 16; i < 64; ++i) { + const std::uint32_t s0 = rotate(words[i - 15], 7) ^ + rotate(words[i - 15], 18) ^ + (words[i - 15] >> 3); + const std::uint32_t s1 = rotate(words[i - 2], 17) ^ + rotate(words[i - 2], 19) ^ + (words[i - 2] >> 10); + words[i] = words[i - 16] + s0 + words[i - 7] + s1; + } + std::uint32_t a = state_[0]; + std::uint32_t b = state_[1]; + std::uint32_t c = state_[2]; + std::uint32_t d = state_[3]; + std::uint32_t e = state_[4]; + std::uint32_t f = state_[5]; + std::uint32_t g = state_[6]; + std::uint32_t h = state_[7]; + for (int i = 0; i < 64; ++i) { + const std::uint32_t sum1 = rotate(e, 6) ^ rotate(e, 11) ^ + rotate(e, 25); + const std::uint32_t choose = (e & f) ^ (~e & g); + const std::uint32_t t1 = h + sum1 + choose + kRound[i] + words[i]; + const std::uint32_t sum0 = rotate(a, 2) ^ rotate(a, 13) ^ + rotate(a, 22); + const std::uint32_t majority = (a & b) ^ (a & c) ^ (b & c); + const std::uint32_t t2 = sum0 + majority; + h = g; + g = f; + f = e; + e = d + t1; + d = c; + c = b; + b = a; + a = t1 + t2; + } + state_[0] += a; + state_[1] += b; + state_[2] += c; + state_[3] += d; + state_[4] += e; + state_[5] += f; + state_[6] += g; + state_[7] += h; + } + + std::array state_ = { + 0x6a09e667u, 0xbb67ae85u, 0x3c6ef372u, 0xa54ff53au, + 0x510e527fu, 0x9b05688cu, 0x1f83d9abu, 0x5be0cd19u, + }; + std::array block_{}; + std::size_t block_bytes_ = 0; + std::uint64_t total_bytes_ = 0; +}; + +} // namespace + +bool sha256_file(const std::string& path, std::string* hex, + std::string* error) { + if (!hex) { + if (error) *error = "SHA-256 output is null"; + return false; + } + std::ifstream file(path, std::ios::binary); + if (!file) { + if (error) *error = "unable to open file for SHA-256: " + path; + return false; + } + Sha256 hash; + std::vector buffer(1 << 20); + while (file) { + file.read(reinterpret_cast(buffer.data()), buffer.size()); + const std::streamsize count = file.gcount(); + if (count > 0) { + hash.update(buffer.data(), static_cast(count)); + } + } + if (!file.eof()) { + if (error) *error = "failed while reading file for SHA-256: " + path; + return false; + } + const auto digest = hash.finish(); + std::ostringstream output; + output << std::hex << std::setfill('0'); + for (unsigned char byte : digest) output << std::setw(2) << unsigned(byte); + *hex = output.str(); + if (error) error->clear(); + return true; +} + +} // namespace loader +} // namespace flashrt diff --git a/cpp/modalities/include/flashrt/cpp/modalities/action.h b/cpp/modalities/include/flashrt/cpp/modalities/action.h index 1596f30e..02d967c5 100644 --- a/cpp/modalities/include/flashrt/cpp/modalities/action.h +++ b/cpp/modalities/include/flashrt/cpp/modalities/action.h @@ -24,6 +24,14 @@ struct ActionPostprocessSpec { bool clamp = false; }; +struct ActionStaging { + void* host_pinned = nullptr; + std::uint64_t bytes = 0; +}; + +Status action_staging_create(ActionStaging* out, std::uint64_t bytes); +void action_staging_destroy(ActionStaging*); + Status postprocess_action_cpu(const ActionPostprocessSpec& spec, TensorView model_output, std::vector* robot_actions); @@ -34,7 +42,8 @@ Status postprocess_action_cpu(const ActionPostprocessSpec& spec, Status postprocess_action(const ActionPostprocessSpec& spec, TensorView model_output, std::vector* robot_actions, - void* stream = nullptr); + void* stream = nullptr, + ActionStaging* staging = nullptr); std::uint64_t required_action_output_bytes(const ActionPostprocessSpec& spec, DType dtype); diff --git a/cpp/modalities/include/flashrt/cpp/modalities/text.h b/cpp/modalities/include/flashrt/cpp/modalities/text.h new file mode 100644 index 00000000..b4a3c39c --- /dev/null +++ b/cpp/modalities/include/flashrt/cpp/modalities/text.h @@ -0,0 +1,44 @@ +#ifndef FLASHRT_MODALITIES_TEXT_H +#define FLASHRT_MODALITIES_TEXT_H + +#include "flashrt/cpp/modalities/types.h" + +#include + +namespace flashrt { +namespace modalities { + +struct EmbeddingGatherSpec { + std::uint64_t vocab_size = 0; + std::uint64_t hidden_dim = 0; + float scale = 1.0f; +}; + +struct TextEmbeddingStaging { + void* device_token_ids = nullptr; + void* device_status = nullptr; + std::uint64_t max_tokens = 0; +}; + +Status gather_token_embeddings_cpu(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output); + +Status text_embedding_staging_create(TextEmbeddingStaging* out, + std::uint64_t max_tokens); +void text_embedding_staging_destroy(TextEmbeddingStaging*); + +Status gather_token_embeddings(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream = nullptr, + TextEmbeddingStaging* staging = nullptr); + +} // namespace modalities +} // namespace flashrt + +#endif // FLASHRT_MODALITIES_TEXT_H diff --git a/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h new file mode 100644 index 00000000..2a4667c0 --- /dev/null +++ b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h @@ -0,0 +1,55 @@ +#ifndef FLASHRT_MODALITIES_TOKENIZER_H +#define FLASHRT_MODALITIES_TOKENIZER_H + +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace modalities { + +struct SentencePieceEncodeOptions { + bool add_bos = false; + bool add_eos = false; + bool pad_to_max_tokens = false; + std::uint64_t max_tokens = 0; + std::int32_t pad_id = 0; +}; + +class SentencePieceTokenizer final { +public: + SentencePieceTokenizer(); + ~SentencePieceTokenizer(); + + SentencePieceTokenizer(SentencePieceTokenizer&&) noexcept; + SentencePieceTokenizer& operator=(SentencePieceTokenizer&&) noexcept; + + SentencePieceTokenizer(const SentencePieceTokenizer&) = delete; + SentencePieceTokenizer& operator=(const SentencePieceTokenizer&) = delete; + + Status load_model(const std::string& model_path); + Status encode(const std::string& text, + const SentencePieceEncodeOptions& options, + std::vector* token_ids); + void reserve(std::uint64_t max_tokens); + std::uint64_t workspace_capacity() const; + + std::int32_t bos_id() const; + std::int32_t eos_id() const; + std::int32_t unk_id() const; + std::int32_t pad_id() const; + std::uint64_t vocab_size() const; + bool loaded() const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace modalities +} // namespace flashrt + +#endif // FLASHRT_MODALITIES_TOKENIZER_H diff --git a/cpp/modalities/include/flashrt/cpp/modalities/vision.h b/cpp/modalities/include/flashrt/cpp/modalities/vision.h index 5bea0881..a5cdc330 100644 --- a/cpp/modalities/include/flashrt/cpp/modalities/vision.h +++ b/cpp/modalities/include/flashrt/cpp/modalities/vision.h @@ -11,12 +11,14 @@ namespace modalities { enum class NormalizeMode { kScaleShift, + kDivideShift, kMeanStd, }; struct NormalizeSpec { NormalizeMode mode = NormalizeMode::kScaleShift; float scale = 1.0f / 127.5f; + float divisor = 127.5f; float shift = -1.0f; float mean[3] = {0.0f, 0.0f, 0.0f}; float inv_std[3] = {1.0f, 1.0f, 1.0f}; diff --git a/cpp/modalities/src/action_cpu.cpp b/cpp/modalities/src/action_cpu.cpp index c020814f..06998fbf 100644 --- a/cpp/modalities/src/action_cpu.cpp +++ b/cpp/modalities/src/action_cpu.cpp @@ -31,6 +31,38 @@ bool has_dim(const std::vector& v, int dim) { } // namespace +Status action_staging_create(ActionStaging* out, std::uint64_t bytes) { + if (!out || !bytes) { + return Status::error(StatusCode::kInvalidArgument, + "invalid action staging capacity"); + } + action_staging_destroy(out); +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return Status::error(StatusCode::kUnsupported, + "action staging requires the CUDA build"); +#else + cudaError_t rc = cudaMallocHost(&out->host_pinned, + static_cast(bytes)); + if (rc != cudaSuccess) { + *out = ActionStaging{}; + return Status::error( + StatusCode::kBackend, + std::string("cuda action host staging allocation failed: ") + + cudaGetErrorString(rc)); + } + out->bytes = bytes; + return Status::ok(); +#endif +} + +void action_staging_destroy(ActionStaging* staging) { + if (!staging) return; +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING + if (staging->host_pinned) cudaFreeHost(staging->host_pinned); +#endif + *staging = ActionStaging{}; +} + std::uint64_t required_action_output_bytes(const ActionPostprocessSpec& spec, DType dtype) { if (spec.chunk <= 0 || spec.model_dim <= 0) return 0; @@ -97,7 +129,8 @@ Status postprocess_action_cpu(const ActionPostprocessSpec& spec, Status postprocess_action(const ActionPostprocessSpec& spec, TensorView model_output, std::vector* robot_actions, - void* stream) { + void* stream, + ActionStaging* persistent_staging) { if (model_output.place == MemoryPlace::kHost || model_output.place == MemoryPlace::kHostPinned) { return postprocess_action_cpu(spec, model_output, robot_actions); @@ -117,15 +150,27 @@ Status postprocess_action(const ActionPostprocessSpec& spec, return Status::error(StatusCode::kInsufficientStorage, "device action output storage is too small"); } - std::vector staging(static_cast(bytes)); + std::vector fallback; + void* staging = nullptr; + if (persistent_staging) { + if (!persistent_staging->host_pinned || + persistent_staging->bytes < bytes) { + return Status::error(StatusCode::kInsufficientStorage, + "action staging capacity is too small"); + } + staging = persistent_staging->host_pinned; + } else { + fallback.resize(static_cast(bytes)); + staging = fallback.data(); + } cudaError_t rc = cudaSuccess; if (stream) { auto* cuda_stream = reinterpret_cast(stream); - rc = cudaMemcpyAsync(staging.data(), model_output.data, bytes, + rc = cudaMemcpyAsync(staging, model_output.data, bytes, cudaMemcpyDeviceToHost, cuda_stream); if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); } else { - rc = cudaMemcpy(staging.data(), model_output.data, bytes, + rc = cudaMemcpy(staging, model_output.data, bytes, cudaMemcpyDeviceToHost); } if (rc != cudaSuccess) { @@ -134,7 +179,7 @@ Status postprocess_action(const ActionPostprocessSpec& spec, cudaGetErrorString(rc)); } TensorView host; - host.data = staging.data(); + host.data = staging; host.bytes = bytes; host.dtype = model_output.dtype; host.place = MemoryPlace::kHost; diff --git a/cpp/modalities/src/sentencepiece_tokenizer.cpp b/cpp/modalities/src/sentencepiece_tokenizer.cpp new file mode 100644 index 00000000..363e61e7 --- /dev/null +++ b/cpp/modalities/src/sentencepiece_tokenizer.cpp @@ -0,0 +1,137 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +#include + +#include +#include + +namespace flashrt { +namespace modalities { + +struct SentencePieceTokenizer::Impl { + sentencepiece::SentencePieceProcessor processor; + std::vector encoded; + bool loaded = false; +}; + +SentencePieceTokenizer::SentencePieceTokenizer() + : impl_(new Impl()) {} + +SentencePieceTokenizer::~SentencePieceTokenizer() = default; + +SentencePieceTokenizer::SentencePieceTokenizer( + SentencePieceTokenizer&&) noexcept = default; + +SentencePieceTokenizer& SentencePieceTokenizer::operator=( + SentencePieceTokenizer&&) noexcept = default; + +Status SentencePieceTokenizer::load_model(const std::string& model_path) { + auto status = impl_->processor.Load(model_path); + if (!status.ok()) { + impl_->loaded = false; + return Status::error(StatusCode::kNotFound, status.ToString()); + } + impl_->loaded = true; + return Status::ok(); +} + +Status SentencePieceTokenizer::encode( + const std::string& text, + const SentencePieceEncodeOptions& options, + std::vector* token_ids) { + if (!token_ids) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids output is null"); + } + token_ids->clear(); + if (!impl_->loaded) { + return Status::error(StatusCode::kInvalidArgument, + "SentencePiece model is not loaded"); + } + + impl_->encoded.clear(); + auto status = impl_->processor.Encode(text, &impl_->encoded); + if (!status.ok()) { + return Status::error(StatusCode::kBackend, status.ToString()); + } + const std::uint64_t extra = + (options.add_bos ? 1u : 0u) + (options.add_eos ? 1u : 0u); + if (options.max_tokens && impl_->encoded.size() + extra > + options.max_tokens) { + return Status::error(StatusCode::kShapeMismatch, + "encoded token sequence exceeds max_tokens"); + } + if (impl_->encoded.size() + extra > + static_cast(std::numeric_limits::max())) { + return Status::error(StatusCode::kInsufficientStorage, + "encoded token sequence is too large"); + } + + if (options.add_bos) { + const int bos = impl_->processor.bos_id(); + if (bos < 0) { + return Status::error(StatusCode::kInvalidArgument, + "tokenizer has no BOS id"); + } + token_ids->push_back(static_cast(bos)); + } + token_ids->reserve(impl_->encoded.size() + extra); + for (int id : impl_->encoded) { + token_ids->push_back(static_cast(id)); + } + if (options.add_eos) { + const int eos = impl_->processor.eos_id(); + if (eos < 0) { + return Status::error(StatusCode::kInvalidArgument, + "tokenizer has no EOS id"); + } + token_ids->push_back(static_cast(eos)); + } + + if (options.max_tokens) { + if (options.pad_to_max_tokens) { + token_ids->resize(options.max_tokens, options.pad_id); + } + } else if (options.pad_to_max_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "pad_to_max_tokens requires max_tokens"); + } + return Status::ok(); +} + +void SentencePieceTokenizer::reserve(std::uint64_t max_tokens) { + impl_->encoded.reserve(static_cast(max_tokens)); +} + +std::uint64_t SentencePieceTokenizer::workspace_capacity() const { + return static_cast(impl_->encoded.capacity()); +} + +std::int32_t SentencePieceTokenizer::bos_id() const { + return impl_->loaded ? impl_->processor.bos_id() : -1; +} + +std::int32_t SentencePieceTokenizer::eos_id() const { + return impl_->loaded ? impl_->processor.eos_id() : -1; +} + +std::int32_t SentencePieceTokenizer::unk_id() const { + return impl_->loaded ? impl_->processor.unk_id() : -1; +} + +std::int32_t SentencePieceTokenizer::pad_id() const { + return impl_->loaded ? impl_->processor.pad_id() : -1; +} + +std::uint64_t SentencePieceTokenizer::vocab_size() const { + return impl_->loaded + ? static_cast(impl_->processor.GetPieceSize()) + : 0; +} + +bool SentencePieceTokenizer::loaded() const { + return impl_->loaded; +} + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/text_cpu.cpp b/cpp/modalities/src/text_cpu.cpp new file mode 100644 index 00000000..454b9acb --- /dev/null +++ b/cpp/modalities/src/text_cpu.cpp @@ -0,0 +1,232 @@ +#include "flashrt/cpp/modalities/text.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace modalities { + +#ifdef FLASHRT_CPP_WITH_CUDA_KERNELS +Status gather_token_embeddings_cuda(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream, + TextEmbeddingStaging* staging); +#endif + +namespace { + +float load_scalar(const void* base, std::uint64_t index, DType dtype) { + switch (dtype) { + case DType::kFloat32: + return static_cast(base)[index]; + case DType::kBFloat16: + return bfloat16_to_float( + static_cast(base)[index]); + case DType::kFloat16: + return float16_to_float( + static_cast(base)[index]); + case DType::kUInt8: + return static_cast( + static_cast(base)[index]); + } + return 0.0f; +} + +void store_scalar(void* base, std::uint64_t index, DType dtype, float value) { + switch (dtype) { + case DType::kFloat32: + static_cast(base)[index] = value; + break; + case DType::kBFloat16: + static_cast(base)[index] = float_to_bfloat16(value); + break; + case DType::kFloat16: + static_cast(base)[index] = float_to_float16(value); + break; + case DType::kUInt8: + static_cast(base)[index] = + static_cast(value); + break; + } +} + +Status validate_matrix(const TensorView& tensor, const char* name, + std::uint64_t rows, std::uint64_t cols) { + Status st = validate_host_tensor(tensor, name); + if (!st.ok_status()) return st; + if (tensor.layout != Layout::kFlat || tensor.shape.rank != 2 || + tensor.shape.dims[0] != rows || tensor.shape.dims[1] != cols) { + return Status::error(StatusCode::kShapeMismatch, + std::string(name) + " shape mismatch"); + } + return Status::ok(); +} + +} // namespace + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +Status text_embedding_staging_create(TextEmbeddingStaging* out, + std::uint64_t max_tokens) { + if (!out || !max_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "invalid text embedding staging capacity"); + } + *out = TextEmbeddingStaging{}; + cudaError_t rc = cudaMalloc(&out->device_token_ids, + max_tokens * sizeof(std::int32_t)); + if (rc != cudaSuccess) { + return Status::error( + StatusCode::kBackend, + std::string("text token staging cudaMalloc failed: ") + + cudaGetErrorString(rc)); + } + rc = cudaMalloc(&out->device_status, sizeof(int)); + if (rc != cudaSuccess) { + cudaFree(out->device_token_ids); + *out = TextEmbeddingStaging{}; + return Status::error( + StatusCode::kBackend, + std::string("text status staging cudaMalloc failed: ") + + cudaGetErrorString(rc)); + } + out->max_tokens = max_tokens; + return Status::ok(); +} + +void text_embedding_staging_destroy(TextEmbeddingStaging* s) { + if (!s) return; + if (s->device_token_ids) cudaFree(s->device_token_ids); + if (s->device_status) cudaFree(s->device_status); + *s = TextEmbeddingStaging{}; +} +#else +Status text_embedding_staging_create(TextEmbeddingStaging* out, + std::uint64_t) { + if (out) *out = TextEmbeddingStaging{}; + return Status::error(StatusCode::kUnsupported, + "text embedding staging requires the CUDA build"); +} + +void text_embedding_staging_destroy(TextEmbeddingStaging* s) { + if (s) *s = TextEmbeddingStaging{}; +} +#endif + +Status gather_token_embeddings_cpu(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output) { + if (!token_ids && n_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids is null"); + } + if (!spec.vocab_size || !spec.hidden_dim) { + return Status::error(StatusCode::kInvalidArgument, + "invalid embedding gather dimensions"); + } + Status st = validate_matrix(embedding_table, "embedding_table", + spec.vocab_size, spec.hidden_dim); + if (!st.ok_status()) return st; + st = validate_matrix(output, "embedding_output", n_tokens, + spec.hidden_dim); + if (!st.ok_status()) return st; + + for (std::uint64_t t = 0; t < n_tokens; ++t) { + const std::int32_t token = token_ids[t]; + if (token < 0 || + static_cast(token) >= spec.vocab_size) { + return Status::error(StatusCode::kInvalidArgument, + "token id is out of vocabulary range"); + } + const std::uint64_t src_base = + static_cast(token) * spec.hidden_dim; + const std::uint64_t dst_base = t * spec.hidden_dim; + for (std::uint64_t d = 0; d < spec.hidden_dim; ++d) { + const float value = load_scalar( + embedding_table.data, src_base + d, embedding_table.dtype); + store_scalar(output.data, dst_base + d, output.dtype, + value * spec.scale); + } + } + return Status::ok(); +} + +Status gather_token_embeddings(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream, + TextEmbeddingStaging* staging) { + if (output.place == MemoryPlace::kHost || + output.place == MemoryPlace::kHostPinned) { + (void)stream; + (void)staging; + return gather_token_embeddings_cpu(spec, token_ids, n_tokens, + embedding_table, output); + } + if (output.place != MemoryPlace::kDevice || + embedding_table.place != MemoryPlace::kDevice) { + return Status::error(StatusCode::kUnsupported, + "device text embedding requires device tensors"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)stream; + (void)staging; + return Status::error(StatusCode::kUnsupported, + "device text embedding was not enabled at build time"); +#else + if (!token_ids && n_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids is null"); + } + if (staging && staging->max_tokens < n_tokens) { + return Status::error(StatusCode::kInsufficientStorage, + "text token staging capacity is too small"); + } +#ifdef FLASHRT_CPP_WITH_CUDA_KERNELS + return gather_token_embeddings_cuda(spec, token_ids, n_tokens, + embedding_table, output, stream, + staging); +#else + std::vector host_bytes( + static_cast(n_tokens * spec.hidden_dim * + dtype_size(output.dtype))); + TensorView host_output{host_bytes.data(), + static_cast(host_bytes.size()), + output.dtype, MemoryPlace::kHost, output.layout, + Shape{n_tokens, spec.hidden_dim}}; + TensorView host_table = embedding_table; + if (embedding_table.place != MemoryPlace::kHost && + embedding_table.place != MemoryPlace::kHostPinned) { + return Status::error(StatusCode::kUnsupported, + "CUDA kernel build is required for device embedding tables"); + } + Status st = gather_token_embeddings_cpu(spec, token_ids, n_tokens, + host_table, host_output); + if (!st.ok_status()) return st; + cudaStream_t cuda_stream = reinterpret_cast(stream); + cudaError_t rc = cudaMemcpyAsync(output.data, host_bytes.data(), + host_bytes.size(), cudaMemcpyHostToDevice, + cuda_stream); + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + if (rc != cudaSuccess) { + return Status::error(StatusCode::kBackend, + std::string("cuda H2D text embedding failed: ") + + cudaGetErrorString(rc)); + } + return Status::ok(); +#endif +#endif +} + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/text_cuda.cu b/cpp/modalities/src/text_cuda.cu new file mode 100644 index 00000000..da8bf63d --- /dev/null +++ b/cpp/modalities/src/text_cuda.cu @@ -0,0 +1,208 @@ +#include "flashrt/cpp/modalities/text.h" + +#include +#include + +#include +#include +#include + +namespace flashrt { +namespace modalities { +namespace { + +__device__ __forceinline__ float bf16_to_f32(std::uint16_t value) { + return __uint_as_float(static_cast(value) << 16); +} + +__device__ __forceinline__ std::uint16_t f32_to_bf16(float value) { + std::uint32_t bits = __float_as_uint(value); + const std::uint32_t lsb = (bits >> 16) & 1u; + bits += 0x7fffu + lsb; + return static_cast(bits >> 16); +} + +__device__ __forceinline__ float load_value(const void* base, + std::uint64_t index, + int dtype) { + if (dtype == 1) return static_cast(base)[index]; + if (dtype == 2) return __half2float(static_cast(base)[index]); + if (dtype == 3) { + return bf16_to_f32(static_cast(base)[index]); + } + return static_cast(static_cast(base)[index]); +} + +__device__ __forceinline__ void store_value(void* base, + std::uint64_t index, + int dtype, + float value) { + if (dtype == 1) { + static_cast(base)[index] = value; + } else if (dtype == 2) { + static_cast<__half*>(base)[index] = __float2half_rn(value); + } else if (dtype == 3) { + static_cast(base)[index] = f32_to_bf16(value); + } else { + static_cast(base)[index] = + static_cast(value); + } +} + +int dtype_code(DType dtype) { + switch (dtype) { + case DType::kFloat32: return 1; + case DType::kFloat16: return 2; + case DType::kBFloat16: return 3; + case DType::kUInt8: return 0; + } + return 0; +} + +Status validate_device_matrix(const TensorView& tensor, const char* name, + std::uint64_t rows, std::uint64_t cols) { + if (!tensor.data) { + return Status::error(StatusCode::kInvalidArgument, + std::string(name) + " has null data"); + } + if (tensor.place != MemoryPlace::kDevice) { + return Status::error(StatusCode::kUnsupported, + std::string(name) + " is not device memory"); + } + if (tensor.layout != Layout::kFlat || tensor.shape.rank != 2 || + tensor.shape.dims[0] != rows || tensor.shape.dims[1] != cols) { + return Status::error(StatusCode::kShapeMismatch, + std::string(name) + " shape mismatch"); + } + const std::uint64_t need = rows * cols * dtype_size(tensor.dtype); + if (tensor.bytes < need) { + return Status::error(StatusCode::kInsufficientStorage, + std::string(name) + " storage is too small"); + } + return Status::ok(); +} + +__global__ void gather_kernel(const std::int32_t* ids, + std::uint64_t n_tokens, + std::uint64_t vocab_size, + std::uint64_t hidden_dim, + const void* table, + int table_dtype, + void* output, + int output_dtype, + float scale, + int* bad_token) { + const std::uint64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const std::uint64_t total = n_tokens * hidden_dim; + if (idx >= total) return; + const std::uint64_t token_index = idx / hidden_dim; + const std::uint64_t dim = idx - token_index * hidden_dim; + const std::int32_t token = ids[token_index]; + if (token < 0 || static_cast(token) >= vocab_size) { + atomicCAS(bad_token, 0, 1); + return; + } + const std::uint64_t src = + static_cast(token) * hidden_dim + dim; + const float value = load_value(table, src, table_dtype) * scale; + store_value(output, idx, output_dtype, value); +} + +const char* cuda_error(cudaError_t rc) { + return cudaGetErrorString(rc); +} + +} // namespace + +Status gather_token_embeddings_cuda(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream, + TextEmbeddingStaging* staging) { + if (!token_ids && n_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids is null"); + } + if (!spec.vocab_size || !spec.hidden_dim) { + return Status::error(StatusCode::kInvalidArgument, + "invalid embedding gather dimensions"); + } + Status st = validate_device_matrix(embedding_table, "embedding_table", + spec.vocab_size, spec.hidden_dim); + if (!st.ok_status()) return st; + st = validate_device_matrix(output, "embedding_output", n_tokens, + spec.hidden_dim); + if (!st.ok_status()) return st; + if (staging && staging->max_tokens < n_tokens) { + return Status::error(StatusCode::kInsufficientStorage, + "text token staging capacity is too small"); + } + + cudaStream_t cuda_stream = reinterpret_cast(stream); + std::int32_t* d_ids = nullptr; + int* d_bad = nullptr; + cudaError_t rc = cudaSuccess; + if (staging) { + d_ids = static_cast(staging->device_token_ids); + d_bad = static_cast(staging->device_status); + } else { + rc = cudaMalloc(&d_ids, n_tokens * sizeof(std::int32_t)); + if (rc != cudaSuccess) { + return Status::error( + StatusCode::kBackend, + std::string("cudaMalloc text token ids failed: ") + + cuda_error(rc)); + } + } + if (!d_bad) rc = cudaMalloc(&d_bad, sizeof(int)); + if (rc == cudaSuccess) { + rc = cudaMemsetAsync(d_bad, 0, sizeof(int), cuda_stream); + } + if (rc == cudaSuccess && n_tokens) { + rc = cudaMemcpyAsync(d_ids, token_ids, n_tokens * sizeof(std::int32_t), + cudaMemcpyHostToDevice, cuda_stream); + } + if (rc != cudaSuccess) { + if (!staging) cudaFree(d_ids); + if (!staging && d_bad) cudaFree(d_bad); + return Status::error(StatusCode::kBackend, + std::string("cuda text token upload failed: ") + + cuda_error(rc)); + } + + const std::uint64_t total = n_tokens * spec.hidden_dim; + if (total) { + const int block = 256; + const int grid = static_cast((total + block - 1) / block); + gather_kernel<<>>( + d_ids, n_tokens, spec.vocab_size, spec.hidden_dim, + embedding_table.data, dtype_code(embedding_table.dtype), + output.data, dtype_code(output.dtype), spec.scale, d_bad); + rc = cudaGetLastError(); + } + + int bad = 0; + if (rc == cudaSuccess) { + rc = cudaMemcpyAsync(&bad, d_bad, sizeof(int), cudaMemcpyDeviceToHost, + cuda_stream); + } + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + if (!staging) cudaFree(d_ids); + if (!staging) cudaFree(d_bad); + if (rc != cudaSuccess) { + return Status::error(StatusCode::kBackend, + std::string("text embedding CUDA failed: ") + + cuda_error(rc)); + } + if (bad) { + return Status::error(StatusCode::kInvalidArgument, + "token id is out of vocabulary range"); + } + return Status::ok(); +} + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/tokenizer_unavailable.cpp b/cpp/modalities/src/tokenizer_unavailable.cpp new file mode 100644 index 00000000..7cb70158 --- /dev/null +++ b/cpp/modalities/src/tokenizer_unavailable.cpp @@ -0,0 +1,52 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +namespace flashrt { +namespace modalities { + +struct SentencePieceTokenizer::Impl {}; + +SentencePieceTokenizer::SentencePieceTokenizer() + : impl_(new Impl()) {} + +SentencePieceTokenizer::~SentencePieceTokenizer() = default; + +SentencePieceTokenizer::SentencePieceTokenizer( + SentencePieceTokenizer&&) noexcept = default; + +SentencePieceTokenizer& SentencePieceTokenizer::operator=( + SentencePieceTokenizer&&) noexcept = default; + +Status SentencePieceTokenizer::load_model(const std::string& model_path) { + (void)model_path; + return Status::error( + StatusCode::kUnsupported, + "native SentencePiece support is not enabled in this build"); +} + +Status SentencePieceTokenizer::encode( + const std::string& text, + const SentencePieceEncodeOptions& options, + std::vector* token_ids) { + (void)text; + (void)options; + if (token_ids) token_ids->clear(); + return Status::error( + StatusCode::kUnsupported, + "native SentencePiece support is not enabled in this build"); +} + +void SentencePieceTokenizer::reserve(std::uint64_t max_tokens) { + (void)max_tokens; +} + +std::uint64_t SentencePieceTokenizer::workspace_capacity() const { return 0; } + +std::int32_t SentencePieceTokenizer::bos_id() const { return -1; } +std::int32_t SentencePieceTokenizer::eos_id() const { return -1; } +std::int32_t SentencePieceTokenizer::unk_id() const { return -1; } +std::int32_t SentencePieceTokenizer::pad_id() const { return -1; } +std::uint64_t SentencePieceTokenizer::vocab_size() const { return 0; } +bool SentencePieceTokenizer::loaded() const { return false; } + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/vision_cpu.cpp b/cpp/modalities/src/vision_cpu.cpp index e6c58e49..cc217d4e 100644 --- a/cpp/modalities/src/vision_cpu.cpp +++ b/cpp/modalities/src/vision_cpu.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #ifdef FLASHRT_CPP_WITH_CUDA_STAGING @@ -111,6 +112,9 @@ float normalize_value(float raw, int c, const NormalizeSpec& spec) { if (spec.mode == NormalizeMode::kScaleShift) { return raw * spec.scale + spec.shift; } + if (spec.mode == NormalizeMode::kDivideShift) { + return raw / spec.divisor + spec.shift; + } return (raw / 255.0f - spec.mean[c]) * spec.inv_std[c]; } @@ -142,6 +146,13 @@ Status validate_frame(const VisionFrame& frame) { return Status::error(StatusCode::kUnsupported, "unsupported pixel format"); } + if (frame.width > std::numeric_limits::max() / ch || + frame.stride_bytes < 0 || + (frame.stride_bytes > 0 && + frame.stride_bytes < frame.width * ch)) { + return Status::error(StatusCode::kShapeMismatch, + "vision frame stride is smaller than one row"); + } const int stride = frame.stride_bytes > 0 ? frame.stride_bytes : frame.width * ch; const std::uint64_t need = static_cast(stride) * static_cast(frame.height); diff --git a/cpp/modalities/src/vision_cuda.cu b/cpp/modalities/src/vision_cuda.cu index b22d441f..50ee7e73 100644 --- a/cpp/modalities/src/vision_cuda.cu +++ b/cpp/modalities/src/vision_cuda.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -61,6 +62,13 @@ Status validate_frame_for_cuda(const VisionFrame& frame) { return Status::error(StatusCode::kUnsupported, "unsupported pixel format"); } + if (frame.width > std::numeric_limits::max() / ch || + frame.stride_bytes < 0 || + (frame.stride_bytes > 0 && + frame.stride_bytes < frame.width * ch)) { + return Status::error(StatusCode::kShapeMismatch, + "vision frame stride is smaller than one row"); + } const int stride = frame.stride_bytes > 0 ? frame.stride_bytes : frame.width * ch; const std::uint64_t need = static_cast(stride) * @@ -108,11 +116,18 @@ __device__ __forceinline__ float normalize_value(float raw, int c, int norm_mode, float scale, + float divisor, float shift, const float* mean, const float* inv_std) { - if (norm_mode == 0) return raw * scale + shift; - return (raw / 255.0f - mean[c]) * inv_std[c]; + if (norm_mode == 0) { + return __fadd_rn(__fmul_rn(raw, scale), shift); + } + if (norm_mode == 1) { + return __fadd_rn(__fdiv_rn(raw, divisor), shift); + } + return __fmul_rn( + __fsub_rn(__fdiv_rn(raw, 255.0f), mean[c]), inv_std[c]); } __device__ __forceinline__ void store_value(void* out, @@ -141,6 +156,7 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, int th, int norm_mode, float scale, + float divisor, float shift, float mean0, float mean1, @@ -152,10 +168,18 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, const int y = blockIdx.y * blockDim.y + threadIdx.y; if (x >= tw || y >= th) return; - const float fx = (static_cast(x) + 0.5f) * - static_cast(sw) / static_cast(tw) - 0.5f; - const float fy = (static_cast(y) + 0.5f) * - static_cast(sh) / static_cast(th) - 0.5f; + const float fx = __fsub_rn( + __fdiv_rn( + __fmul_rn(static_cast(x) + 0.5f, + static_cast(sw)), + static_cast(tw)), + 0.5f); + const float fy = __fsub_rn( + __fdiv_rn( + __fmul_rn(static_cast(y) + 0.5f, + static_cast(sh)), + static_cast(th)), + 0.5f); const int x0 = max(0, min(sw - 1, static_cast(floorf(fx)))); const int y0 = max(0, min(sh - 1, static_cast(floorf(fy)))); const int x1 = max(0, min(sw - 1, x0 + 1)); @@ -172,11 +196,17 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, float mean[3] = {mean0, mean1, mean2}; float inv_std[3] = {inv_std0, inv_std1, inv_std2}; for (int c = 0; c < 3; ++c) { - const float top = p00[c] * (1.0f - wx) + p01[c] * wx; - const float bot = p10[c] * (1.0f - wx) + p11[c] * wx; - const float raw = top * (1.0f - wy) + bot * wy; - const float norm = normalize_value(raw, c, norm_mode, scale, shift, - mean, inv_std); + const float top = __fadd_rn( + __fmul_rn(p00[c], __fsub_rn(1.0f, wx)), + __fmul_rn(p01[c], wx)); + const float bot = __fadd_rn( + __fmul_rn(p10[c], __fsub_rn(1.0f, wx)), + __fmul_rn(p11[c], wx)); + const float raw = __fadd_rn( + __fmul_rn(top, __fsub_rn(1.0f, wy)), + __fmul_rn(bot, wy)); + const float norm = normalize_value( + raw, c, norm_mode, scale, divisor, shift, mean, inv_std); const std::uint64_t out_idx = (((static_cast(view) * th + y) * tw + x) * 3ull) + static_cast(c); @@ -305,8 +335,11 @@ Status preprocess_vision_cuda(const VisionPreprocessSpec& spec, spec.output_dtype == DType::kFloat32 ? 1 : (spec.output_dtype == DType::kFloat16 ? 2 : 0), static_cast(v), spec.target_width, spec.target_height, - spec.normalize.mode == NormalizeMode::kScaleShift ? 0 : 1, - spec.normalize.scale, spec.normalize.shift, + spec.normalize.mode == NormalizeMode::kScaleShift + ? 0 + : (spec.normalize.mode == NormalizeMode::kDivideShift ? 1 : 2), + spec.normalize.scale, spec.normalize.divisor, + spec.normalize.shift, spec.normalize.mean[0], spec.normalize.mean[1], spec.normalize.mean[2], spec.normalize.inv_std[0], spec.normalize.inv_std[1], spec.normalize.inv_std[2]); diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/c_api.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/c_api.h index 9c727036..3eb59ea9 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/c_api.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/c_api.h @@ -55,6 +55,37 @@ typedef struct frt_pi05_runtime_config { * capacity is a per-call error, never a fallback allocation. */ int max_frame_width; int max_frame_height; + + /* Optional ABI extension: native prompt staging source. When all fields + * below are present, frt_pi05_runtime_set_prompt* tokenizes text/state and + * writes embeddings into prompt_embedding_data. The captured graph must + * already copy from this stable source buffer into encoder_x; the model + * runtime does not rebind graph pointers. */ + const char* prompt_tokenizer_model_path; + const void* prompt_embedding_table_data; + uint64_t prompt_embedding_table_bytes; + int prompt_embedding_table_dtype; + uint64_t prompt_embedding_vocab_size; + uint64_t prompt_embedding_hidden_dim; + void* prompt_embedding_data; + uint64_t prompt_embedding_bytes; + int prompt_embedding_dtype; + uint64_t max_prompt_tokens; + float prompt_embedding_scale; + + /* Optional ABI extension: raw proprioception normalization for STATE + * STAGED ports. If present, state is mapped with q01/q99 into [-1, 1] + * before Pi0.5 prompt discretization. */ + const float* state_q01; + uint64_t n_state_q01; + const float* state_q99; + uint64_t n_state_q99; + + /* Optional ABI tail: notify an integrated native graph owner after a + * prompt/state update has produced a new valid token count. */ + int (*prompt_length_update)(void* user, uint64_t prompt_len); + void* prompt_length_update_user; + int prompt_embedding_on_device; } frt_pi05_runtime_config; typedef struct frt_pi05_vision_frame { @@ -75,6 +106,8 @@ int frt_pi05_runtime_create(const frt_runtime_export_v1* exp, void frt_pi05_runtime_destroy(frt_pi05_runtime*); int frt_pi05_runtime_set_prompt(frt_pi05_runtime*, const char* text); +int frt_pi05_runtime_set_prompt_state(frt_pi05_runtime*, const char* text, + const float* state, uint64_t n_state); int frt_pi05_runtime_prepare_vision(frt_pi05_runtime*, const frt_pi05_vision_frame* frames, uint64_t n_frames); diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h index 0ad3f274..66ef4673 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h @@ -23,7 +23,9 @@ class RuntimeIo { int model_action_dim = kModelActionDim, int robot_action_dim = kLiberoActionDim, modalities::DType image_dtype = modalities::DType::kBFloat16, - modalities::VisionStaging* staging = nullptr); + modalities::VisionStaging* staging = nullptr, + modalities::ActionStaging* action_staging = nullptr, + bool strict_rgb8 = true); modalities::Status prepare_vision( const std::vector& frames) const; @@ -42,6 +44,8 @@ class RuntimeIo { modalities::TensorView action_output_; void* stream_ = nullptr; modalities::VisionStaging* staging_ = nullptr; /* borrowed */ + modalities::ActionStaging* action_staging_ = nullptr; /* borrowed */ + bool strict_rgb8_ = true; modalities::VisionPreprocessSpec vision_spec_; modalities::ActionPostprocessSpec action_spec_; }; diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h new file mode 100644 index 00000000..16f4b610 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h @@ -0,0 +1,68 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_BF16_FORWARD_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_BF16_FORWARD_H + +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/native_rtx_attention.h" +#include "flashrt/cpp/models/pi05/native_rtx_attention_driver.h" +#include "flashrt/cpp/models/pi05/native_workspace.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeBf16Forward { +public: + explicit NativeBf16Forward(const NativeKernelDriver* driver) + : driver_(driver) {} + + modalities::Status encoder_qkv( + int layer, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + std::uintptr_t stream) const; +#ifdef FLASHRT_CPP_WITH_FA2 + modalities::Status vision_layer( + int layer, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status vision( + const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status encoder_layer( + int layer, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status encoder( + const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status decoder_layer( + int layer, int step, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status diffusion_step( + int step, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status diffusion( + const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; +#endif + +private: + const NativeKernelDriver* driver_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_BF16_FORWARD_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h new file mode 100644 index 00000000..2cee6402 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h @@ -0,0 +1,59 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_DEVICE_WEIGHTS_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_DEVICE_WEIGHTS_H + +#include "flashrt/cpp/models/pi05/native_weight_ops.h" +#include "flashrt/exec.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeWeightDType { + kBf16, + kFp8E4M3, + kInt8, + kFloat32, +}; + +struct NativeDeviceWeight { + frt_buffer buffer = nullptr; + std::vector shape; + NativeWeightDType dtype = NativeWeightDType::kBf16; +}; + +class NativeDeviceWeightStore { +public: + explicit NativeDeviceWeightStore(frt_ctx ctx) : ctx_(ctx) {} + + NativeDeviceWeightStore(const NativeDeviceWeightStore&) = delete; + NativeDeviceWeightStore& operator=(const NativeDeviceWeightStore&) = delete; + + modalities::Status upload(const std::string& name, + const NativeBf16Tensor& tensor); + modalities::Status upload_bytes( + const std::string& name, + const std::vector& shape, + NativeWeightDType dtype, + const void* data, + std::size_t bytes); + modalities::Status download_bf16( + const std::string& name, + NativeBf16Tensor* out) const; + const NativeDeviceWeight* find(const std::string& name) const; + std::size_t size() const { return weights_.size(); } + +private: + frt_ctx ctx_ = nullptr; // borrowed; the context owns every buffer + std::map weights_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_DEVICE_WEIGHTS_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h new file mode 100644 index 00000000..0bfe18ba --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h @@ -0,0 +1,72 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H + +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeGraphConfig { + int num_views = 2; + int max_prompt_tokens = 200; + int chunk_size = 10; + int num_steps = 10; + int vision_pool_factor = 1; +}; + +class NativeGraphOwner { +public: + static std::unique_ptr create( + const std::string& checkpoint_path, const NativeGraphConfig& config, + modalities::Status* status); + + ~NativeGraphOwner(); + + NativeGraphOwner(const NativeGraphOwner&) = delete; + NativeGraphOwner& operator=(const NativeGraphOwner&) = delete; + + frt_ctx context() const { return ctx_; } + frt_graph infer_graph() const { return infer_graph_; } + int stream_id() const { return stream_id_; } + void* native_stream() const { return replay_stream_; } + const NativeGraphConfig& config() const { return config_; } + NativeDeviceWeightStore& weights() { return weights_; } + const NativeDeviceWeightStore& weights() const { return weights_; } + NativeWorkspace& workspace() { return workspace_; } + const NativeWorkspace& workspace() const { return workspace_; } + NativeRtxAttentionWorkspace& attention() { return attention_; } + const NativeRtxAttentionWorkspace& attention() const { return attention_; } + + modalities::Status set_prompt_length(int prompt_tokens); + int replay() const; + modalities::Status synchronize() const; + +private: + explicit NativeGraphOwner(frt_ctx ctx, const NativeGraphConfig& config); + modalities::Status initialize(const std::string& checkpoint_path); + modalities::Status record(void* stream); + static void record_graph(void* user, void* stream); + + frt_ctx ctx_ = nullptr; + NativeGraphConfig config_; + NativeDeviceWeightStore weights_; + NativeWorkspace workspace_; + NativeRtxAttentionWorkspace attention_; + NativeKernelDriver driver_; + NativeBf16Forward forward_; + std::unique_ptr attention_driver_; + frt_graph infer_graph_ = nullptr; + void* replay_stream_ = nullptr; + int stream_id_ = -1; + modalities::Status capture_status_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h new file mode 100644 index 00000000..c6678413 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h @@ -0,0 +1,86 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_KERNEL_DRIVER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_KERNEL_DRIVER_H + +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeKernelDriver { +public: + NativeKernelDriver() noexcept; + ~NativeKernelDriver(); + + NativeKernelDriver(const NativeKernelDriver&) = delete; + NativeKernelDriver& operator=(const NativeKernelDriver&) = delete; + + modalities::Status status() const; + modalities::Status bf16_nn(void* a, void* b, void* output, + int m, int n, int k, + std::uintptr_t stream) const; + modalities::Status add_bias_bf16(void* values, const void* bias, + int rows, int columns, + std::uintptr_t stream) const; + modalities::Status silu_bf16(void* values, std::size_t elements, + std::uintptr_t stream) const; + modalities::Status gelu_bf16(void* values, std::size_t elements, + std::uintptr_t stream) const; + modalities::Status gate_gelu_bf16(const void* gate, const void* up, + void* output, std::size_t elements, + std::uintptr_t stream) const; + modalities::Status residual_add_bf16(void* residual, const void* values, + std::size_t elements, + std::uintptr_t stream) const; + modalities::Status bias_residual_bf16( + void* residual, const void* values, const void* bias, + int rows, int columns, std::uintptr_t stream) const; + modalities::Status gate_mul_residual_bf16( + void* residual, const void* values, const void* gate, + std::size_t elements, std::uintptr_t stream) const; + modalities::Status rms_norm_bf16( + const void* values, const void* weight, void* output, + int rows, int columns, float epsilon, std::uintptr_t stream) const; + modalities::Status layer_norm_bf16( + const void* values, const void* weight, const void* bias, void* output, + int rows, int columns, float epsilon, std::uintptr_t stream) const; + modalities::Status ada_rms_norm_style_bf16( + const void* values, const void* weight, const void* style, + void* output, void* gate_output, int rows, int columns, + float epsilon, std::uintptr_t stream) const; + modalities::Status qkv_split_bf16( + const void* qkv, void* query, void* key, void* value, + int rows, int query_columns, int key_columns, int value_columns, + std::uintptr_t stream) const; + modalities::Status qkv_split_rope_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + int rows, int query_columns, int key_columns, int value_columns, + int head_dimension, std::uintptr_t stream) const; + modalities::Status qkv_split_rope_devpos_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + const void* device_position, int rows, int query_columns, + int key_columns, int value_columns, int head_dimension, + std::uintptr_t stream) const; + modalities::Status patch_im2col_16bit( + const void* images, void* patches, int num_views, + std::uintptr_t stream) const; + modalities::Status avg_pool_vision_bf16( + const void* values, void* output, int num_views, int height, int width, + int columns, int pool_factor, std::uintptr_t stream) const; + +private: + struct Impl; + std::unique_ptr impl_; + std::string error_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_KERNEL_DRIVER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h new file mode 100644 index 00000000..96b31868 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h @@ -0,0 +1,38 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_QUANTIZATION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_QUANTIZATION_H + +#include "flashrt/cpp/models/pi05/native_weight_ops.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeFp8Tensor { + std::vector shape; + std::vector values; + float scale = 0.0f; +}; + +struct NativeInt8Tensor { + std::vector shape; + std::vector values; + std::vector scales; +}; + +modalities::Status native_quantize_fp8_e4m3( + const NativeFloatTensor& bf16_weight, + bool transpose, + NativeFp8Tensor* out); + +modalities::Status native_quantize_int8_per_output( + const NativeFloatTensor& bf16_weight, + NativeInt8Tensor* out); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_QUANTIZATION_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h new file mode 100644 index 00000000..cec18243 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h @@ -0,0 +1,76 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H + +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeAttentionDType { + kBf16, + kFloat32, + kInt32, +}; + +struct NativeRtxAttentionConfig { + int num_views = 2; + int encoder_sequence = 712; + int encoder_vision_sequence = 512; + int chunk_size = 10; + int encoder_layers = 18; +}; + +struct NativeAttentionBuffer { + frt_buffer buffer = nullptr; + std::vector shape; + NativeAttentionDType dtype = NativeAttentionDType::kBf16; +}; + +class NativeRtxAttentionWorkspace { +public: + explicit NativeRtxAttentionWorkspace(frt_ctx ctx) : ctx_(ctx) {} + + modalities::Status allocate(const NativeRtxAttentionConfig& config); + modalities::Status set_fixed_prompt_length(int prompt_tokens); + const NativeAttentionBuffer* find(const std::string& name) const; + void* encoder_k_layer_dptr(int layer) const; + void* encoder_v_layer_dptr(int layer) const; + + std::size_t size() const { return buffers_.size(); } + std::size_t allocated_bytes() const { return allocated_bytes_; } + std::size_t kv_layer_stride_bytes() const { return kv_layer_stride_bytes_; } + int encoder_splits() const { return encoder_splits_; } + int decoder_splits() const { return decoder_splits_; } + +private: + modalities::Status add(const std::string& name, + std::initializer_list shape, + NativeAttentionDType dtype); + + frt_ctx ctx_ = nullptr; + std::map buffers_; + std::size_t allocated_bytes_ = 0; + std::size_t kv_layer_stride_bytes_ = 0; + int num_views_ = 0; + int encoder_sequence_ = 0; + int encoder_vision_sequence_ = 0; + int chunk_size_ = 0; + int encoder_layers_ = 0; + int encoder_splits_ = 0; + int decoder_splits_ = 0; + frt_buffer prompt_length_buffers_[3] = {nullptr, nullptr, nullptr}; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h new file mode 100644 index 00000000..62259246 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h @@ -0,0 +1,44 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_DRIVER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_DRIVER_H + +#include "flashrt/cpp/models/pi05/native_rtx_attention.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeRtxAttentionDriver { +public: + explicit NativeRtxAttentionDriver( + const NativeRtxAttentionWorkspace* workspace) noexcept; + + modalities::Status status() const; + modalities::Status vision(std::uintptr_t stream) const; + modalities::Status encoder(int layer, std::uintptr_t stream) const; + modalities::Status decoder(int layer, std::uintptr_t stream) const; + + void* vision_output() const; + void* encoder_output() const; + void* decoder_output() const; + int num_sms() const { return num_sms_; } + +private: + const NativeAttentionBuffer* find(const char* name) const; + + const NativeRtxAttentionWorkspace* workspace_ = nullptr; + int num_views_ = 0; + int encoder_sequence_ = 0; + int chunk_size_ = 0; + int total_kv_ = 0; + int num_sms_ = 0; + std::string error_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_DRIVER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_style_precompute.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_style_precompute.h new file mode 100644 index 00000000..df8c99c0 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_style_precompute.h @@ -0,0 +1,28 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_STYLE_PRECOMPUTE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_STYLE_PRECOMPUTE_H + +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/native_workspace.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeStylePrecomputer { +public: + explicit NativeStylePrecomputer(const NativeKernelDriver* driver) + : driver_(driver) {} + + modalities::Status run(const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const; + +private: + const NativeKernelDriver* driver_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_STYLE_PRECOMPUTE_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h new file mode 100644 index 00000000..ece1efb3 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h @@ -0,0 +1,61 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H + +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/models/pi05/native_device_weights.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeMaterializationOptions { + int num_steps = 10; + bool merge_decoder_gate_up = true; + bool include_embedding = true; +}; + +class NativeWeightMaterializer { +public: + NativeWeightMaterializer(const loader::SafetensorsFile& source, + NativeDeviceWeightStore* destination) + : source_(source), destination_(destination) {} + + modalities::Status materialize_encoder_layer(int layer); + modalities::Status materialize_decoder_layer(int layer, + bool merge_gate_up); + modalities::Status materialize_vision_layer(int layer); + modalities::Status materialize_vision_globals(); + modalities::Status materialize_decoder_globals(int num_steps); + modalities::Status materialize_embedding(); + modalities::Status materialize_all( + const NativeMaterializationOptions& options); + +private: + modalities::Status load(const std::string& key, NativeFloatTensor* out); + modalities::Status upload(const std::string& name, + const NativeFloatTensor& tensor); + modalities::Status upload_rounded_transpose( + const std::string& source_key, + const std::string& destination_name); + modalities::Status upload_rounded_copy( + const std::string& source_key, + const std::string& destination_name); + modalities::Status upload_folded_transpose( + const std::string& source_key, + const NativeFloatTensor& norm, + const std::string& destination_name); + modalities::Status upload_rounded_scaled( + const std::string& source_key, + const std::string& destination_name, + float scale, + bool transpose); + + const loader::SafetensorsFile& source_; + NativeDeviceWeightStore* destination_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h new file mode 100644 index 00000000..6e278640 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h @@ -0,0 +1,80 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_OPS_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_OPS_H + +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeFloatTensor { + std::vector shape; + std::vector values; +}; + +struct NativeBf16Tensor { + std::vector shape; + std::vector values; +}; + +modalities::Status load_native_float_tensor( + const loader::SafetensorsFile& file, + const std::string& key, + NativeFloatTensor* out); + +modalities::Status native_to_bf16(const NativeFloatTensor& input, + NativeBf16Tensor* out); + +modalities::Status native_round_to_bf16_float( + const NativeFloatTensor& input, + NativeFloatTensor* out); + +modalities::Status native_transpose_2d(const NativeFloatTensor& input, + NativeFloatTensor* out); + +modalities::Status native_patch_oihw_to_hwio( + const NativeFloatTensor& input, + NativeFloatTensor* out); + +modalities::Status native_interleave_qk_rows( + const NativeFloatTensor& input, + std::uint64_t num_heads, + NativeFloatTensor* out); + +modalities::Status native_fold_rms_columns( + const NativeFloatTensor& weight, + const NativeFloatTensor& norm, + NativeFloatTensor* out); + +modalities::Status native_concat_rows_transpose( + const std::vector& inputs, + NativeFloatTensor* out); + +modalities::Status native_concat_columns( + const NativeFloatTensor& left, + const NativeFloatTensor& right, + NativeFloatTensor* out); + +modalities::Status native_concat_vectors( + const std::vector& inputs, + NativeFloatTensor* out); + +modalities::Status native_scale(const NativeFloatTensor& input, + float scale, + NativeFloatTensor* out); + +modalities::Status native_pi05_time_embeddings( + int num_steps, + std::uint64_t embedding_dim, + NativeFloatTensor* out); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_OPS_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h new file mode 100644 index 00000000..281885d7 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h @@ -0,0 +1,40 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_PACKER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_PACKER_H + +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_quantization.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeWeightPacker { +public: + explicit NativeWeightPacker(NativeDeviceWeightStore* weights) + : weights_(weights) {} + + modalities::Status pack_fp8(const std::string& name, bool transpose); + modalities::Status pack_fp8_as(const std::string& source_name, + const std::string& packed_name, + bool transpose); + modalities::Status pack_int8(const std::string& name); + modalities::Status merge_bf16_columns(const std::string& left_name, + const std::string& right_name, + const std::string& output_name); + modalities::Status pack_all_fp8(bool transpose); + modalities::Status pack_vision_int8(); + modalities::Status pack_encoder_int8(); + modalities::Status pack_decoder_int8(); + +private: + modalities::Status load_bf16(const std::string& name, + NativeFloatTensor* out) const; + + NativeDeviceWeightStore* weights_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_PACKER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weights.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weights.h new file mode 100644 index 00000000..1c93a8ee --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weights.h @@ -0,0 +1,23 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHTS_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHTS_H + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeTensorRequirement { + std::string key; + std::vector shape; +}; + +const std::vector& native_tensor_requirements(); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHTS_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h new file mode 100644 index 00000000..1098bb24 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h @@ -0,0 +1,81 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H + +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/exec.h" + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeWorkspaceConfig { + int num_views = 2; + int max_prompt_tokens = 200; + int chunk_size = 10; + int num_steps = 10; + int vision_pool_factor = 1; +}; + +struct NativeWorkspaceBuffer { + frt_buffer buffer = nullptr; + std::vector shape; + modalities::DType dtype = modalities::DType::kBFloat16; + bool alias = false; +}; + +class NativeWorkspace { +public: + explicit NativeWorkspace(frt_ctx ctx) : ctx_(ctx) {} + + NativeWorkspace(const NativeWorkspace&) = delete; + NativeWorkspace& operator=(const NativeWorkspace&) = delete; + + modalities::Status allocate(const NativeWorkspaceConfig& config); + modalities::Status update_decoder_rope(int prompt_tokens); + modalities::Status expand_vision_position_embedding( + const NativeDeviceWeightStore& weights); + const NativeWorkspaceBuffer* find(const std::string& name) const; + + std::size_t logical_size() const { return buffers_.size(); } + std::size_t allocation_count() const { return allocation_count_; } + std::size_t allocated_bytes() const { return allocated_bytes_; } + int vision_sequence() const { return vision_sequence_; } + int encoder_vision_sequence() const { return encoder_vision_sequence_; } + int encoder_sequence() const { return encoder_sequence_; } + +private: + modalities::Status add(const std::string& name, + std::initializer_list shape, + modalities::DType dtype); + modalities::Status add_alias(const std::string& name, + const std::string& source_name, + std::initializer_list shape); + modalities::Status initialize_rms_ones(); + modalities::Status initialize_rope(); + + frt_ctx ctx_ = nullptr; + std::map buffers_; + std::size_t allocation_count_ = 0; + std::size_t allocated_bytes_ = 0; + int vision_sequence_ = 0; + int encoder_vision_sequence_ = 0; + int encoder_sequence_ = 0; + int num_views_ = 0; + int max_prompt_tokens_ = 0; + int chunk_size_ = 0; + frt_buffer decoder_rope_buffer_ = nullptr; + std::vector rope_table_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h new file mode 100644 index 00000000..f5bd0d36 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h @@ -0,0 +1,53 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_PROMPT_EMBED_H +#define FLASHRT_CPP_MODELS_PI05_PROMPT_EMBED_H + +#include "flashrt/cpp/modalities/text.h" +#include "flashrt/cpp/modalities/tokenizer.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct PromptEmbeddingSpec { + std::uint64_t vocab_size = 0; + std::uint64_t hidden_dim = 0; + std::uint64_t max_tokens = 0; + float scale = 1.0f; + std::int32_t no_state_suffix_token_id = 108; + bool zero_pad_output = true; +}; + +modalities::Status embed_prompt( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len, + void* stream = nullptr, + modalities::TextEmbeddingStaging* staging = nullptr, + std::string* formatted_workspace = nullptr); + +modalities::Status embed_prompt_cpu( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_PROMPT_EMBED_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h new file mode 100644 index 00000000..a5ecc9d7 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h @@ -0,0 +1,32 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H +#define FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +static constexpr int kStatePromptMaxLen = 200; + +std::vector discretize_state_prompt_bins( + const float* state, std::uint64_t n); + +std::string clean_task_prompt(const std::string& prompt); + +std::string format_state_prompt(const std::string& prompt, + const float* state, + std::uint64_t n_state); + +void format_state_prompt_into(const std::string& prompt, + const float* state, + std::uint64_t n_state, + std::string* out); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h index 1f817feb..2c65967d 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -3,8 +3,10 @@ #include "flashrt/cpp/families/vla/runtime.h" #include "flashrt/cpp/models/pi05/io.h" +#include "flashrt/cpp/models/pi05/prompt_embed.h" #include +#include namespace flashrt { namespace models { @@ -12,6 +14,7 @@ namespace pi05 { using ReplayFn = int (*)(frt_graph graph, frt_shape_key key, int stream_id, void* user); +using PromptLengthUpdateFn = int (*)(void* user, std::uint64_t prompt_len); struct RuntimeConfig { int num_views = 3; @@ -33,6 +36,7 @@ struct RuntimeConfig { * construction so prepare_vision never allocates on the hot path. */ int max_frame_width = 1280; int max_frame_height = 720; + bool strict_rgb8 = true; /* Optional host/device overrides. If left null, Runtime derives tensor * views from the export's named buffers. The current CPU reference @@ -41,8 +45,23 @@ struct RuntimeConfig { modalities::TensorView image_input_override; modalities::TensorView action_output_override; + /* Optional native prompt staging. When configured, set_prompt* writes + * token embeddings into prompt_embedding_output. The graph must consume + * this stable source buffer through its own captured copy/update path. */ + std::string prompt_tokenizer_model_path; + modalities::TensorView prompt_embedding_table; + modalities::TensorView prompt_embedding_output; + std::uint64_t prompt_vocab_size = 0; + std::uint64_t prompt_hidden_dim = 0; + std::uint64_t prompt_max_tokens = 0; + float prompt_embedding_scale = 0.0f; + std::vector state_q01; + std::vector state_q99; + ReplayFn replay_fn = nullptr; void* replay_user = nullptr; + PromptLengthUpdateFn prompt_length_update_fn = nullptr; + void* prompt_length_update_user = nullptr; }; class Runtime final : public families::vla::Runtime { @@ -64,6 +83,18 @@ class Runtime final : public families::vla::Runtime { } int set_prompt(const char* text) override; + int set_prompt_state(const char* text, const float* state, + std::uint64_t n_state); + const modalities::Status& prompt_status() const { + return prompt_status_; + } + bool prompt_staging_enabled() const { return prompt_staging_enabled_; } + bool state_normalization_enabled() const { + return !config_.state_q01.empty() && + config_.state_q01.size() == config_.state_q99.size(); + } + std::uint64_t current_prompt_len() const { return current_prompt_len_; } + modalities::Status prepare_vision( const std::vector& frames) override; int replay_tick() override; @@ -73,6 +104,7 @@ class Runtime final : public families::vla::Runtime { void retain_export(); void release_export(); modalities::Status bind(); + modalities::Status bind_prompt_staging(); static int default_replay(frt_graph graph, frt_shape_key key, int stream_id, void* user); @@ -82,7 +114,21 @@ class Runtime final : public families::vla::Runtime { families::vla::Manifest manifest_; modalities::Status status_; modalities::VisionStaging staging_; + modalities::ActionStaging action_staging_; RuntimeIo io_; + modalities::SentencePieceTokenizer prompt_tokenizer_; + modalities::TextEmbeddingStaging prompt_embedding_staging_; + PromptEmbeddingSpec prompt_spec_; + modalities::TensorView prompt_embedding_table_; + modalities::TensorView prompt_embedding_output_; + modalities::Status prompt_status_; + std::vector prompt_token_ids_; + std::vector normalized_state_; + std::string task_prompt_workspace_; + std::string formatted_prompt_workspace_; + std::size_t max_task_prompt_bytes_ = 0; + std::uint64_t current_prompt_len_ = 0; + bool prompt_staging_enabled_ = false; frt_graph graph_ = nullptr; frt_shape_key graph_key_ = 0; int stream_id_ = -1; diff --git a/cpp/models/pi05/src/c_api.cpp b/cpp/models/pi05/src/c_api.cpp index 4582f8aa..ea696b2b 100644 --- a/cpp/models/pi05/src/c_api.cpp +++ b/cpp/models/pi05/src/c_api.cpp @@ -4,23 +4,30 @@ #include #include +#include #include #include #include #include #include +#include #include struct frt_pi05_runtime_s { std::unique_ptr runtime; std::string last_error; + std::vector vision_frames; + std::vector vision_seen; + std::vector action_values; }; namespace { using flashrt::models::pi05::cface::make_config; +using flashrt::models::pi05::cface::pixel_channels; using flashrt::models::pi05::cface::pixel_format; using flashrt::models::pi05::cface::status_code; +using flashrt::models::pi05::cface::valid_pixel_format; } // namespace @@ -38,8 +45,10 @@ extern "C" int frt_pi05_runtime_create( auto* h = new (std::nothrow) frt_pi05_runtime_s(); if (!h) return -5; try { + auto runtime_config = make_config(config); + runtime_config.strict_rgb8 = false; h->runtime.reset( - new flashrt::models::pi05::Runtime(exp, make_config(config))); + new flashrt::models::pi05::Runtime(exp, std::move(runtime_config))); } catch (const std::exception& e) { h->last_error = e.what(); delete h; @@ -54,6 +63,14 @@ extern "C" int frt_pi05_runtime_create( delete h; return rc; } + const auto& manifest = h->runtime->manifest(); + h->vision_frames.resize(manifest.vision.view_order.size()); + h->vision_seen.resize(manifest.vision.view_order.size()); + for (std::size_t i = 0; i < h->vision_frames.size(); ++i) { + h->vision_frames[i].name = manifest.vision.view_order[i]; + } + h->action_values.resize(static_cast( + manifest.action.chunk * manifest.action.robot_dim)); *out = h; return 0; } @@ -66,7 +83,32 @@ extern "C" int frt_pi05_runtime_set_prompt(frt_pi05_runtime* h, const char* text) { if (!h || !h->runtime) return -1; int rc = h->runtime->set_prompt(text); - if (rc != 0) h->last_error = "prompt updates are not supported by adopted-export Pi05 runtime"; + if (rc != 0) { + const auto& st = h->runtime->prompt_status(); + h->last_error = st.message.empty() + ? "prompt updates are not supported by this Pi05 runtime" + : st.message; + } else { + h->last_error.clear(); + } + return rc; +} + +extern "C" int frt_pi05_runtime_set_prompt_state( + frt_pi05_runtime* h, + const char* text, + const float* state, + uint64_t n_state) { + if (!h || !h->runtime || (!state && n_state)) return -1; + int rc = h->runtime->set_prompt_state(text, state, n_state); + if (rc != 0) { + const auto& st = h->runtime->prompt_status(); + h->last_error = st.message.empty() + ? "prompt updates are not supported by this Pi05 runtime" + : st.message; + } else { + h->last_error.clear(); + } return rc; } @@ -75,8 +117,11 @@ extern "C" int frt_pi05_runtime_prepare_vision( const frt_pi05_vision_frame* frames, uint64_t n_frames) { if (!h || !h->runtime || (!frames && n_frames)) return -1; - std::vector v; - v.reserve(static_cast(n_frames)); + if (n_frames != h->vision_frames.size()) { + h->last_error = "Pi05 vision frame count does not match the runtime"; + return -4; + } + std::fill(h->vision_seen.begin(), h->vision_seen.end(), 0); for (uint64_t i = 0; i < n_frames; ++i) { const frt_pi05_vision_frame& in = frames[i]; if (in.struct_size < sizeof(frt_pi05_vision_frame) || @@ -84,8 +129,23 @@ extern "C" int frt_pi05_runtime_prepare_vision( h->last_error = "invalid Pi05 vision frame"; return -1; } - flashrt::modalities::VisionFrame out; - out.name = in.name; + if (!valid_pixel_format(in.pixel_format)) { + h->last_error = "Pi05 vision pixel format is invalid"; + return -4; + } + std::size_t slot = h->vision_frames.size(); + for (std::size_t j = 0; j < h->vision_frames.size(); ++j) { + if (h->vision_frames[j].name == in.name) { + slot = j; + break; + } + } + if (slot == h->vision_frames.size() || h->vision_seen[slot]) { + h->last_error = "Pi05 vision frame name is unknown or duplicated"; + return -4; + } + h->vision_seen[slot] = 1; + auto& out = h->vision_frames[slot]; out.image.data = const_cast(in.data); out.image.bytes = in.bytes; out.image.dtype = flashrt::modalities::DType::kUInt8; @@ -94,15 +154,14 @@ extern "C" int frt_pi05_runtime_prepare_vision( out.image.shape = flashrt::modalities::Shape{ static_cast(std::max(0, in.height)), static_cast(std::max(0, in.width)), - 3}; + pixel_channels(in.pixel_format)}; out.format = pixel_format(in.pixel_format); out.width = in.width; out.height = in.height; out.stride_bytes = in.stride_bytes; out.timestamp_ns = in.timestamp_ns; - v.push_back(std::move(out)); } - auto st = h->runtime->prepare_vision(v); + auto st = h->runtime->prepare_vision(h->vision_frames); if (!st.ok_status()) { h->last_error = st.message; return status_code(st); @@ -123,19 +182,19 @@ extern "C" int frt_pi05_runtime_read_actions(frt_pi05_runtime* h, uint64_t out_capacity, uint64_t* n_written) { if (!h || !h->runtime || !out_actions) return -1; - std::vector actions; - auto st = h->runtime->read_actions(&actions); + auto st = h->runtime->read_actions(&h->action_values); if (!st.ok_status()) { h->last_error = st.message; return status_code(st); } - if (out_capacity < actions.size()) { + if (out_capacity < h->action_values.size()) { h->last_error = "action output buffer is too small"; - if (n_written) *n_written = actions.size(); + if (n_written) *n_written = h->action_values.size(); return -5; } - std::memcpy(out_actions, actions.data(), actions.size() * sizeof(float)); - if (n_written) *n_written = actions.size(); + std::memcpy(out_actions, h->action_values.data(), + h->action_values.size() * sizeof(float)); + if (n_written) *n_written = h->action_values.size(); h->last_error.clear(); return 0; } diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h index 3c085e28..a0bc1d21 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/config_map.h @@ -40,6 +40,19 @@ inline modalities::PixelFormat pixel_format(int value) { } } +inline bool valid_pixel_format(int value) { + return value >= FRT_PI05_PIXEL_RGB8 && value <= FRT_PI05_PIXEL_GRAY8; +} + +inline std::uint64_t pixel_channels(int value) { + switch (value) { + case FRT_PI05_PIXEL_RGBA8: + case FRT_PI05_PIXEL_BGRA8: return 4; + case FRT_PI05_PIXEL_GRAY8: return 1; + default: return 3; + } +} + inline modalities::DType dtype(int value) { using modalities::DType; switch (value) { @@ -90,6 +103,107 @@ inline RuntimeConfig make_config(const frt_pi05_runtime_config* in) { sizeof(in->max_frame_height)) && in->max_frame_height > 0) { cfg.max_frame_height = in->max_frame_height; } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_tokenizer_model_path), + sizeof(in->prompt_tokenizer_model_path)) && + in->prompt_tokenizer_model_path) { + cfg.prompt_tokenizer_model_path = in->prompt_tokenizer_model_path; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_table_data), + sizeof(in->prompt_embedding_table_data)) && + in->prompt_embedding_table_data) { + cfg.prompt_embedding_table.data = + const_cast(in->prompt_embedding_table_data); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_table_bytes), + sizeof(in->prompt_embedding_table_bytes))) { + cfg.prompt_embedding_table.bytes = in->prompt_embedding_table_bytes; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_table_dtype), + sizeof(in->prompt_embedding_table_dtype))) { + cfg.prompt_embedding_table.dtype = + dtype(in->prompt_embedding_table_dtype); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_vocab_size), + sizeof(in->prompt_embedding_vocab_size))) { + cfg.prompt_vocab_size = in->prompt_embedding_vocab_size; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_hidden_dim), + sizeof(in->prompt_embedding_hidden_dim))) { + cfg.prompt_hidden_dim = in->prompt_embedding_hidden_dim; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, prompt_embedding_data), + sizeof(in->prompt_embedding_data)) && + in->prompt_embedding_data) { + cfg.prompt_embedding_output.data = in->prompt_embedding_data; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, prompt_embedding_bytes), + sizeof(in->prompt_embedding_bytes))) { + cfg.prompt_embedding_output.bytes = in->prompt_embedding_bytes; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, prompt_embedding_dtype), + sizeof(in->prompt_embedding_dtype))) { + cfg.prompt_embedding_output.dtype = dtype(in->prompt_embedding_dtype); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, max_prompt_tokens), + sizeof(in->max_prompt_tokens))) { + cfg.prompt_max_tokens = in->max_prompt_tokens; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_scale), + sizeof(in->prompt_embedding_scale))) { + cfg.prompt_embedding_scale = in->prompt_embedding_scale; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, state_q01), + sizeof(in->state_q01)) && + has_field(in, offsetof(frt_pi05_runtime_config, n_state_q01), + sizeof(in->n_state_q01)) && + in->state_q01 && in->n_state_q01) { + cfg.state_q01.assign(in->state_q01, in->state_q01 + in->n_state_q01); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, state_q99), + sizeof(in->state_q99)) && + has_field(in, offsetof(frt_pi05_runtime_config, n_state_q99), + sizeof(in->n_state_q99)) && + in->state_q99 && in->n_state_q99) { + cfg.state_q99.assign(in->state_q99, in->state_q99 + in->n_state_q99); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_length_update), + sizeof(in->prompt_length_update))) { + cfg.prompt_length_update_fn = in->prompt_length_update; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_length_update_user), + sizeof(in->prompt_length_update_user))) { + cfg.prompt_length_update_user = in->prompt_length_update_user; + } + const bool prompt_on_device = + has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_on_device), + sizeof(in->prompt_embedding_on_device)) && + in->prompt_embedding_on_device != 0; + if (cfg.prompt_vocab_size && cfg.prompt_hidden_dim) { + cfg.prompt_embedding_table.place = + prompt_on_device ? modalities::MemoryPlace::kDevice + : modalities::MemoryPlace::kHost; + cfg.prompt_embedding_table.layout = modalities::Layout::kFlat; + cfg.prompt_embedding_table.shape = + modalities::Shape{cfg.prompt_vocab_size, cfg.prompt_hidden_dim}; + } + if (cfg.prompt_max_tokens && cfg.prompt_hidden_dim) { + cfg.prompt_embedding_output.place = + prompt_on_device ? modalities::MemoryPlace::kDevice + : modalities::MemoryPlace::kHost; + cfg.prompt_embedding_output.layout = modalities::Layout::kFlat; + cfg.prompt_embedding_output.shape = + modalities::Shape{cfg.prompt_max_tokens, cfg.prompt_hidden_dim}; + } return cfg; } diff --git a/cpp/models/pi05/src/io.cpp b/cpp/models/pi05/src/io.cpp index fbe95220..27eb5c6e 100644 --- a/cpp/models/pi05/src/io.cpp +++ b/cpp/models/pi05/src/io.cpp @@ -3,6 +3,47 @@ namespace flashrt { namespace models { namespace pi05 { +namespace { + +modalities::Status validate_pi05_frame_contract( + const modalities::VisionFrame& frame, bool strict_rgb8) { + if (strict_rgb8 && frame.format != modalities::PixelFormat::kRGB8) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "Pi05 image input must be RGB8"); + } + if (frame.image.dtype != modalities::DType::kUInt8 || + frame.image.layout != modalities::Layout::kHWC) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "Pi05 image input must be u8 HWC"); + } + std::uint64_t channels = 3; + if (frame.format == modalities::PixelFormat::kRGBA8 || + frame.format == modalities::PixelFormat::kBGRA8) { + channels = 4; + } else if (frame.format == modalities::PixelFormat::kGRAY8) { + channels = 1; + } + if (frame.width <= 0 || frame.height <= 0 || + frame.image.shape.rank != 3 || + frame.image.shape.dims[0] != static_cast(frame.height) || + frame.image.shape.dims[1] != static_cast(frame.width) || + frame.image.shape.dims[2] != channels) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "Pi05 image shape must match HWC dimensions"); + } + if (frame.image.place != modalities::MemoryPlace::kHost && + frame.image.place != modalities::MemoryPlace::kHostPinned) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Pi05 image input must be host memory"); + } + return modalities::Status::ok(); +} + +} // namespace RuntimeIo::RuntimeIo(int num_views, modalities::TensorView image_input, @@ -14,11 +55,15 @@ RuntimeIo::RuntimeIo(int num_views, int model_action_dim, int robot_action_dim, modalities::DType image_dtype, - modalities::VisionStaging* staging) + modalities::VisionStaging* staging, + modalities::ActionStaging* action_staging, + bool strict_rgb8) : image_input_(image_input), action_output_(action_output), stream_(stream), staging_(staging), + action_staging_(action_staging), + strict_rgb8_(strict_rgb8), vision_spec_(vision_preprocess_spec(num_views)), action_spec_(action_postprocess_spec(action_mean, action_stddev, chunk, model_action_dim, robot_action_dim)) { @@ -27,6 +72,10 @@ RuntimeIo::RuntimeIo(int num_views, modalities::Status RuntimeIo::prepare_vision( const std::vector& frames) const { + for (const auto& frame : frames) { + auto st = validate_pi05_frame_contract(frame, strict_rgb8_); + if (!st.ok_status()) return st; + } return modalities::preprocess_vision(vision_spec_, frames, image_input_, stream_, staging_); } @@ -34,7 +83,8 @@ modalities::Status RuntimeIo::prepare_vision( modalities::Status RuntimeIo::read_actions( std::vector* robot_actions) const { return modalities::postprocess_action(action_spec_, action_output_, - robot_actions, stream_); + robot_actions, stream_, + action_staging_); } } // namespace pi05 diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index b91988fd..d2730124 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -31,6 +31,15 @@ struct Adapter { uint32_t images_port = kPortImages; uint32_t noise_port = kPortNoise; uint32_t actions_port = kPortActions; + uint32_t prompt_port = kNoPort; + uint32_t state_port = kNoPort; + bool has_prompt_text = false; + bool has_state = false; + std::size_t prompt_text_limit = 0; + std::string prompt_text; + std::vector state_values; + std::vector vision_frames; + std::vector action_values; int64_t image_shape[4] = {0, 0, 0, 3}; int64_t noise_shape[2] = {0, 0}; @@ -119,18 +128,21 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, } const auto* views = static_cast(data); const uint64_t n = bytes / sizeof(frt_image_view); - std::vector frames; - frames.reserve(n); + if (n != a->vision_frames.size()) { + a->last_error = "image view count does not match the runtime"; + return -4; + } for (uint64_t i = 0; i < n; ++i) { const frt_image_view& in = views[i]; if (in.struct_size < sizeof(frt_image_view) || !in.data) { a->last_error = "invalid image view"; return -1; } - flashrt::modalities::VisionFrame f; - /* generic views carry no names: positional, declared order */ - f.name = i < a->view_order.size() ? a->view_order[i] - : "view" + std::to_string(i); + if (in.pixel_format > FRT_RT_PIXEL_GRAY8) { + a->last_error = "image pixel format is invalid"; + return -4; + } + auto& f = a->vision_frames[static_cast(i)]; f.image.data = const_cast(in.data); f.image.bytes = in.bytes; f.image.dtype = flashrt::modalities::DType::kUInt8; @@ -144,9 +156,8 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, f.height = in.height; f.stride_bytes = in.stride_bytes; f.timestamp_ns = in.timestamp_ns; - frames.push_back(std::move(f)); } - Status st = a->runtime->prepare_vision(frames); + Status st = a->runtime->prepare_vision(a->vision_frames); if (!st.ok_status()) { a->last_error = st.message; return status_code(st); @@ -159,6 +170,69 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, "noise is a SWAP port: write its buffer window directly"; return -3; } + if (port == a->prompt_port) { + if (!data && bytes) { + a->last_error = "prompt payload is null"; + return -1; + } + if (bytes > a->prompt_text_limit) { + a->last_error = "prompt payload exceeds the hot-path capacity"; + return -4; + } + const char* begin = static_cast(data); + if (bytes) { + a->prompt_text.assign(begin, begin + bytes); + } else { + a->prompt_text.clear(); + } + a->has_prompt_text = true; + const int rc = a->has_state + ? a->runtime->set_prompt_state( + a->prompt_text.c_str(), + a->state_values.data(), + a->state_values.size()) + : a->runtime->set_prompt(a->prompt_text.c_str()); + if (rc != 0) { + const auto& st = a->runtime->prompt_status(); + a->last_error = st.message.empty() + ? "prompt staging is not configured" + : st.message; + return status_code(st); + } + a->last_error.clear(); + return 0; + } + if (port == a->state_port) { + if (!data || !bytes || bytes % sizeof(float)) { + a->last_error = "state payload must be f32[]"; + return -1; + } + const uint64_t n = bytes / sizeof(float); + if (n != a->state_values.size()) { + a->last_error = "state dimension does not match the runtime"; + return -4; + } + const auto* values = static_cast(data); + std::memcpy(a->state_values.data(), values, + static_cast(bytes)); + a->has_state = true; + if (!a->has_prompt_text) { + a->last_error.clear(); + return 0; + } + const int rc = a->runtime->set_prompt_state( + a->prompt_text.c_str(), a->state_values.data(), + a->state_values.size()); + if (rc != 0) { + const auto& st = a->runtime->prompt_status(); + a->last_error = st.message.empty() + ? "state staging failed" + : st.message; + return status_code(st); + } + a->last_error.clear(); + return 0; + } a->last_error = "unknown or non-input port"; return -1; } @@ -172,29 +246,40 @@ int get_output(void* self, uint32_t port, void* out, uint64_t capacity, a->last_error = "unknown or non-output port"; return -1; } - std::vector actions; - Status st = a->runtime->read_actions(&actions); + Status st = a->runtime->read_actions(&a->action_values); if (!st.ok_status()) { a->last_error = st.message; return status_code(st); } - const uint64_t need = actions.size() * sizeof(float); + const uint64_t need = a->action_values.size() * sizeof(float); if (written) *written = need; if (capacity < need) { a->last_error = "action output buffer is too small"; return -5; } - std::memcpy(out, actions.data(), need); + std::memcpy(out, a->action_values.data(), need); a->last_error.clear(); return 0; } int prepare(void* self, uint32_t graph, frt_shape_key key) { - (void)graph; - (void)key; auto* a = static_cast(self); - if (a) a->last_error = "adopted-export Pi05 runtime has fixed variants"; - return -3; + if (!a) return -1; + if (!a->source_model || !a->source_model->exp) { + a->last_error = "Pi05 fixed graph variants are captured at setup"; + return -3; + } + const frt_runtime_export_v1* exp = a->source_model->exp; + if (graph >= exp->n_graphs) { + a->last_error = "Pi05 prepare graph index is out of range"; + return -2; + } + if (!frt_graph_has_variant(exp->graphs[graph].handle, key)) { + a->last_error = "Pi05 fixed graph variant was not captured"; + return -2; + } + a->last_error.clear(); + return 0; } int step(void* self) { @@ -306,6 +391,12 @@ extern "C" int frt_pi05_model_runtime_create( const auto& manifest = a->runtime->manifest(); a->view_order = manifest.vision.view_order; + a->vision_frames.resize(a->view_order.size()); + for (std::size_t i = 0; i < a->view_order.size(); ++i) { + a->vision_frames[i].name = a->view_order[i]; + } + a->action_values.resize(static_cast( + manifest.action.chunk * manifest.action.robot_dim)); a->image_shape[0] = static_cast(a->view_order.size()); a->image_shape[1] = manifest.vision.target_height; a->image_shape[2] = manifest.vision.target_width; @@ -335,11 +426,12 @@ extern "C" int frt_pi05_model_runtime_create( FRT_RT_PORT_SWAP, 0, a->noise_shape, 2, 0, action_buf ? action_buf->handle : nullptr, 0, action_buf ? action_buf->bytes : 0}; - ports[kPortActions] = {"actions", FRT_RT_MOD_ACTION, io_dtype, + ports[kPortActions] = {"actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, a->action_shape, 2, 0, - action_buf ? action_buf->handle : nullptr, 0, - action_buf ? action_buf->bytes : 0}; + nullptr, 0, + static_cast(manifest.action.chunk) * + manifest.action.robot_dim * sizeof(float)}; const std::string graph_name = config && config->graph_name ? config->graph_name : "infer"; @@ -393,10 +485,14 @@ extern "C" int frt_pi05_model_runtime_create_over( const uint32_t images = find_port_index(model, "images"); const uint32_t noise = find_port_index(model, "noise"); const uint32_t actions = find_port_index(model, "actions"); + const uint32_t actions_raw = find_port_index(model, "actions_raw"); + const uint32_t prompt = find_port_index(model, "prompt"); + const uint32_t state = find_port_index(model, "state"); if (!compatible_port(model, images, FRT_RT_MOD_IMAGE, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED) || !compatible_port(model, actions, FRT_RT_MOD_ACTION, FRT_RT_PORT_OUT, - FRT_RT_PORT_STAGED)) { + FRT_RT_PORT_STAGED) || + model->ports[actions].dtype != FRT_RT_DTYPE_F32) { return -2; } if (noise != kNoPort && @@ -404,6 +500,21 @@ extern "C" int frt_pi05_model_runtime_create_over( FRT_RT_PORT_SWAP)) { return -2; } + if (actions_raw != kNoPort && + !compatible_port(model, actions_raw, FRT_RT_MOD_TENSOR, + FRT_RT_PORT_OUT, FRT_RT_PORT_SWAP)) { + return -2; + } + if (prompt != kNoPort && + !compatible_port(model, prompt, FRT_RT_MOD_TEXT, FRT_RT_PORT_IN, + FRT_RT_PORT_STAGED)) { + return -2; + } + if (state != kNoPort && + !compatible_port(model, state, FRT_RT_MOD_STATE, FRT_RT_PORT_IN, + FRT_RT_PORT_STAGED)) { + return -2; + } auto cfg = flashrt::models::pi05::cface::make_config(config); if (model->n_stages) { @@ -412,7 +523,8 @@ extern "C" int frt_pi05_model_runtime_create_over( cfg.graph_name = graph->name; } cfg.image_input_override = tensor_from_port(model->ports[images]); - cfg.action_output_override = tensor_from_port(model->ports[actions]); + cfg.action_output_override = tensor_from_port( + model->ports[actions_raw != kNoPort ? actions_raw : actions]); auto a = std::unique_ptr(new (std::nothrow) Adapter()); if (!a) return -5; @@ -420,11 +532,36 @@ extern "C" int frt_pi05_model_runtime_create_over( flashrt::models::pi05::Runtime(model->exp, cfg)); if (!a->runtime) return -5; if (!a->runtime->ok()) return status_code(a->runtime->status()); + if (prompt != kNoPort && !a->runtime->prompt_staging_enabled()) { + return -2; + } + if (state != kNoPort && + (!a->runtime->prompt_staging_enabled() || + !a->runtime->state_normalization_enabled())) { + return -2; + } a->source_model = model; a->images_port = images; a->noise_port = noise; a->actions_port = actions; + a->prompt_port = prompt; + a->state_port = state; a->view_order = a->runtime->manifest().vision.view_order; + a->vision_frames.resize(a->view_order.size()); + for (std::size_t i = 0; i < a->view_order.size(); ++i) { + a->vision_frames[i].name = a->view_order[i]; + } + const auto& action = a->runtime->manifest().action; + a->action_values.resize(static_cast(action.chunk * + action.robot_dim)); + if (cfg.prompt_max_tokens) { + a->prompt_text_limit = static_cast( + cfg.prompt_max_tokens * 8ull); + a->prompt_text.reserve(a->prompt_text_limit); + } + if (state != kNoPort) { + a->state_values.resize(cfg.state_q01.size()); + } frt_model_runtime_verbs verbs{}; verbs.struct_size = sizeof(verbs); diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp new file mode 100644 index 00000000..10d0e9a1 --- /dev/null +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -0,0 +1,675 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +bool shape_is(const NativeWorkspaceBuffer* buffer, + std::initializer_list shape) { + return buffer && buffer->dtype == modalities::DType::kBFloat16 && + buffer->shape == std::vector(shape); +} + +bool shape_is(const NativeAttentionBuffer* buffer, + std::initializer_list shape) { + return buffer && buffer->dtype == NativeAttentionDType::kBf16 && + buffer->shape == std::vector(shape); +} + +#ifdef FLASHRT_CPP_WITH_FA2 +bool shape_is(const NativeDeviceWeight* weight, + std::initializer_list shape) { + return weight && weight->dtype == NativeWeightDType::kBf16 && + weight->shape == std::vector(shape); +} +#endif + +} // namespace + +modalities::Status NativeBf16Forward::encoder_qkv( + int layer, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || layer < 0 || layer >= 18) { + return invalid("native encoder QKV owner is invalid"); + } + const int sequence = workspace->encoder_sequence(); + if (sequence <= 0) { + return invalid("native encoder sequence is invalid"); + } + const NativeWorkspaceBuffer* x = workspace->find("encoder_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("encoder_x_norm"); + const NativeWorkspaceBuffer* qkv = workspace->find("encoder_QKV"); + const NativeWorkspaceBuffer* rms = workspace->find("encoder_rms_ones"); + const NativeWorkspaceBuffer* rope = + workspace->find("encoder_rope_weights"); + const NativeAttentionBuffer* query = attention->find("attn_enc_Q"); + const NativeAttentionBuffer* cache = attention->find("attn_enc_K"); + const NativeAttentionBuffer* value_cache = attention->find("attn_enc_V"); + const NativeDeviceWeight* qkv_weight = + weights.find("encoder_attn_qkv_w_" + std::to_string(layer)); + if (!shape_is(x, {static_cast(sequence), 2048}) || + !shape_is(x_norm, {static_cast(sequence), 2048}) || + !shape_is(qkv, {static_cast(sequence), 2560}) || + !shape_is(rms, {2048}) || + !shape_is(rope, {static_cast(sequence), 256}) || + !shape_is(query, {static_cast(sequence), 8, 256}) || + !cache || cache->dtype != NativeAttentionDType::kBf16 || + cache->shape.size() != 4 || cache->shape[0] != 18 || + cache->shape[1] < static_cast(sequence) || + cache->shape[2] != 1 || cache->shape[3] != 256 || !qkv_weight || + !value_cache || value_cache->dtype != NativeAttentionDType::kBf16 || + value_cache->shape != cache->shape || + qkv_weight->dtype != NativeWeightDType::kBf16 || + qkv_weight->shape != std::vector({2048, 2560})) { + return invalid("native encoder QKV buffers or weight are invalid"); + } + void* key = attention->encoder_k_layer_dptr(layer); + void* value = attention->encoder_v_layer_dptr(layer); + if (!key || !value) { + return invalid("native encoder QKV cache layer is invalid"); + } + modalities::Status st = driver_->rms_norm_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), + frt_buffer_dptr(x_norm->buffer), sequence, 2048, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(qkv_weight->buffer), + frt_buffer_dptr(qkv->buffer), sequence, 2560, 2048, stream); + if (!st.ok_status()) return st; + return driver_->qkv_split_rope_bf16( + frt_buffer_dptr(qkv->buffer), frt_buffer_dptr(rope->buffer), + frt_buffer_dptr(query->buffer), key, value, sequence, 2048, 256, 256, + 256, stream); +} + +#ifdef FLASHRT_CPP_WITH_FA2 +modalities::Status NativeBf16Forward::vision_layer( + int layer, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver || + !attention_driver->status().ok_status() || layer < 0 || layer >= 27) { + return invalid("native vision layer owner is invalid"); + } + const int sequence = workspace->vision_sequence(); + const int num_views = sequence / 256; + const NativeWorkspaceBuffer* x = workspace->find("vision_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("vision_x_norm"); + const NativeWorkspaceBuffer* qkv = workspace->find("vision_QKV"); + const NativeWorkspaceBuffer* hidden = workspace->find("vision_hidden"); + const NativeAttentionBuffer* query = attention->find("attn_vis_Q"); + const NativeAttentionBuffer* key = attention->find("attn_vis_K"); + const NativeAttentionBuffer* value = attention->find("attn_vis_V"); + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* qkv_weight = + weights.find("vision_attn_qkv_w_" + suffix); + const NativeDeviceWeight* qkv_bias = + weights.find("vision_attn_qkv_b_" + suffix); + const NativeDeviceWeight* output_weight = + weights.find("vision_attn_o_w_" + suffix); + const NativeDeviceWeight* output_bias = + weights.find("vision_attn_o_b_" + suffix); + const NativeDeviceWeight* up_weight = + weights.find("vision_ffn_up_w_" + suffix); + const NativeDeviceWeight* up_bias = + weights.find("vision_ffn_up_b_" + suffix); + const NativeDeviceWeight* down_weight = + weights.find("vision_ffn_down_w_" + suffix); + const NativeDeviceWeight* down_bias = + weights.find("vision_ffn_down_b_" + suffix); + const NativeDeviceWeight* ffn_norm_weight = + weights.find("vision_pre_ffn_norm_w_" + suffix); + const NativeDeviceWeight* ffn_norm_bias = + weights.find("vision_pre_ffn_norm_b_" + suffix); + if (sequence <= 0 || sequence % 256 || num_views < 1 || num_views > 3 || + !shape_is(x, {static_cast(sequence), 1152}) || + !shape_is(x_norm, {static_cast(sequence), 1152}) || + !shape_is(qkv, {static_cast(sequence), 3456}) || + !shape_is(hidden, {static_cast(sequence), 4304}) || + !shape_is(query, {static_cast(num_views), 256, 16, + 72}) || + !shape_is(key, {static_cast(num_views), 256, 16, 72}) || + !shape_is(value, + {static_cast(num_views), 256, 16, 72}) || + !shape_is(qkv_weight, {1152, 3456}) || + !shape_is(qkv_bias, {3456}) || + !shape_is(output_weight, {1152, 1152}) || + !shape_is(output_bias, {1152}) || + !shape_is(up_weight, {1152, 4304}) || + !shape_is(up_bias, {4304}) || + !shape_is(down_weight, {4304, 1152}) || + !shape_is(down_bias, {1152}) || + !shape_is(ffn_norm_weight, {1152}) || + !shape_is(ffn_norm_bias, {1152})) { + return invalid("native vision layer buffers or weights are invalid"); + } + modalities::Status st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(qkv_weight->buffer), + frt_buffer_dptr(qkv->buffer), sequence, 3456, 1152, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(qkv->buffer), frt_buffer_dptr(qkv_bias->buffer), + sequence, 3456, stream); + if (!st.ok_status()) return st; + st = driver_->qkv_split_bf16( + frt_buffer_dptr(qkv->buffer), frt_buffer_dptr(query->buffer), + frt_buffer_dptr(key->buffer), frt_buffer_dptr(value->buffer), sequence, + 1152, 1152, 1152, stream); + if (!st.ok_status()) return st; + st = attention_driver->vision(stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + attention_driver->vision_output(), + frt_buffer_dptr(output_weight->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 1152, 1152, stream); + if (!st.ok_status()) return st; + st = driver_->bias_residual_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(output_bias->buffer), sequence, 1152, stream); + if (!st.ok_status()) return st; + st = driver_->layer_norm_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(ffn_norm_weight->buffer), + frt_buffer_dptr(ffn_norm_bias->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 1152, 1e-5f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(up_weight->buffer), + frt_buffer_dptr(hidden->buffer), sequence, 4304, 1152, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(hidden->buffer), frt_buffer_dptr(up_bias->buffer), + sequence, 4304, stream); + if (!st.ok_status()) return st; + st = driver_->gelu_bf16( + frt_buffer_dptr(hidden->buffer), + static_cast(sequence) * 4304, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(hidden->buffer), frt_buffer_dptr(down_weight->buffer), + frt_buffer_dptr(x_norm->buffer), sequence, 1152, 4304, stream); + if (!st.ok_status()) return st; + st = driver_->bias_residual_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(down_bias->buffer), sequence, 1152, stream); + if (!st.ok_status() || layer == 26) return st; + const NativeDeviceWeight* next_norm_weight = weights.find( + "vision_pre_attn_norm_w_" + std::to_string(layer + 1)); + const NativeDeviceWeight* next_norm_bias = weights.find( + "vision_pre_attn_norm_b_" + std::to_string(layer + 1)); + if (!shape_is(next_norm_weight, {1152}) || + !shape_is(next_norm_bias, {1152})) { + return invalid("native next vision norm weights are invalid"); + } + return driver_->layer_norm_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(next_norm_weight->buffer), + frt_buffer_dptr(next_norm_bias->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 1152, 1e-5f, stream); +} + +modalities::Status NativeBf16Forward::vision( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver) { + return invalid("native vision owner is invalid"); + } + const int sequence = workspace->vision_sequence(); + const int encoder_sequence = workspace->encoder_vision_sequence(); + const int num_views = sequence / 256; + const int pool_area = encoder_sequence > 0 ? sequence / encoder_sequence : 0; + const int pool_factor = pool_area == 1 ? 1 : pool_area == 4 ? 2 : + pool_area == 16 ? 4 : 0; + const NativeWorkspaceBuffer* images = + workspace->find("observation_images_normalized"); + const NativeWorkspaceBuffer* patches = workspace->find("vision_patches"); + const NativeWorkspaceBuffer* position = + workspace->find("vision_pos_embed_expanded"); + const NativeWorkspaceBuffer* x = workspace->find("vision_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("vision_x_norm"); + const NativeWorkspaceBuffer* pooled = workspace->find("vision_x_pooled"); + const NativeWorkspaceBuffer* encoder_x = workspace->find("encoder_x"); + const NativeDeviceWeight* patch_weight = + weights.find("vision_patch_embedding_w"); + const NativeDeviceWeight* patch_bias = + weights.find("vision_patch_embedding_b"); + const NativeDeviceWeight* first_norm_weight = + weights.find("vision_pre_attn_norm_w_0"); + const NativeDeviceWeight* first_norm_bias = + weights.find("vision_pre_attn_norm_b_0"); + const NativeDeviceWeight* final_norm_weight = + weights.find("vision_final_norm_w"); + const NativeDeviceWeight* final_norm_bias = + weights.find("vision_final_norm_b"); + const NativeDeviceWeight* projector_weight = + weights.find("encoder_multi_modal_projector_w"); + const NativeDeviceWeight* projector_bias = + weights.find("encoder_multi_modal_projector_b"); + if (sequence <= 0 || sequence % 256 || encoder_sequence <= 0 || + sequence % encoder_sequence || num_views < 1 || num_views > 3 || + !pool_factor || + !shape_is(images, {static_cast(num_views), 224, 224, + 3}) || + !shape_is(patches, {static_cast(sequence), 588}) || + !shape_is(position, {static_cast(sequence), 1152}) || + !shape_is(x, {static_cast(sequence), 1152}) || + !shape_is(x_norm, {static_cast(sequence), 1152}) || + !shape_is(pooled, + {static_cast(encoder_sequence), 1152}) || + !encoder_x || encoder_x->dtype != modalities::DType::kBFloat16 || + encoder_x->shape.size() != 2 || + encoder_x->shape[0] < static_cast(encoder_sequence) || + encoder_x->shape[1] != 2048 || + !shape_is(patch_weight, {14, 14, 3, 1152}) || + !shape_is(patch_bias, {1152}) || + !shape_is(first_norm_weight, {1152}) || + !shape_is(first_norm_bias, {1152}) || + !shape_is(final_norm_weight, {1152}) || + !shape_is(final_norm_bias, {1152}) || + !shape_is(projector_weight, {1152, 2048}) || + !shape_is(projector_bias, {2048})) { + return invalid("native vision buffers or weights are invalid"); + } + modalities::Status st = driver_->patch_im2col_16bit( + frt_buffer_dptr(images->buffer), frt_buffer_dptr(patches->buffer), + num_views, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(patches->buffer), frt_buffer_dptr(patch_weight->buffer), + frt_buffer_dptr(x->buffer), sequence, 1152, 588, stream); + if (!st.ok_status()) return st; + st = driver_->bias_residual_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(position->buffer), + frt_buffer_dptr(patch_bias->buffer), sequence, 1152, stream); + if (!st.ok_status()) return st; + st = driver_->layer_norm_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(first_norm_weight->buffer), + frt_buffer_dptr(first_norm_bias->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 1152, 1e-5f, stream); + if (!st.ok_status()) return st; + for (int layer = 0; layer < 27; ++layer) { + st = vision_layer(layer, weights, workspace, attention, + attention_driver, stream); + if (!st.ok_status()) return st; + } + if (pool_factor > 1) { + st = driver_->avg_pool_vision_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(pooled->buffer), + num_views, 16, 16, 1152, pool_factor, stream); + if (!st.ok_status()) return st; + } + st = driver_->layer_norm_bf16( + frt_buffer_dptr(pooled->buffer), + frt_buffer_dptr(final_norm_weight->buffer), + frt_buffer_dptr(final_norm_bias->buffer), frt_buffer_dptr(x_norm->buffer), + encoder_sequence, 1152, 1e-5f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(projector_weight->buffer), + frt_buffer_dptr(encoder_x->buffer), encoder_sequence, 2048, 1152, + stream); + if (!st.ok_status()) return st; + return driver_->add_bias_bf16( + frt_buffer_dptr(encoder_x->buffer), + frt_buffer_dptr(projector_bias->buffer), encoder_sequence, 2048, + stream); +} + +modalities::Status NativeBf16Forward::encoder_layer( + int layer, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + modalities::Status st = + encoder_qkv(layer, weights, workspace, attention, stream); + if (!st.ok_status() || layer == 17) return st; + if (!attention_driver || !attention_driver->status().ok_status()) { + return invalid("native encoder attention driver is invalid"); + } + const int sequence = workspace->encoder_sequence(); + const NativeWorkspaceBuffer* x = workspace->find("encoder_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("encoder_x_norm"); + const NativeWorkspaceBuffer* gate = + workspace->find("encoder_gate_merged"); + const NativeWorkspaceBuffer* hidden = workspace->find("encoder_hidden"); + const NativeWorkspaceBuffer* rms = workspace->find("encoder_rms_ones"); + const NativeDeviceWeight* output_weight = + weights.find("encoder_attn_o_w_" + std::to_string(layer)); + const NativeDeviceWeight* gate_weight = + weights.find("encoder_ffn_gate_w_" + std::to_string(layer)); + const NativeDeviceWeight* up_weight = + weights.find("encoder_ffn_up_w_" + std::to_string(layer)); + const NativeDeviceWeight* down_weight = + weights.find("encoder_ffn_down_w_" + std::to_string(layer)); + if (!shape_is(x, {static_cast(sequence), 2048}) || + !shape_is(x_norm, {static_cast(sequence), 2048}) || + !shape_is(gate, {static_cast(sequence), 32768}) || + !shape_is(hidden, {static_cast(sequence), 16384}) || + !shape_is(rms, {2048}) || + !shape_is(output_weight, {2048, 2048}) || + !shape_is(gate_weight, {2048, 16384}) || + !shape_is(up_weight, {2048, 16384}) || + !shape_is(down_weight, {16384, 2048})) { + return invalid("native encoder layer buffers or weights are invalid"); + } + st = attention_driver->encoder(layer, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + attention_driver->encoder_output(), + frt_buffer_dptr(output_weight->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 2048, 2048, stream); + if (!st.ok_status()) return st; + st = driver_->residual_add_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + static_cast(sequence) * 2048, stream); + if (!st.ok_status()) return st; + st = driver_->rms_norm_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), + frt_buffer_dptr(x_norm->buffer), sequence, 2048, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate_weight->buffer), + frt_buffer_dptr(gate->buffer), sequence, 16384, 2048, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(up_weight->buffer), + frt_buffer_dptr(hidden->buffer), sequence, 16384, 2048, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_bf16( + frt_buffer_dptr(gate->buffer), frt_buffer_dptr(hidden->buffer), + frt_buffer_dptr(hidden->buffer), + static_cast(sequence) * 16384, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(hidden->buffer), frt_buffer_dptr(down_weight->buffer), + frt_buffer_dptr(x_norm->buffer), sequence, 2048, 16384, stream); + if (!st.ok_status()) return st; + return driver_->residual_add_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + static_cast(sequence) * 2048, stream); +} + +modalities::Status NativeBf16Forward::encoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + for (int layer = 0; layer < 18; ++layer) { + modalities::Status st = encoder_layer( + layer, weights, workspace, attention, attention_driver, stream); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +modalities::Status NativeBf16Forward::decoder_layer( + int layer, + int step, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver || + !attention_driver->status().ok_status() || layer < 0 || layer >= 18) { + return invalid("native decoder layer owner is invalid"); + } + const NativeWorkspaceBuffer* x = workspace->find("decoder_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("x_normed_buf"); + const NativeWorkspaceBuffer* gate = workspace->find("gate_buf"); + const NativeWorkspaceBuffer* qkv = workspace->find("decoder_QKV"); + const NativeWorkspaceBuffer* hidden = workspace->find("decoder_hidden"); + const NativeWorkspaceBuffer* gate_projection = + workspace->find("decoder_gate_merged"); + const NativeWorkspaceBuffer* rms = workspace->find("decoder_rms_ones"); + const NativeWorkspaceBuffer* rope = + workspace->find("decoder_rope_weights"); + const NativeWorkspaceBuffer* style_attn = + workspace->find("decoder_style_attn"); + const NativeWorkspaceBuffer* style_ffn = + workspace->find("decoder_style_ffn"); + if (!x || x->shape.size() != 2) { + return invalid("native decoder workspace is invalid"); + } + const int sequence = static_cast(x->shape[0]); + const NativeAttentionBuffer* query = attention->find("attn_dec_Q"); + const NativeAttentionBuffer* devpos = attention->find("attn_dec_devpos"); + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* qkv_weight = + weights.find("decoder_attn_qkv_w_" + suffix); + const NativeDeviceWeight* output_weight = + weights.find("decoder_attn_o_w_" + suffix); + const NativeDeviceWeight* gate_weight = + weights.find("decoder_ffn_gate_w_" + suffix); + const NativeDeviceWeight* up_weight = + weights.find("decoder_ffn_up_w_" + suffix); + const NativeDeviceWeight* down_weight = + weights.find("decoder_ffn_down_w_" + suffix); + if (sequence <= 0 || step < 0 || + !shape_is(x, {static_cast(sequence), 1024}) || + !shape_is(x_norm, {static_cast(sequence), 1024}) || + !shape_is(gate, {static_cast(sequence), 1024}) || + !shape_is(qkv, {static_cast(sequence), 2560}) || + !shape_is(hidden, {static_cast(sequence), 4096}) || + !shape_is(gate_projection, + {static_cast(sequence), 8192}) || + !shape_is(rms, {1024}) || + !shape_is(rope, {static_cast(sequence), 256}) || + !style_attn || style_attn->dtype != modalities::DType::kBFloat16 || + style_attn->shape.size() != 4 || + style_attn->shape[0] <= static_cast(step) || + style_attn->shape[1] != 18 || + style_attn->shape[2] != static_cast(sequence) || + style_attn->shape[3] != 3072 || !style_ffn || + style_ffn->dtype != modalities::DType::kBFloat16 || + style_ffn->shape != style_attn->shape || + !shape_is(query, {static_cast(sequence), 8, 256}) || + !devpos || devpos->dtype != NativeAttentionDType::kInt32 || + devpos->shape != std::vector({1}) || + !shape_is(qkv_weight, {1024, 2560}) || + !shape_is(output_weight, {2048, 1024}) || + !shape_is(gate_weight, {1024, 4096}) || + !shape_is(up_weight, {1024, 4096}) || + !shape_is(down_weight, {4096, 1024})) { + return invalid("native decoder layer buffers or weights are invalid"); + } + const std::size_t style_offset = + (static_cast(step) * 18 + layer) * sequence * 3072 * + sizeof(std::uint16_t); + const auto* attn_style = + static_cast(frt_buffer_dptr(style_attn->buffer)) + + style_offset; + const auto* ffn_style = + static_cast(frt_buffer_dptr(style_ffn->buffer)) + + style_offset; + modalities::Status st = driver_->ada_rms_norm_style_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), attn_style, + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), + sequence, 1024, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(qkv_weight->buffer), + frt_buffer_dptr(qkv->buffer), sequence, 2560, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->qkv_split_rope_devpos_bf16( + frt_buffer_dptr(qkv->buffer), frt_buffer_dptr(rope->buffer), + frt_buffer_dptr(query->buffer), attention->encoder_k_layer_dptr(layer), + attention->encoder_v_layer_dptr(layer), + frt_buffer_dptr(devpos->buffer), sequence, 2048, 256, 256, 256, + stream); + if (!st.ok_status()) return st; + st = attention_driver->decoder(layer, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + attention_driver->decoder_output(), + frt_buffer_dptr(output_weight->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 1024, 2048, stream); + if (!st.ok_status()) return st; + st = driver_->gate_mul_residual_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(gate->buffer), + static_cast(sequence) * 1024, stream); + if (!st.ok_status()) return st; + st = driver_->ada_rms_norm_style_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), ffn_style, + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), + sequence, 1024, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate_weight->buffer), + frt_buffer_dptr(gate_projection->buffer), sequence, 4096, 1024, + stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(up_weight->buffer), + frt_buffer_dptr(hidden->buffer), sequence, 4096, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_bf16( + frt_buffer_dptr(gate_projection->buffer), + frt_buffer_dptr(hidden->buffer), frt_buffer_dptr(hidden->buffer), + static_cast(sequence) * 4096, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(hidden->buffer), frt_buffer_dptr(down_weight->buffer), + frt_buffer_dptr(x_norm->buffer), sequence, 1024, 4096, stream); + if (!st.ok_status()) return st; + return driver_->gate_mul_residual_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(gate->buffer), + static_cast(sequence) * 1024, stream); +} + +modalities::Status NativeBf16Forward::diffusion_step( + int step, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver || step < 0) { + return invalid("native diffusion step owner is invalid"); + } + const NativeWorkspaceBuffer* noise = workspace->find("diffusion_noise"); + const NativeWorkspaceBuffer* x = workspace->find("decoder_x"); + const NativeWorkspaceBuffer* action = + workspace->find("decoder_action_buf"); + const NativeWorkspaceBuffer* x_norm = workspace->find("x_normed_buf"); + const NativeWorkspaceBuffer* gate = workspace->find("gate_buf"); + const NativeWorkspaceBuffer* rms = workspace->find("decoder_rms_ones"); + const NativeWorkspaceBuffer* style = + workspace->find("decoder_style_final"); + if (!noise || noise->shape.size() != 2) { + return invalid("native diffusion workspace is invalid"); + } + const int sequence = static_cast(noise->shape[0]); + const NativeDeviceWeight* input_weight = + weights.find("decoder_action_in_proj_w"); + const NativeDeviceWeight* input_bias = + weights.find("decoder_action_in_proj_b"); + const NativeDeviceWeight* output_weight = + weights.find("decoder_action_out_proj_w"); + const NativeDeviceWeight* output_bias = + weights.find("decoder_action_out_proj_b"); + if (sequence <= 0 || + !shape_is(noise, {static_cast(sequence), 32}) || + !shape_is(x, {static_cast(sequence), 1024}) || + !shape_is(action, {static_cast(sequence), 32}) || + !shape_is(x_norm, {static_cast(sequence), 1024}) || + !shape_is(gate, {static_cast(sequence), 1024}) || + !shape_is(rms, {1024}) || !style || + style->dtype != modalities::DType::kBFloat16 || + style->shape.size() != 3 || + style->shape[0] <= static_cast(step) || + style->shape[1] != static_cast(sequence) || + style->shape[2] != 3072 || + !shape_is(input_weight, {32, 1024}) || + !shape_is(input_bias, {1024}) || + !shape_is(output_weight, {1024, 32}) || + !shape_is(output_bias, {32})) { + return invalid("native diffusion buffers or weights are invalid"); + } + modalities::Status st = driver_->bf16_nn( + frt_buffer_dptr(noise->buffer), frt_buffer_dptr(input_weight->buffer), + frt_buffer_dptr(x->buffer), sequence, 1024, 32, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(input_bias->buffer), + sequence, 1024, stream); + if (!st.ok_status()) return st; + for (int layer = 0; layer < 18; ++layer) { + st = decoder_layer(layer, step, weights, workspace, attention, + attention_driver, stream); + if (!st.ok_status()) return st; + } + const std::size_t style_offset = + static_cast(step) * sequence * 3072 * + sizeof(std::uint16_t); + const auto* final_style = + static_cast(frt_buffer_dptr(style->buffer)) + + style_offset; + st = driver_->ada_rms_norm_style_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), final_style, + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), + sequence, 1024, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(output_weight->buffer), frt_buffer_dptr(action->buffer), + sequence, 32, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(action->buffer), frt_buffer_dptr(output_bias->buffer), + sequence, 32, stream); + if (!st.ok_status()) return st; + return driver_->residual_add_bf16( + frt_buffer_dptr(noise->buffer), frt_buffer_dptr(action->buffer), + static_cast(sequence) * 32, stream); +} + +modalities::Status NativeBf16Forward::diffusion( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!workspace) return invalid("native diffusion workspace is invalid"); + const NativeWorkspaceBuffer* style = + workspace->find("decoder_style_final"); + if (!style || style->shape.size() != 3 || !style->shape[0] || + style->shape[0] > static_cast(INT_MAX)) { + return invalid("native diffusion step count is invalid"); + } + const int steps = static_cast(style->shape[0]); + for (int step = 0; step < steps; ++step) { + modalities::Status st = diffusion_step( + step, weights, workspace, attention, attention_driver, stream); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} +#endif + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp new file mode 100644 index 00000000..a6289b55 --- /dev/null +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -0,0 +1,150 @@ +#include "flashrt/cpp/models/pi05/native_device_weights.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +bool element_count(const std::vector& shape, + std::size_t* out) { + std::size_t count = 1; + for (std::uint64_t dim : shape) { + if (dim > std::numeric_limits::max() || + (dim && count > std::numeric_limits::max() / + static_cast(dim))) { + return false; + } + count *= static_cast(dim); + } + if (out) *out = count; + return true; +} + +std::size_t element_bytes(NativeWeightDType dtype) { + switch (dtype) { + case NativeWeightDType::kBf16: return sizeof(std::uint16_t); + case NativeWeightDType::kFp8E4M3: return sizeof(std::uint8_t); + case NativeWeightDType::kInt8: return sizeof(std::int8_t); + case NativeWeightDType::kFloat32: return sizeof(float); + } + return 0; +} + +} // namespace + +modalities::Status NativeDeviceWeightStore::upload( + const std::string& name, + const NativeBf16Tensor& tensor) { + return upload_bytes(name, tensor.shape, NativeWeightDType::kBf16, + tensor.values.data(), + tensor.values.size() * sizeof(std::uint16_t)); +} + +modalities::Status NativeDeviceWeightStore::upload_bytes( + const std::string& name, + const std::vector& shape, + NativeWeightDType dtype, + const void* data, + std::size_t bytes) { + if (!ctx_ || name.empty()) return invalid("invalid device weight store"); + if (weights_.find(name) != weights_.end()) { + return invalid("duplicate device weight name"); + } + std::size_t elements = 0; + const std::size_t width = element_bytes(dtype); + if (!data || !width || !element_count(shape, &elements) || + elements > std::numeric_limits::max() / width || + elements * width != bytes) { + return invalid("device weight shape does not match typed payload"); + } + if (!bytes) return invalid("device weight payload is empty"); + +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)data; + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "device weight upload requires the CUDA build"); +#else + frt_buffer buffer = frt_buffer_alloc(ctx_, name.c_str(), bytes); + if (!buffer) { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + cudaMemGetInfo(&free_bytes, &total_bytes); + std::ostringstream message; + message << "device weight allocation failed: " << name + << " requested=" << bytes << " free=" << free_bytes + << " total=" << total_bytes; + return modalities::Status::error(modalities::StatusCode::kBackend, + message.str()); + } + const cudaError_t rc = cudaMemcpy(frt_buffer_dptr(buffer), + data, bytes, + cudaMemcpyHostToDevice); + if (rc != cudaSuccess) { + return modalities::Status::error( + modalities::StatusCode::kBackend, + std::string("device weight upload failed: ") + + cudaGetErrorString(rc)); + } + weights_.emplace(name, NativeDeviceWeight{buffer, shape, dtype}); + return modalities::Status::ok(); +#endif +} + +const NativeDeviceWeight* NativeDeviceWeightStore::find( + const std::string& name) const { + const auto it = weights_.find(name); + return it == weights_.end() ? nullptr : &it->second; +} + +modalities::Status NativeDeviceWeightStore::download_bf16( + const std::string& name, + NativeBf16Tensor* out) const { + if (!out) return invalid("BF16 download destination is null"); + const NativeDeviceWeight* weight = find(name); + if (!weight || !weight->buffer || + weight->dtype != NativeWeightDType::kBf16) { + return invalid("BF16 device weight was not found"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "device weight download requires the CUDA build"); +#else + const std::size_t bytes = frt_buffer_bytes(weight->buffer); + if (!bytes || bytes % sizeof(std::uint16_t) != 0) { + return invalid("BF16 device weight has an invalid byte size"); + } + NativeBf16Tensor result; + result.shape = weight->shape; + result.values.resize(bytes / sizeof(std::uint16_t)); + const cudaError_t rc = cudaMemcpy(result.values.data(), + frt_buffer_dptr(weight->buffer), bytes, + cudaMemcpyDeviceToHost); + if (rc != cudaSuccess) { + return modalities::Status::error( + modalities::StatusCode::kBackend, + std::string("device weight download failed: ") + + cudaGetErrorString(rc)); + } + *out = std::move(result); + return modalities::Status::ok(); +#endif +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp new file mode 100644 index 00000000..ad449eb5 --- /dev/null +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -0,0 +1,239 @@ +#include "flashrt/cpp/models/pi05/native_graph_owner.h" + +#include "flashrt/cpp/models/pi05/native_style_precompute.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +} // namespace + +NativeGraphOwner::NativeGraphOwner( + frt_ctx ctx, + const NativeGraphConfig& config) + : ctx_(ctx), + config_(config), + weights_(ctx), + workspace_(ctx), + attention_(ctx), + forward_(&driver_), + capture_status_(modalities::Status::ok()) {} + +NativeGraphOwner::~NativeGraphOwner() { + if (replay_stream_) { + cudaStreamSynchronize(static_cast(replay_stream_)); + cudaStreamDestroy(static_cast(replay_stream_)); + replay_stream_ = nullptr; + } + if (ctx_) { + frt_ctx_destroy(ctx_); + ctx_ = nullptr; + } +} + +std::unique_ptr NativeGraphOwner::create( + const std::string& checkpoint_path, + const NativeGraphConfig& config, + modalities::Status* status) { + if (config.num_views < 1 || config.num_views > 3 || + config.max_prompt_tokens < 1 || config.chunk_size < 1 || + config.num_steps < 1 || + (config.vision_pool_factor != 1 && + config.vision_pool_factor != 2 && config.vision_pool_factor != 4)) { + if (status) *status = invalid("native graph configuration is invalid"); + return nullptr; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) { + if (status) *status = backend("native graph context creation failed"); + return nullptr; + } + std::unique_ptr owner( + new (std::nothrow) NativeGraphOwner(ctx, config)); + if (!owner) { + frt_ctx_destroy(ctx); + if (status) *status = backend("native graph owner allocation failed"); + return nullptr; + } + modalities::Status st = owner->initialize(checkpoint_path); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + if (status) *status = modalities::Status::ok(); + return owner; +} + +modalities::Status NativeGraphOwner::initialize( + const std::string& checkpoint_path) { + loader::SafetensorsFile source; + if (!source.open(checkpoint_path + "/model.safetensors")) { + return modalities::Status::error(modalities::StatusCode::kNotFound, + source.error()); + } + NativeWeightMaterializer materializer(source, &weights_); + NativeMaterializationOptions options; + options.num_steps = config_.num_steps; + options.merge_decoder_gate_up = false; + options.include_embedding = true; + modalities::Status st = materializer.materialize_all(options); + if (!st.ok_status()) return st; + + NativeWorkspaceConfig workspace_config; + workspace_config.num_views = config_.num_views; + workspace_config.max_prompt_tokens = config_.max_prompt_tokens; + workspace_config.chunk_size = config_.chunk_size; + workspace_config.num_steps = config_.num_steps; + workspace_config.vision_pool_factor = config_.vision_pool_factor; + st = workspace_.allocate(workspace_config); + if (!st.ok_status()) return st; + st = workspace_.expand_vision_position_embedding(weights_); + if (!st.ok_status()) return st; + + NativeRtxAttentionConfig attention_config; + attention_config.num_views = config_.num_views; + attention_config.encoder_sequence = workspace_.encoder_sequence(); + attention_config.encoder_vision_sequence = + workspace_.encoder_vision_sequence(); + attention_config.chunk_size = config_.chunk_size; + st = attention_.allocate(attention_config); + if (!st.ok_status()) return st; + st = set_prompt_length(0); + if (!st.ok_status()) return st; + + NativeStylePrecomputer precomputer(&driver_); + st = precomputer.run(weights_, &workspace_, 0); + if (!st.ok_status()) return st; + attention_driver_.reset(new (std::nothrow) + NativeRtxAttentionDriver(&attention_)); + if (!attention_driver_) { + return backend("native attention driver allocation failed"); + } + st = attention_driver_->status(); + if (!st.ok_status()) return st; + + for (const char* name : {"observation_images_normalized", + "prompt_embedding", "diffusion_noise"}) { + const NativeWorkspaceBuffer* buffer = workspace_.find(name); + if (!buffer || + cudaMemset(frt_buffer_dptr(buffer->buffer), 0, + frt_buffer_bytes(buffer->buffer)) != cudaSuccess) { + return backend("native graph input initialization failed"); + } + } + if (cudaDeviceSynchronize() != cudaSuccess) { + return backend("native graph setup synchronization failed"); + } + + infer_graph_ = frt_graph_create(ctx_, "infer", 1); + const NativeWorkspaceBuffer* images = + workspace_.find("observation_images_normalized"); + const NativeWorkspaceBuffer* prompt = workspace_.find("prompt_embedding"); + const NativeWorkspaceBuffer* encoder = workspace_.find("encoder_x"); + const NativeWorkspaceBuffer* noise = workspace_.find("diffusion_noise"); + if (!infer_graph_ || !images || !prompt || !encoder || !noise || + frt_graph_bind(infer_graph_, "images", images->buffer) != FRT_OK || + frt_graph_bind(infer_graph_, "prompt", prompt->buffer) != FRT_OK || + frt_graph_bind(infer_graph_, "encoder_x", encoder->buffer) != FRT_OK || + frt_graph_bind(infer_graph_, "noise", noise->buffer) != FRT_OK) { + return backend("native graph binding failed"); + } + capture_status_ = modalities::Status::ok(); + if (frt_graph_capture(infer_graph_, 0, record_graph, this) != FRT_OK) { + return capture_status_.ok_status() + ? backend("native full graph capture failed") + : capture_status_; + } + if (!capture_status_.ok_status() || + frt_graph_variant_count(infer_graph_) != 1) { + return capture_status_.ok_status() + ? backend("native full graph variant is missing") + : capture_status_; + } + + cudaStream_t stream = nullptr; + if (cudaStreamCreate(&stream) != cudaSuccess) { + return backend("native replay stream creation failed"); + } + replay_stream_ = stream; + stream_id_ = frt_ctx_wrap_stream(ctx_, replay_stream_); + if (stream_id_ < 0) return backend("native replay stream wrapping failed"); + return modalities::Status::ok(); +} + +modalities::Status NativeGraphOwner::record(void* stream) { + const NativeWorkspaceBuffer* prompt = workspace_.find("prompt_embedding"); + const NativeWorkspaceBuffer* encoder = workspace_.find("encoder_x"); + if (!prompt || !encoder) return invalid("native prompt buffers are missing"); + const std::size_t prompt_bytes = frt_buffer_bytes(prompt->buffer); + const std::size_t prompt_offset = + static_cast(workspace_.encoder_vision_sequence()) * 2048 * + sizeof(std::uint16_t); + if (prompt_offset > frt_buffer_bytes(encoder->buffer) || + prompt_bytes > frt_buffer_bytes(encoder->buffer) - prompt_offset) { + return invalid("native prompt window exceeds encoder storage"); + } + auto* destination = + static_cast(frt_buffer_dptr(encoder->buffer)) + + prompt_offset; + if (cudaMemcpyAsync(destination, frt_buffer_dptr(prompt->buffer), + prompt_bytes, cudaMemcpyDeviceToDevice, + static_cast(stream)) != cudaSuccess) { + return backend("native prompt graph copy failed"); + } + modalities::Status st = forward_.vision( + weights_, &workspace_, &attention_, attention_driver_.get(), + reinterpret_cast(stream)); + if (!st.ok_status()) return st; + st = forward_.encoder(weights_, &workspace_, &attention_, + attention_driver_.get(), + reinterpret_cast(stream)); + if (!st.ok_status()) return st; + return forward_.diffusion(weights_, &workspace_, &attention_, + attention_driver_.get(), + reinterpret_cast(stream)); +} + +void NativeGraphOwner::record_graph(void* user, void* stream) { + auto* owner = static_cast(user); + owner->capture_status_ = owner->record(stream); +} + +modalities::Status NativeGraphOwner::set_prompt_length(int prompt_tokens) { + modalities::Status st = attention_.set_fixed_prompt_length(prompt_tokens); + if (!st.ok_status()) return st; + return workspace_.update_decoder_rope(prompt_tokens); +} + +int NativeGraphOwner::replay() const { + if (!infer_graph_ || stream_id_ < 0) return FRT_ERR_INVALID; + return frt_graph_replay(infer_graph_, 0, stream_id_); +} + +modalities::Status NativeGraphOwner::synchronize() const { + if (!replay_stream_) return invalid("native replay stream is missing"); + const cudaError_t rc = + cudaStreamSynchronize(static_cast(replay_stream_)); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_kernel_driver.cu b/cpp/models/pi05/src/native_kernel_driver.cu new file mode 100644 index 00000000..cff13040 --- /dev/null +++ b/cpp/models/pi05/src/native_kernel_driver.cu @@ -0,0 +1,345 @@ +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" + +#include "activation.cuh" +#include "elementwise.cuh" +#include "gemm_runner.h" +#include "norm.cuh" +#include "patch_embed.cuh" +#include "rope.cuh" + +#include +#include + +#include +#include + +void add_bias_bf16(__nv_bfloat16* x, const __nv_bfloat16* b, + int rows, int columns, cudaStream_t stream); + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +__global__ void native_silu_bf16_kernel(__nv_bfloat16* values, + std::size_t elements) { + const std::size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < elements) { + const float value = __bfloat162float(values[index]); + values[index] = + __float2bfloat16(value / (1.0f + expf(-value))); + } +} + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +modalities::Status launch_status() { + const cudaError_t rc = cudaGetLastError(); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +} // namespace + +struct NativeKernelDriver::Impl { + GemmRunner gemm; +}; + +NativeKernelDriver::NativeKernelDriver() noexcept { + try { + impl_.reset(new Impl()); + } catch (const std::exception& e) { + error_ = e.what(); + } catch (...) { + error_ = "native kernel driver initialization failed"; + } +} + +NativeKernelDriver::~NativeKernelDriver() = default; + +modalities::Status NativeKernelDriver::status() const { + return impl_ ? modalities::Status::ok() : backend(error_); +} + +modalities::Status NativeKernelDriver::bf16_nn( + void* a, + void* b, + void* output, + int m, + int n, + int k, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!a || !b || !output || m <= 0 || n <= 0 || k <= 0) { + return invalid("native BF16 GEMM arguments are invalid"); + } + try { + impl_->gemm.bf16_nn(a, b, output, m, n, k, + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("native BF16 GEMM launch failed"); + } +} + +modalities::Status NativeKernelDriver::add_bias_bf16( + void* values, + const void* bias, + int rows, + int columns, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !bias || rows <= 0 || columns <= 0) { + return invalid("native BF16 bias arguments are invalid"); + } + ::add_bias_bf16(static_cast<__nv_bfloat16*>(values), + static_cast(bias), rows, columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::silu_bf16( + void* values, + std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !elements) { + return invalid("native BF16 SiLU arguments are invalid"); + } + native_silu_bf16_kernel<<<(elements + 255) / 256, 256, 0, + reinterpret_cast(stream)>>>( + static_cast<__nv_bfloat16*>(values), elements); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gelu_bf16( + void* values, std::size_t elements, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !elements || (elements & 1) || + elements > static_cast(INT_MAX)) { + return invalid("native BF16 GELU arguments are invalid"); + } + ::gelu_inplace(static_cast<__nv_bfloat16*>(values), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gate_gelu_bf16( + const void* gate, const void* up, void* output, std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!gate || !up || !output || !elements || + elements > static_cast(INT_MAX)) { + return invalid("native BF16 gated GELU arguments are invalid"); + } + ::gate_silu_mul(static_cast(gate), + static_cast(up), + static_cast<__nv_bfloat16*>(output), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::residual_add_bf16( + void* residual, const void* values, std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!residual || !values || !elements || (elements & 1) || + elements > static_cast(INT_MAX)) { + return invalid("native BF16 residual arguments are invalid"); + } + ::residual_add(static_cast<__nv_bfloat16*>(residual), + static_cast(values), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::bias_residual_bf16( + void* residual, const void* values, const void* bias, int rows, + int columns, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!residual || !values || !bias || rows <= 0 || columns <= 0 || + (columns & 1)) { + return invalid("native BF16 bias residual arguments are invalid"); + } + ::bias_residual(static_cast<__nv_bfloat16*>(residual), + static_cast(values), + static_cast(bias), rows, columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gate_mul_residual_bf16( + void* residual, const void* values, const void* gate, + std::size_t elements, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!residual || !values || !gate || !elements || (elements & 1) || + elements > static_cast(INT_MAX)) { + return invalid("native BF16 gated residual arguments are invalid"); + } + ::gate_mul_residual(static_cast<__nv_bfloat16*>(residual), + static_cast(values), + static_cast(gate), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::rms_norm_bf16( + const void* values, const void* weight, void* output, int rows, + int columns, float epsilon, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !weight || !output || rows <= 0 || columns <= 0 || + (columns & 1) || !(epsilon > 0.0f)) { + return invalid("native BF16 RMSNorm arguments are invalid"); + } + ::rms_norm(static_cast(values), + static_cast(weight), + static_cast<__nv_bfloat16*>(output), rows, columns, epsilon, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::layer_norm_bf16( + const void* values, const void* weight, const void* bias, void* output, + int rows, int columns, float epsilon, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !weight || !bias || !output || rows <= 0 || columns <= 0 || + (columns & 1) || !(epsilon > 0.0f)) { + return invalid("native BF16 LayerNorm arguments are invalid"); + } + ::layer_norm(static_cast(values), + static_cast(weight), + static_cast(bias), + static_cast<__nv_bfloat16*>(output), rows, columns, epsilon, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::ada_rms_norm_style_bf16( + const void* values, const void* weight, const void* style, void* output, + void* gate_output, int rows, int columns, float epsilon, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !weight || !style || !output || !gate_output || rows <= 0 || + columns <= 0 || (columns & 1) || !(epsilon > 0.0f)) { + return invalid("native BF16 AdaRMSNorm arguments are invalid"); + } + ::ada_rms_norm_style( + static_cast(values), + static_cast(weight), + static_cast(style), + static_cast<__nv_bfloat16*>(output), + static_cast<__nv_bfloat16*>(gate_output), rows, columns, epsilon, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::qkv_split_bf16( + const void* qkv, void* query, void* key, void* value, int rows, + int query_columns, int key_columns, int value_columns, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!qkv || !query || !key || !value || rows <= 0 || query_columns <= 0 || + key_columns <= 0 || value_columns <= 0) { + return invalid("native BF16 QKV split arguments are invalid"); + } + ::qkv_split(static_cast(qkv), + static_cast<__nv_bfloat16*>(query), + static_cast<__nv_bfloat16*>(key), + static_cast<__nv_bfloat16*>(value), rows, query_columns, + key_columns, value_columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::qkv_split_rope_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + int rows, int query_columns, int key_columns, int value_columns, + int head_dimension, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!qkv || !rope || !query || !key || !value || rows <= 0 || + query_columns <= 0 || key_columns <= 0 || value_columns <= 0 || + head_dimension <= 0 || (head_dimension & 1) || + query_columns % head_dimension || key_columns % head_dimension) { + return invalid("native BF16 QKV RoPE arguments are invalid"); + } + ::qkv_split_rope( + static_cast(qkv), + static_cast(rope), + static_cast<__nv_bfloat16*>(query), + static_cast<__nv_bfloat16*>(key), + static_cast<__nv_bfloat16*>(value), rows, query_columns, key_columns, + value_columns, head_dimension, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::qkv_split_rope_devpos_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + const void* device_position, int rows, int query_columns, int key_columns, + int value_columns, int head_dimension, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!qkv || !rope || !query || !key || !value || !device_position || + rows <= 0 || query_columns <= 0 || key_columns <= 0 || + value_columns <= 0 || head_dimension <= 0 || (head_dimension & 1) || + query_columns % head_dimension || key_columns % head_dimension) { + return invalid("native BF16 QKV devpos arguments are invalid"); + } + ::qkv_split_rope_devpos( + static_cast(qkv), + static_cast(rope), + static_cast<__nv_bfloat16*>(query), + static_cast<__nv_bfloat16*>(key), + static_cast<__nv_bfloat16*>(value), + static_cast(device_position), rows, query_columns, + key_columns, value_columns, head_dimension, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::patch_im2col_16bit( + const void* images, void* patches, int num_views, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!images || !patches || num_views <= 0) { + return invalid("native patch im2col arguments are invalid"); + } + ::patch_im2col(static_cast(images), + static_cast<__half*>(patches), num_views, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::avg_pool_vision_bf16( + const void* values, void* output, int num_views, int height, int width, + int columns, int pool_factor, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !output || num_views <= 0 || height <= 0 || width <= 0 || + columns <= 0 || pool_factor <= 0 || height % pool_factor || + width % pool_factor) { + return invalid("native vision pooling arguments are invalid"); + } + ::avg_pool_vision_tokens( + static_cast(values), + static_cast<__nv_bfloat16*>(output), num_views, height, width, columns, + pool_factor, reinterpret_cast(stream)); + return launch_status(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp new file mode 100644 index 00000000..a334f32d --- /dev/null +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -0,0 +1,351 @@ +#include "native_open_internal.h" + +#if defined(FLASHRT_CPP_WITH_FA2) && defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + +#include "config_map.h" +#include "flashrt/cpp/loader/sha256.h" +#include "flashrt/cpp/models/pi05/model_runtime.h" +#include "flashrt/cpp/models/pi05/native_graph_owner.h" + +#include + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +const NativeWorkspaceBuffer* workspace_buffer( + const NativeGraphOwner& owner, + const char* name) { + return owner.workspace().find(name); +} + +void release_graph_owner(void* owner) { + delete static_cast(owner); +} + +int update_prompt_length(void* owner, std::uint64_t prompt_len) { + auto* graph = static_cast(owner); + if (!graph || prompt_len > static_cast(INT_MAX)) return -1; + return cface::status_code( + graph->set_prompt_length(static_cast(prompt_len))); +} + +bool add_identity(frt_runtime_builder builder, const char* key, + const std::string& value) { + return frt_runtime_builder_add_identity(builder, key, value.c_str()) == 0; +} + +int unpublished_set_input(void*, uint32_t, const void*, uint64_t, int) { + return -3; +} +int unpublished_get_output(void*, uint32_t, void*, uint64_t, uint64_t*, int) { + return -3; +} + +frt_model_runtime_verbs unpublished_verbs() { + frt_model_runtime_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.set_input = unpublished_set_input; + verbs.get_output = unpublished_get_output; + return verbs; +} + +int fail_builder(frt_runtime_builder builder, std::string* error, + const char* message) { + frt_model_runtime_verbs discard_verbs = unpublished_verbs(); + frt_model_runtime_v1* discarded = frt_runtime_builder_finish_model( + builder, &discard_verbs, nullptr, nullptr, nullptr, nullptr); + if (discarded) discarded->release(discarded->owner); + if (error) *error = message; + return -6; +} + +} // namespace + +int build_native_model_runtime(const NativeOpenConfig& config, + frt_model_runtime_v1** out, + std::string* error) { + if (!out) return -1; + *out = nullptr; + int device = 0; + cudaDeviceProp properties{}; + cudaError_t cuda_rc = cudaGetDevice(&device); + if (cuda_rc == cudaSuccess) { + cuda_rc = cudaGetDeviceProperties(&properties, device); + } + if (cuda_rc != cudaSuccess) { + if (error) *error = cudaGetErrorString(cuda_rc); + return -6; + } + if (properties.major != 12 || properties.minor != 0) { + if (error) *error = "Pi0.5 native_v2 requires RTX SM120"; + return -3; + } + const std::string hardware_id = + "sm" + std::to_string(properties.major * 10 + properties.minor); + + std::string weights_sha256; + std::string tokenizer_sha256; + std::string hash_error; + if (!loader::sha256_file(config.checkpoint_path + "/model.safetensors", + &weights_sha256, &hash_error) || + !loader::sha256_file(config.tokenizer_model_path, &tokenizer_sha256, + &hash_error)) { + if (error) *error = hash_error; + return -2; + } + + NativeGraphConfig graph_config; + graph_config.num_views = config.num_views; + graph_config.max_prompt_tokens = config.max_prompt_tokens; + graph_config.chunk_size = config.chunk; + graph_config.num_steps = config.num_steps; + graph_config.vision_pool_factor = config.vision_pool_factor; + modalities::Status st; + std::unique_ptr graph = NativeGraphOwner::create( + config.checkpoint_path, graph_config, &st); + if (!graph) { + if (error) *error = st.message; + return cface::status_code(st); + } + + const NativeWorkspaceBuffer* images = + workspace_buffer(*graph, "observation_images_normalized"); + const NativeWorkspaceBuffer* noise = + workspace_buffer(*graph, "diffusion_noise"); + const NativeWorkspaceBuffer* encoder = + workspace_buffer(*graph, "encoder_x"); + const NativeWorkspaceBuffer* previous = + workspace_buffer(*graph, "rtc_prev_action_chunk"); + const NativeWorkspaceBuffer* prefix_weights = + workspace_buffer(*graph, "rtc_prefix_weights"); + const NativeWorkspaceBuffer* guidance = + workspace_buffer(*graph, "rtc_guidance_weight"); + const NativeWorkspaceBuffer* prompt = + workspace_buffer(*graph, "prompt_embedding"); + const NativeDeviceWeight* embedding = graph->weights().find( + "embedding_weight"); + if (!images || !noise || !encoder || !previous || !prefix_weights || + !guidance || !prompt || !embedding || + embedding->dtype != NativeWeightDType::kBf16 || + embedding->shape.size() != 2 || embedding->shape[1] != 2048) { + if (error) *error = "native graph export buffers are incomplete"; + return -6; + } + + frt_runtime_builder builder = frt_runtime_builder_create(graph->context()); + if (!builder) { + if (error) *error = "native runtime builder creation failed"; + return -6; + } + const frt_shape_key keys[] = {0}; + bool ok = + frt_runtime_builder_add_stream( + builder, "main", graph->stream_id(), 0, + graph->native_stream()) == 0 && + frt_runtime_builder_add_graph( + builder, "infer", graph->infer_graph(), 0, keys, 1, + graph->stream_id()) == 0 && + frt_runtime_builder_add_buffer( + builder, "observation_images_normalized", images->buffer, + frt_buffer_bytes(images->buffer), FRT_RT_ROLE_INPUT) == 0 && + frt_runtime_builder_add_buffer( + builder, "diffusion_noise", noise->buffer, + frt_buffer_bytes(noise->buffer), + FRT_RT_ROLE_INPUT | FRT_RT_ROLE_OUTPUT) == 0 && + frt_runtime_builder_add_buffer( + builder, "encoder_x", encoder->buffer, + frt_buffer_bytes(encoder->buffer), + FRT_RT_ROLE_INPUT | FRT_RT_ROLE_STATE) == 0 && + frt_runtime_builder_add_buffer( + builder, "rtc_prev_action_chunk", previous->buffer, + frt_buffer_bytes(previous->buffer), FRT_RT_ROLE_INPUT) == 0 && + frt_runtime_builder_add_buffer( + builder, "rtc_prefix_weights", prefix_weights->buffer, + frt_buffer_bytes(prefix_weights->buffer), FRT_RT_ROLE_INPUT) == 0 && + frt_runtime_builder_add_buffer( + builder, "rtc_guidance_weight", guidance->buffer, + frt_buffer_bytes(guidance->buffer), FRT_RT_ROLE_INPUT) == 0 && + frt_runtime_builder_add_buffer( + builder, "prompt_embedding", prompt->buffer, + frt_buffer_bytes(prompt->buffer), + FRT_RT_ROLE_INPUT | FRT_RT_ROLE_STATE) == 0; + if (!ok) return fail_builder(builder, error, "native descriptor build failed"); + + ok = frt_runtime_builder_add_region( + builder, "rollout_boundary", noise->buffer, 0, + frt_buffer_bytes(noise->buffer), + FRT_RT_REGION_SNAPSHOT | FRT_RT_REGION_RESTORE) == 0; + if (!ok) return fail_builder(builder, error, "native region build failed"); + + ok = add_identity(builder, "model", "pi05") && + add_identity(builder, "producer", "native") && + add_identity(builder, "pipeline", "NativeBf16") && + add_identity(builder, "hardware", hardware_id) && + add_identity(builder, "tensor_dtype", "bf16") && + add_identity(builder, "weights_sha256", weights_sha256) && + add_identity(builder, "tokenizer_sha256", tokenizer_sha256) && + add_identity(builder, "io", "native_v2") && + add_identity(builder, "state_prompt_mode", "fixed") && + add_identity(builder, "num_views", std::to_string(config.num_views)) && + add_identity(builder, "max_prompt_len", + std::to_string(config.max_prompt_tokens)) && + add_identity(builder, "state_dim", std::to_string(config.state_dim)) && + add_identity(builder, "chunk_size", std::to_string(config.chunk)) && + add_identity(builder, "num_steps", std::to_string(config.num_steps)) && + add_identity(builder, "vision_pool_factor", + std::to_string(config.vision_pool_factor)) && + add_identity(builder, "model_action_dim", "32") && + add_identity(builder, "robot_action_dim", + std::to_string(config.action_q01.size())); + if (!ok) return fail_builder(builder, error, "native identity build failed"); + + std::ostringstream manifest; + manifest << "{\"model\":\"pi05\",\"producer\":\"native\"," + << "\"hardware\":\"" << hardware_id + << "\",\"io\":\"native_v2\"," + << "\"stage_plan\":{\"name\":\"full\"," + << "\"stages\":[{\"name\":\"infer\"," + << "\"graph\":\"infer\",\"after\":[]}]}}"; + if (frt_runtime_builder_set_manifest(builder, manifest.str().c_str()) != 0) { + return fail_builder(builder, error, "native manifest build failed"); + } + + const int64_t prompt_shape[] = {-1}; + const int64_t state_shape[] = {config.state_dim}; + const int64_t image_shape[] = {config.num_views, 224, 224, 3}; + const int64_t raw_action_shape[] = {config.chunk, 32}; + const int64_t action_shape[] = { + config.chunk, static_cast(config.action_q01.size())}; + const std::uint64_t action_bytes = + static_cast(config.chunk) * + config.action_q01.size() * sizeof(float); + ok = frt_runtime_builder_add_port( + builder, "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) == 0 && + frt_runtime_builder_add_port( + builder, "state", FRT_RT_MOD_STATE, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, + state_shape, 1, 0, nullptr, 0, 0) == 0 && + frt_runtime_builder_add_port( + builder, "images", FRT_RT_MOD_IMAGE, FRT_RT_DTYPE_BF16, + FRT_RT_LAYOUT_NHWC, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, + image_shape, 4, 30, images->buffer, 0, + frt_buffer_bytes(images->buffer)) == 0 && + frt_runtime_builder_add_port( + builder, "noise", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_BF16, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_SWAP, 0, + raw_action_shape, 2, 0, noise->buffer, 0, + frt_buffer_bytes(noise->buffer)) == 0 && + frt_runtime_builder_add_port( + builder, "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, action_bytes) == 0 && + frt_runtime_builder_add_port( + builder, "actions_raw", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_BF16, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_SWAP, 0, + raw_action_shape, 2, 0, noise->buffer, 0, + frt_buffer_bytes(noise->buffer)) == 0 && + frt_runtime_builder_add_stage(builder, 0, nullptr, 0) == 0; + if (!ok) return fail_builder(builder, error, "native port/stage build failed"); + + NativeGraphOwner* raw_graph = graph.release(); + /* This base is retained only by the verb override below and is never + * returned to a consumer. The published object always has real verbs. */ + frt_model_runtime_verbs base_verbs = unpublished_verbs(); + frt_model_runtime_v1* base = frt_runtime_builder_finish_model( + builder, &base_verbs, nullptr, raw_graph, nullptr, + release_graph_owner); + if (!base) { + delete raw_graph; + if (error) *error = "native integrated runtime finish failed"; + return -6; + } + + std::vector action_mean(config.action_q01.size()); + std::vector action_stddev(config.action_q01.size()); + for (std::size_t i = 0; i < action_mean.size(); ++i) { + action_stddev[i] = + (config.action_q99[i] - config.action_q01[i] + 1e-6f) * 0.5f; + action_mean[i] = config.action_q01[i] + action_stddev[i]; + } + frt_pi05_runtime_config runtime_config{}; + runtime_config.struct_size = sizeof(runtime_config); + runtime_config.num_views = config.num_views; + runtime_config.chunk = config.chunk; + runtime_config.model_action_dim = 32; + runtime_config.robot_action_dim = static_cast(action_mean.size()); + runtime_config.action_mean = action_mean.data(); + runtime_config.n_action_mean = action_mean.size(); + runtime_config.action_stddev = action_stddev.data(); + runtime_config.n_action_stddev = action_stddev.size(); + runtime_config.graph_name = "infer"; + runtime_config.image_buffer_name = "observation_images_normalized"; + runtime_config.action_buffer_name = "diffusion_noise"; + runtime_config.image_dtype = FRT_PI05_DTYPE_BFLOAT16; + runtime_config.action_dtype = FRT_PI05_DTYPE_BFLOAT16; + runtime_config.prompt_tokenizer_model_path = + config.tokenizer_model_path.c_str(); + runtime_config.prompt_embedding_table_data = + frt_buffer_dptr(embedding->buffer); + runtime_config.prompt_embedding_table_bytes = + frt_buffer_bytes(embedding->buffer); + runtime_config.prompt_embedding_table_dtype = FRT_PI05_DTYPE_BFLOAT16; + runtime_config.prompt_embedding_vocab_size = embedding->shape[0]; + runtime_config.prompt_embedding_hidden_dim = 2048; + runtime_config.prompt_embedding_data = frt_buffer_dptr(prompt->buffer); + runtime_config.prompt_embedding_bytes = frt_buffer_bytes(prompt->buffer); + runtime_config.prompt_embedding_dtype = FRT_PI05_DTYPE_BFLOAT16; + runtime_config.max_prompt_tokens = config.max_prompt_tokens; + runtime_config.prompt_embedding_scale = std::sqrt(2048.0f); + runtime_config.state_q01 = config.state_q01.data(); + runtime_config.n_state_q01 = config.state_q01.size(); + runtime_config.state_q99 = config.state_q99.data(); + runtime_config.n_state_q99 = config.state_q99.size(); + runtime_config.prompt_length_update = update_prompt_length; + runtime_config.prompt_length_update_user = raw_graph; + runtime_config.prompt_embedding_on_device = 1; + + frt_model_runtime_v1* model = nullptr; + const int rc = frt_pi05_model_runtime_create_over( + base, &runtime_config, &model); + base->release(base->owner); + if (rc != 0 || !model) { + if (error) *error = "native Pi0.5 verb overlay failed"; + return rc != 0 ? rc : -6; + } + *out = model; + if (error) error->clear(); + return 0; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#else + +namespace flashrt { +namespace models { +namespace pi05 { + +int build_native_model_runtime(const NativeOpenConfig&, + frt_model_runtime_v1** out, + std::string* error) { + if (out) *out = nullptr; + if (error) *error = "native FA2 and SentencePiece are unavailable"; + return -3; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp new file mode 100644 index 00000000..fb078cf2 --- /dev/null +++ b/cpp/models/pi05/src/native_open.cpp @@ -0,0 +1,715 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/models/pi05/native_weights.h" +#include "native_open_internal.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +thread_local std::string g_last_error; + +struct JsonValue { + enum class Type { kString, kInteger, kBool, kNull }; + Type type = Type::kNull; + std::string text; + int64_t integer = 0; + bool boolean = false; +}; + +class JsonParser { +public: + explicit JsonParser(const char* src) : cur_(src ? src : "") {} + + bool parse_object(std::map* out) { + skip_ws(); + if (!consume('{')) return fail("config_json must be a JSON object"); + skip_ws(); + if (consume('}')) return finish(out); + while (true) { + std::string key; + if (!parse_string(&key)) return false; + skip_ws(); + if (!consume(':')) return fail("expected ':' after JSON key"); + skip_ws(); + JsonValue value; + if (!parse_value(&value)) return false; + values_[key] = value; + skip_ws(); + if (consume('}')) return finish(out); + if (!consume(',')) return fail("expected ',' or '}' in object"); + skip_ws(); + } + } + + const std::string& error() const { return error_; } + +private: + bool finish(std::map* out) { + skip_ws(); + if (*cur_) return fail("unexpected trailing data after JSON object"); + if (out) *out = std::move(values_); + return true; + } + + void skip_ws() { + while (*cur_ && std::isspace(static_cast(*cur_))) ++cur_; + } + + bool consume(char c) { + if (*cur_ != c) return false; + ++cur_; + return true; + } + + bool parse_value(JsonValue* value) { + if (!value) return fail("internal parser error"); + if (*cur_ == '"') { + value->type = JsonValue::Type::kString; + return parse_string(&value->text); + } + if (*cur_ == '-' || std::isdigit(static_cast(*cur_))) { + value->type = JsonValue::Type::kInteger; + return parse_integer(&value->integer); + } + if (match_literal("true")) { + value->type = JsonValue::Type::kBool; + value->boolean = true; + return true; + } + if (match_literal("false")) { + value->type = JsonValue::Type::kBool; + value->boolean = false; + return true; + } + if (match_literal("null")) { + value->type = JsonValue::Type::kNull; + return true; + } + return fail("unsupported JSON value"); + } + + bool parse_string(std::string* out) { + if (!consume('"')) return fail("expected JSON string"); + std::string s; + while (*cur_ && *cur_ != '"') { + unsigned char c = static_cast(*cur_++); + if (c < 0x20) return fail("control character in JSON string"); + if (c != '\\') { + s.push_back(static_cast(c)); + continue; + } + char esc = *cur_++; + switch (esc) { + case '"': s.push_back('"'); break; + case '\\': s.push_back('\\'); break; + case '/': s.push_back('/'); break; + case 'b': s.push_back('\b'); break; + case 'f': s.push_back('\f'); break; + case 'n': s.push_back('\n'); break; + case 'r': s.push_back('\r'); break; + case 't': s.push_back('\t'); break; + default: + return fail("unsupported JSON string escape"); + } + } + if (!consume('"')) return fail("unterminated JSON string"); + if (out) *out = std::move(s); + return true; + } + + bool parse_integer(int64_t* out) { + const char* begin = cur_; + if (*cur_ == '-') ++cur_; + if (!std::isdigit(static_cast(*cur_))) { + return fail("expected JSON integer"); + } + if (*cur_ == '0') { + ++cur_; + } else { + while (std::isdigit(static_cast(*cur_))) ++cur_; + } + if (*cur_ == '.' || *cur_ == 'e' || *cur_ == 'E') { + return fail("JSON number must be an integer"); + } + errno = 0; + char* end = nullptr; + const long long value = std::strtoll(begin, &end, 10); + if (errno || end != cur_) return fail("integer value is out of range"); + if (out) *out = static_cast(value); + return true; + } + + bool match_literal(const char* text) { + const std::size_t n = std::strlen(text); + if (std::strncmp(cur_, text, n) != 0) return false; + cur_ += n; + return true; + } + + bool fail(const char* msg) { + error_ = msg; + return false; + } + + const char* cur_; + std::string error_; + std::map values_; +}; + +bool path_exists(const std::string& path) { + struct stat st {}; + return !path.empty() && ::stat(path.c_str(), &st) == 0; +} + +bool regular_file_exists(const std::string& path) { + struct stat st {}; + return !path.empty() && ::stat(path.c_str(), &st) == 0 && + S_ISREG(st.st_mode); +} + +std::string join_path(const std::string& dir, const char* leaf) { + if (dir.empty() || dir.back() == '/') return dir + leaf; + return dir + "/" + leaf; +} + +std::string quoted_key(const std::string& key) { + std::string out = "\""; + for (char c : key) { + if (c == '"' || c == '\\') out.push_back('\\'); + out.push_back(c); + } + out.push_back('"'); + return out; +} + +bool object_for_key(const std::string& json, + const std::string& key, + std::string* object) { + const std::string q = quoted_key(key); + size_t pos = json.find(q); + while (pos != std::string::npos) { + size_t p = pos + q.size(); + while (p < json.size() && + std::isspace(static_cast(json[p]))) { + ++p; + } + if (p < json.size() && json[p] == ':') { + ++p; + while (p < json.size() && + std::isspace(static_cast(json[p]))) { + ++p; + } + if (p < json.size() && json[p] == '{') { + int depth = 0; + bool in_string = false; + bool escaped = false; + for (size_t i = p; i < json.size(); ++i) { + const char c = json[i]; + if (in_string) { + if (escaped) { + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + in_string = false; + } + continue; + } + if (c == '"') { + in_string = true; + } else if (c == '{') { + ++depth; + } else if (c == '}') { + --depth; + if (depth == 0) { + if (object) *object = json.substr(p, i - p + 1); + return true; + } + } + } + } + } + pos = json.find(q, pos + 1); + } + return false; +} + +bool parse_f64_array_property(const std::string& object, + const char* name, + std::vector* out) { + const std::string q = quoted_key(name); + size_t p = object.find(q); + if (p == std::string::npos) return false; + p += q.size(); + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p >= object.size() || object[p++] != ':') return false; + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p >= object.size() || object[p++] != '[') return false; + std::vector values; + while (p < object.size()) { + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p < object.size() && object[p] == ']') { + ++p; + if (out) *out = std::move(values); + return true; + } + errno = 0; + char* end = nullptr; + const double value = std::strtod(object.c_str() + p, &end); + if (errno || end == object.c_str() + p) return false; + values.push_back(value); + p = static_cast(end - object.c_str()); + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p < object.size() && object[p] == ',') { + ++p; + continue; + } + if (p < object.size() && object[p] == ']') continue; + return false; + } + return false; +} + +bool read_safetensors_f32_vector(const std::string& path, + const char* key, + std::vector* out) { + if (!out) return false; + flashrt::loader::SafetensorsFile file; + if (!file.open(path)) { + g_last_error = file.error(); + return false; + } + const auto* tensor = file.find(key); + if (!tensor || tensor->dtype != "F32" || tensor->shape.size() != 1) { + g_last_error = "safetensors F32 vector metadata is invalid"; + return false; + } + const uint64_t n = tensor->shape[0]; + if (n > (1ull << 20)) { + g_last_error = "safetensors vector is too large"; + return false; + } + std::vector values(static_cast(n)); + std::memcpy(values.data(), file.data(*tensor), + static_cast(tensor->bytes)); + *out = std::move(values); + return true; +} + +bool sane_quantile_pair(const std::vector& q01, + const std::vector& q99) { + if (q01.empty() || q01.size() != q99.size()) return false; + for (size_t i = 0; i < q01.size(); ++i) { + if (!std::isfinite(q01[i]) || !std::isfinite(q99[i]) || + q99[i] <= q01[i]) { + return false; + } + } + return true; +} + +bool sane_quantile_pair(const std::vector& q01, + const std::vector& q99) { + if (q01.empty() || q01.size() != q99.size()) return false; + for (size_t i = 0; i < q01.size(); ++i) { + if (!std::isfinite(q01[i]) || !std::isfinite(q99[i]) || + q99[i] <= q01[i]) { + return false; + } + } + return true; +} + +bool read_text_file(const std::string& path, std::string* out) { + if (!out) return false; + std::ifstream f(path); + if (!f) return false; + out->assign((std::istreambuf_iterator(f)), + std::istreambuf_iterator()); + return f.good() || f.eof(); +} + +std::string dirname(const std::string& path) { + const size_t p = path.find_last_of('/'); + if (p == std::string::npos) return "."; + if (p == 0) return "/"; + return path.substr(0, p); +} + +bool norm_block_values(const std::string& json, + const char* block_name, + std::vector* q01_out, + std::vector* q99_out) { + std::string block; + if (!object_for_key(json, block_name, &block)) return false; + std::vector q01; + std::vector q99; + if (!parse_f64_array_property(block, "q01", &q01) || + !parse_f64_array_property(block, "q99", &q99) || + !sane_quantile_pair(q01, q99)) { + return false; + } + if (q01_out) q01_out->assign(q01.begin(), q01.end()); + if (q99_out) q99_out->assign(q99.begin(), q99.end()); + return true; +} + +bool validate_norm_stats_file(const std::string& path, + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* config) { + std::string json; + if (!read_text_file(path, &json)) return false; + std::vector action_q01; + std::vector action_q99; + std::vector state_q01; + std::vector state_q99; + if (!norm_block_values(json, "actions", &action_q01, &action_q99) || + !norm_block_values(json, "state", &state_q01, &state_q99)) { + g_last_error = "norm_stats.json is missing actions/state q01/q99"; + return false; + } + if (action_q01.empty() || action_q01.size() > 32) { + g_last_error = "norm_stats action dimension is invalid"; + return false; + } + if (state_q01.size() != static_cast(state_dim)) { + g_last_error = "norm_stats state dimension does not match config"; + return false; + } + if (config) { + config->state_q01 = std::move(state_q01); + config->state_q99 = std::move(state_q99); + config->action_q01 = std::move(action_q01); + config->action_q99 = std::move(action_q99); + } + g_last_error.clear(); + return true; +} + +bool has_prefix(const std::string& s, const char* prefix) { + const size_t n = std::strlen(prefix); + return s.size() >= n && s.compare(0, n, prefix) == 0; +} + +bool has_suffix(const std::string& s, const char* suffix) { + const size_t n = std::strlen(suffix); + return s.size() >= n && s.compare(s.size() - n, n, suffix) == 0; +} + +std::string find_child(const std::string& dir, + const char* prefix, + const char* suffix) { + DIR* d = ::opendir(dir.c_str()); + if (!d) return ""; + std::string found; + while (dirent* ent = ::readdir(d)) { + const std::string name = ent->d_name; + if (has_prefix(name, prefix) && has_suffix(name, suffix)) { + found = join_path(dir, name.c_str()); + break; + } + } + ::closedir(d); + return found; +} + +bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* + config) { + const std::string pre = find_child( + checkpoint_path, "policy_preprocessor_step_", + "_normalizer_processor.safetensors"); + const std::string post = find_child( + checkpoint_path, "policy_postprocessor_step_", + "_unnormalizer_processor.safetensors"); + if (pre.empty() || post.empty()) return false; + + std::vector state_q01; + std::vector state_q99; + std::vector action_q01; + std::vector action_q99; + if (!read_safetensors_f32_vector(pre, "observation.state.q01", + &state_q01) || + !read_safetensors_f32_vector(pre, "observation.state.q99", + &state_q99) || + !read_safetensors_f32_vector(post, "action.q01", &action_q01) || + !read_safetensors_f32_vector(post, "action.q99", &action_q99)) { + g_last_error = + "lerobot policy stats are missing action/state q01/q99"; + return false; + } + if (state_q01.size() != static_cast(state_dim) || + !sane_quantile_pair(state_q01, state_q99)) { + g_last_error = + "lerobot policy state dimension does not match config"; + return false; + } + if (action_q01.size() > 32 || + !sane_quantile_pair(action_q01, action_q99)) { + g_last_error = "lerobot policy action dimension is invalid"; + return false; + } + if (config) { + config->state_q01 = std::move(state_q01); + config->state_q99 = std::move(state_q99); + config->action_q01 = std::move(action_q01); + config->action_q99 = std::move(action_q99); + } + g_last_error.clear(); + return true; +} + +bool validate_norm_stats(const std::string& checkpoint_path, + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* config) { + const std::string parent = dirname(checkpoint_path); + const std::string candidates[] = { + join_path(checkpoint_path, + "assets/physical-intelligence/libero/norm_stats.json"), + join_path(checkpoint_path, "assets/droid/norm_stats.json"), + join_path(checkpoint_path, "norm_stats.json"), + join_path(parent, + "pi05_libero/assets/physical-intelligence/libero/" + "norm_stats.json"), + join_path(parent, "pi05_droid/assets/droid/norm_stats.json"), + join_path(parent, "pi05_droid_pytorch/assets/droid/norm_stats.json"), + }; + bool saw_malformed = false; + std::string malformed_error; + for (const std::string& path : candidates) { + if (!regular_file_exists(path)) continue; + if (validate_norm_stats_file(path, state_dim, config)) return true; + saw_malformed = true; + malformed_error = g_last_error; + } + if (validate_lerobot_policy_norm_stats(checkpoint_path, state_dim, + config)) { + return true; + } + g_last_error = saw_malformed + ? malformed_error + : "norm_stats.json not found for Pi0.5 native_v2"; + return false; +} + +bool validate_pi05_safetensors(const std::string& checkpoint_path) { + const std::string path = join_path(checkpoint_path, "model.safetensors"); + if (!regular_file_exists(path)) { + g_last_error = "checkpoint_path must contain model.safetensors"; + return false; + } + flashrt::loader::SafetensorsFile file; + if (!file.open(path)) { + g_last_error = file.error(); + return false; + } + + for (const auto& req : + flashrt::models::pi05::native_tensor_requirements()) { + std::string key = req.key; + const flashrt::loader::SafetensorInfo* meta = file.find(key); + if (!meta) { + key = std::string("model.") + req.key; + meta = file.find(key); + if (!meta) { + g_last_error = std::string("model.safetensors is missing ") + + req.key; + return false; + } + } + if (meta->dtype != "BF16" && meta->dtype != "F16" && + meta->dtype != "F32") { + g_last_error = std::string("Pi0.5 tensor dtype is unsupported: ") + + req.key; + return false; + } + if (meta->shape != req.shape) { + g_last_error = std::string("Pi0.5 tensor shape mismatch: ") + + req.key; + return false; + } + } + g_last_error.clear(); + return true; +} + +bool string_field(const std::map& obj, + const char* key, + std::string* out, + bool required) { + auto it = obj.find(key); + if (it == obj.end()) { + if (!required) return true; + g_last_error = std::string("missing required field: ") + key; + return false; + } + if (it->second.type != JsonValue::Type::kString || + it->second.text.empty()) { + g_last_error = std::string("field must be a non-empty string: ") + key; + return false; + } + if (out) *out = it->second.text; + return true; +} + +bool integer_field(const std::map& obj, + const char* key, + int64_t* out) { + auto it = obj.find(key); + if (it == obj.end()) return true; + if (it->second.type != JsonValue::Type::kInteger) { + g_last_error = std::string("field must be an integer: ") + key; + return false; + } + if (out) *out = it->second.integer; + return true; +} + +int validate_config( + const char* config_json, + flashrt::models::pi05::NativeOpenConfig* parsed) { + if (!config_json) { + g_last_error = "config_json is null"; + return -1; + } + std::map obj; + JsonParser parser(config_json); + if (!parser.parse_object(&obj)) { + g_last_error = parser.error(); + return -1; + } + + std::string io; + std::string checkpoint_path; + std::string tokenizer_model_path; + std::string state_prompt_mode; + if (!string_field(obj, "io", &io, true) || + !string_field(obj, "checkpoint_path", &checkpoint_path, true) || + !string_field(obj, "tokenizer_model_path", &tokenizer_model_path, + true) || + !string_field(obj, "state_prompt_mode", &state_prompt_mode, true)) { + return -1; + } + if (io != "native_v2") { + g_last_error = "Pi0.5 native open requires io='native_v2'"; + return -1; + } + if (state_prompt_mode != "fixed") { + g_last_error = + "Pi0.5 native_v2 requires state_prompt_mode='fixed'"; + return -1; + } + if (!path_exists(checkpoint_path)) { + g_last_error = "checkpoint_path does not exist"; + return -2; + } + if (!regular_file_exists(tokenizer_model_path)) { + g_last_error = "tokenizer_model_path does not name a file"; + return -2; + } + if (!validate_pi05_safetensors(checkpoint_path)) { + return -2; + } + + int64_t max_prompt_tokens = 0; + int64_t state_dim = 0; + int64_t num_views = 0; + int64_t chunk = 0; + int64_t num_steps = 10; + int64_t vision_pool_factor = 1; + if (!integer_field(obj, "max_prompt_tokens", &max_prompt_tokens) || + !integer_field(obj, "state_dim", &state_dim) || + !integer_field(obj, "num_views", &num_views) || + !integer_field(obj, "chunk", &chunk) || + !integer_field(obj, "num_steps", &num_steps) || + !integer_field(obj, "vision_pool_factor", &vision_pool_factor)) { + return -1; + } + if (max_prompt_tokens < 200 || max_prompt_tokens > INT_MAX) { + g_last_error = "max_prompt_tokens must be in [200, INT_MAX]"; + return -1; + } + if (state_dim <= 0 || state_dim > INT_MAX) { + g_last_error = "state_dim must be in [1, INT_MAX]"; + return -1; + } + if (num_views && (num_views < 1 || num_views > 3)) { + g_last_error = "num_views must be in [1, 3]"; + return -1; + } + if (chunk && (chunk <= 0 || chunk > INT_MAX)) { + g_last_error = "chunk must be in [1, INT_MAX]"; + return -1; + } + if (num_steps <= 0 || num_steps > INT_MAX) { + g_last_error = "num_steps must be in [1, INT_MAX]"; + return -1; + } + if (vision_pool_factor != 1 && vision_pool_factor != 2 && + vision_pool_factor != 4) { + g_last_error = "vision_pool_factor must be one of 1, 2, or 4"; + return -1; + } + flashrt::models::pi05::NativeOpenConfig config; + config.checkpoint_path = checkpoint_path; + config.tokenizer_model_path = tokenizer_model_path; + config.max_prompt_tokens = static_cast(max_prompt_tokens); + config.state_dim = static_cast(state_dim); + config.num_views = static_cast(num_views ? num_views : 2); + config.chunk = static_cast(chunk ? chunk : 10); + config.num_steps = static_cast(num_steps); + config.vision_pool_factor = static_cast(vision_pool_factor); + if (!validate_norm_stats(checkpoint_path, state_dim, &config)) { + return -2; + } + + if (parsed) *parsed = std::move(config); + g_last_error.clear(); + return 0; +} + +} // namespace + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out) { + if (!out) { + g_last_error = "out is null"; + return -1; + } + *out = nullptr; + flashrt::models::pi05::NativeOpenConfig config; + const int rc = validate_config(config_json, &config); + if (rc != 0) return rc; +#if defined(FLASHRT_CPP_WITH_FA2) && defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + return flashrt::models::pi05::build_native_model_runtime( + config, out, &g_last_error); +#else + g_last_error = + "Pi0.5 native open validated config; this build requires native " + "FA2 and SentencePiece for graph capture"; + return -3; +#endif +} + +extern "C" const char* frt_pi05_native_open_last_error() { + return g_last_error.c_str(); +} diff --git a/cpp/models/pi05/src/native_open_internal.h b/cpp/models/pi05/src/native_open_internal.h new file mode 100644 index 00000000..21683dfe --- /dev/null +++ b/cpp/models/pi05/src/native_open_internal.h @@ -0,0 +1,36 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_OPEN_INTERNAL_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_OPEN_INTERNAL_H + +#include "flashrt/model_runtime.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeOpenConfig { + std::string checkpoint_path; + std::string tokenizer_model_path; + int max_prompt_tokens = 200; + int state_dim = 0; + int num_views = 2; + int chunk = 10; + int num_steps = 10; + int vision_pool_factor = 1; + std::vector state_q01; + std::vector state_q99; + std::vector action_q01; + std::vector action_q99; +}; + +int build_native_model_runtime(const NativeOpenConfig& config, + frt_model_runtime_v1** out, + std::string* error); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_OPEN_INTERNAL_H diff --git a/cpp/models/pi05/src/native_quantization.cu b/cpp/models/pi05/src/native_quantization.cu new file mode 100644 index 00000000..e53ec458 --- /dev/null +++ b/cpp/models/pi05/src/native_quantization.cu @@ -0,0 +1,111 @@ +#include "flashrt/cpp/models/pi05/native_quantization.h" + +#include + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +bool valid_matrix(const NativeFloatTensor& tensor) { + if (tensor.shape.size() != 2 || !tensor.shape[0] || !tensor.shape[1]) { + return false; + } + const std::uint64_t rows = tensor.shape[0]; + const std::uint64_t columns = tensor.shape[1]; + return rows <= SIZE_MAX / columns && + rows * columns == tensor.values.size(); +} + +bool finite_values(const NativeFloatTensor& tensor) { + for (float value : tensor.values) { + if (!std::isfinite(value)) return false; + } + return true; +} + +} // namespace + +modalities::Status native_quantize_fp8_e4m3( + const NativeFloatTensor& bf16_weight, + bool transpose, + NativeFp8Tensor* out) { + if (!out || !valid_matrix(bf16_weight) || !finite_values(bf16_weight)) { + return invalid("FP8 weight must be a finite BF16 matrix"); + } + NativeFloatTensor arranged; + if (transpose) { + const modalities::Status st = + native_transpose_2d(bf16_weight, &arranged); + if (!st.ok_status()) return st; + } else { + arranged = bf16_weight; + } + float amax = 0.0f; + for (float value : arranged.values) { + amax = std::max(amax, std::fabs(value)); + } + NativeFp8Tensor result; + result.shape = arranged.shape; + result.scale = std::max(amax / 448.0f, 1.0e-12f); + result.values.resize(arranged.values.size()); + for (std::size_t i = 0; i < arranged.values.size(); ++i) { + const float value = std::max( + -448.0f, + std::min(448.0f, arranged.values[i] / result.scale)); + result.values[i] = __nv_fp8_e4m3(value).__x; + } + *out = std::move(result); + return modalities::Status::ok(); +} + +modalities::Status native_quantize_int8_per_output( + const NativeFloatTensor& bf16_weight, + NativeInt8Tensor* out) { + if (!out || !valid_matrix(bf16_weight) || !finite_values(bf16_weight)) { + return invalid("INT8 weight must be a finite BF16 matrix"); + } + NativeFloatTensor transposed; + modalities::Status st = native_transpose_2d(bf16_weight, &transposed); + if (!st.ok_status()) return st; + const std::size_t rows = static_cast(transposed.shape[0]); + const std::size_t columns = + static_cast(transposed.shape[1]); + NativeInt8Tensor result; + result.shape = transposed.shape; + result.values.resize(transposed.values.size()); + result.scales.resize(rows); + const float inv_int8_max = 1.0f / 127.0f; + for (std::size_t row = 0; row < rows; ++row) { + float amax = 0.0f; + for (std::size_t column = 0; column < columns; ++column) { + amax = std::max( + amax, std::fabs(transposed.values[row * columns + column])); + } + const float scale = std::max(amax * inv_int8_max, 1.0e-12f); + result.scales[row] = scale; + for (std::size_t column = 0; column < columns; ++column) { + const float scaled = + transposed.values[row * columns + column] / scale; + const float rounded = std::nearbyint(scaled); + result.values[row * columns + column] = static_cast( + std::max(-127.0f, std::min(127.0f, rounded))); + } + } + *out = std::move(result); + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_quantization_unavailable.cpp b/cpp/models/pi05/src/native_quantization_unavailable.cpp new file mode 100644 index 00000000..94d6957c --- /dev/null +++ b/cpp/models/pi05/src/native_quantization_unavailable.cpp @@ -0,0 +1,31 @@ +#include "flashrt/cpp/models/pi05/native_quantization.h" + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status unavailable() { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "native weight quantization requires the CUDA kernels build"); +} + +} // namespace + +modalities::Status native_quantize_fp8_e4m3( + const NativeFloatTensor&, + bool, + NativeFp8Tensor*) { + return unavailable(); +} + +modalities::Status native_quantize_int8_per_output( + const NativeFloatTensor&, + NativeInt8Tensor*) { + return unavailable(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_rtx_attention.cpp b/cpp/models/pi05/src/native_rtx_attention.cpp new file mode 100644 index 00000000..f77a6b91 --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_attention.cpp @@ -0,0 +1,205 @@ +#include "flashrt/cpp/models/pi05/native_rtx_attention.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +std::size_t dtype_size(NativeAttentionDType dtype) { + switch (dtype) { + case NativeAttentionDType::kBf16: return sizeof(std::uint16_t); + case NativeAttentionDType::kFloat32: return sizeof(float); + case NativeAttentionDType::kInt32: return sizeof(std::int32_t); + } + return 0; +} + +bool element_count(std::initializer_list shape, + std::size_t* out) { + std::size_t count = 1; + for (std::uint64_t dim : shape) { + if (!dim || dim > std::numeric_limits::max() || + count > std::numeric_limits::max() / + static_cast(dim)) { + return false; + } + count *= static_cast(dim); + } + if (out) *out = count; + return true; +} + +std::uint64_t round_up_128(std::uint64_t value) { + return ((value + 127) / 128) * 128; +} + +} // namespace + +modalities::Status NativeRtxAttentionWorkspace::add( + const std::string& name, + std::initializer_list shape, + NativeAttentionDType dtype) { + if (!ctx_ || name.empty() || buffers_.find(name) != buffers_.end()) { + return invalid("native attention buffer definition is invalid"); + } + std::size_t elements = 0; + const std::size_t width = dtype_size(dtype); + if (!width || !element_count(shape, &elements) || + elements > std::numeric_limits::max() / width) { + return invalid("native attention buffer shape is invalid"); + } + const std::size_t bytes = elements * width; + frt_buffer buffer = frt_buffer_alloc(ctx_, name.c_str(), bytes); + if (!buffer) return backend("native attention allocation failed"); + buffers_.emplace(name, NativeAttentionBuffer{ + buffer, std::vector(shape), + dtype}); + allocated_bytes_ += bytes; + return modalities::Status::ok(); +} + +modalities::Status NativeRtxAttentionWorkspace::allocate( + const NativeRtxAttentionConfig& config) { + if (!ctx_ || !buffers_.empty() || config.num_views < 1 || + config.num_views > 3 || config.encoder_sequence <= 0 || + config.encoder_vision_sequence <= 0 || + config.encoder_vision_sequence > config.encoder_sequence || + config.chunk_size <= 0 || config.encoder_layers != 18) { + return invalid("Pi0.5 RTX attention configuration is invalid"); + } + num_views_ = config.num_views; + encoder_sequence_ = config.encoder_sequence; + encoder_vision_sequence_ = config.encoder_vision_sequence; + chunk_size_ = config.chunk_size; + encoder_layers_ = config.encoder_layers; + const std::uint64_t nv = static_cast(num_views_); + const std::uint64_t es = static_cast(encoder_sequence_); + const std::uint64_t ds = static_cast(chunk_size_); + const std::uint64_t layers = static_cast(encoder_layers_); + const std::uint64_t total_kv = es + ds; + encoder_splits_ = std::min(128, (encoder_sequence_ + 63) / 64); + decoder_splits_ = + std::min(128, (encoder_sequence_ + chunk_size_ + 63) / 64); + kv_layer_stride_bytes_ = + static_cast(total_kv) * 256 * sizeof(std::uint16_t); + modalities::Status st; +#define FRT_ADD(...) \ + do { \ + st = add(__VA_ARGS__); \ + if (!st.ok_status()) return st; \ + } while (false) + FRT_ADD("attn_vis_Q", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_vis_K", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_vis_V", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_Q", {es, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_K", {layers, total_kv, 1, 256}, + NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_V", {layers, total_kv, 1, 256}, + NativeAttentionDType::kBf16); + FRT_ADD("attn_dec_Q", {ds, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_seqused", {1}, NativeAttentionDType::kInt32); + FRT_ADD("attn_dec_seqused", {1}, NativeAttentionDType::kInt32); + FRT_ADD("attn_dec_devpos", {1}, NativeAttentionDType::kInt32); + + FRT_ADD("attn_vis_O", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_vis_lse", {nv, 16, 256}, NativeAttentionDType::kFloat32); + FRT_ADD("attn_vis_lse_accum", {2, nv, 16, 256}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_vis_o_accum", {2, nv, 16, 256, 96}, + NativeAttentionDType::kFloat32); + + FRT_ADD("attn_enc_O", {1, es, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_lse", {1, 8, round_up_128(es)}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_enc_lse_accum", + {static_cast(encoder_splits_), 1, 8, es}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_enc_o_accum", + {static_cast(encoder_splits_), 1, 8, es, 256}, + NativeAttentionDType::kFloat32); + + FRT_ADD("attn_dec_O", {1, ds, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_dec_lse", {1, 8, round_up_128(ds)}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_dec_lse_accum", + {static_cast(decoder_splits_), 1, 8, ds}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_dec_o_accum", + {static_cast(decoder_splits_), 1, 8, ds, 256}, + NativeAttentionDType::kFloat32); +#undef FRT_ADD + const char* prompt_length_names[] = { + "attn_enc_seqused", "attn_dec_seqused", "attn_dec_devpos"}; + for (int i = 0; i < 3; ++i) { + const NativeAttentionBuffer* target = find(prompt_length_names[i]); + if (!target) return invalid("prompt length buffer was not allocated"); + prompt_length_buffers_[i] = target->buffer; + } + return set_fixed_prompt_length(0); +} + +modalities::Status NativeRtxAttentionWorkspace::set_fixed_prompt_length( + int prompt_tokens) { + const int max_prompt = encoder_sequence_ - encoder_vision_sequence_; + if (prompt_tokens < 0 || prompt_tokens > max_prompt || buffers_.empty()) { + return invalid("Pi0.5 fixed attention prompt length is invalid"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "fixed attention update requires the CUDA build"); +#else + const std::int32_t valid = encoder_vision_sequence_ + prompt_tokens; + const std::int32_t values[] = {valid, valid + chunk_size_, valid}; + for (int i = 0; i < 3; ++i) { + if (!prompt_length_buffers_[i] || + cudaMemcpy(frt_buffer_dptr(prompt_length_buffers_[i]), &values[i], + sizeof(values[i]), cudaMemcpyHostToDevice) != + cudaSuccess) { + return backend("fixed attention length upload failed"); + } + } + return modalities::Status::ok(); +#endif +} + +const NativeAttentionBuffer* NativeRtxAttentionWorkspace::find( + const std::string& name) const { + const auto it = buffers_.find(name); + return it == buffers_.end() ? nullptr : &it->second; +} + +void* NativeRtxAttentionWorkspace::encoder_k_layer_dptr(int layer) const { + const NativeAttentionBuffer* cache = find("attn_enc_K"); + if (!cache || layer < 0 || layer >= encoder_layers_) return nullptr; + return static_cast(frt_buffer_dptr(cache->buffer)) + + static_cast(layer) * kv_layer_stride_bytes_; +} + +void* NativeRtxAttentionWorkspace::encoder_v_layer_dptr(int layer) const { + const NativeAttentionBuffer* cache = find("attn_enc_V"); + if (!cache || layer < 0 || layer >= encoder_layers_) return nullptr; + return static_cast(frt_buffer_dptr(cache->buffer)) + + static_cast(layer) * kv_layer_stride_bytes_; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_rtx_attention_driver.cu b/cpp/models/pi05/src/native_rtx_attention_driver.cu new file mode 100644 index 00000000..c43cce7b --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_attention_driver.cu @@ -0,0 +1,208 @@ +#include "flashrt/cpp/models/pi05/native_rtx_attention_driver.h" + +#include "attention/fa2_wrapper.h" + +#include + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +modalities::Status launch_status() { + const cudaError_t rc = cudaGetLastError(); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +__global__ void fill_negative_infinity(float* values, std::size_t count) { + const std::size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) values[index] = __int_as_float(0xff800000); +} + +bool exact_shape(const NativeAttentionBuffer* buffer, + std::initializer_list expected) { + return buffer && buffer->shape == std::vector(expected); +} + +} // namespace + +NativeRtxAttentionDriver::NativeRtxAttentionDriver( + const NativeRtxAttentionWorkspace* workspace) noexcept + : workspace_(workspace) { + const NativeAttentionBuffer* vis = find("attn_vis_Q"); + const NativeAttentionBuffer* enc = find("attn_enc_Q"); + const NativeAttentionBuffer* dec = find("attn_dec_Q"); + const NativeAttentionBuffer* kv = find("attn_enc_K"); + if (!vis || !enc || !dec || !kv || vis->shape.size() != 4 || + enc->shape.size() != 3 || dec->shape.size() != 3 || + kv->shape.size() != 4 || vis->shape[1] != 256 || + vis->shape[2] != 16 || vis->shape[3] != 72 || + enc->shape[1] != 8 || enc->shape[2] != 256 || + dec->shape[1] != 8 || dec->shape[2] != 256 || + kv->shape[0] != 18 || kv->shape[2] != 1 || kv->shape[3] != 256) { + error_ = "Pi0.5 RTX attention workspace is not allocated"; + return; + } + num_views_ = static_cast(vis->shape[0]); + encoder_sequence_ = static_cast(enc->shape[0]); + chunk_size_ = static_cast(dec->shape[0]); + total_kv_ = static_cast(kv->shape[1]); + if (total_kv_ != encoder_sequence_ + chunk_size_ || + !exact_shape(find("attn_vis_O"), + {static_cast(num_views_), 256, 16, 72}) || + !exact_shape(find("attn_enc_O"), + {1, static_cast(encoder_sequence_), 8, + 256}) || + !exact_shape(find("attn_dec_O"), + {1, static_cast(chunk_size_), 8, 256})) { + error_ = "Pi0.5 RTX attention workspace shapes are inconsistent"; + return; + } + int device = 0; + cudaDeviceProp properties{}; + cudaError_t rc = cudaGetDevice(&device); + if (rc == cudaSuccess) rc = cudaGetDeviceProperties(&properties, device); + if (rc != cudaSuccess) { + error_ = cudaGetErrorString(rc); + return; + } + if (properties.major < 8) { + error_ = "native BF16 FA2 requires compute capability 8.0 or newer"; + return; + } + num_sms_ = properties.multiProcessorCount; +} + +const NativeAttentionBuffer* NativeRtxAttentionDriver::find( + const char* name) const { + return workspace_ ? workspace_->find(name) : nullptr; +} + +modalities::Status NativeRtxAttentionDriver::status() const { + return error_.empty() ? modalities::Status::ok() : backend(error_); +} + +modalities::Status NativeRtxAttentionDriver::vision( + std::uintptr_t stream) const { + if (!error_.empty()) return backend(error_); + const NativeAttentionBuffer* q = find("attn_vis_Q"); + const NativeAttentionBuffer* k = find("attn_vis_K"); + const NativeAttentionBuffer* v = find("attn_vis_V"); + const NativeAttentionBuffer* o = find("attn_vis_O"); + const NativeAttentionBuffer* lse = find("attn_vis_lse"); + const NativeAttentionBuffer* lse_accum = find("attn_vis_lse_accum"); + const NativeAttentionBuffer* o_accum = find("attn_vis_o_accum"); + if (!q || !k || !v || !o || !lse || !lse_accum || !o_accum) { + return invalid("native vision attention buffers are incomplete"); + } + constexpr int row_stride = 16 * 72; + constexpr int batch_stride = 256 * row_stride; + fvk_attention_fa2_fwd_bf16( + frt_buffer_dptr(q->buffer), frt_buffer_dptr(k->buffer), + frt_buffer_dptr(v->buffer), frt_buffer_dptr(o->buffer), + frt_buffer_dptr(lse->buffer), frt_buffer_dptr(lse_accum->buffer), + frt_buffer_dptr(o_accum->buffer), num_views_, 256, 256, 16, 16, 72, + batch_stride, row_stride, 72, batch_stride, row_stride, 72, + batch_stride, row_stride, 72, batch_stride, row_stride, 72, + 1.0f / std::sqrt(72.0f), num_sms_, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeRtxAttentionDriver::encoder( + int layer, + std::uintptr_t stream) const { + if (!error_.empty()) return backend(error_); + void* k = workspace_->encoder_k_layer_dptr(layer); + void* v = workspace_->encoder_v_layer_dptr(layer); + const NativeAttentionBuffer* q = find("attn_enc_Q"); + const NativeAttentionBuffer* o = find("attn_enc_O"); + const NativeAttentionBuffer* lse = find("attn_enc_lse"); + const NativeAttentionBuffer* seqused = find("attn_enc_seqused"); + if (!q || !k || !v || !o || !lse || !seqused) { + return invalid("native encoder attention arguments are invalid"); + } + const int q_row_stride = 8 * 256; + const int q_batch_stride = encoder_sequence_ * q_row_stride; + const int kv_batch_stride = total_kv_ * 256; + fvk_attention_fa2_fwd_bf16_seqused( + frt_buffer_dptr(q->buffer), k, v, frt_buffer_dptr(o->buffer), + frt_buffer_dptr(lse->buffer), frt_buffer_dptr(seqused->buffer), 1, + encoder_sequence_, encoder_sequence_, 8, 1, 256, q_batch_stride, + q_row_stride, 256, kv_batch_stride, 256, 256, kv_batch_stride, 256, + 256, q_batch_stride, q_row_stride, 256, + 1.0f / std::sqrt(256.0f), num_sms_, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeRtxAttentionDriver::decoder( + int layer, + std::uintptr_t stream) const { + if (!error_.empty()) return backend(error_); + void* k = workspace_->encoder_k_layer_dptr(layer); + void* v = workspace_->encoder_v_layer_dptr(layer); + const NativeAttentionBuffer* q = find("attn_dec_Q"); + const NativeAttentionBuffer* o = find("attn_dec_O"); + const NativeAttentionBuffer* lse = find("attn_dec_lse"); + const NativeAttentionBuffer* seqused = find("attn_dec_seqused"); + const NativeAttentionBuffer* lse_accum = find("attn_dec_lse_accum"); + const NativeAttentionBuffer* o_accum = find("attn_dec_o_accum"); + if (!q || !k || !v || !o || !lse || !seqused || !lse_accum || + !o_accum) { + return invalid("native decoder attention arguments are invalid"); + } + const std::size_t accum_count = + frt_buffer_bytes(lse_accum->buffer) / sizeof(float); + fill_negative_infinity<<<(accum_count + 255) / 256, 256, 0, + reinterpret_cast(stream)>>>( + static_cast(frt_buffer_dptr(lse_accum->buffer)), accum_count); + cudaError_t rc = cudaGetLastError(); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + + const int q_row_stride = 8 * 256; + const int q_batch_stride = chunk_size_ * q_row_stride; + const int kv_batch_stride = total_kv_ * 256; + fvk_attention_fa2_fwd_bf16_seqused_splitkv( + frt_buffer_dptr(q->buffer), k, v, frt_buffer_dptr(o->buffer), + frt_buffer_dptr(lse->buffer), frt_buffer_dptr(seqused->buffer), + frt_buffer_dptr(lse_accum->buffer), frt_buffer_dptr(o_accum->buffer), + 1, chunk_size_, total_kv_, 8, 1, 256, q_batch_stride, q_row_stride, + 256, kv_batch_stride, 256, 256, kv_batch_stride, 256, 256, + q_batch_stride, q_row_stride, 256, 1.0f / std::sqrt(256.0f), + num_sms_, reinterpret_cast(stream)); + return launch_status(); +} + +void* NativeRtxAttentionDriver::vision_output() const { + const NativeAttentionBuffer* output = find("attn_vis_O"); + return output ? frt_buffer_dptr(output->buffer) : nullptr; +} + +void* NativeRtxAttentionDriver::encoder_output() const { + const NativeAttentionBuffer* output = find("attn_enc_O"); + return output ? frt_buffer_dptr(output->buffer) : nullptr; +} + +void* NativeRtxAttentionDriver::decoder_output() const { + const NativeAttentionBuffer* output = find("attn_dec_O"); + return output ? frt_buffer_dptr(output->buffer) : nullptr; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_style_precompute.cu b/cpp/models/pi05/src/native_style_precompute.cu new file mode 100644 index 00000000..07c496ab --- /dev/null +++ b/cpp/models/pi05/src/native_style_precompute.cu @@ -0,0 +1,197 @@ +#include "flashrt/cpp/models/pi05/native_style_precompute.h" + +#include + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +bool weight_shape(const NativeDeviceWeightStore& weights, + const std::string& name, + std::initializer_list shape, + const NativeDeviceWeight** out) { + const NativeDeviceWeight* weight = weights.find(name); + if (!weight || weight->dtype != NativeWeightDType::kBf16 || + weight->shape != std::vector(shape)) { + return false; + } + if (out) *out = weight; + return true; +} + +void* offset(void* base, std::size_t elements) { + return static_cast(base) + + elements * sizeof(std::uint16_t); +} + +} // namespace + +modalities::Status NativeStylePrecomputer::run( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace) { + return invalid("native style precomputer is invalid"); + } + const NativeWorkspaceBuffer* time_output = + workspace->find("decoder_time_emb"); + const NativeWorkspaceBuffer* style_attn = + workspace->find("decoder_style_attn"); + const NativeWorkspaceBuffer* style_ffn = + workspace->find("decoder_style_ffn"); + const NativeWorkspaceBuffer* style_final = + workspace->find("decoder_style_final"); + const NativeWorkspaceBuffer* scratch_a = workspace->find("decoder_x"); + const NativeWorkspaceBuffer* scratch_b = workspace->find("x_normed_buf"); + if (!time_output || !style_attn || !style_ffn || !style_final || + !scratch_a || !scratch_b || time_output->shape.size() != 3 || + style_attn->shape.size() != 4 || style_ffn->shape != style_attn->shape || + style_final->shape.size() != 3) { + return invalid("native style workspace layout is invalid"); + } + const int steps = static_cast(time_output->shape[0]); + const int chunk = static_cast(time_output->shape[1]); + if (time_output->shape[2] != 1024 || style_attn->shape[0] != steps || + style_attn->shape[1] != 18 || style_attn->shape[2] != chunk || + style_attn->shape[3] != 3072 || + style_final->shape != + std::vector( + {static_cast(steps), + static_cast(chunk), 3072})) { + return invalid("native style workspace shape is invalid"); + } + + const NativeDeviceWeight* time_source = nullptr; + const NativeDeviceWeight* time_in_w = nullptr; + const NativeDeviceWeight* time_in_b = nullptr; + const NativeDeviceWeight* time_out_w = nullptr; + const NativeDeviceWeight* time_out_b = nullptr; + const NativeDeviceWeight* final_w = nullptr; + const NativeDeviceWeight* final_b = nullptr; + if (!weight_shape(weights, "decoder_time_embeds", + {static_cast(steps), 1024}, + &time_source) || + !weight_shape(weights, "decoder_time_mlp_in_w", {1024, 1024}, &time_in_w) || + !weight_shape(weights, "decoder_time_mlp_in_b", {1024}, &time_in_b) || + !weight_shape(weights, "decoder_time_mlp_out_w", {1024, 1024}, &time_out_w) || + !weight_shape(weights, "decoder_time_mlp_out_b", {1024}, &time_out_b) || + !weight_shape(weights, "decoder_final_norm_mod_w", {1024, 3072}, &final_w) || + !weight_shape(weights, "decoder_final_norm_mod_b", {3072}, &final_b)) { + return invalid("native style global weights are incomplete"); + } + const cudaStream_t cuda_stream = reinterpret_cast(stream); + const std::uintptr_t native_stream = stream; + for (int step = 0; step < steps; ++step) { + void* time_row = offset(frt_buffer_dptr(time_source->buffer), + static_cast(step) * 1024); + modalities::Status st = driver_->bf16_nn( + time_row, frt_buffer_dptr(time_in_w->buffer), + frt_buffer_dptr(scratch_a->buffer), 1, 1024, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(scratch_a->buffer), + frt_buffer_dptr(time_in_b->buffer), 1, 1024, native_stream); + if (!st.ok_status()) return st; + st = driver_->silu_bf16(frt_buffer_dptr(scratch_a->buffer), 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(scratch_a->buffer), + frt_buffer_dptr(time_out_w->buffer), + frt_buffer_dptr(scratch_b->buffer), 1, 1024, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(scratch_b->buffer), + frt_buffer_dptr(time_out_b->buffer), 1, 1024, native_stream); + if (!st.ok_status()) return st; + st = driver_->silu_bf16(frt_buffer_dptr(scratch_b->buffer), 1024, + native_stream); + if (!st.ok_status()) return st; + + void* expanded = offset( + frt_buffer_dptr(time_output->buffer), + static_cast(step) * chunk * 1024); + for (int row = 0; row < chunk; ++row) { + const cudaError_t rc = cudaMemcpyAsync( + offset(expanded, static_cast(row) * 1024), + frt_buffer_dptr(scratch_b->buffer), + 1024 * sizeof(std::uint16_t), cudaMemcpyDeviceToDevice, + cuda_stream); + if (rc != cudaSuccess) return backend("time style expansion failed"); + } + + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* attn_w = nullptr; + const NativeDeviceWeight* attn_b = nullptr; + const NativeDeviceWeight* ffn_w = nullptr; + const NativeDeviceWeight* ffn_b = nullptr; + if (!weight_shape(weights, "decoder_pre_attn_norm_mod_w_" + suffix, + {1024, 3072}, &attn_w) || + !weight_shape(weights, "decoder_pre_attn_norm_mod_b_" + suffix, + {3072}, &attn_b) || + !weight_shape(weights, "decoder_pre_ffn_norm_mod_w_" + suffix, + {1024, 3072}, &ffn_w) || + !weight_shape(weights, "decoder_pre_ffn_norm_mod_b_" + suffix, + {3072}, &ffn_b)) { + return invalid("native style layer weights are incomplete"); + } + const std::size_t style_offset = + (static_cast(step) * 18 + layer) * chunk * 3072; + void* attn_target = + offset(frt_buffer_dptr(style_attn->buffer), style_offset); + void* ffn_target = + offset(frt_buffer_dptr(style_ffn->buffer), style_offset); + st = driver_->bf16_nn(expanded, frt_buffer_dptr(attn_w->buffer), + attn_target, chunk, 3072, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16(attn_target, + frt_buffer_dptr(attn_b->buffer), + chunk, 3072, native_stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn(expanded, frt_buffer_dptr(ffn_w->buffer), + ffn_target, chunk, 3072, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16(ffn_target, + frt_buffer_dptr(ffn_b->buffer), + chunk, 3072, native_stream); + if (!st.ok_status()) return st; + } + void* final_target = offset( + frt_buffer_dptr(style_final->buffer), + static_cast(step) * chunk * 3072); + st = driver_->bf16_nn(expanded, frt_buffer_dptr(final_w->buffer), + final_target, chunk, 3072, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16(final_target, + frt_buffer_dptr(final_b->buffer), + chunk, 3072, native_stream); + if (!st.ok_status()) return st; + } + const cudaError_t rc = cudaStreamSynchronize(cuda_stream); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("native style precompute synchronization failed"); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp new file mode 100644 index 00000000..251f7420 --- /dev/null +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -0,0 +1,467 @@ +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +std::string encoder_prefix(int layer) { + return "paligemma_with_expert.paligemma.model.language_model.layers." + + std::to_string(layer); +} + +std::string decoder_prefix(int layer) { + return "paligemma_with_expert.gemma_expert.model.layers." + + std::to_string(layer); +} + +const std::string& vision_prefix() { + static const std::string prefix = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; + return prefix; +} + +std::string layer_name(const char* stem, int layer) { + return std::string(stem) + std::to_string(layer); +} + +} // namespace + +modalities::Status NativeWeightMaterializer::load( + const std::string& key, + NativeFloatTensor* out) { + return load_native_float_tensor(source_, key, out); +} + +modalities::Status NativeWeightMaterializer::upload( + const std::string& name, + const NativeFloatTensor& tensor) { + if (!destination_) return invalid("native weight destination is null"); + NativeBf16Tensor bf16; + modalities::Status st = native_to_bf16(tensor, &bf16); + if (!st.ok_status()) return st; + return destination_->upload(name, bf16); +} + +modalities::Status NativeWeightMaterializer::upload_rounded_transpose( + const std::string& source_key, + const std::string& destination_name) { + NativeFloatTensor source; + NativeFloatTensor rounded; + NativeFloatTensor transposed; + modalities::Status st = load(source_key, &source); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(source, &rounded); + if (!st.ok_status()) return st; + st = native_transpose_2d(rounded, &transposed); + if (!st.ok_status()) return st; + return upload(destination_name, transposed); +} + +modalities::Status NativeWeightMaterializer::upload_rounded_copy( + const std::string& source_key, + const std::string& destination_name) { + NativeFloatTensor source; + NativeFloatTensor rounded; + modalities::Status st = load(source_key, &source); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(source, &rounded); + if (!st.ok_status()) return st; + return upload(destination_name, rounded); +} + +modalities::Status NativeWeightMaterializer::upload_folded_transpose( + const std::string& source_key, + const NativeFloatTensor& norm, + const std::string& destination_name) { + NativeFloatTensor source; + NativeFloatTensor folded; + NativeFloatTensor transposed; + modalities::Status st = load(source_key, &source); + if (!st.ok_status()) return st; + st = native_fold_rms_columns(source, norm, &folded); + if (!st.ok_status()) return st; + st = native_transpose_2d(folded, &transposed); + if (!st.ok_status()) return st; + return upload(destination_name, transposed); +} + +modalities::Status NativeWeightMaterializer::upload_rounded_scaled( + const std::string& source_key, + const std::string& destination_name, + float scale, + bool transpose) { + NativeFloatTensor source; + NativeFloatTensor rounded; + NativeFloatTensor arranged; + NativeFloatTensor scaled; + modalities::Status st = load(source_key, &source); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(source, &rounded); + if (!st.ok_status()) return st; + if (transpose) { + st = native_transpose_2d(rounded, &arranged); + if (!st.ok_status()) return st; + } else { + arranged = std::move(rounded); + } + st = native_scale(arranged, scale, &scaled); + if (!st.ok_status()) return st; + return upload(destination_name, scaled); +} + +modalities::Status NativeWeightMaterializer::materialize_encoder_layer( + int layer) { + if (layer < 0 || layer >= 18 || !destination_) { + return invalid("Pi0.5 encoder layer index is invalid"); + } + const std::string prefix = encoder_prefix(layer); + NativeFloatTensor norm; + modalities::Status st = load(prefix + ".input_layernorm.weight", &norm); + if (!st.ok_status()) return st; + + NativeFloatTensor q; + NativeFloatTensor k; + NativeFloatTensor v; + NativeFloatTensor qi; + NativeFloatTensor ki; + NativeFloatTensor qf; + NativeFloatTensor kf; + NativeFloatTensor vf; + NativeFloatTensor qkv; + st = load(prefix + ".self_attn.q_proj.weight", &q); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.k_proj.weight", &k); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.v_proj.weight", &v); + if (!st.ok_status()) return st; + st = native_interleave_qk_rows(q, 8, &qi); + if (!st.ok_status()) return st; + st = native_interleave_qk_rows(k, 1, &ki); + if (!st.ok_status()) return st; + st = native_fold_rms_columns(qi, norm, &qf); + if (!st.ok_status()) return st; + st = native_fold_rms_columns(ki, norm, &kf); + if (!st.ok_status()) return st; + st = native_fold_rms_columns(v, norm, &vf); + if (!st.ok_status()) return st; + st = native_concat_rows_transpose({&qf, &kf, &vf}, &qkv); + if (!st.ok_status()) return st; + st = upload(layer_name("encoder_attn_qkv_w_", layer), qkv); + if (!st.ok_status()) return st; + + st = upload_rounded_transpose( + prefix + ".self_attn.o_proj.weight", + layer_name("encoder_attn_o_w_", layer)); + if (!st.ok_status()) return st; + + st = load(prefix + ".post_attention_layernorm.weight", &norm); + if (!st.ok_status()) return st; + st = upload_folded_transpose( + prefix + ".mlp.gate_proj.weight", norm, + layer_name("encoder_ffn_gate_w_", layer)); + if (!st.ok_status()) return st; + st = upload_folded_transpose( + prefix + ".mlp.up_proj.weight", norm, + layer_name("encoder_ffn_up_w_", layer)); + if (!st.ok_status()) return st; + return upload_rounded_transpose( + prefix + ".mlp.down_proj.weight", + layer_name("encoder_ffn_down_w_", layer)); +} + +modalities::Status NativeWeightMaterializer::materialize_decoder_layer( + int layer, + bool merge_gate_up) { + if (layer < 0 || layer >= 18 || !destination_) { + return invalid("Pi0.5 decoder layer index is invalid"); + } + const std::string prefix = decoder_prefix(layer); + NativeFloatTensor q; + NativeFloatTensor k; + NativeFloatTensor v; + NativeFloatTensor qr; + NativeFloatTensor kr; + NativeFloatTensor vr; + NativeFloatTensor qi; + NativeFloatTensor ki; + NativeFloatTensor qkv; + modalities::Status st = load(prefix + ".self_attn.q_proj.weight", &q); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.k_proj.weight", &k); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.v_proj.weight", &v); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(q, &qr); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(k, &kr); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(v, &vr); + if (!st.ok_status()) return st; + st = native_interleave_qk_rows(qr, 8, &qi); + if (!st.ok_status()) return st; + st = native_interleave_qk_rows(kr, 1, &ki); + if (!st.ok_status()) return st; + st = native_concat_rows_transpose({&qi, &ki, &vr}, &qkv); + if (!st.ok_status()) return st; + st = upload(layer_name("decoder_attn_qkv_w_", layer), qkv); + if (!st.ok_status()) return st; + + st = upload_rounded_transpose( + prefix + ".self_attn.o_proj.weight", + layer_name("decoder_attn_o_w_", layer)); + if (!st.ok_status()) return st; + + NativeFloatTensor gate; + NativeFloatTensor up; + NativeFloatTensor gate_rounded; + NativeFloatTensor up_rounded; + NativeFloatTensor gate_t; + NativeFloatTensor up_t; + st = load(prefix + ".mlp.gate_proj.weight", &gate); + if (!st.ok_status()) return st; + st = load(prefix + ".mlp.up_proj.weight", &up); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(gate, &gate_rounded); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(up, &up_rounded); + if (!st.ok_status()) return st; + st = native_transpose_2d(gate_rounded, &gate_t); + if (!st.ok_status()) return st; + st = native_transpose_2d(up_rounded, &up_t); + if (!st.ok_status()) return st; + st = upload(layer_name("decoder_ffn_gate_w_", layer), gate_t); + if (!st.ok_status()) return st; + st = upload(layer_name("decoder_ffn_up_w_", layer), up_t); + if (!st.ok_status()) return st; + if (merge_gate_up) { + NativeFloatTensor gate_up; + st = native_concat_columns(gate_t, up_t, &gate_up); + if (!st.ok_status()) return st; + st = upload(layer_name("decoder_ffn_gate_up_w_", layer), gate_up); + if (!st.ok_status()) return st; + } + st = upload_rounded_transpose( + prefix + ".mlp.down_proj.weight", + layer_name("decoder_ffn_down_w_", layer)); + if (!st.ok_status()) return st; + + st = upload_rounded_transpose( + prefix + ".input_layernorm.dense.weight", + layer_name("decoder_pre_attn_norm_mod_w_", layer)); + if (!st.ok_status()) return st; + st = upload_rounded_copy( + prefix + ".input_layernorm.dense.bias", + layer_name("decoder_pre_attn_norm_mod_b_", layer)); + if (!st.ok_status()) return st; + st = upload_rounded_transpose( + prefix + ".post_attention_layernorm.dense.weight", + layer_name("decoder_pre_ffn_norm_mod_w_", layer)); + if (!st.ok_status()) return st; + return upload_rounded_copy( + prefix + ".post_attention_layernorm.dense.bias", + layer_name("decoder_pre_ffn_norm_mod_b_", layer)); +} + +modalities::Status NativeWeightMaterializer::materialize_vision_layer( + int layer) { + if (layer < 0 || layer >= 27 || !destination_) { + return invalid("Pi0.5 vision layer index is invalid"); + } + const std::string prefix = vision_prefix() + ".encoder.layers." + + std::to_string(layer); + NativeFloatTensor q; + NativeFloatTensor k; + NativeFloatTensor v; + NativeFloatTensor qr; + NativeFloatTensor kr; + NativeFloatTensor vr; + NativeFloatTensor qkv; + modalities::Status st = load(prefix + ".self_attn.q_proj.weight", &q); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.k_proj.weight", &k); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.v_proj.weight", &v); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(q, &qr); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(k, &kr); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(v, &vr); + if (!st.ok_status()) return st; + st = native_concat_rows_transpose({&qr, &kr, &vr}, &qkv); + if (!st.ok_status()) return st; + st = upload(layer_name("vision_attn_qkv_w_", layer), qkv); + if (!st.ok_status()) return st; + + st = load(prefix + ".self_attn.q_proj.bias", &q); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.k_proj.bias", &k); + if (!st.ok_status()) return st; + st = load(prefix + ".self_attn.v_proj.bias", &v); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(q, &qr); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(k, &kr); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(v, &vr); + if (!st.ok_status()) return st; + st = native_concat_vectors({&qr, &kr, &vr}, &qkv); + if (!st.ok_status()) return st; + st = upload(layer_name("vision_attn_qkv_b_", layer), qkv); + if (!st.ok_status()) return st; + + const struct { + const char* source; + const char* destination; + bool transpose; + } entries[] = { + {"self_attn.out_proj.weight", "vision_attn_o_w_", true}, + {"self_attn.out_proj.bias", "vision_attn_o_b_", false}, + {"mlp.fc1.weight", "vision_ffn_up_w_", true}, + {"mlp.fc1.bias", "vision_ffn_up_b_", false}, + {"mlp.fc2.weight", "vision_ffn_down_w_", true}, + {"mlp.fc2.bias", "vision_ffn_down_b_", false}, + {"layer_norm1.weight", "vision_pre_attn_norm_w_", false}, + {"layer_norm1.bias", "vision_pre_attn_norm_b_", false}, + {"layer_norm2.weight", "vision_pre_ffn_norm_w_", false}, + {"layer_norm2.bias", "vision_pre_ffn_norm_b_", false}, + }; + for (const auto& entry : entries) { + st = entry.transpose + ? upload_rounded_transpose( + prefix + "." + entry.source, + layer_name(entry.destination, layer)) + : upload_rounded_copy( + prefix + "." + entry.source, + layer_name(entry.destination, layer)); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +modalities::Status NativeWeightMaterializer::materialize_vision_globals() { + if (!destination_) return invalid("native weight destination is null"); + const std::string prefix = vision_prefix(); + NativeFloatTensor patch; + NativeFloatTensor rounded; + NativeFloatTensor permuted; + modalities::Status st = load( + prefix + ".embeddings.patch_embedding.weight", &patch); + if (!st.ok_status()) return st; + st = native_round_to_bf16_float(patch, &rounded); + if (!st.ok_status()) return st; + st = native_patch_oihw_to_hwio(rounded, &permuted); + if (!st.ok_status()) return st; + st = upload("vision_patch_embedding_w", permuted); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".embeddings.patch_embedding.bias", + "vision_patch_embedding_b"); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".embeddings.position_embedding.weight", + "vision_position_embedding"); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".post_layernorm.weight", + "vision_final_norm_w"); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".post_layernorm.bias", + "vision_final_norm_b"); + if (!st.ok_status()) return st; + + const std::string projector = + "paligemma_with_expert.paligemma.model.multi_modal_projector.linear"; + st = upload_rounded_transpose(projector + ".weight", + "encoder_multi_modal_projector_w"); + if (!st.ok_status()) return st; + return upload_rounded_copy(projector + ".bias", + "encoder_multi_modal_projector_b"); +} + +modalities::Status NativeWeightMaterializer::materialize_decoder_globals( + int num_steps) { + if (!destination_ || num_steps <= 0) { + return invalid("Pi0.5 decoder global configuration is invalid"); + } + const struct { + const char* source; + const char* destination; + bool transpose; + } entries[] = { + {"paligemma_with_expert.gemma_expert.model.norm.dense.weight", + "decoder_final_norm_mod_w", true}, + {"paligemma_with_expert.gemma_expert.model.norm.dense.bias", + "decoder_final_norm_mod_b", false}, + {"time_mlp_in.weight", "decoder_time_mlp_in_w", true}, + {"time_mlp_in.bias", "decoder_time_mlp_in_b", false}, + {"time_mlp_out.weight", "decoder_time_mlp_out_w", true}, + {"time_mlp_out.bias", "decoder_time_mlp_out_b", false}, + {"action_in_proj.weight", "decoder_action_in_proj_w", true}, + {"action_in_proj.bias", "decoder_action_in_proj_b", false}, + }; + for (const auto& entry : entries) { + const modalities::Status st = + entry.transpose + ? upload_rounded_transpose(entry.source, entry.destination) + : upload_rounded_copy(entry.source, entry.destination); + if (!st.ok_status()) return st; + } + + NativeFloatTensor time_embeddings; + modalities::Status st = + native_pi05_time_embeddings(num_steps, 1024, &time_embeddings); + if (!st.ok_status()) return st; + st = upload("decoder_time_embeds", time_embeddings); + if (!st.ok_status()) return st; + + const float step_scale = -1.0f / static_cast(num_steps); + st = upload_rounded_scaled( + "action_out_proj.weight", "decoder_action_out_proj_w", step_scale, + true); + if (!st.ok_status()) return st; + return upload_rounded_scaled( + "action_out_proj.bias", "decoder_action_out_proj_b", step_scale, + false); +} + +modalities::Status NativeWeightMaterializer::materialize_embedding() { + if (!destination_) return invalid("native weight destination is null"); + return upload_rounded_copy( + "paligemma_with_expert.paligemma.lm_head.weight", + "embedding_weight"); +} + +modalities::Status NativeWeightMaterializer::materialize_all( + const NativeMaterializationOptions& options) { + if (!destination_ || options.num_steps <= 0) { + return invalid("Pi0.5 materialization options are invalid"); + } + modalities::Status st = materialize_vision_globals(); + if (!st.ok_status()) return st; + for (int layer = 0; layer < 27; ++layer) { + st = materialize_vision_layer(layer); + if (!st.ok_status()) return st; + } + for (int layer = 0; layer < 18; ++layer) { + st = materialize_encoder_layer(layer); + if (!st.ok_status()) return st; + } + for (int layer = 0; layer < 18; ++layer) { + st = materialize_decoder_layer( + layer, options.merge_decoder_gate_up); + if (!st.ok_status()) return st; + } + st = materialize_decoder_globals(options.num_steps); + if (!st.ok_status() || !options.include_embedding) return st; + return materialize_embedding(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_weight_ops.cpp b/cpp/models/pi05/src/native_weight_ops.cpp new file mode 100644 index 00000000..d3dd5134 --- /dev/null +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -0,0 +1,372 @@ +#include "flashrt/cpp/models/pi05/native_weight_ops.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +bool element_count(const std::vector& shape, + std::size_t* out) { + std::size_t count = 1; + for (std::uint64_t dim : shape) { + if (dim > std::numeric_limits::max() || + (dim && count > std::numeric_limits::max() / + static_cast(dim))) { + return false; + } + count *= static_cast(dim); + } + if (out) *out = count; + return true; +} + +bool valid_tensor(const NativeFloatTensor& tensor) { + std::size_t expected = 0; + return element_count(tensor.shape, &expected) && + expected == tensor.values.size(); +} + +const loader::SafetensorInfo* find_source_tensor( + const loader::SafetensorsFile& file, + const std::string& key) { + const loader::SafetensorInfo* tensor = file.find(key); + if (!tensor) tensor = file.find(std::string("model.") + key); + return tensor; +} + +} // namespace + +modalities::Status load_native_float_tensor( + const loader::SafetensorsFile& file, + const std::string& key, + NativeFloatTensor* out) { + if (!file.is_open() || !out) return invalid("invalid native tensor load"); + const loader::SafetensorInfo* tensor = find_source_tensor(file, key); + if (!tensor) { + return modalities::Status::error(modalities::StatusCode::kNotFound, + "native tensor not found: " + key); + } + std::size_t count = 0; + if (!element_count(tensor->shape, &count)) { + return invalid("native tensor shape overflows size_t"); + } + const void* data = file.data(*tensor); + if (!data && tensor->bytes) return invalid("native tensor has no payload"); + + NativeFloatTensor loaded; + loaded.shape = tensor->shape; + loaded.values.resize(count); + if (tensor->dtype == "F32") { + std::memcpy(loaded.values.data(), data, count * sizeof(float)); + } else if (tensor->dtype == "BF16") { + const auto* src = static_cast(data); + for (std::size_t i = 0; i < count; ++i) { + loaded.values[i] = modalities::bfloat16_to_float(src[i]); + } + } else if (tensor->dtype == "F16") { + const auto* src = static_cast(data); + for (std::size_t i = 0; i < count; ++i) { + loaded.values[i] = modalities::float16_to_float(src[i]); + } + } else { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "native tensor dtype is not a floating-point weight: " + + tensor->dtype); + } + *out = std::move(loaded); + return modalities::Status::ok(); +} + +modalities::Status native_to_bf16(const NativeFloatTensor& input, + NativeBf16Tensor* out) { + if (!out || !valid_tensor(input)) return invalid("invalid BF16 input"); + NativeBf16Tensor converted; + converted.shape = input.shape; + converted.values.resize(input.values.size()); + for (std::size_t i = 0; i < input.values.size(); ++i) { + converted.values[i] = modalities::float_to_bfloat16(input.values[i]); + } + *out = std::move(converted); + return modalities::Status::ok(); +} + +modalities::Status native_round_to_bf16_float( + const NativeFloatTensor& input, + NativeFloatTensor* out) { + if (!out || !valid_tensor(input)) { + return invalid("invalid BF16 round-trip input"); + } + NativeFloatTensor rounded = input; + for (float& value : rounded.values) { + value = modalities::bfloat16_to_float( + modalities::float_to_bfloat16(value)); + } + *out = std::move(rounded); + return modalities::Status::ok(); +} + +modalities::Status native_transpose_2d(const NativeFloatTensor& input, + NativeFloatTensor* out) { + if (!out || !valid_tensor(input) || input.shape.size() != 2) { + return invalid("transpose requires a valid rank-2 tensor"); + } + const std::size_t rows = static_cast(input.shape[0]); + const std::size_t cols = static_cast(input.shape[1]); + NativeFloatTensor transposed; + transposed.shape = {input.shape[1], input.shape[0]}; + transposed.values.resize(input.values.size()); + for (std::size_t row = 0; row < rows; ++row) { + for (std::size_t col = 0; col < cols; ++col) { + transposed.values[col * rows + row] = + input.values[row * cols + col]; + } + } + *out = std::move(transposed); + return modalities::Status::ok(); +} + +modalities::Status native_patch_oihw_to_hwio( + const NativeFloatTensor& input, + NativeFloatTensor* out) { + if (!out || !valid_tensor(input) || input.shape.size() != 4) { + return invalid("patch permutation requires a valid rank-4 tensor"); + } + const std::size_t outputs = static_cast(input.shape[0]); + const std::size_t channels = static_cast(input.shape[1]); + const std::size_t height = static_cast(input.shape[2]); + const std::size_t width = static_cast(input.shape[3]); + NativeFloatTensor permuted; + permuted.shape = {input.shape[2], input.shape[3], input.shape[1], + input.shape[0]}; + permuted.values.resize(input.values.size()); + for (std::size_t o = 0; o < outputs; ++o) { + for (std::size_t c = 0; c < channels; ++c) { + for (std::size_t h = 0; h < height; ++h) { + for (std::size_t w = 0; w < width; ++w) { + const std::size_t src = + ((o * channels + c) * height + h) * width + w; + const std::size_t dst = + ((h * width + w) * channels + c) * outputs + o; + permuted.values[dst] = input.values[src]; + } + } + } + } + *out = std::move(permuted); + return modalities::Status::ok(); +} + +modalities::Status native_interleave_qk_rows( + const NativeFloatTensor& input, + std::uint64_t num_heads, + NativeFloatTensor* out) { + if (!out || !valid_tensor(input) || input.shape.size() != 2 || + !num_heads || input.shape[0] % num_heads != 0) { + return invalid("Q/K interleave requires divisible rank-2 rows"); + } + const std::uint64_t head_dim = input.shape[0] / num_heads; + if (head_dim % 2 != 0) { + return invalid("Q/K interleave requires an even head dimension"); + } + const std::size_t cols = static_cast(input.shape[1]); + NativeFloatTensor interleaved; + interleaved.shape = input.shape; + interleaved.values.resize(input.values.size()); + for (std::uint64_t head = 0; head < num_heads; ++head) { + for (std::uint64_t pair = 0; pair < head_dim / 2; ++pair) { + for (std::uint64_t half = 0; half < 2; ++half) { + const std::uint64_t src_row = + head * head_dim + half * (head_dim / 2) + pair; + const std::uint64_t dst_row = + head * head_dim + pair * 2 + half; + std::memcpy(interleaved.values.data() + dst_row * cols, + input.values.data() + src_row * cols, + cols * sizeof(float)); + } + } + } + *out = std::move(interleaved); + return modalities::Status::ok(); +} + +modalities::Status native_fold_rms_columns( + const NativeFloatTensor& weight, + const NativeFloatTensor& norm, + NativeFloatTensor* out) { + if (!out || !valid_tensor(weight) || !valid_tensor(norm) || + weight.shape.size() != 2 || norm.shape.size() != 1 || + weight.shape[1] != norm.shape[0]) { + return invalid("RMS fold requires weight[out,in] and norm[in]"); + } + NativeFloatTensor folded = weight; + const std::size_t rows = static_cast(weight.shape[0]); + const std::size_t cols = static_cast(weight.shape[1]); + for (std::size_t row = 0; row < rows; ++row) { + for (std::size_t col = 0; col < cols; ++col) { + folded.values[row * cols + col] *= 1.0f + norm.values[col]; + } + } + *out = std::move(folded); + return modalities::Status::ok(); +} + +modalities::Status native_concat_rows_transpose( + const std::vector& inputs, + NativeFloatTensor* out) { + if (!out || inputs.empty() || !inputs[0] || + !valid_tensor(*inputs[0]) || inputs[0]->shape.size() != 2) { + return invalid("row concat requires rank-2 tensors"); + } + const std::uint64_t cols = inputs[0]->shape[1]; + std::uint64_t total_rows = 0; + for (const NativeFloatTensor* input : inputs) { + if (!input || !valid_tensor(*input) || input->shape.size() != 2 || + input->shape[1] != cols || + total_rows > std::numeric_limits::max() - + input->shape[0]) { + return invalid("row concat tensors have incompatible shapes"); + } + total_rows += input->shape[0]; + } + NativeFloatTensor joined; + joined.shape = {cols, total_rows}; + std::size_t joined_count = 0; + if (!element_count(joined.shape, &joined_count)) { + return invalid("row concat output shape overflows size_t"); + } + joined.values.resize(joined_count); + std::uint64_t row_offset = 0; + for (const NativeFloatTensor* input : inputs) { + for (std::uint64_t row = 0; row < input->shape[0]; ++row) { + for (std::uint64_t col = 0; col < cols; ++col) { + joined.values[static_cast(col * total_rows + + row_offset + row)] = + input->values[static_cast(row * cols + col)]; + } + } + row_offset += input->shape[0]; + } + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_concat_columns( + const NativeFloatTensor& left, + const NativeFloatTensor& right, + NativeFloatTensor* out) { + if (!out || !valid_tensor(left) || !valid_tensor(right) || + left.shape.size() != 2 || right.shape.size() != 2 || + left.shape[0] != right.shape[0]) { + return invalid("column concat tensors have incompatible shapes"); + } + const std::size_t rows = static_cast(left.shape[0]); + const std::size_t left_cols = static_cast(left.shape[1]); + const std::size_t right_cols = static_cast(right.shape[1]); + if (left.shape[1] > std::numeric_limits::max() - + right.shape[1]) { + return invalid("column concat output shape overflows uint64"); + } + NativeFloatTensor joined; + joined.shape = {left.shape[0], left.shape[1] + right.shape[1]}; + std::size_t joined_count = 0; + if (!element_count(joined.shape, &joined_count)) { + return invalid("column concat output shape overflows size_t"); + } + joined.values.resize(joined_count); + for (std::size_t row = 0; row < rows; ++row) { + float* dst = joined.values.data() + row * (left_cols + right_cols); + std::memcpy(dst, left.values.data() + row * left_cols, + left_cols * sizeof(float)); + std::memcpy(dst + left_cols, + right.values.data() + row * right_cols, + right_cols * sizeof(float)); + } + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_concat_vectors( + const std::vector& inputs, + NativeFloatTensor* out) { + if (!out || inputs.empty()) return invalid("vector concat has no inputs"); + std::size_t total = 0; + for (const NativeFloatTensor* input : inputs) { + if (!input || !valid_tensor(*input) || input->shape.size() != 1 || + input->values.size() > + std::numeric_limits::max() - total) { + return invalid("vector concat tensors have incompatible shapes"); + } + total += input->values.size(); + } + NativeFloatTensor joined; + joined.shape = {static_cast(total)}; + joined.values.reserve(total); + for (const NativeFloatTensor* input : inputs) { + joined.values.insert(joined.values.end(), input->values.begin(), + input->values.end()); + } + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_scale(const NativeFloatTensor& input, + float scale, + NativeFloatTensor* out) { + if (!out || !valid_tensor(input)) return invalid("invalid scale input"); + NativeFloatTensor scaled = input; + for (float& value : scaled.values) value *= scale; + *out = std::move(scaled); + return modalities::Status::ok(); +} + +modalities::Status native_pi05_time_embeddings( + int num_steps, + std::uint64_t embedding_dim, + NativeFloatTensor* out) { + if (!out || num_steps <= 0 || embedding_dim < 2 || + embedding_dim % 2 != 0) { + return invalid("Pi0.5 time embedding shape is invalid"); + } + const std::uint64_t half = embedding_dim / 2; + NativeFloatTensor result; + result.shape = {static_cast(num_steps), embedding_dim}; + result.values.resize(static_cast(num_steps) * embedding_dim); + const float dt = -1.0f / static_cast(num_steps); + const float min_period = 4.0e-3f; + const float period_ratio = 1000.0f; + const float pi = static_cast(3.14159265358979323846); + const float fraction_step = + half == 1 ? 0.0f : 1.0f / static_cast(half - 1); + float t = 1.0f; + for (int step = 0; step < num_steps; ++step) { + const std::size_t row = static_cast(step) * embedding_dim; + for (std::uint64_t i = 0; i < half; ++i) { + const float fraction = static_cast(i) * fraction_step; + const float period = + min_period * std::pow(period_ratio, fraction); + float angle = t * (1.0f / period); + angle *= 2.0f; + angle *= pi; + result.values[row + i] = std::sin(angle); + result.values[row + half + i] = std::cos(angle); + } + t += dt; + } + *out = std::move(result); + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_weight_packer.cpp b/cpp/models/pi05/src/native_weight_packer.cpp new file mode 100644 index 00000000..54e3d9b3 --- /dev/null +++ b/cpp/models/pi05/src/native_weight_packer.cpp @@ -0,0 +1,194 @@ +#include "flashrt/cpp/models/pi05/native_weight_packer.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +} // namespace + +modalities::Status NativeWeightPacker::load_bf16( + const std::string& name, + NativeFloatTensor* out) const { + if (!weights_ || !out) return invalid("native weight packer is invalid"); + NativeBf16Tensor source; + modalities::Status st = weights_->download_bf16(name, &source); + if (!st.ok_status()) return st; + NativeFloatTensor result; + result.shape = source.shape; + result.values.resize(source.values.size()); + for (std::size_t i = 0; i < source.values.size(); ++i) { + result.values[i] = + modalities::bfloat16_to_float(source.values[i]); + } + *out = std::move(result); + return modalities::Status::ok(); +} + +modalities::Status NativeWeightPacker::pack_fp8( + const std::string& name, + bool transpose) { + return pack_fp8_as(name, name, transpose); +} + +modalities::Status NativeWeightPacker::pack_fp8_as( + const std::string& source_name, + const std::string& packed_name, + bool transpose) { + NativeFloatTensor source; + modalities::Status st = load_bf16(source_name, &source); + if (!st.ok_status()) return st; + NativeFp8Tensor packed; + st = native_quantize_fp8_e4m3(source, transpose, &packed); + if (!st.ok_status()) return st; + const std::string prefix = "fp8." + packed_name; + st = weights_->upload_bytes(prefix, packed.shape, + NativeWeightDType::kFp8E4M3, + packed.values.data(), packed.values.size()); + if (!st.ok_status()) return st; + return weights_->upload_bytes(prefix + ".scale", {1}, + NativeWeightDType::kFloat32, + &packed.scale, sizeof(packed.scale)); +} + +modalities::Status NativeWeightPacker::pack_int8(const std::string& name) { + NativeFloatTensor source; + modalities::Status st = load_bf16(name, &source); + if (!st.ok_status()) return st; + NativeInt8Tensor packed; + st = native_quantize_int8_per_output(source, &packed); + if (!st.ok_status()) return st; + const std::string prefix = "int8." + name; + st = weights_->upload_bytes(prefix, packed.shape, + NativeWeightDType::kInt8, + packed.values.data(), packed.values.size()); + if (!st.ok_status()) return st; + return weights_->upload_bytes( + prefix + ".scale", {static_cast(packed.scales.size())}, + NativeWeightDType::kFloat32, packed.scales.data(), + packed.scales.size() * sizeof(float)); +} + +modalities::Status NativeWeightPacker::merge_bf16_columns( + const std::string& left_name, + const std::string& right_name, + const std::string& output_name) { + NativeFloatTensor left; + NativeFloatTensor right; + NativeFloatTensor merged; + modalities::Status st = load_bf16(left_name, &left); + if (!st.ok_status()) return st; + st = load_bf16(right_name, &right); + if (!st.ok_status()) return st; + st = native_concat_columns(left, right, &merged); + if (!st.ok_status()) return st; + NativeBf16Tensor bf16; + st = native_to_bf16(merged, &bf16); + if (!st.ok_status()) return st; + return weights_->upload(output_name, bf16); +} + +modalities::Status NativeWeightPacker::pack_all_fp8(bool transpose) { + if (!weights_) return invalid("native weight packer is invalid"); + modalities::Status st; + for (int layer = 0; layer < 27; ++layer) { + for (const char* stem : {"vision_attn_qkv_w_", "vision_attn_o_w_", + "vision_ffn_up_w_", + "vision_ffn_down_w_"}) { + st = pack_fp8(std::string(stem) + std::to_string(layer), + transpose); + if (!st.ok_status()) return st; + } + } + st = pack_fp8_as("encoder_multi_modal_projector_w", + "vision_projector_w", transpose); + if (!st.ok_status()) return st; + + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + const std::string gate_up = "encoder_ffn_gate_up_w_" + suffix; + st = merge_bf16_columns("encoder_ffn_gate_w_" + suffix, + "encoder_ffn_up_w_" + suffix, gate_up); + if (!st.ok_status()) return st; + for (const std::string& name : { + "encoder_attn_qkv_w_" + suffix, + "encoder_attn_o_w_" + suffix, + gate_up, + "encoder_ffn_down_w_" + suffix}) { + st = pack_fp8(name, transpose); + if (!st.ok_status()) return st; + } + } + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + for (const std::string& name : { + "decoder_attn_qkv_w_" + suffix, + "decoder_attn_o_w_" + suffix, + "decoder_ffn_gate_up_w_" + suffix, + "decoder_ffn_down_w_" + suffix}) { + st = pack_fp8(name, transpose); + if (!st.ok_status()) return st; + } + } + return modalities::Status::ok(); +} + +modalities::Status NativeWeightPacker::pack_vision_int8() { + if (!weights_) return invalid("native weight packer is invalid"); + for (int layer = 0; layer < 27; ++layer) { + for (const char* stem : {"vision_attn_qkv_w_", "vision_attn_o_w_", + "vision_ffn_up_w_", + "vision_ffn_down_w_"}) { + const modalities::Status st = + pack_int8(std::string(stem) + std::to_string(layer)); + if (!st.ok_status()) return st; + } + } + return modalities::Status::ok(); +} + +modalities::Status NativeWeightPacker::pack_encoder_int8() { + if (!weights_) return invalid("native weight packer is invalid"); + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + for (const std::string& name : { + "encoder_attn_qkv_w_" + suffix, + "encoder_attn_o_w_" + suffix, + "encoder_ffn_gate_w_" + suffix, + "encoder_ffn_up_w_" + suffix, + "encoder_ffn_down_w_" + suffix}) { + const modalities::Status st = pack_int8(name); + if (!st.ok_status()) return st; + } + } + return modalities::Status::ok(); +} + +modalities::Status NativeWeightPacker::pack_decoder_int8() { + if (!weights_) return invalid("native weight packer is invalid"); + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + for (const std::string& name : { + "decoder_attn_qkv_w_" + suffix, + "decoder_attn_o_w_" + suffix, + "decoder_ffn_gate_w_" + suffix, + "decoder_ffn_up_w_" + suffix, + "decoder_ffn_down_w_" + suffix}) { + const modalities::Status st = pack_int8(name); + if (!st.ok_status()) return st; + } + } + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_weights.cpp b/cpp/models/pi05/src/native_weights.cpp new file mode 100644 index 00000000..5abedc13 --- /dev/null +++ b/cpp/models/pi05/src/native_weights.cpp @@ -0,0 +1,115 @@ +#include "flashrt/cpp/models/pi05/native_weights.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +using Requirement = NativeTensorRequirement; + +void add(std::vector* out, const std::string& key, + std::initializer_list shape) { + out->push_back(Requirement{key, shape}); +} + +std::vector build_requirements() { + std::vector out; + out.reserve(820); + + const std::string vision = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; + add(&out, vision + ".embeddings.patch_embedding.weight", {1152, 3, 14, 14}); + add(&out, vision + ".embeddings.patch_embedding.bias", {1152}); + add(&out, vision + ".embeddings.position_embedding.weight", {256, 1152}); + for (int layer = 0; layer < 27; ++layer) { + const std::string p = vision + ".encoder.layers." + + std::to_string(layer); + for (const char* projection : {"q_proj", "k_proj", "v_proj", + "out_proj"}) { + add(&out, p + ".self_attn." + projection + ".weight", + {1152, 1152}); + add(&out, p + ".self_attn." + projection + ".bias", {1152}); + } + add(&out, p + ".mlp.fc1.weight", {4304, 1152}); + add(&out, p + ".mlp.fc1.bias", {4304}); + add(&out, p + ".mlp.fc2.weight", {1152, 4304}); + add(&out, p + ".mlp.fc2.bias", {1152}); + for (const char* norm : {"layer_norm1", "layer_norm2"}) { + add(&out, p + "." + norm + ".weight", {1152}); + add(&out, p + "." + norm + ".bias", {1152}); + } + } + add(&out, vision + ".post_layernorm.weight", {1152}); + add(&out, vision + ".post_layernorm.bias", {1152}); + + const std::string projector = + "paligemma_with_expert.paligemma.model.multi_modal_projector.linear"; + add(&out, projector + ".weight", {2048, 1152}); + add(&out, projector + ".bias", {2048}); + + const std::string encoder = + "paligemma_with_expert.paligemma.model.language_model.layers."; + for (int layer = 0; layer < 18; ++layer) { + const std::string p = encoder + std::to_string(layer); + add(&out, p + ".self_attn.q_proj.weight", {2048, 2048}); + add(&out, p + ".self_attn.k_proj.weight", {256, 2048}); + add(&out, p + ".self_attn.v_proj.weight", {256, 2048}); + add(&out, p + ".self_attn.o_proj.weight", {2048, 2048}); + add(&out, p + ".mlp.gate_proj.weight", {16384, 2048}); + add(&out, p + ".mlp.up_proj.weight", {16384, 2048}); + add(&out, p + ".mlp.down_proj.weight", {2048, 16384}); + add(&out, p + ".input_layernorm.weight", {2048}); + add(&out, p + ".post_attention_layernorm.weight", {2048}); + } + add(&out, "paligemma_with_expert.paligemma.model.language_model.norm.weight", + {2048}); + add(&out, "paligemma_with_expert.paligemma.lm_head.weight", + {257152, 2048}); + + const std::string decoder = + "paligemma_with_expert.gemma_expert.model.layers."; + for (int layer = 0; layer < 18; ++layer) { + const std::string p = decoder + std::to_string(layer); + add(&out, p + ".self_attn.q_proj.weight", {2048, 1024}); + add(&out, p + ".self_attn.k_proj.weight", {256, 1024}); + add(&out, p + ".self_attn.v_proj.weight", {256, 1024}); + add(&out, p + ".self_attn.o_proj.weight", {1024, 2048}); + add(&out, p + ".mlp.gate_proj.weight", {4096, 1024}); + add(&out, p + ".mlp.up_proj.weight", {4096, 1024}); + add(&out, p + ".mlp.down_proj.weight", {1024, 4096}); + for (const char* norm : {"input_layernorm", "post_attention_layernorm"}) { + add(&out, p + "." + norm + ".dense.weight", {3072, 1024}); + add(&out, p + "." + norm + ".dense.bias", {3072}); + } + } + add(&out, "paligemma_with_expert.gemma_expert.model.norm.dense.weight", + {3072, 1024}); + add(&out, "paligemma_with_expert.gemma_expert.model.norm.dense.bias", + {3072}); + add(&out, "paligemma_with_expert.gemma_expert.lm_head.weight", + {257152, 1024}); + + add(&out, "action_in_proj.weight", {1024, 32}); + add(&out, "action_in_proj.bias", {1024}); + add(&out, "action_out_proj.weight", {32, 1024}); + add(&out, "action_out_proj.bias", {32}); + add(&out, "time_mlp_in.weight", {1024, 1024}); + add(&out, "time_mlp_in.bias", {1024}); + add(&out, "time_mlp_out.weight", {1024, 1024}); + add(&out, "time_mlp_out.bias", {1024}); + return out; +} + +} // namespace + +const std::vector& native_tensor_requirements() { + static const std::vector requirements = + build_requirements(); + return requirements; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_workspace.cpp b/cpp/models/pi05/src/native_workspace.cpp new file mode 100644 index 00000000..81a0004b --- /dev/null +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -0,0 +1,323 @@ +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +bool element_count(std::initializer_list shape, + std::size_t* out) { + std::size_t count = 1; + for (std::uint64_t dim : shape) { + if (!dim || dim > std::numeric_limits::max() || + count > std::numeric_limits::max() / + static_cast(dim)) { + return false; + } + count *= static_cast(dim); + } + if (out) *out = count; + return true; +} + +} // namespace + +modalities::Status NativeWorkspace::add( + const std::string& name, + std::initializer_list shape, + modalities::DType dtype) { + if (!ctx_ || name.empty() || buffers_.find(name) != buffers_.end()) { + return invalid("native workspace buffer definition is invalid"); + } + std::size_t elements = 0; + const std::size_t width = modalities::dtype_size(dtype); + if (!width || !element_count(shape, &elements) || + elements > std::numeric_limits::max() / width) { + return invalid("native workspace buffer shape is invalid"); + } + const std::size_t bytes = elements * width; + frt_buffer buffer = frt_buffer_alloc(ctx_, name.c_str(), bytes); + if (!buffer) return backend("native workspace allocation failed"); + buffers_.emplace(name, NativeWorkspaceBuffer{ + buffer, std::vector(shape), + dtype, false}); + ++allocation_count_; + allocated_bytes_ += bytes; + return modalities::Status::ok(); +} + +modalities::Status NativeWorkspace::add_alias( + const std::string& name, + const std::string& source_name, + std::initializer_list shape) { + if (name.empty() || buffers_.find(name) != buffers_.end()) { + return invalid("native workspace alias definition is invalid"); + } + const auto source = buffers_.find(source_name); + if (source == buffers_.end() || !source->second.buffer) { + return invalid("native workspace alias source was not found"); + } + std::size_t elements = 0; + const std::size_t width = modalities::dtype_size(source->second.dtype); + if (!width || !element_count(shape, &elements) || + elements > std::numeric_limits::max() / width || + elements * width != + frt_buffer_bytes(source->second.buffer)) { + return invalid("native workspace alias shape does not match source"); + } + buffers_.emplace(name, NativeWorkspaceBuffer{ + source->second.buffer, + std::vector(shape), + source->second.dtype, true}); + return modalities::Status::ok(); +} + +modalities::Status NativeWorkspace::initialize_rms_ones() { +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "native workspace initialization requires the CUDA build"); +#else + for (const char* name : {"encoder_rms_ones", "decoder_rms_ones"}) { + const NativeWorkspaceBuffer* target = find(name); + if (!target) return invalid("native RMS buffer was not allocated"); + std::vector ones( + target->shape[0], modalities::float_to_bfloat16(1.0f)); + const cudaError_t rc = cudaMemcpy( + frt_buffer_dptr(target->buffer), ones.data(), + ones.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice); + if (rc != cudaSuccess) return backend("native RMS upload failed"); + } + return modalities::Status::ok(); +#endif +} + +modalities::Status NativeWorkspace::initialize_rope() { +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "native RoPE initialization requires the CUDA build"); +#else + const int max_positions = encoder_sequence_ + chunk_size_; + rope_table_.resize(static_cast(max_positions) * 256); + for (int position = 0; position < max_positions; ++position) { + const std::size_t row = static_cast(position) * 256; + for (int i = 0; i < 128; ++i) { + const double exponent = static_cast(2 * i) / 256.0; + const double inverse_frequency = + 1.0 / std::pow(10000.0, exponent); + const double phase = + static_cast(position) * inverse_frequency; + rope_table_[row + 2 * i] = + modalities::float_to_bfloat16( + static_cast(std::cos(phase))); + rope_table_[row + 2 * i + 1] = + modalities::float_to_bfloat16( + static_cast(std::sin(phase))); + } + } + const NativeWorkspaceBuffer* encoder = find("encoder_rope_weights"); + if (!encoder) return invalid("encoder RoPE buffer was not allocated"); + const std::size_t encoder_bytes = + static_cast(encoder_sequence_) * 256 * + sizeof(std::uint16_t); + const cudaError_t rc = cudaMemcpy( + frt_buffer_dptr(encoder->buffer), rope_table_.data(), encoder_bytes, + cudaMemcpyHostToDevice); + if (rc != cudaSuccess) return backend("encoder RoPE upload failed"); + return update_decoder_rope(0); +#endif +} + +modalities::Status NativeWorkspace::update_decoder_rope(int prompt_tokens) { + if (prompt_tokens < 0 || prompt_tokens > max_prompt_tokens_ || + rope_table_.empty()) { + return invalid("Pi0.5 decoder RoPE prompt length is invalid"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "decoder RoPE update requires the CUDA build"); +#else + if (!decoder_rope_buffer_) + return invalid("decoder RoPE buffer was not allocated"); + const std::size_t start = + static_cast(encoder_vision_sequence_ + prompt_tokens) * + 256; + const std::size_t elements = + static_cast(chunk_size_) * 256; + if (start > rope_table_.size() || + elements > rope_table_.size() - start) { + return invalid("decoder RoPE slice exceeds the generated table"); + } + const cudaError_t rc = cudaMemcpy( + frt_buffer_dptr(decoder_rope_buffer_), rope_table_.data() + start, + elements * sizeof(std::uint16_t), cudaMemcpyHostToDevice); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("decoder RoPE upload failed"); +#endif +} + +modalities::Status NativeWorkspace::expand_vision_position_embedding( + const NativeDeviceWeightStore& weights) { + const NativeDeviceWeight* source = + weights.find("vision_position_embedding"); + const NativeWorkspaceBuffer* destination = + find("vision_pos_embed_expanded"); + if (!source || !destination || + source->dtype != NativeWeightDType::kBf16 || + source->shape != std::vector({256, 1152})) { + return invalid("vision position embedding source is invalid"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "position embedding expansion requires the CUDA build"); +#else + const std::size_t view_bytes = 256 * 1152 * sizeof(std::uint16_t); + if (frt_buffer_bytes(destination->buffer) != + static_cast(num_views_) * view_bytes) { + return invalid("expanded position embedding buffer size is invalid"); + } + for (int view = 0; view < num_views_; ++view) { + auto* target = static_cast( + frt_buffer_dptr(destination->buffer)) + + static_cast(view) * view_bytes; + const cudaError_t rc = cudaMemcpy( + target, frt_buffer_dptr(source->buffer), view_bytes, + cudaMemcpyDeviceToDevice); + if (rc != cudaSuccess) { + return backend("vision position embedding expansion failed"); + } + } + return modalities::Status::ok(); +#endif +} + +modalities::Status NativeWorkspace::allocate( + const NativeWorkspaceConfig& config) { + if (!ctx_ || !buffers_.empty() || config.num_views < 1 || + config.num_views > 3 || config.max_prompt_tokens <= 0 || + config.max_prompt_tokens > std::numeric_limits::max() - 768 || + config.chunk_size <= 0 || config.num_steps <= 0 || + (config.vision_pool_factor != 1 && + config.vision_pool_factor != 2 && + config.vision_pool_factor != 4)) { + return invalid("Pi0.5 native workspace configuration is invalid"); + } + const int pool_area = + config.vision_pool_factor * config.vision_pool_factor; + num_views_ = config.num_views; + max_prompt_tokens_ = config.max_prompt_tokens; + chunk_size_ = config.chunk_size; + vision_sequence_ = config.num_views * 256; + encoder_vision_sequence_ = vision_sequence_ / pool_area; + encoder_sequence_ = + encoder_vision_sequence_ + config.max_prompt_tokens; + const std::uint64_t nv = static_cast(config.num_views); + const std::uint64_t vs = static_cast(vision_sequence_); + const std::uint64_t vs_enc = + static_cast(encoder_vision_sequence_); + const std::uint64_t es = static_cast(encoder_sequence_); + const std::uint64_t ds = static_cast(config.chunk_size); + const std::uint64_t steps = static_cast(config.num_steps); + modalities::Status st; +#define FRT_ADD(...) \ + do { \ + st = add(__VA_ARGS__); \ + if (!st.ok_status()) return st; \ + } while (false) + FRT_ADD("observation_images_normalized", {nv, 224, 224, 3}, + modalities::DType::kBFloat16); + FRT_ADD("vision_x", {vs, 1152}, modalities::DType::kBFloat16); + FRT_ADD("vision_x_norm", {vs, 1152}, modalities::DType::kBFloat16); + if (config.vision_pool_factor == 1) { + st = add_alias("vision_x_pooled", "vision_x", {vs_enc, 1152}); + if (!st.ok_status()) return st; + } else { + FRT_ADD("vision_x_pooled", {vs_enc, 1152}, + modalities::DType::kBFloat16); + } + FRT_ADD("vision_QKV", {vs, 3456}, modalities::DType::kBFloat16); + FRT_ADD("vision_hidden", {vs, 4304}, modalities::DType::kBFloat16); + FRT_ADD("vision_pos_embed_expanded", {vs, 1152}, + modalities::DType::kBFloat16); + FRT_ADD("vision_patches", {vs, 588}, modalities::DType::kBFloat16); + + FRT_ADD("encoder_rope_weights", {es, 256}, + modalities::DType::kBFloat16); + FRT_ADD("prompt_embedding", + {static_cast(max_prompt_tokens_), 2048}, + modalities::DType::kBFloat16); + FRT_ADD("encoder_x", {es, 2048}, modalities::DType::kBFloat16); + FRT_ADD("encoder_x_norm", {es, 2048}, modalities::DType::kBFloat16); + FRT_ADD("encoder_QKV", {es, 2560}, modalities::DType::kBFloat16); + FRT_ADD("encoder_hidden", {es, 16384}, modalities::DType::kBFloat16); + FRT_ADD("encoder_gate_merged", {es, 32768}, + modalities::DType::kBFloat16); + FRT_ADD("encoder_gate_buf", {es, 16384}, + modalities::DType::kBFloat16); + FRT_ADD("encoder_rms_ones", {2048}, modalities::DType::kBFloat16); + + FRT_ADD("decoder_rope_weights", {ds, 256}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_x", {ds, 1024}, modalities::DType::kBFloat16); + FRT_ADD("decoder_action_buf", {ds, 32}, modalities::DType::kBFloat16); + FRT_ADD("decoder_time_emb", {steps, ds, 1024}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_style_attn", {steps, 18, ds, 3072}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_style_ffn", {steps, 18, ds, 3072}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_style_final", {steps, ds, 3072}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_QKV", {ds, 2560}, modalities::DType::kBFloat16); + FRT_ADD("decoder_hidden", {ds, 4096}, modalities::DType::kBFloat16); + FRT_ADD("decoder_gate_merged", {ds, 8192}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_gate_buf", {ds, 4096}, + modalities::DType::kBFloat16); + FRT_ADD("diffusion_noise", {ds, 32}, modalities::DType::kBFloat16); + FRT_ADD("rtc_prev_action_chunk", {ds, 32}, + modalities::DType::kBFloat16); + FRT_ADD("rtc_prefix_weights", {ds}, modalities::DType::kFloat32); + FRT_ADD("rtc_guidance_weight", {1}, modalities::DType::kFloat32); + FRT_ADD("x_normed_buf", {ds, 1024}, modalities::DType::kBFloat16); + FRT_ADD("gate_buf", {ds, 1024}, modalities::DType::kBFloat16); + FRT_ADD("decoder_rms_ones", {1024}, modalities::DType::kBFloat16); +#undef FRT_ADD + const NativeWorkspaceBuffer* decoder = find("decoder_rope_weights"); + if (!decoder) return invalid("decoder RoPE buffer was not allocated"); + decoder_rope_buffer_ = decoder->buffer; + st = initialize_rms_ones(); + if (!st.ok_status()) return st; + return initialize_rope(); +} + +const NativeWorkspaceBuffer* NativeWorkspace::find( + const std::string& name) const { + const auto it = buffers_.find(name); + return it == buffers_.end() ? nullptr : &it->second; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/prompt_embed.cpp b/cpp/models/pi05/src/prompt_embed.cpp new file mode 100644 index 00000000..50a9c046 --- /dev/null +++ b/cpp/models/pi05/src/prompt_embed.cpp @@ -0,0 +1,171 @@ +#include "flashrt/cpp/models/pi05/prompt_embed.h" + +#include "flashrt/cpp/models/pi05/prompt_format.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status validate_output_capacity( + const PromptEmbeddingSpec& spec, + const modalities::TensorView& output) { + if (!spec.vocab_size || !spec.hidden_dim || !spec.max_tokens) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "invalid prompt embedding dimensions"); + } + if (!output.data) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt_embedding has null data"); + } + if (output.place != modalities::MemoryPlace::kHost && + output.place != modalities::MemoryPlace::kHostPinned && + output.place != modalities::MemoryPlace::kDevice) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "prompt_embedding memory place is unsupported"); + } + if (output.layout != modalities::Layout::kFlat || + output.shape.rank != 2 || + output.shape.dims[0] != spec.max_tokens || + output.shape.dims[1] != spec.hidden_dim) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "prompt_embedding shape mismatch"); + } + const std::uint64_t need = + spec.max_tokens * spec.hidden_dim * modalities::dtype_size(output.dtype); + if (output.bytes < need) { + return modalities::Status::error( + modalities::StatusCode::kInsufficientStorage, + "prompt_embedding storage is too small"); + } + return modalities::Status::ok(); +} + +modalities::Status zero_prompt_output(const modalities::TensorView& output, + void* stream) { + if (output.place == modalities::MemoryPlace::kHost || + output.place == modalities::MemoryPlace::kHostPinned) { + std::memset(output.data, 0, static_cast(output.bytes)); + return modalities::Status::ok(); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)stream; + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "device prompt zeroing requires the CUDA build"); +#else + cudaStream_t cuda_stream = reinterpret_cast(stream); + cudaError_t rc = cudaMemsetAsync(output.data, 0, output.bytes, + cuda_stream); + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + if (rc != cudaSuccess) { + return modalities::Status::error( + modalities::StatusCode::kBackend, + std::string("cuda prompt zeroing failed: ") + + cudaGetErrorString(rc)); + } + return modalities::Status::ok(); +#endif +} + +} // namespace + +modalities::Status embed_prompt( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len, + void* stream, + modalities::TextEmbeddingStaging* staging, + std::string* formatted_workspace) { + if (!token_ids || !prompt_len) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt embedding outputs are null"); + } + token_ids->clear(); + *prompt_len = 0; + auto st = validate_output_capacity(spec, output); + if (!st.ok_status()) return st; + if (!tokenizer.loaded()) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "SentencePiece model is not loaded"); + } + + modalities::SentencePieceEncodeOptions options; + options.add_bos = true; + options.max_tokens = spec.max_tokens; + if (state) { + std::string local; + std::string* formatted = formatted_workspace ? formatted_workspace + : &local; + format_state_prompt_into(prompt, state, n_state, formatted); + st = tokenizer.encode(*formatted, options, token_ids); + } else { + st = tokenizer.encode(prompt, options, token_ids); + if (st.ok_status() && spec.no_state_suffix_token_id >= 0) { + token_ids->push_back(spec.no_state_suffix_token_id); + } + } + if (!st.ok_status()) return st; + if (token_ids->size() > spec.max_tokens) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "prompt token count exceeds max_tokens"); + } + + if (spec.zero_pad_output) { + st = zero_prompt_output(output, stream); + if (!st.ok_status()) return st; + } + modalities::TensorView prefix = output; + prefix.shape = modalities::Shape{static_cast( + token_ids->size()), + spec.hidden_dim}; + prefix.bytes = static_cast(token_ids->size()) * + spec.hidden_dim * modalities::dtype_size(output.dtype); + + modalities::EmbeddingGatherSpec gather{spec.vocab_size, spec.hidden_dim, + spec.scale}; + st = modalities::gather_token_embeddings( + gather, token_ids->data(), token_ids->size(), embedding_table, prefix, + stream, staging); + if (!st.ok_status()) return st; + *prompt_len = static_cast(token_ids->size()); + return modalities::Status::ok(); +} + +modalities::Status embed_prompt_cpu( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len) { + return embed_prompt(tokenizer, spec, prompt, state, n_state, + embedding_table, output, token_ids, prompt_len); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/prompt_format.cpp b/cpp/models/pi05/src/prompt_format.cpp new file mode 100644 index 00000000..5e1c36a1 --- /dev/null +++ b/cpp/models/pi05/src/prompt_format.cpp @@ -0,0 +1,95 @@ +#include "flashrt/cpp/models/pi05/prompt_format.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +std::vector make_openpi_bins() { + std::vector bins; + bins.reserve(256); + for (int i = 0; i < 256; ++i) { + bins.push_back(-1.0f + static_cast(i) * (2.0f / 256.0f)); + } + return bins; +} + +bool ascii_space(char c) { + return std::isspace(static_cast(c)) != 0; +} + +} // namespace + +std::vector discretize_state_prompt_bins( + const float* state, std::uint64_t n) { + static const std::vector bins = make_openpi_bins(); + std::vector out; + out.reserve(static_cast(n)); + for (std::uint64_t i = 0; i < n; ++i) { + const auto it = std::upper_bound(bins.begin(), bins.end(), state[i]); + out.push_back(static_cast(it - bins.begin()) - 1); + } + return out; +} + +std::string clean_task_prompt(const std::string& prompt) { + auto begin = prompt.begin(); + auto end = prompt.end(); + while (begin != end && ascii_space(*begin)) ++begin; + while (begin != end && ascii_space(*(end - 1))) --end; + + std::string cleaned(begin, end); + for (char& c : cleaned) { + if (c == '_' || c == '\n') c = ' '; + } + return cleaned; +} + +std::string format_state_prompt(const std::string& prompt, + const float* state, + std::uint64_t n_state) { + std::string out; + out.reserve(prompt.size() + static_cast(n_state) * 5 + 32); + format_state_prompt_into(prompt, state, n_state, &out); + return out; +} + +void format_state_prompt_into(const std::string& prompt, + const float* state, + std::uint64_t n_state, + std::string* out) { + if (!out) return; + out->clear(); + auto begin = prompt.begin(); + auto end = prompt.end(); + while (begin != end && ascii_space(*begin)) ++begin; + while (begin != end && ascii_space(*(end - 1))) --end; + + if (state) out->append("Task: "); + for (auto it = begin; it != end; ++it) { + out->push_back(*it == '_' || *it == '\n' ? ' ' : *it); + } + if (!state) return; + + static const std::vector bins = make_openpi_bins(); + out->append(", State: "); + char number[24]; + for (std::uint64_t i = 0; i < n_state; ++i) { + if (i) out->push_back(' '); + const auto bin = static_cast( + std::upper_bound(bins.begin(), bins.end(), state[i]) - + bins.begin()) - + 1; + const auto result = std::to_chars(number, number + sizeof(number), bin); + out->append(number, result.ptr); + } + out->append(";\nAction: "); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index 37b48b2c..ecdbae61 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -1,6 +1,8 @@ #include "flashrt/cpp/models/pi05/runtime.h" +#include #include +#include namespace flashrt { namespace models { @@ -70,6 +72,8 @@ Runtime::Runtime(const frt_runtime_export_v1* exp, RuntimeConfig config) } Runtime::~Runtime() { + modalities::text_embedding_staging_destroy(&prompt_embedding_staging_); + modalities::action_staging_destroy(&action_staging_); modalities::vision_staging_destroy(&staging_); release_export(); } @@ -166,18 +170,85 @@ modalities::Status Runtime::bind() { staging = &staging_; } + modalities::ActionStaging* action_staging = nullptr; + if (action.place == modalities::MemoryPlace::kDevice) { + modalities::ActionPostprocessSpec action_spec = action_postprocess_spec( + config_.action_mean, config_.action_stddev, config_.chunk, + config_.model_action_dim, config_.robot_action_dim); + modalities::Status st = modalities::action_staging_create( + &action_staging_, + modalities::required_action_output_bytes(action_spec, + action.dtype)); + if (!st.ok_status()) return st; + action_staging = &action_staging_; + } + io_ = RuntimeIo(config_.num_views, image, action, config_.action_mean, config_.action_stddev, find_native_stream(exp_, stream_id_), config_.chunk, config_.model_action_dim, - config_.robot_action_dim, config_.image_dtype, staging); - return modalities::Status::ok(); + config_.robot_action_dim, config_.image_dtype, staging, + action_staging, config_.strict_rgb8); + return bind_prompt_staging(); } int Runtime::set_prompt(const char* text) { - /* The adopted-export path assumes prompt/token embedding was prepared by - * the producer before capture/export. A native Pi0.5 producer will replace - * this with tokenizer + prompt-region binding without changing Nexus. */ - return (text == nullptr || text[0] == '\0') ? 0 : -1; + return set_prompt_state(text, nullptr, 0); +} + +int Runtime::set_prompt_state(const char* text, const float* state, + std::uint64_t n_state) { + if (!prompt_staging_enabled_) { + return (text == nullptr || text[0] == '\0') ? 0 : -1; + } + if (!text) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt text is null"); + return -1; + } + const std::size_t text_bytes = std::strlen(text); + if (text_bytes > max_task_prompt_bytes_) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "prompt text exceeds the configured hot-path capacity"); + return -1; + } + task_prompt_workspace_.assign(text, text_bytes); + const float* state_for_prompt = state; + if (state && state_normalization_enabled()) { + if (n_state != config_.state_q01.size()) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "state dimension does not match norm stats"); + return -1; + } + for (std::uint64_t i = 0; i < n_state; ++i) { + const float lo = config_.state_q01[i]; + const float hi = config_.state_q99[i]; + normalized_state_[i] = + ((state[i] - lo) / (hi - lo + 1e-6f)) * 2.0f - 1.0f; + } + state_for_prompt = normalized_state_.data(); + } + prompt_status_ = embed_prompt( + prompt_tokenizer_, prompt_spec_, task_prompt_workspace_, + state_for_prompt, n_state, + prompt_embedding_table_, prompt_embedding_output_, &prompt_token_ids_, + ¤t_prompt_len_, find_native_stream(exp_, stream_id_), + prompt_embedding_output_.place == modalities::MemoryPlace::kDevice + ? &prompt_embedding_staging_ + : nullptr, + &formatted_prompt_workspace_); + if (prompt_status_.ok_status() && config_.prompt_length_update_fn) { + const int rc = config_.prompt_length_update_fn( + config_.prompt_length_update_user, current_prompt_len_); + if (rc != 0) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kBackend, + "prompt length device update failed"); + } + } + return prompt_status_.ok_status() ? 0 : -1; } modalities::Status Runtime::prepare_vision( @@ -203,6 +274,70 @@ int Runtime::default_replay(frt_graph graph, frt_shape_key key, return frt_graph_replay(graph, key, stream_id); } +modalities::Status Runtime::bind_prompt_staging() { + const bool any = + !config_.prompt_tokenizer_model_path.empty() || + config_.prompt_embedding_table.data || + config_.prompt_embedding_output.data || + config_.prompt_vocab_size || config_.prompt_hidden_dim || + config_.prompt_max_tokens; + if (!any) { + prompt_status_ = modalities::Status::ok(); + return modalities::Status::ok(); + } + if (config_.prompt_tokenizer_model_path.empty() || + !config_.prompt_embedding_table.data || + !config_.prompt_embedding_output.data || + !config_.prompt_vocab_size || !config_.prompt_hidden_dim || + !config_.prompt_max_tokens) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "incomplete Pi05 prompt staging config"); + } + prompt_status_ = + prompt_tokenizer_.load_model(config_.prompt_tokenizer_model_path); + if (!prompt_status_.ok_status()) return prompt_status_; + + prompt_embedding_table_ = config_.prompt_embedding_table; + prompt_embedding_output_ = config_.prompt_embedding_output; + prompt_spec_.vocab_size = config_.prompt_vocab_size; + prompt_spec_.hidden_dim = config_.prompt_hidden_dim; + prompt_spec_.max_tokens = config_.prompt_max_tokens; + prompt_spec_.scale = config_.prompt_embedding_scale > 0.0f + ? config_.prompt_embedding_scale + : std::sqrt(static_cast( + config_.prompt_hidden_dim)); + if (config_.state_q01.size() > + (std::numeric_limits::max() - 32ull) / 5ull) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "state workspace capacity overflows size_t"); + } + const std::size_t state_bytes = config_.state_q01.size() * 5ull + 32ull; + if (config_.prompt_max_tokens > + (std::numeric_limits::max() - state_bytes) / 8ull) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt workspace capacity overflows size_t"); + } + const std::size_t max_prompt_bytes = + static_cast(config_.prompt_max_tokens * 8ull); + max_task_prompt_bytes_ = max_prompt_bytes; + task_prompt_workspace_.reserve(max_prompt_bytes); + formatted_prompt_workspace_.reserve(max_prompt_bytes + state_bytes); + prompt_token_ids_.reserve(static_cast( + config_.prompt_max_tokens + 1ull)); + normalized_state_.resize(config_.state_q01.size()); + prompt_tokenizer_.reserve(config_.prompt_max_tokens); + if (prompt_embedding_output_.place == modalities::MemoryPlace::kDevice) { + prompt_status_ = modalities::text_embedding_staging_create( + &prompt_embedding_staging_, config_.prompt_max_tokens); + if (!prompt_status_.ok_status()) return prompt_status_; + } + prompt_staging_enabled_ = true; + return modalities::Status::ok(); +} + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/models/pi05/src/spec.cpp b/cpp/models/pi05/src/spec.cpp index 4d3efea2..f6eb0067 100644 --- a/cpp/models/pi05/src/spec.cpp +++ b/cpp/models/pi05/src/spec.cpp @@ -16,8 +16,8 @@ modalities::VisionPreprocessSpec vision_preprocess_spec(int num_views) { spec.target_height = kImageSize; spec.output_dtype = modalities::DType::kBFloat16; spec.output_layout = modalities::Layout::kNHWC; - spec.normalize.mode = modalities::NormalizeMode::kScaleShift; - spec.normalize.scale = 1.0f / 127.5f; + spec.normalize.mode = modalities::NormalizeMode::kDivideShift; + spec.normalize.divisor = 127.5f; spec.normalize.shift = -1.0f; spec.require_exact_views = true; return spec; diff --git a/cpp/tests/data/pi05_native_v2_schema.records b/cpp/tests/data/pi05_native_v2_schema.records new file mode 100644 index 00000000..de6fb986 --- /dev/null +++ b/cpp/tests/data/pi05_native_v2_schema.records @@ -0,0 +1,8 @@ +region:0:rollout_boundary:0:640:3 +port:0:prompt:2:0:0:0:1:1:-1:-1:0:0 +port:1:state:3:1:0:0:1:1:8:-1:0:0 +port:2:images:1:3:2:0:1:1:2,224,224,3:0:0:602112 +port:3:noise:0:3:0:0:0:0:10,32:1:0:640 +port:4:actions:4:1:0:1:1:0:10,7:-1:0:280 +port:5:actions_raw:0:3:0:1:0:0:10,32:1:0:640 +stage:0:0: diff --git a/cpp/tests/gate_pi05_hot_allocator.py b/cpp/tests/gate_pi05_hot_allocator.py new file mode 100644 index 00000000..8d809049 --- /dev/null +++ b/cpp/tests/gate_pi05_hot_allocator.py @@ -0,0 +1,86 @@ +"""Reject CUDA allocation APIs in a replay-only Pi0.5 Nsight trace.""" + +from __future__ import annotations + +import argparse +from collections import Counter +import csv +from pathlib import Path + + +FORBIDDEN = ( + "cudaMalloc", + "cudaFree", + "cudaHostAlloc", + "cudaHostRegister", + "cudaHostUnregister", + "cudaMemPoolCreate", + "cudaMemPoolDestroy", + "cuMemAlloc", + "cuMemFree", + "cuMemHostAlloc", + "cuMemHostRegister", + "cuMemHostUnregister", + "cuMemCreate", + "cuMemMap", + "cuMemUnmap", + "cuMemAddressReserve", + "cuMemAddressFree", +) + + +def _read_rows(path: Path) -> list[dict[str, str]]: + lines = path.read_text(encoding="utf-8").splitlines(keepends=True) + try: + header = next( + index for index, line in enumerate(lines) + if line.startswith("Start (ns),") + ) + except StopIteration as exc: + raise ValueError(f"{path}: cuda_api_trace CSV header is missing") from exc + rows = list(csv.DictReader(lines[header:])) + if not rows: + raise ValueError(f"{path}: CUDA API trace is empty") + return rows + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--trace", type=Path, required=True) + parser.add_argument("--expected-replays", type=int, default=1000) + args = parser.parse_args() + if args.expected_replays <= 0: + parser.error("--expected-replays must be positive") + + rows = _read_rows(args.trace) + failed = [row for row in rows if int(row["Result"]) != 0] + if failed: + raise AssertionError(f"CUDA API calls failed: {failed[:4]}") + names = Counter(row["Name"] for row in rows) + graph_launches = sum( + count for name, count in names.items() + if name.startswith("cudaGraphLaunch") or name.startswith("cuGraphLaunch") + ) + if graph_launches != args.expected_replays: + raise AssertionError( + f"graph replay count differs: actual={graph_launches} " + f"expected={args.expected_replays}" + ) + allocator_calls = { + name: count for name, count in names.items() + if any(marker in name for marker in FORBIDDEN) + } + if allocator_calls: + raise AssertionError(f"hot path called CUDA allocators: {allocator_calls}") + print({ + "ok": True, + "graph_replays": graph_launches, + "allocator_calls": 0, + "cuda_api_calls": sum(names.values()), + }) + print("PASS Pi0.5 hot replay performed no CUDA allocation calls") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/gate_pi05_kernel_sequence.py b/cpp/tests/gate_pi05_kernel_sequence.py new file mode 100644 index 00000000..1c8ae859 --- /dev/null +++ b/cpp/tests/gate_pi05_kernel_sequence.py @@ -0,0 +1,119 @@ +"""Compare replay-only Pi0.5 native and Python Nsight kernel traces.""" + +from __future__ import annotations + +import argparse +from collections import Counter +import csv +from pathlib import Path + + +def _read_names(path: Path) -> list[str]: + lines = path.read_text(encoding="utf-8").splitlines(keepends=True) + try: + header = next( + index for index, line in enumerate(lines) + if line.startswith("Start (ns),") + ) + except StopIteration as exc: + raise ValueError(f"{path}: cuda_gpu_trace CSV header is missing") from exc + names = [row["Name"] for row in csv.DictReader(lines[header:])] + if not names: + raise ValueError(f"{path}: CUDA trace is empty") + return names + + +def _classify(name: str) -> tuple[str | None, str | None]: + # These two nodes are an implementation detail of the selected GEMM + # algorithm. A split-K algorithm substitutes a reduction for workspace + # initialization without changing the surrounding logical GEMM sequence. + if name == "[CUDA memset]": + return None, "gemm_workspace_init" + if "cublasLt::splitKreduce_kernel" in name: + return None, "gemm_splitk_reduce" + + patterns = ( + ("copy", "[CUDA memcpy"), + ("attention_fill", "FillFunctor"), + ("attention_fill", "fill_negative_infinity"), + ("gemm", "cutlass::Kernel2"), + ("gemm", "gemmSN_NN_kernel"), + ("attention_combine", "flash_fwd_splitkv_combine_kernel"), + ("attention_split", "flash_fwd_splitkv_kernel"), + ("attention", "flash_fwd_kernel"), + ("ada_norm", "ada_rms_norm_style_kernel"), + ("gate_residual", "gate_mul_res_kernel"), + ("gate_silu", "gate_silu_mul_kernel"), + ("qkv_devpos", "qkv_split_rope_devpos_kernel"), + ("qkv_rope", "qkv_split_rope_kernel"), + ("qkv", "qkv_split_kernel"), + ("bias", "bias_res_kernel"), + ("bias", "add_bias"), + ("layer_norm", "layer_norm_kernel"), + ("rms_norm", "rms_norm_kernel"), + ("gelu", "gelu_kernel"), + ("residual", "res_add_kernel"), + ("patch", "patch_im2col_kernel"), + ) + for logical, marker in patterns: + if marker in name: + return logical, None + raise ValueError(f"kernel is not in the explicit Pi0.5 whitelist: {name}") + + +def _normalize(names: list[str]) -> tuple[list[str], Counter[str]]: + result = [] + ignored: Counter[str] = Counter() + for name in names: + logical, ignored_kind = _classify(name) + if ignored_kind is not None: + ignored[ignored_kind] += 1 + else: + result.append(logical) + return result, ignored + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--native", type=Path, required=True) + parser.add_argument("--python", type=Path, required=True) + args = parser.parse_args() + + native_names = _read_names(args.native) + python_names = _read_names(args.python) + if len(native_names) != len(python_names): + raise AssertionError( + f"raw event count differs: native={len(native_names)} " + f"python={len(python_names)}" + ) + native, native_ignored = _normalize(native_names) + python, python_ignored = _normalize(python_names) + if sum(native_ignored.values()) != sum(python_ignored.values()): + raise AssertionError( + "allowlisted GEMM helper count differs: " + f"native={dict(native_ignored)} python={dict(python_ignored)}" + ) + if native != python: + mismatch = next( + (index for index, pair in enumerate(zip(native, python)) + if pair[0] != pair[1]), + min(len(native), len(python)), + ) + raise AssertionError( + f"logical kernel sequence differs at {mismatch}: " + f"native={native[mismatch:mismatch + 8]} " + f"python={python[mismatch:mismatch + 8]}" + ) + print({ + "ok": True, + "raw_events": len(native_names), + "logical_events": len(native), + "native_gemm_helpers": dict(native_ignored), + "python_gemm_helpers": dict(python_ignored), + }) + print("PASS Pi0.5 native/Python logical kernel sequences are identical") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/gate_pi05_model_runtime_export.py b/cpp/tests/gate_pi05_model_runtime_export.py index 7af2a09c..d0fd8335 100644 --- a/cpp/tests/gate_pi05_model_runtime_export.py +++ b/cpp/tests/gate_pi05_model_runtime_export.py @@ -2,10 +2,10 @@ Run inside the CUDA container from the repo root: - PYTHONPATH=.:./exec/build-container:./runtime/build-container \ + FLASHRT_BUILD_DIR=cpp/build-sm120-debug \ python cpp/tests/gate_pi05_model_runtime_export.py \ --checkpoint "${PI05_CHECKPOINT:-/path/to/pi05_libero_pytorch}" --fp8 \ - --lib cpp/build-container/libflashrt_cpp_pi05_c.so + --lib cpp/build-sm120-debug/libflashrt_cpp_pi05_c.so The gate compares three surfaces: 1. Python frontend staging/replay/postprocess. @@ -19,6 +19,7 @@ import argparse import ctypes +import os import statistics import sys import time @@ -34,6 +35,13 @@ p = str(ROOT / rel) if rel else str(ROOT) if p not in sys.path: sys.path.insert(0, p) +configured_build = os.environ.get("FLASHRT_BUILD_DIR") +if configured_build: + for subdir in ("exec", "runtime"): + path = str(Path(configured_build).resolve() / subdir) + if path in sys.path: + sys.path.remove(path) + sys.path.insert(0, path) import flash_rt # noqa: E402 from flash_rt.core.utils.actions import LIBERO_ACTION_DIM, unnormalize_actions # noqa: E402 @@ -200,6 +208,16 @@ def _cos(a: np.ndarray, b: np.ndarray, dtype: torch.dtype) -> float: return float(np.dot(af, bf) / (na * nb)) +def _cos_f32(a: np.ndarray, b: np.ndarray) -> float: + af = np.asarray(a, dtype=np.float64).reshape(-1) + bf = np.asarray(b, dtype=np.float64).reshape(-1) + na = float(np.linalg.norm(af)) + nb = float(np.linalg.norm(bf)) + if na == 0.0 or nb == 0.0: + return float("nan") + return float(np.dot(af, bf) / (na * nb)) + + def _make_images(num_views: int, seed: int) -> list[np.ndarray]: rng = np.random.default_rng(seed) return [ @@ -331,6 +349,13 @@ def _python_replay(pipe, pl, obs, start_noise: np.ndarray) -> np.ndarray: return _read(pl.input_noise_buf) +def _python_replay_staged(pl, start_noise: np.ndarray) -> np.ndarray: + """Replay without touching the already-staged image buffer.""" + _upload_bytes(pl.input_noise_buf, start_noise) + pl.forward() + return _read(pl.input_noise_buf) + + def _python_replay_actions(pipe, pl, obs, start_noise: np.ndarray, dtype: torch.dtype) -> np.ndarray: raw = _python_replay(pipe, pl, obs, start_noise) @@ -395,6 +420,9 @@ def main() -> None: ap.add_argument("--lib", default=str( ROOT / "cpp/build-container/libflashrt_cpp_pi05_c.so")) args = ap.parse_args() + if configured_build and Path(args.lib).resolve().parent != Path( + configured_build).resolve(): + ap.error("--lib must come from FLASHRT_BUILD_DIR") images = _make_images(args.num_views, args.seed) obs = _make_obs(images) @@ -470,6 +498,28 @@ def main() -> None: assert m_split.n_stages == 2, m_split.n_stages assert m_rtc.n_ports == 5, m_rtc.n_ports assert m_rtc.n_stages == 2, m_rtc.n_stages + for runtime in (mr_full, mr_split, mr_rtc): + action_port = next( + port for port in runtime.ports() + if port["name"] == "actions" + ) + assert action_port["dtype"] == 1, action_port + assert action_port["update"] == 1, action_port + assert action_port["buffer"] == 0, action_port + assert action_port["bytes"] == ( + int(pl.chunk_size) * LIBERO_ACTION_DIM * 4 + ), action_port + rtc_raw_port = next( + port for port in mr_rtc.ports() + if port["name"] == "actions_raw" + ) + rtc_noise_port = next( + port for port in mr_rtc.ports() + if port["name"] == "noise" + ) + assert rtc_raw_port["dtype"] == rtc_noise_port["dtype"], rtc_raw_port + assert rtc_raw_port["update"] == 0, rtc_raw_port + assert rtc_raw_port["buffer"] != 0, rtc_raw_port assert mr_full.fingerprint != mr_split.fingerprint assert mr_rtc.fingerprint not in ( mr_full.fingerprint, mr_split.fingerprint) @@ -484,6 +534,18 @@ def main() -> None: _raw_to_float(cxx_img, torch_dtype)))) start_noise = _seed_noise(pipe, pl, args.seed + 1009, torch_dtype) + + # Mechanism lane: hold the exact image/noise bytes fixed and compare + # the Python graph entry with the model-runtime graph entry. Numerical + # differences from the two image preprocessors do not belong here. + _upload_bytes(pl.input_images_buf, cxx_img) + mechanism_py_raw = _python_replay_staged(pl, start_noise) + _upload_bytes(pl.input_images_buf, cxx_img) + _upload_bytes(pl.input_noise_buf, start_noise) + mechanism_full_actions = _model_step_get_actions( + m_full, int(pl.chunk_size)) + mechanism_full_raw = _read(pl.input_noise_buf) + py_raw = _python_replay(pipe, pl, obs, start_noise) _upload_bytes(pl.input_noise_buf, start_noise) @@ -517,21 +579,41 @@ def main() -> None: py_actions = unnormalize_actions(py_raw_f, pipe.norm_stats)[ :, :LIBERO_ACTION_DIM].astype(np.float32) + mechanism_raw_exact = bool(np.array_equal( + mechanism_py_raw, mechanism_full_raw)) + mechanism_raw_max = float(np.max(np.abs( + _raw_to_float(mechanism_py_raw, torch_dtype) - + _raw_to_float(mechanism_full_raw, torch_dtype)))) + mechanism_py_actions = unnormalize_actions( + _raw_to_float(mechanism_py_raw, torch_dtype).reshape( + int(pl.chunk_size), 32), + pipe.norm_stats, + )[:, :LIBERO_ACTION_DIM].astype(np.float32) + mechanism_action_max = float(np.max(np.abs( + mechanism_py_actions - mechanism_full_actions))) + mechanism_action_close = bool(np.allclose( + mechanism_py_actions, mechanism_full_actions, + rtol=1e-6, atol=1e-6, + )) + raw_exact = bool(np.array_equal(py_raw, full_raw)) raw_cos = _cos(py_raw, full_raw, torch_dtype) raw_max = float(np.max(np.abs( _raw_to_float(py_raw, torch_dtype) - _raw_to_float(full_raw, torch_dtype)))) act_max = float(np.max(np.abs(py_actions - full_actions))) - act_ok = bool(np.allclose(py_actions, full_actions, rtol=1e-4, atol=1e-3)) + act_cos = _cos_f32(py_actions, full_actions) + act_close = bool(np.allclose( + py_actions, full_actions, rtol=1e-4, atol=1e-3 + )) + act_ok = act_cos >= 0.999 and act_close split_raw_exact = bool(np.array_equal(full_raw, split_raw)) split_raw_cos = _cos(full_raw, split_raw, torch_dtype) split_raw_max = float(np.max(np.abs( _raw_to_float(full_raw, torch_dtype) - _raw_to_float(split_raw, torch_dtype)))) + split_act_exact = bool(np.array_equal(full_actions, split_actions)) split_act_max = float(np.max(np.abs(full_actions - split_actions))) - split_act_ok = bool(np.allclose( - full_actions, split_actions, rtol=1e-4, atol=1e-3)) print("\n===== REAL PI0.5 MODEL-RUNTIME EXPORT GATE =====") print(f"full fingerprint : 0x{mr_full.fingerprint:016x}") @@ -541,20 +623,49 @@ def main() -> None: print(f"split runtime : ports={m_split.n_ports} stages={m_split.n_stages}") print(f"rtc runtime : ports={m_rtc.n_ports} stages={m_rtc.n_stages} prefix={prefix_len}") print(f"image buffer exact : {img_exact} cos={img_cos:.8f} max_abs={img_max:.6g}") + print( + "same-bytes mechanism raw: " + f"exact={mechanism_raw_exact} max_abs={mechanism_raw_max:.6g}" + ) + print( + "same-bytes action stage : " + f"allclose={mechanism_action_close} " + f"max_abs={mechanism_action_max:.6g}" + ) print(f"py vs full raw exact : {raw_exact} cos={raw_cos:.8f} max_abs={raw_max:.6g}") - print(f"py vs full action : {act_ok} max_abs={act_max:.6g}") + print( + f"py vs full action : accepted={act_ok} " + f"cos={act_cos:.8f} allclose={act_close} max_abs={act_max:.6g}" + ) print(f"full vs split raw exact: {split_raw_exact} cos={split_raw_cos:.8f} max_abs={split_raw_max:.6g}") - print(f"full vs split action : {split_act_ok} max_abs={split_act_max:.6g}") + print( + f"full vs split action : exact={split_act_exact} " + f"max_abs={split_act_max:.6g}" + ) print(f"rtc prefix exact : {rtc_prefix_exact} max_abs={rtc_prefix_max:.6g}") assert img_cos >= 0.999, f"image preprocess cosine too low: {img_cos}" - assert raw_cos >= 0.999, f"raw replay cosine too low: {raw_cos}" - assert act_ok, f"robot actions differ: max_abs={act_max}" + assert mechanism_raw_exact, ( + "same image/noise bytes must replay bit-exactly through Python " + f"and model-runtime entries; max_abs={mechanism_raw_max:.6g}" + ) + assert mechanism_action_close, ( + "logical action staging differs for identical raw bytes: " + f"max_abs={mechanism_action_max:.6g}" + ) + assert raw_exact, ( + "Python and model-runtime replay must be bit-exact after image " + f"staging; cos={raw_cos:.8f} max_abs={raw_max:.6g}" + ) + assert act_ok, ( + f"robot actions differ: cos={act_cos:.8f} max_abs={act_max}" + ) assert split_raw_exact, ( "split replay must be bit-exact against full replay; " f"cos={split_raw_cos:.8f} max_abs={split_raw_max:.6g}") - assert split_act_ok, ( - f"split robot actions differ: max_abs={split_act_max}") + assert split_act_exact, ( + f"split robot actions are not bit-exact: max_abs={split_act_max}" + ) assert rtc_prefix_exact, ( "RTC-prefix action graph did not preserve prev_action_chunk " f"prefix; max_abs={rtc_prefix_max:.6g}") diff --git a/cpp/tests/gate_pi05_native_diffusion.py b/cpp/tests/gate_pi05_native_diffusion.py new file mode 100644 index 00000000..ebd4be25 --- /dev/null +++ b/cpp/tests/gate_pi05_native_diffusion.py @@ -0,0 +1,273 @@ +#!/usr/bin/env python3 +import argparse +import gc +import math +import pathlib +import subprocess +import tempfile + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors import safe_open + + +CHUNK = 10 +PREFIX = 712 + + +def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: + output, inputs = weight.shape + head_dim = output // heads + return ( + weight.reshape(heads, head_dim, inputs) + .reshape(heads, 2, head_dim // 2, inputs) + .permute(0, 2, 1, 3) + .reshape(output, inputs) + ) + + +def ada_rms(values: torch.Tensor, style: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + source = values.float() + normalized = source * torch.rsqrt(source.square().mean(-1, keepdim=True) + 1e-6) + scale, shift, gate = style.float().chunk(3, dim=-1) + output = (normalized * (1.0 + scale) + shift).to(torch.bfloat16) + return output, gate.to(torch.bfloat16) + + +def rotate(tensor: torch.Tensor, heads: int, rope: torch.Tensor) -> torch.Tensor: + pairs = tensor.reshape(CHUNK, heads, 128, 2).float() + cosine = rope[:, None, :, 0].float() + sine = rope[:, None, :, 1].float() + return torch.stack( + [ + pairs[..., 0] * cosine - pairs[..., 1] * sine, + pairs[..., 1] * cosine + pairs[..., 0] * sine, + ], + -1, + ).to(torch.bfloat16).reshape(CHUNK, heads, 256) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + parser.add_argument("--steps", type=int, default=1) + parser.add_argument("--start-step", type=int, default=0) + args = parser.parse_args() + if args.steps < 1 or args.start_step < 0 or args.start_step + args.steps > 10: + raise ValueError("start-step and steps must select a subset of [0, 10)") + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + root = "model." if "model.action_in_proj.weight" in keys else "" + + def raw(name: str) -> torch.Tensor: + return file.get_tensor(root + name) + + def bf16(name: str) -> torch.Tensor: + return raw(name).to(device="cuda", dtype=torch.bfloat16) + + decoder = "paligemma_with_expert.gemma_expert.model.layers" + time_in_w = bf16("time_mlp_in.weight").t().contiguous() + time_in_b = bf16("time_mlp_in.bias") + time_out_w = bf16("time_mlp_out.weight").t().contiguous() + time_out_b = bf16("time_mlp_out.bias") + attn_mod_w = [ + bf16(f"{decoder}.{i}.input_layernorm.dense.weight").t().contiguous() + for i in range(18) + ] + attn_mod_b = [ + bf16(f"{decoder}.{i}.input_layernorm.dense.bias") for i in range(18) + ] + ffn_mod_w = [ + bf16(f"{decoder}.{i}.post_attention_layernorm.dense.weight") + .t() + .contiguous() + for i in range(18) + ] + ffn_mod_b = [ + bf16(f"{decoder}.{i}.post_attention_layernorm.dense.bias") + for i in range(18) + ] + final_mod_w = bf16( + "paligemma_with_expert.gemma_expert.model.norm.dense.weight" + ).t().contiguous() + final_mod_b = bf16( + "paligemma_with_expert.gemma_expert.model.norm.dense.bias" + ) + fraction = torch.linspace(0.0, 1.0, 512) + period = 4e-3 * (4.0 / 4e-3) ** fraction + current = torch.tensor(1.0, dtype=torch.float32) + schedule = [] + for _ in range(10): + angle = current * (1.0 / period) * 2 * math.pi + schedule.append( + torch.cat([torch.sin(angle), torch.cos(angle)]).to( + device="cuda", dtype=torch.bfloat16 + ) + ) + current = current - 0.1 + styles_attn = torch.empty( + 10, 18, CHUNK, 3072, device="cuda", dtype=torch.bfloat16 + ) + styles_ffn = torch.empty_like(styles_attn) + styles_final = torch.empty( + 10, CHUNK, 3072, device="cuda", dtype=torch.bfloat16 + ) + step_range = range(args.start_step, args.start_step + args.steps) + for step in step_range: + value = schedule[step][None, :] @ time_in_w + value = (value.float() + time_in_b.float()).to(torch.bfloat16) + value_float = value.float() + value = (value_float * torch.sigmoid(value_float)).to(torch.bfloat16) + value = value @ time_out_w + value = (value.float() + time_out_b.float()).to(torch.bfloat16) + value_float = value.float() + value = (value_float * torch.sigmoid(value_float)).to(torch.bfloat16) + expanded = value.expand(CHUNK, -1).contiguous() + for layer in range(18): + styles_attn[step, layer] = ( + (expanded @ attn_mod_w[layer]).float() + + attn_mod_b[layer].float() + ).to(torch.bfloat16) + styles_ffn[step, layer] = ( + (expanded @ ffn_mod_w[layer]).float() + + ffn_mod_b[layer].float() + ).to(torch.bfloat16) + styles_final[step] = ( + (expanded @ final_mod_w).float() + final_mod_b.float() + ).to(torch.bfloat16) + + layers = [] + for index in range(18): + prefix = f"{decoder}.{index}" + q = interleave_qk(raw(f"{prefix}.self_attn.q_proj.weight"), 8) + k = interleave_qk(raw(f"{prefix}.self_attn.k_proj.weight"), 1) + v = raw(f"{prefix}.self_attn.v_proj.weight") + layers.append( + { + "qkv": torch.cat([q, k, v], dim=0) + .t() + .to(device="cuda", dtype=torch.bfloat16) + .contiguous(), + "output": bf16(f"{prefix}.self_attn.o_proj.weight") + .t() + .contiguous(), + "gate": bf16(f"{prefix}.mlp.gate_proj.weight").t().contiguous(), + "up": bf16(f"{prefix}.mlp.up_proj.weight").t().contiguous(), + "down": bf16(f"{prefix}.mlp.down_proj.weight").t().contiguous(), + } + ) + + positions = torch.arange(PREFIX, PREFIX + CHUNK, dtype=torch.float64)[:, None] + pair = torch.arange(128, dtype=torch.float64)[None, :] + phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) + rope = torch.stack([torch.cos(phase), torch.sin(phase)], -1).to( + device="cuda", dtype=torch.bfloat16 + ) + layer_index = torch.arange(18, device="cuda")[:, None, None] + row_index = torch.arange(722, device="cuda")[None, :, None] + column_index = torch.arange(256, device="cuda")[None, None, :] + cache_k = ( + ((layer_index + row_index + column_index) % 17 - 8).float() / 16.0 + ).to(torch.bfloat16) + cache_v = ( + ((2 * layer_index + row_index + 3 * column_index) % 19 - 9).float() + / 16.0 + ).to(torch.bfloat16) + flat = torch.arange(CHUNK * 32, device="cuda") + noise = (((flat % 23) - 11).float() / 12.0).to(torch.bfloat16).reshape( + CHUNK, 32 + ) + input_weight = bf16("action_in_proj.weight").t().contiguous() + input_bias = bf16("action_in_proj.bias") + output_weight = ( + bf16("action_out_proj.weight").float() * -0.1 + ).to(torch.bfloat16).t().contiguous() + output_bias = ( + bf16("action_out_proj.bias").float() * -0.1 + ).to(torch.bfloat16) + + for step in step_range: + x = noise @ input_weight + x = (x.float() + input_bias.float()).to(torch.bfloat16) + for index, weights in enumerate(layers): + x_norm, residual_gate = ada_rms(x, styles_attn[step, index]) + qkv = x_norm @ weights["qkv"] + query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) + query = rotate(query, 8, rope) + key = rotate(key, 1, rope).reshape(CHUNK, 256) + value = value.reshape(CHUNK, 256) + cache_k[index, PREFIX : PREFIX + CHUNK] = key + cache_v[index, PREFIX : PREFIX + CHUNK] = value + attended = F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + cache_k[index].reshape(722, 1, 256).transpose(0, 1).unsqueeze(0), + cache_v[index].reshape(722, 1, 256).transpose(0, 1).unsqueeze(0), + scale=1.0 / 16.0, + enable_gqa=True, + ).squeeze(0).transpose(0, 1).reshape(CHUNK, 2048) + projected = attended @ weights["output"] + x = ( + x.float() + projected.float() * residual_gate.float() + ).to(torch.bfloat16) + x_norm, residual_gate = ada_rms(x, styles_ffn[step, index]) + gate = x_norm @ weights["gate"] + up = x_norm @ weights["up"] + gate_float = gate.float() + activated = gate_float / ( + 1.0 + + torch.exp( + -1.5957691216057308 + * gate_float + * (1.0 + 0.044715 * gate_float.square()) + ) + ) + hidden = (activated * up.float()).to(torch.bfloat16) + down = hidden @ weights["down"] + x = (x.float() + down.float() * residual_gate.float()).to( + torch.bfloat16 + ) + x_norm, _ = ada_rms(x, styles_final[step]) + action = x_norm @ output_weight + action = (action.float() + output_bias.float()).to(torch.bfloat16) + noise = (noise.float() + action.float()).to(torch.bfloat16) + + expected = noise.cpu().float() + del cache_k, cache_v, noise, x, x_norm, action + gc.collect() + torch.cuda.empty_cache() + with tempfile.TemporaryDirectory() as directory: + output = str(pathlib.Path(directory) / "diffusion.bin") + subprocess.check_call( + [ + args.probe, + args.checkpoint, + output, + str(args.steps), + str(args.start_step), + ] + ) + bits = np.fromfile(output, dtype=np.uint16) + if bits.size != CHUNK * 32: + raise AssertionError(f"diffusion probe output elements={bits.size}") + actual = torch.from_numpy(bits.copy()).view(torch.bfloat16).float().reshape( + CHUNK, 32 + ) + cosine = float( + F.cosine_similarity( + actual.flatten().double(), expected.flatten().double(), dim=0 + ) + ) + maximum = float((actual - expected).abs().max()) + if cosine < 0.9999: + raise AssertionError(f"cosine={cosine:.8f} max={maximum:.6f}") + print( + f"PASS diffusion steps {args.start_step}.." + f"{args.start_step + args.steps - 1} " + f"cosine={cosine:.8f} max={maximum:.6f}" + ) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_e2e.py b/cpp/tests/gate_pi05_native_e2e.py new file mode 100644 index 00000000..5a012c6f --- /dev/null +++ b/cpp/tests/gate_pi05_native_e2e.py @@ -0,0 +1,275 @@ +"""Compare native SM120 Pi0.5 against the official OpenPI PyTorch policy.""" + +from __future__ import annotations + +import argparse +import io +import json +import os +from pathlib import Path +import subprocess +import sys +import tempfile + +import ml_dtypes +import numpy as np +from PIL import Image +import pyarrow.compute as pc +import pyarrow.parquet as pq + + +def _cosine(a: np.ndarray, b: np.ndarray) -> float: + x = np.asarray(a, dtype=np.float64).reshape(-1) + y = np.asarray(b, dtype=np.float64).reshape(-1) + nx = np.linalg.norm(x) + ny = np.linalg.norm(y) + return float(x @ y / (nx * ny)) if nx and ny else float("nan") + + +def _decode_image(cell) -> np.ndarray: + raw = cell["bytes"] if isinstance(cell, dict) else cell + image = Image.open(io.BytesIO(raw)).convert("RGB") + image = image.resize((224, 224), Image.Resampling.BILINEAR) + return np.ascontiguousarray(image, dtype=np.uint8) + + +def _task(root: Path, task_index: int) -> str: + with (root / "meta" / "tasks.jsonl").open(encoding="utf-8") as stream: + for line in stream: + item = json.loads(line) + if int(item["task_index"]) == task_index: + return str(item["task"]) + raise KeyError(f"task_index={task_index} is missing") + + +def _make_fixture(args, fixture: Path) -> None: + info = json.loads((args.dataset / "meta" / "info.json").read_text()) + chunk = args.episode // int(info.get("chunks_size", 1000)) + relative = info.get( + "data_path", + "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + ).format(episode_chunk=chunk, episode_index=args.episode) + table = pq.read_table(args.dataset / relative) + rows = table.filter(pc.equal(table["frame_index"], args.frame)).to_pylist() + if len(rows) != 1: + raise RuntimeError( + f"expected one episode={args.episode} frame={args.frame} row" + ) + row = rows[0] + _decode_image(row["image"]).tofile(fixture / "image_0.rgb") + _decode_image(row["wrist_image"]).tofile(fixture / "image_1.rgb") + np.asarray(row["state"], dtype=np.float32).tofile(fixture / "state.f32") + (fixture / "prompt.txt").write_text( + _task(args.dataset, int(row["task_index"])), encoding="utf-8" + ) + values = np.random.default_rng(args.seed).standard_normal(10 * 32) + np.asarray(values, dtype=np.float32).astype(ml_dtypes.bfloat16).tofile( + fixture / "noise.bf16" + ) + + +def _official_baseline(checkpoint: Path, fixture: Path, output: Path) -> None: + import torch + + from openpi.models import model as model_api + from openpi.models import tokenizer as tokenizer_api + from openpi.training import config as training_config + + def image(name: str) -> np.ndarray: + return np.fromfile(fixture / name, dtype=np.uint8).reshape(224, 224, 3) + + state = np.fromfile(fixture / "state.f32", dtype=np.float32) + prompt = (fixture / "prompt.txt").read_text(encoding="utf-8") + noise = np.fromfile(fixture / "noise.bf16", dtype=ml_dtypes.bfloat16) + noise = noise.astype(np.float32).reshape(10, 32) + stats = json.loads( + (checkpoint / "assets" / "physical-intelligence" / "libero" / + "norm_stats.json").read_text() + )["norm_stats"] + state_q01 = np.asarray(stats["state"]["q01"], dtype=np.float32) + state_q99 = np.asarray(stats["state"]["q99"], dtype=np.float32) + normalized_state = ( + (state - state_q01) / (state_q99 - state_q01 + 1e-6) * 2.0 - 1.0 + ) + tokens, token_mask = tokenizer_api.PaligemmaTokenizer(200).tokenize( + prompt, normalized_state + ) + padded_state = np.zeros(32, dtype=np.float32) + padded_state[:state.size] = normalized_state + base = image("image_0.rgb") + wrist = image("image_1.rgb") + device_inputs = { + "image": { + "base_0_rgb": torch.from_numpy(base).to("cuda")[None, ...], + "left_wrist_0_rgb": torch.from_numpy(wrist).to("cuda")[None, ...], + "right_wrist_0_rgb": torch.zeros_like( + torch.from_numpy(base).to("cuda")[None, ...] + ), + }, + "image_mask": { + "base_0_rgb": torch.ones(1, dtype=torch.bool, device="cuda"), + "left_wrist_0_rgb": torch.ones(1, dtype=torch.bool, device="cuda"), + "right_wrist_0_rgb": torch.zeros(1, dtype=torch.bool, device="cuda"), + }, + "state": torch.from_numpy(padded_state).to("cuda")[None, ...], + "tokenized_prompt": torch.from_numpy(tokens).to("cuda")[None, ...], + "tokenized_prompt_mask": torch.from_numpy(token_mask).to("cuda")[None, ...], + } + model_observation = model_api.Observation.from_dict(device_inputs) + train_config = training_config.get_config("pi05_libero") + model = train_config.model.load_pytorch( + train_config, str(checkpoint / "model.safetensors") + ) + model.paligemma_with_expert.to_bfloat16_for_selected_params("bfloat16") + model.to("cuda").eval() + noise_tensor = torch.from_numpy(noise).to("cuda")[None, ...] + with torch.inference_mode(): + raw = model.sample_actions( + "cuda", model_observation, noise=noise_tensor, num_steps=10 + )[0].float().cpu().numpy() + action_q01 = np.asarray(stats["actions"]["q01"], dtype=np.float32) + action_q99 = np.asarray(stats["actions"]["q99"], dtype=np.float32) + clipped = np.clip(raw[:, :action_q01.size], -1.0, 1.0) + actions = ((clipped + 1.0) * 0.5 * + (action_q99 - action_q01 + 1e-6) + action_q01) + np.asarray(raw, dtype=np.float32).tofile(output / "openpi_raw.f32") + np.asarray(actions, dtype=np.float32).tofile( + output / "openpi_actions.f32" + ) + + +def _run(args) -> None: + with tempfile.TemporaryDirectory(prefix="pi05_native_e2e_") as temp: + root = Path(temp) + _make_fixture(args, root) + env = dict(os.environ) + env["TORCH_COMPILE_DISABLE"] = "1" + baseline_prefix = env.get("OPENPI_BASELINE_SITE_PACKAGES") + if baseline_prefix: + baseline_packages = Path(baseline_prefix) + if not baseline_packages.is_dir(): + raise FileNotFoundError(baseline_packages) + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = str(baseline_packages) + ( + os.pathsep + existing if existing else "" + ) + subprocess.run( + [ + sys.executable, + __file__, + "--baseline-fixture", + str(root), + "--checkpoint", + str(args.checkpoint), + ], + check=True, + env=env, + ) + subprocess.run( + [ + str(args.probe), + str(args.checkpoint), + str(args.tokenizer), + str(root), + str(root), + ], + check=True, + ) + openpi_raw = np.fromfile(root / "openpi_raw.f32", dtype=np.float32) + native_raw = np.fromfile( + root / "native_raw.bf16", dtype=ml_dtypes.bfloat16 + ).astype(np.float32) + openpi_actions = np.fromfile( + root / "openpi_actions.f32", dtype=np.float32 + ) + native_actions = np.fromfile( + root / "native_actions.f32", dtype=np.float32 + ) + sizes = { + "openpi_raw": openpi_raw.size, + "native_raw": native_raw.size, + "openpi_actions": openpi_actions.size, + "native_actions": native_actions.size, + } + expected_sizes = { + "openpi_raw": 10 * 32, + "native_raw": 10 * 32, + "openpi_actions": 10 * 7, + "native_actions": 10 * 7, + } + if sizes != expected_sizes: + raise RuntimeError(f"unexpected E2E output sizes: {sizes}") + raw_cos = _cosine(openpi_raw, native_raw) + action_cos = _cosine(openpi_actions, native_actions) + raw_max = float(np.max(np.abs(openpi_raw - native_raw))) + action_max = float(np.max(np.abs(openpi_actions - native_actions))) + stats = json.loads( + (args.checkpoint / "assets" / "physical-intelligence" / + "libero" / "norm_stats.json").read_text() + )["norm_stats"]["actions"] + q01 = np.asarray(stats["q01"], dtype=np.float32) + q99 = np.asarray(stats["q99"], dtype=np.float32) + native_model = native_raw.reshape(10, 32)[:, :q01.size] + native_contract_actions = ( + (np.clip(native_model, -1.0, 1.0) + 1.0) * 0.5 * + (q99 - q01 + 1e-6) + q01 + ) + contract_max = float(np.max(np.abs( + native_contract_actions.reshape(-1) - native_actions + ))) + contract_close = bool( + np.allclose(native_contract_actions.reshape(-1), native_actions, + rtol=1e-6, atol=1e-6) + ) + print("\n===== PI0.5 NATIVE VS OFFICIAL OPENPI =====") + print(f"episode/frame : {args.episode}/{args.frame}") + print(f"raw action cosine : {raw_cos:.8f} max_abs={raw_max:.6g}") + print( + f"robot action : cos={action_cos:.8f} " + f"max_abs_vs_fp32={action_max:.6g}" + ) + print( + f"action postprocess: allclose={contract_close} " + f"max_abs={contract_max:.6g}" + ) + if raw_cos < 0.9999: + raise AssertionError(f"raw action cosine {raw_cos:.8f} < 0.9999") + if action_cos < 0.9999: + raise AssertionError( + f"robot action cosine {action_cos:.8f} < 0.9999" + ) + if not contract_close: + raise AssertionError( + f"native action postprocess differs; max_abs={contract_max:.6g}" + ) + print("PASS native Pi0.5 real-episode E2E") + + +def _parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path) + parser.add_argument("--dataset", type=Path) + parser.add_argument("--probe", type=Path) + parser.add_argument("--episode", type=int, default=0) + parser.add_argument("--frame", type=int, default=0) + parser.add_argument("--seed", type=int, default=20260709) + parser.add_argument("--baseline-fixture", type=Path, help=argparse.SUPPRESS) + args = parser.parse_args() + if args.baseline_fixture is None: + for name in ("tokenizer", "dataset", "probe"): + if getattr(args, name) is None: + parser.error(f"--{name} is required") + return args + + +if __name__ == "__main__": + parsed = _parse_args() + if parsed.baseline_fixture is not None: + _official_baseline( + parsed.checkpoint, + parsed.baseline_fixture, + parsed.baseline_fixture, + ) + else: + _run(parsed) diff --git a/cpp/tests/gate_pi05_native_encoder.py b/cpp/tests/gate_pi05_native_encoder.py new file mode 100644 index 00000000..c7d04ba6 --- /dev/null +++ b/cpp/tests/gate_pi05_native_encoder.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +import argparse +import gc +import pathlib +import subprocess +import tempfile + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors import safe_open + + +SEQUENCE = 712 +WIDTH = 2048 + + +def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: + output, inputs = weight.shape + head_dim = output // heads + return ( + weight.reshape(heads, head_dim, inputs) + .reshape(heads, 2, head_dim // 2, inputs) + .permute(0, 2, 1, 3) + .reshape(output, inputs) + ) + + +def rms(values: torch.Tensor) -> torch.Tensor: + source = values.float() + return (source * torch.rsqrt(source.square().mean(-1, keepdim=True) + 1e-6)).to( + torch.bfloat16 + ) + + +def rotate(tensor: torch.Tensor, heads: int, rope: torch.Tensor) -> torch.Tensor: + pairs = tensor.reshape(SEQUENCE, heads, 128, 2).float() + cosine = rope[:, None, :, 0].float() + sine = rope[:, None, :, 1].float() + return torch.stack( + [ + pairs[..., 0] * cosine - pairs[..., 1] * sine, + pairs[..., 1] * cosine + pairs[..., 0] * sine, + ], + -1, + ).to(torch.bfloat16).reshape(SEQUENCE, heads, 256) + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor) -> str: + cosine = float( + F.cosine_similarity( + actual.flatten().double(), expected.flatten().double(), dim=0 + ) + ) + maximum = float((actual - expected).abs().max()) + if cosine < 0.9999: + raise AssertionError(f"{name}: cosine={cosine:.8f} max={maximum:.6f}") + return f"{name} cosine={cosine:.8f} max={maximum:.6f}" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + root = "model." if "model.action_in_proj.weight" in keys else "" + + def raw(name: str) -> torch.Tensor: + return file.get_tensor(root + name) + + x = torch.zeros((SEQUENCE, WIDTH), device="cuda", dtype=torch.bfloat16) + rows = torch.arange(SEQUENCE, device="cuda")[:, None] + columns = torch.arange(512, device="cuda")[None, :] + x[:, :512] = (((rows + columns) % 15 - 7).float() / 8).to(torch.bfloat16) + positions = torch.arange(SEQUENCE, dtype=torch.float64)[:, None] + pair = torch.arange(128, dtype=torch.float64)[None, :] + phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) + rope = torch.stack([torch.cos(phase), torch.sin(phase)], -1).to( + device="cuda", dtype=torch.bfloat16 + ) + + final_q = final_k = final_v = None + for index in range(18): + layer = ( + "paligemma_with_expert.paligemma.model.language_model.layers." + f"{index}" + ) + input_norm = 1.0 + raw(f"{layer}.input_layernorm.weight").float() + q = interleave_qk(raw(f"{layer}.self_attn.q_proj.weight").float(), 8) + k = interleave_qk(raw(f"{layer}.self_attn.k_proj.weight").float(), 1) + v = raw(f"{layer}.self_attn.v_proj.weight").float() + qkv_weight = torch.cat( + [ + q * input_norm[None, :], + k * input_norm[None, :], + v * input_norm[None, :], + ], + dim=0, + ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() + qkv = rms(x) @ qkv_weight + query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) + query = rotate(query, 8, rope) + key = rotate(key, 1, rope) + value = value.reshape(SEQUENCE, 1, 256) + if index == 17: + final_q = query.reshape(SEQUENCE, 2048).cpu().float() + final_k = key.reshape(SEQUENCE, 256).cpu().float() + final_v = value.reshape(SEQUENCE, 256).cpu().float() + break + + attended = F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + scale=1.0 / 16.0, + enable_gqa=True, + ).squeeze(0).transpose(0, 1).reshape(SEQUENCE, 2048) + output_weight = raw(f"{layer}.self_attn.o_proj.weight").to( + device="cuda", dtype=torch.bfloat16 + ).t().contiguous() + x = (x.float() + (attended @ output_weight).float()).to(torch.bfloat16) + post_norm = 1.0 + raw( + f"{layer}.post_attention_layernorm.weight" + ).float() + gate_weight = ( + raw(f"{layer}.mlp.gate_proj.weight").float() * post_norm[None, :] + ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() + up_weight = ( + raw(f"{layer}.mlp.up_proj.weight").float() * post_norm[None, :] + ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() + down_weight = raw(f"{layer}.mlp.down_proj.weight").to( + device="cuda", dtype=torch.bfloat16 + ).t().contiguous() + normalized = rms(x) + gate = normalized @ gate_weight + up = normalized @ up_weight + gate_float = gate.float() + activated = gate_float / ( + 1.0 + + torch.exp( + -1.5957691216057308 + * gate_float + * (1.0 + 0.044715 * gate_float.square()) + ) + ) + hidden = (activated * up.float()).to(torch.bfloat16) + x = (x.float() + (hidden @ down_weight).float()).to(torch.bfloat16) + del q, k, v, qkv_weight, qkv, query, key, value + del attended, output_weight, gate_weight, up_weight, down_weight + del normalized, gate, up, gate_float, activated, hidden + gc.collect() + + expected_x = x.cpu().float() + del x, rope + torch.cuda.empty_cache() + with tempfile.TemporaryDirectory() as directory: + output = str(pathlib.Path(directory) / "encoder.bin") + subprocess.check_call([args.probe, args.checkpoint, output]) + bits = np.fromfile(output, dtype=np.uint16) + sizes = [SEQUENCE * 2048, SEQUENCE * 2048, SEQUENCE * 256, SEQUENCE * 256] + if bits.size != sum(sizes): + raise AssertionError(f"encoder probe output elements={bits.size}") + tensors = [] + offset = 0 + for size in sizes: + tensors.append( + torch.from_numpy(bits[offset : offset + size].copy()) + .view(torch.bfloat16) + .float() + ) + offset += size + messages = [ + compare("x", tensors[0].reshape(SEQUENCE, 2048), expected_x), + compare("q", tensors[1].reshape(SEQUENCE, 2048), final_q), + compare("k", tensors[2].reshape(SEQUENCE, 256), final_k), + compare("v", tensors[3].reshape(SEQUENCE, 256), final_v), + ] + print("PASS encoder 18 layers " + "; ".join(messages)) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_encoder_layer.py b/cpp/tests/gate_pi05_native_encoder_layer.py new file mode 100644 index 00000000..cfabab74 --- /dev/null +++ b/cpp/tests/gate_pi05_native_encoder_layer.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +import argparse +import pathlib +import subprocess +import tempfile + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors import safe_open + + +def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: + output, inputs = weight.shape + head_dim = output // heads + return ( + weight.reshape(heads, head_dim, inputs) + .reshape(heads, 2, head_dim // 2, inputs) + .permute(0, 2, 1, 3) + .reshape(output, inputs) + ) + + +def rms(values: torch.Tensor) -> torch.Tensor: + source = values.float() + return (source * torch.rsqrt(source.square().mean(-1, keepdim=True) + 1e-6)).to( + torch.bfloat16 + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + root = "model." if "model.action_in_proj.weight" in keys else "" + layer = "paligemma_with_expert.paligemma.model.language_model.layers.0" + + def raw(name: str) -> torch.Tensor: + return file.get_tensor(root + name) + + input_norm = 1.0 + raw(f"{layer}.input_layernorm.weight").float() + q = interleave_qk(raw(f"{layer}.self_attn.q_proj.weight").float(), 8) + k = interleave_qk(raw(f"{layer}.self_attn.k_proj.weight").float(), 1) + v = raw(f"{layer}.self_attn.v_proj.weight").float() + qkv_weight = torch.cat( + [q * input_norm[None, :], k * input_norm[None, :], + v * input_norm[None, :]], dim=0 + ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() + output_weight = raw(f"{layer}.self_attn.o_proj.weight").to( + device="cuda", dtype=torch.bfloat16 + ).t().contiguous() + post_norm = 1.0 + raw(f"{layer}.post_attention_layernorm.weight").float() + gate_weight = (raw(f"{layer}.mlp.gate_proj.weight").float() * + post_norm[None, :]).t().to( + device="cuda", dtype=torch.bfloat16 + ).contiguous() + up_weight = (raw(f"{layer}.mlp.up_proj.weight").float() * + post_norm[None, :]).t().to( + device="cuda", dtype=torch.bfloat16 + ).contiguous() + down_weight = raw(f"{layer}.mlp.down_proj.weight").to( + device="cuda", dtype=torch.bfloat16 + ).t().contiguous() + + x = torch.zeros((712, 2048), device="cuda", dtype=torch.bfloat16) + rows = torch.arange(712, device="cuda")[:, None] + columns = torch.arange(512, device="cuda")[None, :] + x[:, :512] = (((rows + columns) % 15 - 7).float() / 8).to(torch.bfloat16) + qkv = rms(x) @ qkv_weight + query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) + positions = torch.arange(712, dtype=torch.float64)[:, None] + pair = torch.arange(128, dtype=torch.float64)[None, :] + phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) + rope = torch.stack([torch.cos(phase), torch.sin(phase)], -1).to( + device="cuda", dtype=torch.bfloat16 + ) + + def rotate(tensor: torch.Tensor, heads: int) -> torch.Tensor: + pairs = tensor.reshape(712, heads, 128, 2).float() + cosine = rope[:, None, :, 0].float() + sine = rope[:, None, :, 1].float() + return torch.stack( + [pairs[..., 0] * cosine - pairs[..., 1] * sine, + pairs[..., 1] * cosine + pairs[..., 0] * sine], -1 + ).to(torch.bfloat16).reshape(712, heads, 256) + + query = rotate(query, 8).transpose(0, 1).unsqueeze(0) + key = rotate(key, 1).transpose(0, 1).unsqueeze(0) + value = value.reshape(712, 1, 256).transpose(0, 1).unsqueeze(0) + attended = F.scaled_dot_product_attention( + query, key, value, scale=1.0 / 16.0, enable_gqa=True + ).squeeze(0).transpose(0, 1).reshape(712, 2048) + projected = attended @ output_weight + x = (x.float() + projected.float()).to(torch.bfloat16) + normalized = rms(x) + gate = normalized @ gate_weight + up = normalized @ up_weight + gate_float = gate.float() + activated = gate_float / ( + 1.0 + torch.exp(-1.5957691216057308 * gate_float * + (1.0 + 0.044715 * gate_float.square())) + ) + hidden = (activated * up.float()).to(torch.bfloat16) + down = hidden @ down_weight + expected = (x.float() + down.float()).to(torch.bfloat16).cpu().float() + + with tempfile.TemporaryDirectory() as directory: + output = str(pathlib.Path(directory) / "encoder.bin") + subprocess.check_call([args.probe, args.checkpoint, output]) + bits = np.fromfile(output, dtype=np.uint16).reshape(712, 2048) + actual = torch.from_numpy(bits.copy()).view(torch.bfloat16).float() + cosine = float(F.cosine_similarity( + actual.flatten().double(), expected.flatten().double(), dim=0 + )) + maximum = float((actual - expected).abs().max()) + if cosine < 0.9999: + raise AssertionError(f"cosine={cosine:.8f} max={maximum:.6f}") + print(f"PASS encoder layer0 cosine={cosine:.8f} max={maximum:.6f}") + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_encoder_qkv.py b/cpp/tests/gate_pi05_native_encoder_qkv.py new file mode 100644 index 00000000..4bf57dca --- /dev/null +++ b/cpp/tests/gate_pi05_native_encoder_qkv.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +import argparse +import pathlib +import subprocess +import tempfile + +import numpy as np +import torch +from safetensors import safe_open + + +def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: + output, inputs = weight.shape + head_dim = output // heads + return ( + weight.reshape(heads, head_dim, inputs) + .reshape(heads, 2, head_dim // 2, inputs) + .permute(0, 2, 1, 3) + .reshape(output, inputs) + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the encoder QKV gate") + + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + prefix = "model." if "model.action_in_proj.weight" in keys else "" + layer = "paligemma_with_expert.paligemma.model.language_model.layers.17" + + def raw(name: str) -> torch.Tensor: + return file.get_tensor(prefix + name) + + q = interleave_qk(raw(f"{layer}.self_attn.q_proj.weight").float(), 8) + k = interleave_qk(raw(f"{layer}.self_attn.k_proj.weight").float(), 1) + v = raw(f"{layer}.self_attn.v_proj.weight").float() + norm = 1.0 + raw(f"{layer}.input_layernorm.weight").float() + weight = torch.cat( + [q * norm[None, :], k * norm[None, :], v * norm[None, :]], dim=0 + ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() + + x = torch.zeros((712, 2048), device="cuda", dtype=torch.bfloat16) + rows = torch.arange(712, device="cuda")[:, None] + columns = torch.arange(512, device="cuda")[None, :] + x[:, :512] = (((rows + columns) % 15 - 7).float() / 8).to(torch.bfloat16) + x_float = x.float() + normalized = ( + x_float * torch.rsqrt(x_float.square().mean(dim=-1, keepdim=True) + 1e-6) + ).to(torch.bfloat16) + qkv = normalized @ weight + query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) + + positions = torch.arange(712, dtype=torch.float64)[:, None] + pair = torch.arange(128, dtype=torch.float64)[None, :] + phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) + rope = torch.stack([torch.cos(phase), torch.sin(phase)], dim=-1).to( + device="cuda", dtype=torch.bfloat16 + ) + + def apply_rope(tensor: torch.Tensor, heads: int) -> torch.Tensor: + pairs = tensor.reshape(712, heads, 128, 2).float() + cosine = rope[:, None, :, 0].float() + sine = rope[:, None, :, 1].float() + even = pairs[..., 0] * cosine - pairs[..., 1] * sine + odd = pairs[..., 1] * cosine + pairs[..., 0] * sine + return torch.stack([even, odd], dim=-1).to(torch.bfloat16).reshape( + 712, heads * 256 + ) + + expected = { + "q": apply_rope(query, 8), + "k": apply_rope(key, 1), + "v": value.contiguous(), + } + with tempfile.TemporaryDirectory() as directory: + output = str(pathlib.Path(directory) / "encoder") + subprocess.check_call([args.probe, args.checkpoint, output]) + for name, reference in expected.items(): + actual_bits = np.fromfile(f"{output}.{name}.bin", dtype=np.uint16) + actual = torch.from_numpy(actual_bits.copy()).view(torch.bfloat16) + actual = actual.reshape(reference.shape).float() + target = reference.cpu().float() + cosine = float(torch.nn.functional.cosine_similarity( + actual.flatten().double(), target.flatten().double(), dim=0 + )) + maximum = float((actual - target).abs().max()) + if cosine < 0.9999: + raise AssertionError( + f"{name}: cosine={cosine:.8f} max={maximum:.6f}" + ) + print(f"PASS encoder17 {name} cosine={cosine:.8f} max={maximum:.6f}") + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_quantization.py b/cpp/tests/gate_pi05_native_quantization.py new file mode 100644 index 00000000..50d424ad --- /dev/null +++ b/cpp/tests/gate_pi05_native_quantization.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +import argparse +import subprocess + +import torch +from safetensors import safe_open + + +DECODER = "paligemma_with_expert.gemma_expert.model.layers.0" + + +def interleave_qk(weight: torch.Tensor, num_heads: int) -> torch.Tensor: + out_dim, in_dim = weight.shape + head_dim = out_dim // num_heads + return ( + weight.reshape(num_heads, head_dim, in_dim) + .reshape(num_heads, 2, head_dim // 2, in_dim) + .permute(0, 2, 1, 3) + .reshape(out_dim, in_dim) + ) + + +def fnv1a(data: bytes) -> int: + value = 14695981039346656037 + for byte in data: + value ^= byte + value = (value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return value + + +def digest(tensor: torch.Tensor) -> int: + return fnv1a(tensor.contiguous().cpu().numpy().tobytes()) + + +def parse_probe(text: str) -> tuple[tuple[int, ...], int, int, int]: + fields = dict(field.split("=", 1) for field in text.strip().split()) + return ( + tuple(int(dim) for dim in fields["shape"].split(",")), + int(fields["values_fnv"], 16), + int(fields["scale_shape"]), + int(fields["scales_fnv"], 16), + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the producer quantization gate") + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + prefix = "model." if "model.action_in_proj.weight" in keys else "" + + def bf16(key: str) -> torch.Tensor: + return file.get_tensor(prefix + key).to(torch.bfloat16) + + q = interleave_qk(bf16(f"{DECODER}.self_attn.q_proj.weight").float(), 8) + k = interleave_qk(bf16(f"{DECODER}.self_attn.k_proj.weight").float(), 1) + v = bf16(f"{DECODER}.self_attn.v_proj.weight") + weight = torch.cat([q, k, v], dim=0).t().to( + device="cuda", dtype=torch.bfloat16 + ).contiguous() + + expected = {} + for layout in ("kn", "nk"): + arranged = weight.t().contiguous() if layout == "nk" else weight + scale = max(arranged.float().abs().max().item() / 448.0, 1e-12) + quantized = (arranged.float() / scale).clamp(-448.0, 448.0).to( + torch.float8_e4m3fn + ) + scale_tensor = torch.tensor([scale], dtype=torch.float32, device="cuda") + expected[f"decoder_qkv0_fp8_{layout}"] = ( + tuple(quantized.shape), + digest(quantized.view(torch.uint8)), + 1, + digest(scale_tensor), + ) + + transposed = weight.float().transpose(0, 1).contiguous() + scales = torch.clamp( + transposed.abs().amax(dim=1) / 127.0, min=1e-12 + ).to(dtype=torch.float32).contiguous() + quantized = torch.clamp( + torch.round(transposed / scales[:, None]), -127, 127 + ).to(torch.int8).contiguous() + expected["decoder_qkv0_int8"] = ( + tuple(quantized.shape), + digest(quantized), + scales.numel(), + digest(scales), + ) + + for operation, reference in expected.items(): + output = subprocess.check_output( + [args.probe, args.checkpoint, operation], text=True + ) + actual = parse_probe(output) + if actual != reference: + raise AssertionError( + f"{operation}: C++ {actual} != PyTorch {reference}" + ) + print( + f"PASS {operation} shape={actual[0]} " + f"values_fnv={actual[1]:016x} scales_fnv={actual[3]:016x}" + ) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_rope.py b/cpp/tests/gate_pi05_native_rope.py new file mode 100644 index 00000000..8afbd013 --- /dev/null +++ b/cpp/tests/gate_pi05_native_rope.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +import argparse +import subprocess + +import ml_dtypes +import numpy as np + + +def fnv1a(data: bytes) -> int: + value = 14695981039346656037 + for byte in data: + value ^= byte + value = (value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return value + + +def parse_probe(text: str) -> dict[str, str]: + return dict(field.split("=", 1) for field in text.strip().split()) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--probe", required=True) + args = parser.parse_args() + cases = [(2, 200, 10, 1, 37), (3, 256, 50, 2, 256)] + for views, max_prompt, chunk, pool, prompt in cases: + vision = views * 256 // (pool * pool) + encoder_length = vision + max_prompt + max_positions = encoder_length + chunk + inverse_frequency = 1.0 / ( + 10000 ** (np.arange(0, 256, 2, dtype=np.float64) / 256) + ) + positions = np.arange(max_positions, dtype=np.float64) + phase = positions[:, None] * inverse_frequency[None, :] + cosine = np.cos(phase).astype(ml_dtypes.bfloat16) + sine = np.sin(phase).astype(ml_dtypes.bfloat16) + table = np.stack([cosine, sine], axis=-1).reshape(max_positions, 256) + encoder = np.ascontiguousarray(table[:encoder_length]) + decoder = np.ascontiguousarray( + table[vision + prompt : vision + prompt + chunk] + ) + output = subprocess.check_output( + [ + args.probe, + str(views), + str(max_prompt), + str(chunk), + str(pool), + str(prompt), + ], + text=True, + ) + actual = parse_probe(output) + expected = { + "encoder_shape": f"{encoder_length},256", + "encoder_fnv": f"{fnv1a(encoder.tobytes()):016x}", + "decoder_shape": f"{chunk},256", + "decoder_fnv": f"{fnv1a(decoder.tobytes()):016x}", + } + if actual != expected: + raise AssertionError(f"C++ {actual} != NumPy {expected}") + print( + f"PASS views={views} pool={pool} prompt={prompt} " + f"encoder_fnv={actual['encoder_fnv']} " + f"decoder_fnv={actual['decoder_fnv']}" + ) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_schema_parity.py b/cpp/tests/gate_pi05_native_schema_parity.py new file mode 100644 index 00000000..845cf698 --- /dev/null +++ b/cpp/tests/gate_pi05_native_schema_parity.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Compare Python and C++ native-v2 port/stage/region declarations.""" + +from __future__ import annotations + +import argparse +import difflib +import os +from pathlib import Path +import subprocess +import sys +import tempfile + +import numpy as np + + +ROOT = Path(__file__).resolve().parents[2] +GOLDEN = Path(__file__).with_name("data") / "pi05_native_v2_schema.records" +sys.path.insert(0, str(ROOT)) +configured_build = os.environ.get("FLASHRT_BUILD_DIR") +if not configured_build: + raise SystemExit("Set FLASHRT_BUILD_DIR to the Python producer CMake build") +for subdir in ("exec", "runtime"): + sys.path.insert(0, str(Path(configured_build).resolve() / subdir)) + +import _flashrt_runtime as runtime_abi # noqa: E402 +import flash_rt # noqa: E402 + + +def canonical_records(identity: str) -> list[str]: + prefixes = ("region:", "port:", "stage:") + return [line for line in identity.splitlines() + if line.startswith(prefixes)] + + +def assert_records(label: str, actual: list[str], expected: list[str]) -> None: + if actual == expected: + return + diff = "\n".join(difflib.unified_diff( + expected, actual, fromfile="golden", tofile=label, lineterm="", + )) + raise AssertionError(f"{label} schema mismatch:\n{diff}") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--native-probe", type=Path, required=True) + args = parser.parse_args() + for name in ("checkpoint", "tokenizer", "native_probe"): + path = getattr(args, name).resolve() + if not path.exists(): + parser.error(f"--{name.replace('_', '-')} does not exist: {path}") + setattr(args, name, path) + golden_records = GOLDEN.read_text(encoding="utf-8").splitlines() + expected_ports = sum(line.startswith("port:") for line in golden_records) + expected_stages = sum(line.startswith("stage:") for line in golden_records) + expected_regions = sum(line.startswith("region:") for line in golden_records) + + rng = np.random.default_rng(20260710) + images = [ + np.ascontiguousarray( + rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + ) + for _ in range(2) + ] + state = np.linspace(-0.25, 0.25, 8, dtype=np.float32) + model = flash_rt.load_model( + str(args.checkpoint), framework="torch", config="pi05", + hardware="auto", num_views=2, num_steps=10, cache_frames=1, + use_fp8=True, use_fp16=False, state_prompt_mode="fixed", + ) + model.predict(images, prompt="pick up the red block", state=state) + producer = model._pipe.pipeline.export_model_runtime( + identity={"gate": "native_v2_schema_parity"}, + stage_plan="full", io="native_v2", + ) + try: + counts = dict(runtime_abi.export_counts(producer.export_ptr)) + if len(producer.ports()) != expected_ports or \ + len(producer.stages()) != expected_stages or \ + counts.get("capsule_regions") != expected_regions: + raise RuntimeError( + f"unexpected Python native-v2 counts: ports=" + f"{len(producer.ports())} stages={len(producer.stages())} " + f"regions={counts.get('capsule_regions')}" + ) + python_records = canonical_records(producer.identity) + assert_records("python-native-v2", python_records, golden_records) + with tempfile.TemporaryDirectory(prefix="pi05_schema_parity_") as tmp: + native_path = Path(tmp) / "native.schema" + env = dict(os.environ) + env["FLASHRT_SCHEMA_OUTPUT"] = str(native_path) + env["FLASHRT_SCHEMA_ONLY"] = "1" + subprocess.run( + [str(args.native_probe), str(args.checkpoint), + str(args.tokenizer)], + check=True, env=env, + ) + native_records = native_path.read_text().splitlines() + + assert_records("cpp-native-v2", native_records, golden_records) + + print("\n===== PI0.5 NATIVE-V2 SCHEMA PARITY =====") + print(f"records : {len(python_records)}") + print(f"ports/stage : {expected_ports} / {expected_stages}") + print(f"regions : {expected_regions}") + print("PASS - Python and C++ native-v2 schemas match the golden records") + return 0 + finally: + producer.release() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/gate_pi05_native_style.py b/cpp/tests/gate_pi05_native_style.py new file mode 100644 index 00000000..87bef640 --- /dev/null +++ b/cpp/tests/gate_pi05_native_style.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +import argparse +import math +import pathlib +import subprocess +import tempfile + +import numpy as np +import torch +from safetensors import safe_open + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the style precompute gate") + + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + prefix = "model." if "model.action_in_proj.weight" in keys else "" + + def bf16(key: str) -> torch.Tensor: + return file.get_tensor(prefix + key).to( + device="cuda", dtype=torch.bfloat16 + ) + + decoder = "paligemma_with_expert.gemma_expert.model.layers" + time_in_w = bf16("time_mlp_in.weight").t().contiguous() + time_in_b = bf16("time_mlp_in.bias") + time_out_w = bf16("time_mlp_out.weight").t().contiguous() + time_out_b = bf16("time_mlp_out.bias") + attn_w = torch.stack( + [bf16(f"{decoder}.{i}.input_layernorm.dense.weight").t() for i in range(18)] + ) + attn_b = torch.stack( + [bf16(f"{decoder}.{i}.input_layernorm.dense.bias") for i in range(18)] + ) + ffn_w = torch.stack( + [ + bf16(f"{decoder}.{i}.post_attention_layernorm.dense.weight").t() + for i in range(18) + ] + ) + ffn_b = torch.stack( + [ + bf16(f"{decoder}.{i}.post_attention_layernorm.dense.bias") + for i in range(18) + ] + ) + final_w = bf16( + "paligemma_with_expert.gemma_expert.model.norm.dense.weight" + ).t() + final_b = bf16("paligemma_with_expert.gemma_expert.model.norm.dense.bias") + + fraction = torch.linspace(0.0, 1.0, 512) + period = 4e-3 * (4.0 / 4e-3) ** fraction + t = torch.tensor(1.0, dtype=torch.float32) + rows = [] + for _ in range(10): + angle = t * (1.0 / period) * 2 * math.pi + rows.append( + torch.cat([torch.sin(angle), torch.cos(angle)]).to( + device="cuda", dtype=torch.bfloat16 + ) + ) + t = t - 0.1 + schedule = torch.stack(rows) + expected = { + "decoder_time_emb": torch.empty( + 10, 10, 1024, dtype=torch.bfloat16, device="cuda" + ), + "decoder_style_attn": torch.empty( + 10, 18, 10, 3072, dtype=torch.bfloat16, device="cuda" + ), + "decoder_style_ffn": torch.empty_like( + torch.empty(10, 18, 10, 3072, dtype=torch.bfloat16, device="cuda") + ), + "decoder_style_final": torch.empty( + 10, 10, 3072, dtype=torch.bfloat16, device="cuda" + ), + } + for step in range(10): + value = schedule[step : step + 1] + value = (value @ time_in_w + time_in_b[None, :]).float() + value = (value * torch.sigmoid(value)).to(torch.bfloat16) + value = (value @ time_out_w + time_out_b[None, :]).float() + value = (value * torch.sigmoid(value)).to(torch.bfloat16) + expanded = value.expand(10, -1).contiguous() + expected["decoder_time_emb"][step] = expanded + for layer in range(18): + expected["decoder_style_attn"][step, layer] = ( + expanded @ attn_w[layer] + attn_b[layer][None, :] + ) + expected["decoder_style_ffn"][step, layer] = ( + expanded @ ffn_w[layer] + ffn_b[layer][None, :] + ) + expected["decoder_style_final"][step] = ( + expanded @ final_w + final_b[None, :] + ) + + with tempfile.TemporaryDirectory() as directory: + output_prefix = str(pathlib.Path(directory) / "styles") + subprocess.check_call([args.probe, args.checkpoint, output_prefix]) + for name, reference in expected.items(): + actual_bits = np.fromfile( + f"{output_prefix}.{name}.bin", dtype=np.uint16 + ).reshape(tuple(reference.shape)) + reference_bits = reference.contiguous().view(torch.uint16).cpu().numpy() + exact = float(np.mean(actual_bits == reference_bits)) + actual = torch.from_numpy(actual_bits.copy()).view(torch.bfloat16).float() + target = reference.cpu().float() + maximum = float((actual - target).abs().max()) + cosine = float( + torch.nn.functional.cosine_similarity( + actual.flatten().double(), target.flatten().double(), dim=0 + ) + ) + if cosine < 0.9999: + raise AssertionError( + f"{name}: exact={exact} max={maximum} cosine={cosine}" + ) + print( + f"PASS {name} exact={exact:.6f} " + f"max={maximum:.6f} cosine={cosine:.8f}" + ) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_vision.py b/cpp/tests/gate_pi05_native_vision.py new file mode 100644 index 00000000..c40eb17b --- /dev/null +++ b/cpp/tests/gate_pi05_native_vision.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +import argparse +import gc +import pathlib +import subprocess +import tempfile + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors import safe_open + + +NUM_VIEWS = 2 +SEQUENCE = NUM_VIEWS * 256 +WIDTH = 1152 +HIDDEN = 4304 + + +def layer_norm( + values: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor +) -> torch.Tensor: + source = values.float() + mean = source.mean(-1, keepdim=True) + variance = (source - mean).square().mean(-1, keepdim=True) + return ( + (source - mean) + * torch.rsqrt(variance + 1e-5) + * weight.float() + + bias.float() + ).to(torch.bfloat16) + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor) -> str: + cosine = float( + F.cosine_similarity( + actual.flatten().double(), expected.flatten().double(), dim=0 + ) + ) + maximum = float((actual - expected).abs().max()) + if cosine < 0.9999: + raise AssertionError(f"{name}: cosine={cosine:.8f} max={maximum:.6f}") + return f"{name} cosine={cosine:.8f} max={maximum:.6f}" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + root = "model." if "model.action_in_proj.weight" in keys else "" + vision = "paligemma_with_expert.paligemma.model.vision_tower.vision_model" + + def raw(name: str) -> torch.Tensor: + return file.get_tensor(root + name) + + def bf16(name: str) -> torch.Tensor: + return raw(name).to(device="cuda", dtype=torch.bfloat16) + + flat = torch.arange(NUM_VIEWS * 224 * 224 * 3, device="cuda") + images = (((flat % 257) - 128).float() / 128.0).to(torch.bfloat16).reshape( + NUM_VIEWS, 224, 224, 3 + ) + patches = ( + images.reshape(NUM_VIEWS, 16, 14, 16, 14, 3) + .permute(0, 1, 3, 2, 4, 5) + .reshape(SEQUENCE, 588) + ) + patch_weight = bf16(f"{vision}.embeddings.patch_embedding.weight") + patch_weight = patch_weight.permute(2, 3, 1, 0).reshape(588, WIDTH) + patch_bias = bf16(f"{vision}.embeddings.patch_embedding.bias") + position = bf16(f"{vision}.embeddings.position_embedding.weight").repeat( + NUM_VIEWS, 1 + ) + x = (patches @ patch_weight).to(torch.bfloat16) + x = (x.float() + position.float() + patch_bias.float()).to(torch.bfloat16) + first = f"{vision}.encoder.layers.0" + x_norm = layer_norm( + x, bf16(f"{first}.layer_norm1.weight"), bf16(f"{first}.layer_norm1.bias") + ) + + for index in range(27): + layer = f"{vision}.encoder.layers.{index}" + q_weight = bf16(f"{layer}.self_attn.q_proj.weight") + k_weight = bf16(f"{layer}.self_attn.k_proj.weight") + v_weight = bf16(f"{layer}.self_attn.v_proj.weight") + qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0).t().contiguous() + qkv_bias = torch.cat( + [ + bf16(f"{layer}.self_attn.q_proj.bias"), + bf16(f"{layer}.self_attn.k_proj.bias"), + bf16(f"{layer}.self_attn.v_proj.bias"), + ] + ) + qkv = x_norm @ qkv_weight + qkv = (qkv.float() + qkv_bias.float()).to(torch.bfloat16) + query, key, value = qkv.reshape(NUM_VIEWS, 256, 3, 16, 72).unbind(2) + attended = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + scale=1.0 / np.sqrt(72.0), + ).transpose(1, 2).reshape(SEQUENCE, WIDTH) + output_weight = bf16(f"{layer}.self_attn.out_proj.weight").t().contiguous() + output_bias = bf16(f"{layer}.self_attn.out_proj.bias") + projected = attended @ output_weight + x = (x.float() + projected.float() + output_bias.float()).to(torch.bfloat16) + x_norm = layer_norm( + x, + bf16(f"{layer}.layer_norm2.weight"), + bf16(f"{layer}.layer_norm2.bias"), + ) + up_weight = bf16(f"{layer}.mlp.fc1.weight").t().contiguous() + up_bias = bf16(f"{layer}.mlp.fc1.bias") + hidden = x_norm @ up_weight + hidden = (hidden.float() + up_bias.float()).to(torch.bfloat16) + hidden_float = hidden.float() + hidden = ( + hidden_float + * 0.5 + * ( + 1.0 + + torch.tanh( + 0.7978845608 + * (hidden_float + 0.044715 * hidden_float.pow(3)) + ) + ) + ).to(torch.bfloat16) + down_weight = bf16(f"{layer}.mlp.fc2.weight").t().contiguous() + down_bias = bf16(f"{layer}.mlp.fc2.bias") + down = hidden @ down_weight + x = (x.float() + down.float() + down_bias.float()).to(torch.bfloat16) + if index != 26: + next_layer = f"{vision}.encoder.layers.{index + 1}" + x_norm = layer_norm( + x, + bf16(f"{next_layer}.layer_norm1.weight"), + bf16(f"{next_layer}.layer_norm1.bias"), + ) + del q_weight, k_weight, v_weight, qkv_weight, qkv_bias, qkv + del query, key, value, attended, output_weight, output_bias, projected + del up_weight, up_bias, hidden, hidden_float, down_weight, down_bias, down + gc.collect() + + expected_vision = x.cpu().float() + final_norm = layer_norm( + x, + bf16(f"{vision}.post_layernorm.weight"), + bf16(f"{vision}.post_layernorm.bias"), + ) + projector = ( + "paligemma_with_expert.paligemma.model.multi_modal_projector.linear" + ) + projected = final_norm @ bf16(f"{projector}.weight").t().contiguous() + expected_encoder = ( + projected.float() + bf16(f"{projector}.bias").float() + ).to(torch.bfloat16).cpu().float() + del x, x_norm, final_norm, projected, images, patches + torch.cuda.empty_cache() + + with tempfile.TemporaryDirectory() as directory: + output = str(pathlib.Path(directory) / "vision.bin") + subprocess.check_call([args.probe, args.checkpoint, output]) + bits = np.fromfile(output, dtype=np.uint16) + sizes = [SEQUENCE * WIDTH, SEQUENCE * 2048] + if bits.size != sum(sizes): + raise AssertionError(f"vision probe output elements={bits.size}") + vision_bits = bits[: sizes[0]].copy() + encoder_bits = bits[sizes[0] :].copy() + actual_vision = ( + torch.from_numpy(vision_bits).view(torch.bfloat16).float().reshape(SEQUENCE, WIDTH) + ) + actual_encoder = ( + torch.from_numpy(encoder_bits) + .view(torch.bfloat16) + .float() + .reshape(SEQUENCE, 2048) + ) + print( + "PASS vision 27 layers " + + compare("vision", actual_vision, expected_vision) + + "; " + + compare("encoder", actual_encoder, expected_encoder) + ) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_native_weight_ops.py b/cpp/tests/gate_pi05_native_weight_ops.py new file mode 100644 index 00000000..ad75773a --- /dev/null +++ b/cpp/tests/gate_pi05_native_weight_ops.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +import argparse +import math +import subprocess + +import torch +from safetensors import safe_open + + +VISION = "paligemma_with_expert.paligemma.model.vision_tower.vision_model" +ENCODER = "paligemma_with_expert.paligemma.model.language_model.layers.0" +DECODER = "paligemma_with_expert.gemma_expert.model.layers.0" + + +def interleave_qk(weight: torch.Tensor, num_heads: int) -> torch.Tensor: + out_dim, in_dim = weight.shape + head_dim = out_dim // num_heads + return ( + weight.reshape(num_heads, head_dim, in_dim) + .reshape(num_heads, 2, head_dim // 2, in_dim) + .permute(0, 2, 1, 3) + .reshape(out_dim, in_dim) + ) + + +def fnv1a(data: bytes) -> int: + value = 14695981039346656037 + for byte in data: + value ^= byte + value = (value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return value + + +def digest(tensor: torch.Tensor) -> tuple[tuple[int, ...], int]: + tensor = tensor.contiguous().view(torch.uint16).cpu() + return tuple(tensor.shape), fnv1a(tensor.numpy().tobytes()) + + +def parse_probe(text: str) -> tuple[tuple[int, ...], int]: + fields = dict(field.split("=", 1) for field in text.strip().split()) + shape = tuple(int(dim) for dim in fields["shape"].split(",")) + return shape, int(fields["fnv"], 16) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--probe", required=True) + args = parser.parse_args() + + file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") + keys = set(file.keys()) + prefix = "model." if "model.action_in_proj.weight" in keys else "" + + def raw(key: str) -> torch.Tensor: + return file.get_tensor(prefix + key) + + def bf16(key: str) -> torch.Tensor: + return raw(key).to(torch.bfloat16) + + patch = bf16(f"{VISION}.embeddings.patch_embedding.weight") + expected = { + "patch": patch.permute(2, 3, 1, 0).contiguous(), + } + + q = interleave_qk(raw(f"{ENCODER}.self_attn.q_proj.weight").float(), 8) + k = interleave_qk(raw(f"{ENCODER}.self_attn.k_proj.weight").float(), 1) + v = raw(f"{ENCODER}.self_attn.v_proj.weight").float() + norm = 1.0 + raw(f"{ENCODER}.input_layernorm.weight").float() + expected["encoder_qkv0"] = torch.cat( + [q * norm.unsqueeze(0), k * norm.unsqueeze(0), v * norm.unsqueeze(0)], + dim=0, + ).t().to(torch.bfloat16).contiguous() + + q = interleave_qk(bf16(f"{DECODER}.self_attn.q_proj.weight").float(), 8) + k = interleave_qk(bf16(f"{DECODER}.self_attn.k_proj.weight").float(), 1) + v = bf16(f"{DECODER}.self_attn.v_proj.weight") + expected["decoder_qkv0"] = torch.cat([q, k, v], dim=0).t().to( + torch.bfloat16 + ).contiguous() + + gate = bf16(f"{DECODER}.mlp.gate_proj.weight").t() + up = bf16(f"{DECODER}.mlp.up_proj.weight").t() + expected["decoder_gate_up0"] = torch.cat([gate, up], dim=1).contiguous() + + def time_embeds(num_steps: int) -> torch.Tensor: + fraction = torch.linspace(0.0, 1.0, 512) + period = 4e-3 * (4.0 / 4e-3) ** fraction + t = torch.tensor(1.0, dtype=torch.float32) + rows = [] + for _ in range(num_steps): + angle = ( + t.unsqueeze(-1) + * (1.0 / period).unsqueeze(0) + * 2 + * math.pi + ) + rows.append( + torch.cat([torch.sin(angle), torch.cos(angle)], dim=-1).to( + torch.bfloat16 + ) + ) + t = t - 1.0 / num_steps + return torch.cat(rows, dim=0).contiguous() + + action_out = bf16("action_out_proj.weight").t().to(torch.bfloat16) + for num_steps in (10, 5): + expected[f"action_out{num_steps}"] = ( + action_out * (-1.0 / num_steps) + ).contiguous() + expected[f"time_embeds{num_steps}"] = time_embeds(num_steps) + + for operation, tensor in expected.items(): + output = subprocess.check_output( + [args.probe, args.checkpoint, operation], text=True + ) + actual = parse_probe(output) + reference = digest(tensor) + if actual != reference: + raise AssertionError( + f"{operation}: C++ {actual} != PyTorch {reference}" + ) + print(f"PASS {operation} shape={actual[0]} fnv={actual[1]:016x}") + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/gate_pi05_tokenizer_corpus.py b/cpp/tests/gate_pi05_tokenizer_corpus.py new file mode 100644 index 00000000..9f9c3c6c --- /dev/null +++ b/cpp/tests/gate_pi05_tokenizer_corpus.py @@ -0,0 +1,140 @@ +"""Token-exact Pi0.5 gate over real LIBERO prompt/state records.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import struct +import subprocess +import sys +import tempfile + +import numpy as np +import pyarrow.parquet as pq + + +CORPUS_MAGIC = 0x50303554 +OUTPUT_MAGIC = 0x50303549 + + +def _load_openpi_tokenizer(): + prefix = os.environ.get("OPENPI_BASELINE_SITE_PACKAGES") + if prefix: + path = Path(prefix) + if not path.is_dir(): + raise FileNotFoundError(path) + sys.path.insert(0, str(path)) + from openpi.models import tokenizer as tokenizer_api + + return tokenizer_api.PaligemmaTokenizer(200) + + +def _tasks(dataset: Path) -> dict[int, str]: + result = {} + with (dataset / "meta" / "tasks.jsonl").open(encoding="utf-8") as stream: + for line in stream: + item = json.loads(line) + result[int(item["task_index"])] = str(item["task"]) + return result + + +def _records(dataset: Path, limit: int): + info = json.loads((dataset / "meta" / "info.json").read_text()) + template = info.get( + "data_path", + "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + ) + chunk_size = int(info.get("chunks_size", 1000)) + total_episodes = int(info["total_episodes"]) + count = 0 + for episode in range(total_episodes): + path = dataset / template.format( + episode_chunk=episode // chunk_size, + episode_index=episode, + ) + table = pq.read_table(path, columns=["state", "task_index"]) + for row in table.to_pylist(): + yield int(row["task_index"]), np.asarray( + row["state"], dtype=np.float32 + ) + count += 1 + if count >= limit: + return + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--probe", type=Path, required=True) + parser.add_argument("--count", type=int, default=10000) + args = parser.parse_args() + if args.count <= 0: + parser.error("--count must be positive") + tasks = _tasks(args.dataset) + stats = json.loads( + (args.checkpoint / "assets" / "physical-intelligence" / "libero" / + "norm_stats.json").read_text() + )["norm_stats"]["state"] + q01 = np.asarray(stats["q01"], dtype=np.float32) + q99 = np.asarray(stats["q99"], dtype=np.float32) + official = _load_openpi_tokenizer() + with tempfile.TemporaryDirectory(prefix="pi05_tokenizer_gate_") as temp: + corpus = Path(temp) / "corpus.bin" + output = Path(temp) / "ids.bin" + expected = [] + lengths = set() + with corpus.open("wb") as stream: + stream.write(struct.pack(" + +#include +#include +#include +#include + +namespace { + +struct CaptureArgs { + const flashrt::models::pi05::NativeBf16Forward* forward = nullptr; + const flashrt::models::pi05::NativeDeviceWeightStore* weights = nullptr; + flashrt::models::pi05::NativeWorkspace* workspace = nullptr; + flashrt::models::pi05::NativeRtxAttentionWorkspace* attention = nullptr; + const flashrt::models::pi05::NativeRtxAttentionDriver* attention_driver = + nullptr; + int start_step = 0; + int steps = 10; + bool recorded = false; + std::string error; +}; + +void record_diffusion(void* user, void* stream) { + auto* args = static_cast(user); + flashrt::modalities::Status st = flashrt::modalities::Status::ok(); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + if (args->start_step == 0 && args->steps == 10) { + st = args->forward->diffusion( + *args->weights, args->workspace, args->attention, + args->attention_driver, native_stream); + } else { + for (int offset = 0; offset < args->steps && st.ok_status(); ++offset) { + const int step = args->start_step + offset; + st = args->forward->diffusion_step( + step, *args->weights, args->workspace, args->attention, + args->attention_driver, native_stream); + } + } + args->recorded = st.ok_status(); + args->error = st.message; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 3 || argc > 5) { + std::cerr << "usage: pi05_native_diffusion_probe CHECKPOINT OUTPUT " + "[STEPS [START_STEP]]\n"; + return 2; + } + const int steps = argc >= 4 ? std::stoi(argv[3]) : 10; + const int start_step = argc == 5 ? std::stoi(argv[4]) : 0; + if (steps < 1 || start_step < 0 || start_step + steps > 10) return 2; + using namespace flashrt::models::pi05; + flashrt::loader::SafetensorsFile source; + if (!source.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << source.error() << '\n'; + return 2; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + NativeDeviceWeightStore weights(ctx); + NativeWeightMaterializer materializer(source, &weights); + flashrt::modalities::Status st = materializer.materialize_decoder_globals(10); + for (int layer = 0; layer < 18 && st.ok_status(); ++layer) { + st = materializer.materialize_decoder_layer(layer, false); + } + NativeWorkspace workspace(ctx); + NativeRtxAttentionWorkspace attention(ctx); + if (!st.ok_status() || + !workspace.allocate(NativeWorkspaceConfig{}).ok_status() || + !workspace.update_decoder_rope(200).ok_status() || + !attention.allocate(NativeRtxAttentionConfig{}).ok_status() || + !attention.set_fixed_prompt_length(200).ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + NativeKernelDriver driver; + NativeStylePrecomputer precomputer(&driver); + st = precomputer.run(weights, &workspace, 0); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + const auto* noise = workspace.find("diffusion_noise"); + const auto* cache_k = attention.find("attn_enc_K"); + const auto* cache_v = attention.find("attn_enc_V"); + std::vector host_noise(10 * 32); + for (std::size_t i = 0; i < host_noise.size(); ++i) { + const float value = static_cast(static_cast(i % 23) - 11) / + 12.0f; + host_noise[i] = flashrt::modalities::float_to_bfloat16(value); + } + std::vector host_k(18 * 722 * 256); + std::vector host_v(host_k.size()); + for (int layer = 0; layer < 18; ++layer) { + for (int row = 0; row < 722; ++row) { + for (int column = 0; column < 256; ++column) { + const std::size_t offset = + (static_cast(layer) * 722 + row) * 256 + + column; + host_k[offset] = flashrt::modalities::float_to_bfloat16( + static_cast((layer + row + column) % 17 - 8) / + 16.0f); + host_v[offset] = flashrt::modalities::float_to_bfloat16( + static_cast((2 * layer + row + 3 * column) % 19 - + 9) / + 16.0f); + } + } + } + if (!noise || !cache_k || !cache_v || + cudaMemcpy(frt_buffer_dptr(noise->buffer), host_noise.data(), + host_noise.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess || + cudaMemcpy(frt_buffer_dptr(cache_k->buffer), host_k.data(), + host_k.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess || + cudaMemcpy(frt_buffer_dptr(cache_v->buffer), host_v.data(), + host_v.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess) { + frt_ctx_destroy(ctx); + return 1; + } + NativeRtxAttentionDriver attention_driver(&attention); + NativeBf16Forward forward(&driver); + frt_graph graph = frt_graph_create(ctx, "native_diffusion", 10); + cudaStream_t stream = nullptr; + if (!graph || cudaStreamCreate(&stream) != cudaSuccess || + frt_graph_bind(graph, "noise", noise->buffer) != FRT_OK || + frt_graph_bind(graph, "encoder_k", cache_k->buffer) != FRT_OK || + frt_graph_bind(graph, "encoder_v", cache_v->buffer) != FRT_OK) { + frt_ctx_destroy(ctx); + return 1; + } + CaptureArgs capture{&forward, &weights, &workspace, &attention, + &attention_driver, start_step, steps, false, {}}; + const int capture_rc = frt_graph_capture( + graph, 10, record_diffusion, &capture); + if (capture_rc != FRT_OK || !capture.recorded) { + std::cerr << "diffusion capture failed: rc=" << capture_rc + << " status=" << capture.error << '\n'; + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + for (int i = 0; i < 100; ++i) { + if (cudaMemcpyAsync(frt_buffer_dptr(noise->buffer), host_noise.data(), + host_noise.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, stream) != cudaSuccess || + frt_graph_replay(graph, 10, stream_id) != FRT_OK) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + } + if (frt_graph_variant_count(graph) != 1 || + cudaStreamSynchronize(stream) != cudaSuccess) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + std::vector output(host_noise.size()); + if (cudaMemcpy(output.data(), frt_buffer_dptr(noise->buffer), + output.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) != cudaSuccess) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + std::ofstream file(argv[2], std::ios::binary | std::ios::trunc); + file.write(reinterpret_cast(output.data()), + static_cast(output.size() * + sizeof(std::uint16_t))); + const bool ok = file.good(); + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + if (!ok) return 1; + std::cout << "PASS native diffusion steps " << start_step << ".." + << start_step + steps - 1 << '\n'; + return 0; +} diff --git a/cpp/tests/pi05_native_dlopen_probe.cpp b/cpp/tests/pi05_native_dlopen_probe.cpp new file mode 100644 index 00000000..88d19964 --- /dev/null +++ b/cpp/tests/pi05_native_dlopen_probe.cpp @@ -0,0 +1,82 @@ +#include "flashrt/model_runtime.h" + +#include + +#include +#include +#include +#include + +namespace { + +std::string json_string(const std::string& value) { + std::string output = "\""; + for (char c : value) { + if (c == '\\' || c == '"') output.push_back('\\'); + output.push_back(c); + } + output.push_back('"'); + return output; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 5) { + std::cerr << "usage: pi05_native_dlopen_probe SO CHECKPOINT " + "TOKENIZER CYCLES\n"; + return 2; + } + const int cycles = std::atoi(argv[4]); + if (cycles <= 0) return 2; + std::ostringstream config; + config << "{\"io\":\"native_v2\",\"checkpoint_path\":" + << json_string(argv[2]) << ",\"tokenizer_model_path\":" + << json_string(argv[3]) + << ",\"state_prompt_mode\":\"fixed\"," + "\"max_prompt_tokens\":200,\"state_dim\":8," + "\"num_views\":2,\"chunk\":10,\"num_steps\":10," + "\"vision_pool_factor\":1}"; + for (int cycle = 0; cycle < cycles; ++cycle) { + void* library = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL); + if (!library) { + std::cerr << "dlopen failed: " << dlerror() << '\n'; + return 1; + } + auto open = reinterpret_cast( + dlsym(library, FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL)); + auto last_error = reinterpret_cast( + dlsym(library, "frt_pi05_native_open_last_error")); + if (!open || !last_error) { + std::cerr << "native factory symbols are missing\n"; + dlclose(library); + return 1; + } + frt_model_runtime_v1* model = nullptr; + const int rc = open(config.str().c_str(), &model); + if (rc != 0 || !model) { + std::cerr << "native open failed: rc=" << rc << " error=" + << last_error() << '\n'; + dlclose(library); + return 1; + } + if (model->abi_version != FRT_MODEL_RUNTIME_ABI_VERSION || + model->struct_size < sizeof(*model) || !model->retain || + !model->release) { + std::cerr << "native model ABI is invalid\n"; + if (model->release) model->release(model->owner); + dlclose(library); + return 1; + } + model->retain(model->owner); + model->release(model->owner); + model->release(model->owner); + if (dlclose(library) != 0) { + std::cerr << "dlclose failed: " << dlerror() << '\n'; + return 1; + } + std::cout << "cycle " << (cycle + 1) << " released\n"; + } + std::cout << "PASS native model dlopen/release/dlclose lifecycle\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_e2e_probe.cpp b/cpp/tests/pi05_native_e2e_probe.cpp new file mode 100644 index 00000000..0b3e5249 --- /dev/null +++ b/cpp/tests/pi05_native_e2e_probe.cpp @@ -0,0 +1,177 @@ +#include "flashrt/model_runtime.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +namespace { + +bool read_file(const std::string& path, std::vector* out) { + std::ifstream input(path, std::ios::binary); + if (!input) return false; + input.seekg(0, std::ios::end); + const std::streamoff size = input.tellg(); + if (size < 0) return false; + input.seekg(0, std::ios::beg); + std::vector data(static_cast(size)); + if (size && !input.read(reinterpret_cast(data.data()), size)) { + return false; + } + *out = std::move(data); + return true; +} + +bool write_file(const std::string& path, const void* data, std::size_t bytes) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) return false; + output.write(static_cast(data), + static_cast(bytes)); + return output.good(); +} + +std::string json_string(const std::string& value) { + std::string output = "\""; + for (char c : value) { + if (c == '\\' || c == '"') output.push_back('\\'); + output.push_back(c); + } + output.push_back('"'); + return output; +} + +int model_error(frt_model_runtime_v1* model, const char* message) { + std::cerr << message; + if (model && model->verbs.last_error) { + const char* detail = model->verbs.last_error(model->self); + if (detail && *detail) std::cerr << ": " << detail; + } + std::cerr << '\n'; + if (model) model->release(model->owner); + return 1; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 5) { + std::cerr << "usage: pi05_native_e2e_probe CHECKPOINT TOKENIZER " + "FIXTURE_DIR OUTPUT_DIR\n"; + return 2; + } + const std::string checkpoint = argv[1]; + const std::string tokenizer = argv[2]; + const std::string fixture = argv[3]; + const std::string output = argv[4]; + + std::ostringstream json; + json << "{\"io\":\"native_v2\",\"checkpoint_path\":" + << json_string(checkpoint) << ",\"tokenizer_model_path\":" + << json_string(tokenizer) + << ",\"state_prompt_mode\":\"fixed\"," + "\"max_prompt_tokens\":200,\"state_dim\":8," + "\"num_views\":2,\"chunk\":10,\"num_steps\":10," + "\"vision_pool_factor\":1}"; + frt_model_runtime_v1* model = nullptr; + const int open_rc = frt_model_runtime_open_v1(json.str().c_str(), &model); + if (open_rc != 0 || !model) { + std::cerr << "native open failed: rc=" << open_rc << " error=" + << frt_pi05_native_open_last_error() << '\n'; + return 1; + } + const char* names[] = { + "prompt", "state", "images", "noise", "actions", "actions_raw"}; + if (model->n_ports != 6) return model_error(model, "unexpected port count"); + for (std::uint64_t i = 0; i < model->n_ports; ++i) { + if (!model->ports[i].name || + std::strcmp(model->ports[i].name, names[i]) != 0) { + return model_error(model, "unexpected port schema"); + } + } + if (model->ports[4].dtype != FRT_RT_DTYPE_F32 || + model->ports[4].update != FRT_RT_PORT_STAGED || + model->ports[4].buffer != nullptr || + model->ports[4].bytes != 10 * 7 * sizeof(float) || + model->ports[5].dtype != FRT_RT_DTYPE_BF16 || + model->ports[5].update != FRT_RT_PORT_SWAP || + model->ports[5].buffer != model->ports[3].buffer) { + return model_error(model, "native action port contract mismatch"); + } + + std::vector prompt; + std::vector state; + std::vector image0; + std::vector image1; + std::vector noise; + if (!read_file(fixture + "/prompt.txt", &prompt) || + !read_file(fixture + "/state.f32", &state) || + !read_file(fixture + "/image_0.rgb", &image0) || + !read_file(fixture + "/image_1.rgb", &image1) || + !read_file(fixture + "/noise.bf16", &noise) || + state.size() != 8 * sizeof(float) || + image0.size() != 224 * 224 * 3 || image1.size() != image0.size() || + noise.size() != 10 * 32 * sizeof(std::uint16_t)) { + return model_error(model, "invalid E2E fixture"); + } + if (model->verbs.set_input(model->self, 0, prompt.data(), prompt.size(), + -1) != 0 || + model->verbs.set_input(model->self, 1, state.data(), state.size(), + -1) != 0) { + return model_error(model, "prompt/state staging failed"); + } + + frt_image_view views[2]{}; + const std::vector* images[] = {&image0, &image1}; + for (int i = 0; i < 2; ++i) { + views[i].struct_size = sizeof(frt_image_view); + views[i].pixel_format = FRT_RT_PIXEL_RGB8; + views[i].data = images[i]->data(); + views[i].bytes = images[i]->size(); + views[i].width = 224; + views[i].height = 224; + views[i].stride_bytes = 224 * 3; + } + if (model->verbs.set_input(model->self, 2, views, sizeof(views), -1) != 0) { + return model_error(model, "image staging failed"); + } + frt_buffer noise_buffer = model->ports[3].buffer; + if (!noise_buffer || frt_buffer_bytes(noise_buffer) != noise.size() || + cudaMemcpy(frt_buffer_dptr(noise_buffer), noise.data(), noise.size(), + cudaMemcpyHostToDevice) != cudaSuccess) { + return model_error(model, "noise upload failed"); + } + if (model->verbs.step(model->self) != 0) { + return model_error(model, "native infer failed"); + } + + std::vector actions(10 * 7); + std::uint64_t written = 0; + if (model->verbs.get_output(model->self, 4, actions.data(), + actions.size() * sizeof(float), &written, + -1) != 0 || + written != actions.size() * sizeof(float)) { + return model_error(model, "action output failed"); + } + std::vector raw(noise.size()); + if (cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), + raw.size(), cudaMemcpyDeviceToHost) != cudaSuccess) { + return model_error(model, "raw action download failed"); + } + if (!write_file(output + "/native_raw.bf16", raw.data(), raw.size()) || + !write_file(output + "/native_actions.f32", actions.data(), + actions.size() * sizeof(float))) { + return model_error(model, "native output write failed"); + } + model->release(model->owner); + std::cout << "PASS native real-episode fixture\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_encoder_layer_probe.cpp b/cpp/tests/pi05_native_encoder_layer_probe.cpp new file mode 100644 index 00000000..0aff2e74 --- /dev/null +++ b/cpp/tests/pi05_native_encoder_layer_probe.cpp @@ -0,0 +1,133 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +#include +#include +#include +#include + +namespace { + +struct CaptureArgs { + const flashrt::models::pi05::NativeBf16Forward* forward = nullptr; + const flashrt::models::pi05::NativeDeviceWeightStore* weights = nullptr; + flashrt::models::pi05::NativeWorkspace* workspace = nullptr; + flashrt::models::pi05::NativeRtxAttentionWorkspace* attention = nullptr; + const flashrt::models::pi05::NativeRtxAttentionDriver* attention_driver = + nullptr; + bool recorded = false; +}; + +void record_layer(void* user, void* stream) { + auto* args = static_cast(user); + args->recorded = args->forward + ->encoder_layer(0, *args->weights, args->workspace, args->attention, + args->attention_driver, + reinterpret_cast(stream)) + .ok_status(); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_encoder_layer_probe CHECKPOINT OUTPUT\n"; + return 2; + } + using namespace flashrt::models::pi05; + flashrt::loader::SafetensorsFile source; + if (!source.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << source.error() << '\n'; + return 2; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + NativeDeviceWeightStore weights(ctx); + NativeWeightMaterializer materializer(source, &weights); + flashrt::modalities::Status st = materializer.materialize_encoder_layer(0); + NativeWorkspace workspace(ctx); + NativeRtxAttentionWorkspace attention(ctx); + if (!st.ok_status() || + !workspace.allocate(NativeWorkspaceConfig{}).ok_status() || + !attention.allocate(NativeRtxAttentionConfig{}).ok_status() || + !attention.set_fixed_prompt_length(200).ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + const auto* encoder_x = workspace.find("encoder_x"); + std::vector host_x(712 * 2048, 0); + for (int row = 0; row < 712; ++row) { + for (int column = 0; column < 512; ++column) { + const float value = float((row + column) % 15 - 7) / 8.0f; + host_x[static_cast(row) * 2048 + column] = + flashrt::modalities::float_to_bfloat16(value); + } + } + if (!encoder_x || + cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), + host_x.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess) { + frt_ctx_destroy(ctx); + return 1; + } + NativeKernelDriver driver; + NativeRtxAttentionDriver attention_driver(&attention); + NativeBf16Forward forward(&driver); + frt_graph graph = frt_graph_create(ctx, "native_encoder_layer", 712); + cudaStream_t stream = nullptr; + if (!graph || cudaStreamCreate(&stream) != cudaSuccess || + frt_graph_bind(graph, "encoder_x", encoder_x->buffer) != FRT_OK) { + frt_ctx_destroy(ctx); + return 1; + } + CaptureArgs capture{&forward, &weights, &workspace, &attention, + &attention_driver, false}; + if (frt_graph_capture(graph, 712, record_layer, &capture) != FRT_OK || + !capture.recorded) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + for (int i = 0; i < 100; ++i) { + if (cudaMemcpyAsync(frt_buffer_dptr(encoder_x->buffer), host_x.data(), + host_x.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, stream) != cudaSuccess || + frt_graph_replay(graph, 712, stream_id) != FRT_OK) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + } + if (frt_graph_variant_count(graph) != 1 || + cudaStreamSynchronize(stream) != cudaSuccess) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + std::vector output(712 * 2048); + if (cudaMemcpy(output.data(), frt_buffer_dptr(encoder_x->buffer), + output.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) != cudaSuccess) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + std::ofstream file(argv[2], std::ios::binary | std::ios::trunc); + file.write(reinterpret_cast(output.data()), + static_cast(output.size() * sizeof(std::uint16_t))); + const bool ok = file.good(); + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + if (!ok) return 1; + std::cout << "PASS native encoder layer 0\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_encoder_probe.cpp b/cpp/tests/pi05_native_encoder_probe.cpp new file mode 100644 index 00000000..c3930e6b --- /dev/null +++ b/cpp/tests/pi05_native_encoder_probe.cpp @@ -0,0 +1,143 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +#include +#include +#include +#include + +namespace { + +struct CaptureArgs { + const flashrt::models::pi05::NativeBf16Forward* forward = nullptr; + const flashrt::models::pi05::NativeDeviceWeightStore* weights = nullptr; + flashrt::models::pi05::NativeWorkspace* workspace = nullptr; + flashrt::models::pi05::NativeRtxAttentionWorkspace* attention = nullptr; + const flashrt::models::pi05::NativeRtxAttentionDriver* attention_driver = + nullptr; + bool recorded = false; +}; + +void record_encoder(void* user, void* stream) { + auto* args = static_cast(user); + args->recorded = args->forward + ->encoder(*args->weights, args->workspace, args->attention, + args->attention_driver, + reinterpret_cast(stream)) + .ok_status(); +} + +bool write_buffer(std::ofstream* file, const void* device, std::size_t elements) { + std::vector host(elements); + if (cudaMemcpy(host.data(), device, elements * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + file->write(reinterpret_cast(host.data()), + static_cast(host.size() * + sizeof(std::uint16_t))); + return file->good(); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_encoder_probe CHECKPOINT OUTPUT\n"; + return 2; + } + using namespace flashrt::models::pi05; + flashrt::loader::SafetensorsFile source; + if (!source.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << source.error() << '\n'; + return 2; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + NativeDeviceWeightStore weights(ctx); + NativeWeightMaterializer materializer(source, &weights); + flashrt::modalities::Status st = flashrt::modalities::Status::ok(); + for (int layer = 0; layer < 18 && st.ok_status(); ++layer) { + st = materializer.materialize_encoder_layer(layer); + } + NativeWorkspace workspace(ctx); + NativeRtxAttentionWorkspace attention(ctx); + if (!st.ok_status() || + !workspace.allocate(NativeWorkspaceConfig{}).ok_status() || + !attention.allocate(NativeRtxAttentionConfig{}).ok_status() || + !attention.set_fixed_prompt_length(200).ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + const auto* encoder_x = workspace.find("encoder_x"); + const auto* encoder_q = attention.find("attn_enc_Q"); + std::vector host_x(712 * 2048, 0); + for (int row = 0; row < 712; ++row) { + for (int column = 0; column < 512; ++column) { + const float value = float((row + column) % 15 - 7) / 8.0f; + host_x[static_cast(row) * 2048 + column] = + flashrt::modalities::float_to_bfloat16(value); + } + } + if (!encoder_x || !encoder_q || + cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), + host_x.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess) { + frt_ctx_destroy(ctx); + return 1; + } + NativeKernelDriver driver; + NativeRtxAttentionDriver attention_driver(&attention); + NativeBf16Forward forward(&driver); + frt_graph graph = frt_graph_create(ctx, "native_encoder", 712); + cudaStream_t stream = nullptr; + if (!graph || cudaStreamCreate(&stream) != cudaSuccess || + frt_graph_bind(graph, "encoder_x", encoder_x->buffer) != FRT_OK || + frt_graph_bind(graph, "encoder_q", encoder_q->buffer) != FRT_OK) { + frt_ctx_destroy(ctx); + return 1; + } + CaptureArgs capture{&forward, &weights, &workspace, &attention, + &attention_driver, false}; + if (frt_graph_capture(graph, 712, record_encoder, &capture) != FRT_OK || + !capture.recorded) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + for (int i = 0; i < 100; ++i) { + if (cudaMemcpyAsync(frt_buffer_dptr(encoder_x->buffer), host_x.data(), + host_x.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, stream) != cudaSuccess || + frt_graph_replay(graph, 712, stream_id) != FRT_OK) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + } + if (frt_graph_variant_count(graph) != 1 || + cudaStreamSynchronize(stream) != cudaSuccess) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + std::ofstream file(argv[2], std::ios::binary | std::ios::trunc); + const bool ok = file && + write_buffer(&file, frt_buffer_dptr(encoder_x->buffer), 712 * 2048) && + write_buffer(&file, frt_buffer_dptr(encoder_q->buffer), 712 * 2048) && + write_buffer(&file, attention.encoder_k_layer_dptr(17), 712 * 256) && + write_buffer(&file, attention.encoder_v_layer_dptr(17), 712 * 256); + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + if (!ok) return 1; + std::cout << "PASS native encoder 18 layers\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_encoder_qkv_probe.cpp b/cpp/tests/pi05_native_encoder_qkv_probe.cpp new file mode 100644 index 00000000..ce6b22b6 --- /dev/null +++ b/cpp/tests/pi05_native_encoder_qkv_probe.cpp @@ -0,0 +1,99 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +#include +#include +#include +#include + +namespace { + +bool write_device(const std::string& path, const void* device, + std::size_t elements) { + std::vector host(elements); + if (cudaMemcpy(host.data(), device, host.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + std::ofstream file(path, std::ios::binary | std::ios::trunc); + file.write(reinterpret_cast(host.data()), + static_cast(host.size() * sizeof(std::uint16_t))); + return file.good(); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_encoder_qkv_probe CHECKPOINT OUTPUT\n"; + return 2; + } + using namespace flashrt::models::pi05; + flashrt::loader::SafetensorsFile source; + if (!source.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << source.error() << '\n'; + return 2; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + NativeDeviceWeightStore weights(ctx); + NativeWeightMaterializer materializer(source, &weights); + flashrt::modalities::Status st = materializer.materialize_encoder_layer(17); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + NativeWorkspace workspace(ctx); + if (!workspace.allocate(NativeWorkspaceConfig{}).ok_status()) { + frt_ctx_destroy(ctx); + return 1; + } + NativeRtxAttentionWorkspace attention(ctx); + if (!attention.allocate(NativeRtxAttentionConfig{}).ok_status()) { + frt_ctx_destroy(ctx); + return 1; + } + const auto* encoder_x = workspace.find("encoder_x"); + if (!encoder_x) { + frt_ctx_destroy(ctx); + return 1; + } + std::vector host_x(712 * 2048, 0); + for (int row = 0; row < 712; ++row) { + for (int column = 0; column < 512; ++column) { + const float value = float((row + column) % 15 - 7) / 8.0f; + host_x[static_cast(row) * 2048 + column] = + flashrt::modalities::float_to_bfloat16(value); + } + } + if (cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), + host_x.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess) { + frt_ctx_destroy(ctx); + return 1; + } + NativeKernelDriver driver; + NativeBf16Forward forward(&driver); + st = forward.encoder_qkv(17, weights, &workspace, &attention, 0); + if (!st.ok_status() || cudaDeviceSynchronize() != cudaSuccess) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + const auto* query = attention.find("attn_enc_Q"); + const std::string prefix = argv[2]; + const bool ok = query && + write_device(prefix + ".q.bin", frt_buffer_dptr(query->buffer), + 712 * 2048) && + write_device(prefix + ".k.bin", attention.encoder_k_layer_dptr(17), + 712 * 256) && + write_device(prefix + ".v.bin", attention.encoder_v_layer_dptr(17), + 712 * 256); + frt_ctx_destroy(ctx); + if (!ok) return 1; + std::cout << "PASS native encoder QKV layer 17\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_graph_probe.cpp b/cpp/tests/pi05_native_graph_probe.cpp new file mode 100644 index 00000000..c561bb50 --- /dev/null +++ b/cpp/tests/pi05_native_graph_probe.cpp @@ -0,0 +1,103 @@ +#include "flashrt/cpp/models/pi05/native_graph_owner.h" +#include "flashrt/cpp/modalities/types.h" + +#include + +#include +#include +#include + +namespace { + +std::vector download(frt_buffer buffer) { + std::vector host(frt_buffer_bytes(buffer) / + sizeof(std::uint16_t)); + if (cudaMemcpy(host.data(), frt_buffer_dptr(buffer), + frt_buffer_bytes(buffer), cudaMemcpyDeviceToHost) != + cudaSuccess) { + host.clear(); + } + return host; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 2) { + std::cerr << "usage: pi05_native_graph_probe CHECKPOINT\n"; + return 2; + } + using namespace flashrt::models::pi05; + flashrt::modalities::Status st; + std::unique_ptr owner = NativeGraphOwner::create( + argv[1], NativeGraphConfig{}, &st); + if (!owner) { + std::cerr << st.message << '\n'; + return 1; + } + const NativeWorkspaceBuffer* images = + owner->workspace().find("observation_images_normalized"); + const NativeWorkspaceBuffer* prompt = + owner->workspace().find("prompt_embedding"); + const NativeWorkspaceBuffer* noise = + owner->workspace().find("diffusion_noise"); + if (!images || !prompt || !noise || !owner->infer_graph() || + frt_graph_variant_count(owner->infer_graph()) != 1 || + owner->stream_id() < 0 || !owner->native_stream()) { + return 1; + } + std::vector host_images( + frt_buffer_bytes(images->buffer) / sizeof(std::uint16_t)); + std::vector host_prompt( + frt_buffer_bytes(prompt->buffer) / sizeof(std::uint16_t)); + std::vector host_noise( + frt_buffer_bytes(noise->buffer) / sizeof(std::uint16_t)); + for (std::size_t i = 0; i < host_images.size(); ++i) { + host_images[i] = flashrt::modalities::float_to_bfloat16( + static_cast(static_cast(i % 257) - 128) / 128.0f); + } + for (std::size_t i = 0; i < host_prompt.size(); ++i) { + host_prompt[i] = flashrt::modalities::float_to_bfloat16( + static_cast(static_cast(i % 31) - 15) / 32.0f); + } + for (std::size_t i = 0; i < host_noise.size(); ++i) { + host_noise[i] = flashrt::modalities::float_to_bfloat16( + static_cast(static_cast(i % 23) - 11) / 12.0f); + } + if (cudaMemcpy(frt_buffer_dptr(images->buffer), host_images.data(), + frt_buffer_bytes(images->buffer), + cudaMemcpyHostToDevice) != cudaSuccess || + cudaMemcpy(frt_buffer_dptr(prompt->buffer), host_prompt.data(), + frt_buffer_bytes(prompt->buffer), + cudaMemcpyHostToDevice) != cudaSuccess || + !owner->set_prompt_length(37).ok_status()) { + return 1; + } + const std::size_t allocation_count = owner->workspace().allocation_count(); + if (cudaMemcpy(frt_buffer_dptr(noise->buffer), host_noise.data(), + frt_buffer_bytes(noise->buffer), + cudaMemcpyHostToDevice) != cudaSuccess || + owner->replay() != FRT_OK || !owner->synchronize().ok_status()) { + return 1; + } + const std::vector expected = download(noise->buffer); + if (expected.empty()) return 1; + for (int replay = 0; replay < 100; ++replay) { + if (cudaMemcpyAsync( + frt_buffer_dptr(noise->buffer), host_noise.data(), + frt_buffer_bytes(noise->buffer), cudaMemcpyHostToDevice, + static_cast(owner->native_stream())) != + cudaSuccess || + owner->replay() != FRT_OK) { + return 1; + } + } + if (!owner->synchronize().ok_status() || + frt_graph_variant_count(owner->infer_graph()) != 1 || + owner->workspace().allocation_count() != allocation_count || + download(noise->buffer) != expected) { + return 1; + } + std::cout << "PASS native full graph 100 replays\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp new file mode 100644 index 00000000..edb0043d --- /dev/null +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -0,0 +1,337 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/modalities/types.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_open_probe CHECKPOINT TOKENIZER\n"; + return 2; + } + int replay_count = 1; + const char* replay_env = std::getenv("FLASHRT_PROFILE_REPLAYS"); + if (replay_env) { + char* end = nullptr; + const long parsed = std::strtol(replay_env, &end, 10); + if (!end || *end != '\0' || parsed <= 0 || parsed > 100000) { + std::cerr << "FLASHRT_PROFILE_REPLAYS must be in [1, 100000]\n"; + return 2; + } + replay_count = static_cast(parsed); + } + const bool profile_service_loop = + std::getenv("FLASHRT_PROFILE_SERVICE_LOOP") != nullptr; + if (profile_service_loop && !replay_env) { + std::cerr << "FLASHRT_PROFILE_SERVICE_LOOP requires " + "FLASHRT_PROFILE_REPLAYS\n"; + return 2; + } + int hot_state_updates = 0; + const char* hot_updates_env = std::getenv("FLASHRT_HOT_STATE_UPDATES"); + if (hot_updates_env) { + char* end = nullptr; + const long parsed = std::strtol(hot_updates_env, &end, 10); + if (!end || *end != '\0' || parsed <= 0 || parsed > 100000) { + std::cerr << "FLASHRT_HOT_STATE_UPDATES must be in [1, 100000]\n"; + return 2; + } + hot_state_updates = static_cast(parsed); + } + double hot_state_p99_limit_us = 0.0; + const char* hot_limit_env = std::getenv("FLASHRT_HOT_STATE_P99_US"); + if (hot_limit_env) { + char* end = nullptr; + hot_state_p99_limit_us = std::strtod(hot_limit_env, &end); + if (!end || *end != '\0' || !std::isfinite(hot_state_p99_limit_us) || + hot_state_p99_limit_us <= 0.0) { + std::cerr << "FLASHRT_HOT_STATE_P99_US must be positive\n"; + return 2; + } + } + std::ostringstream json; + json << "{\"io\":\"native_v2\",\"checkpoint_path\":\"" + << argv[1] << "\",\"tokenizer_model_path\":\"" << argv[2] + << "\",\"state_prompt_mode\":\"fixed\"," + << "\"max_prompt_tokens\":200,\"state_dim\":8," + << "\"num_views\":2,\"chunk\":10,\"num_steps\":10," + << "\"vision_pool_factor\":1}"; + frt_model_runtime_v1* model = nullptr; + const int open_rc = frt_model_runtime_open_v1(json.str().c_str(), &model); + if (open_rc != 0 || !model) { + std::cerr << "native open failed: rc=" << open_rc << " error=" + << frt_pi05_native_open_last_error() << '\n'; + return 1; + } + const char* port_names[] = { + "prompt", "state", "images", "noise", "actions", "actions_raw"}; + const frt_runtime_export_v1* exp = model->exp; + int active_device = 0; + cudaDeviceProp active_properties{}; + const bool device_identity_ok = + cudaGetDevice(&active_device) == cudaSuccess && + cudaGetDeviceProperties(&active_properties, active_device) == + cudaSuccess; + const std::string hardware_id = device_identity_ok + ? "sm" + std::to_string(active_properties.major * 10 + + active_properties.minor) + : std::string(); + const std::string hardware_identity = "hardware=" + hardware_id; + const std::string hardware_manifest = + "\"hardware\":\"" + hardware_id + "\""; + bool ok = model->abi_version == FRT_MODEL_RUNTIME_ABI_VERSION && + model->struct_size == sizeof(frt_model_runtime_v1) && exp && + exp->abi_version == FRT_RUNTIME_ABI_VERSION && + exp->struct_size == sizeof(frt_runtime_export_v1) && + model->n_ports == 6 && model->n_stages == 1 && + exp->n_graphs == 1 && exp->n_streams == 1 && + exp->n_capsule_regions == 1 && exp->n_buffers == 7 && + exp->fingerprint != 0 && exp->identity && + std::strstr(exp->identity, "producer=native") && + device_identity_ok && + std::strstr(exp->identity, hardware_identity.c_str()) && + exp->manifest_json && + std::strstr(exp->manifest_json, hardware_manifest.c_str()) && + std::strstr(exp->identity, "weights_sha256=") && + std::strstr(exp->identity, "tokenizer_sha256=") && + model->stages[0].graph == 0 && + frt_graph_variant_count(exp->graphs[0].handle) == 1; + for (std::uint64_t i = 0; i < model->n_ports; ++i) { + ok = ok && std::strcmp(model->ports[i].name, port_names[i]) == 0; + } + ok = ok && + std::strcmp(exp->capsule_regions[0].name, "rollout_boundary") == 0 && + model->ports[0].modality == FRT_RT_MOD_TEXT && + model->ports[0].update == FRT_RT_PORT_STAGED && + model->ports[1].modality == FRT_RT_MOD_STATE && + model->ports[2].modality == FRT_RT_MOD_IMAGE && + model->ports[3].update == FRT_RT_PORT_SWAP && + model->ports[4].direction == FRT_RT_PORT_OUT && + model->ports[4].dtype == FRT_RT_DTYPE_F32 && + model->ports[4].update == FRT_RT_PORT_STAGED && + model->ports[4].buffer == nullptr && + model->ports[4].bytes == 10 * 7 * sizeof(float) && + model->ports[5].dtype == FRT_RT_DTYPE_BF16 && + model->ports[5].update == FRT_RT_PORT_SWAP && + model->ports[5].buffer == model->ports[3].buffer && + model->ports[5].offset == model->ports[3].offset && + model->ports[5].bytes == model->ports[3].bytes; + if (!ok) { + std::cerr << "native schema validation failed\n"; + model->release(model->owner); + return 1; + } + const char* schema_output = std::getenv("FLASHRT_SCHEMA_OUTPUT"); + if (schema_output && schema_output[0] != '\0') { + std::ofstream output(schema_output); + std::istringstream identity(exp->identity ? exp->identity : ""); + std::string line; + while (std::getline(identity, line)) { + if (line.compare(0, 7, "region:") == 0 || + line.compare(0, 5, "port:") == 0 || + line.compare(0, 6, "stage:") == 0) { + output << line << '\n'; + } + } + if (!output) { + std::cerr << "native schema output failed\n"; + model->release(model->owner); + return 1; + } + if (std::getenv("FLASHRT_SCHEMA_ONLY")) { + model->release(model->owner); + std::cout << "PASS native schema export\n"; + return 0; + } + } + if (model->verbs.prepare(model->self, 0, 0) != 0 || + model->verbs.prepare(model->self, 0, 1) != -2) { + std::cerr << "native prepare validation failed\n"; + model->release(model->owner); + return 1; + } + + const std::string prompt = "pick up the black bowl"; + float state[8] = {0.1f, -0.2f, 0.3f, -0.4f, + 0.5f, -0.6f, 0.7f, -0.8f}; + if (model->verbs.set_input(model->self, 0, prompt.data(), prompt.size(), + -1) != 0 || + model->verbs.set_input(model->self, 1, state, sizeof(state), -1) != 0) { + std::cerr << "native prompt/state staging failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + if (hot_state_updates) { + constexpr int kWarmUpdates = 20; + std::vector hot_state_latencies; + hot_state_latencies.reserve(hot_state_updates); + for (int update = -kWarmUpdates; update < hot_state_updates; ++update) { + for (int dim = 0; dim < 8; ++dim) { + state[dim] = std::sin( + static_cast((update + kWarmUpdates) * 8 + dim) * + 0.017f); + } + const auto begin = std::chrono::steady_clock::now(); + const int rc = model->verbs.set_input( + model->self, 1, state, sizeof(state), -1); + const auto end = std::chrono::steady_clock::now(); + if (rc != 0) { + std::cerr << "native hot state update failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + if (update >= 0) { + hot_state_latencies.push_back( + std::chrono::duration(end - begin) + .count()); + } + } + std::sort(hot_state_latencies.begin(), hot_state_latencies.end()); + const std::size_t p99_index = + (hot_state_latencies.size() * 99 + 99) / 100 - 1; + const double p50 = hot_state_latencies[hot_state_latencies.size() / 2]; + const double p99 = hot_state_latencies[p99_index]; + const double maximum = hot_state_latencies.back(); + std::cout << "hot state updates: n=" << hot_state_latencies.size() + << " p50_us=" << p50 << " p99_us=" << p99 + << " max_us=" << maximum << '\n'; + if ((hot_state_p99_limit_us > 0.0 && + p99 > hot_state_p99_limit_us) || + frt_graph_variant_count(exp->graphs[0].handle) != 1) { + std::cerr << "native hot state update gate failed\n"; + model->release(model->owner); + return 1; + } + } + std::vector rgb0(224 * 224 * 3); + std::vector rgb1(rgb0.size()); + for (std::size_t i = 0; i < rgb0.size(); ++i) { + rgb0[i] = static_cast(i % 251); + rgb1[i] = static_cast((3 * i + 17) % 251); + } + frt_image_view views[2]{}; + for (int i = 0; i < 2; ++i) { + views[i].struct_size = sizeof(frt_image_view); + views[i].pixel_format = FRT_RT_PIXEL_RGB8; + views[i].data = i ? static_cast(rgb1.data()) + : static_cast(rgb0.data()); + views[i].bytes = rgb0.size(); + views[i].width = 224; + views[i].height = 224; + views[i].stride_bytes = 224 * 3; + } + if (model->verbs.set_input(model->self, 2, views, sizeof(views), -1) != 0) { + std::cerr << "native image staging failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + frt_buffer noise = model->ports[3].buffer; + std::vector host_noise(10 * 32); + for (std::size_t i = 0; i < host_noise.size(); ++i) { + host_noise[i] = flashrt::modalities::float_to_bfloat16( + static_cast(static_cast(i % 23) - 11) / 12.0f); + } + const bool profile_range = replay_env || + std::getenv("FLASHRT_PROFILE_RANGE") != nullptr; + float actions[10 * 7]{}; + std::uint64_t written = 0; + if (!noise || model->verbs.set_input(model->self, 3, host_noise.data(), + host_noise.size() * 2, -1) != -3 || + cudaMemcpy(frt_buffer_dptr(noise), host_noise.data(), + host_noise.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess || + (profile_range && cudaProfilerStart() != cudaSuccess)) { + std::cerr << "native step failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + int step_rc = 0; + cudaError_t upload_rc = cudaSuccess; + for (int replay = 0; replay < replay_count; ++replay) { + if (profile_service_loop) { + for (int dim = 0; dim < 8; ++dim) { + state[dim] = std::sin( + static_cast(replay * 8 + dim) * 0.017f); + } + const char* live_prompt = replay % 2 == 0 + ? "pick up the black bowl" + : "move the black bowl to the plate"; + if (model->verbs.set_input( + model->self, 0, live_prompt, std::strlen(live_prompt), + -1) != 0 || + model->verbs.set_input( + model->self, 1, state, sizeof(state), -1) != 0 || + model->verbs.set_input( + model->self, 2, views, sizeof(views), -1) != 0) { + step_rc = -1; + break; + } + } + if (replay != 0 || profile_service_loop) { + upload_rc = cudaMemcpy( + frt_buffer_dptr(noise), host_noise.data(), + host_noise.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice); + if (upload_rc != cudaSuccess) break; + } + step_rc = model->verbs.step(model->self); + if (step_rc != 0) break; + if (profile_service_loop && + (model->verbs.get_output( + model->self, 4, actions, sizeof(actions), &written, -1) != 0 || + written != sizeof(actions))) { + step_rc = -1; + break; + } + } + const cudaError_t sync_rc = cudaDeviceSynchronize(); + const cudaError_t profiler_rc = + profile_range ? cudaProfilerStop() : cudaSuccess; + if (step_rc != 0 || upload_rc != cudaSuccess || sync_rc != cudaSuccess || + profiler_rc != cudaSuccess || + frt_graph_variant_count(exp->graphs[0].handle) != 1) { + std::cerr << "native step failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + if (model->verbs.get_output(model->self, 4, actions, sizeof(actions), + &written, -1) != 0 || + written != sizeof(actions)) { + std::cerr << "native action output failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + for (float value : actions) { + if (!std::isfinite(value)) { + std::cerr << "native action output is not finite\n"; + model->release(model->owner); + return 1; + } + } + model->retain(model->owner); + model->release(model->owner); + model->release(model->owner); + std::cout << "PASS native open_v1 full lifecycle\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_quant_probe.cpp b/cpp/tests/pi05_native_quant_probe.cpp new file mode 100644 index 00000000..8a5e12f7 --- /dev/null +++ b/cpp/tests/pi05_native_quant_probe.cpp @@ -0,0 +1,130 @@ +#include "flashrt/cpp/models/pi05/native_quantization.h" + +#include +#include +#include +#include +#include + +namespace { + +using flashrt::loader::SafetensorsFile; +using flashrt::models::pi05::NativeFloatTensor; +using flashrt::models::pi05::NativeFp8Tensor; +using flashrt::models::pi05::NativeInt8Tensor; +using flashrt::modalities::Status; + +constexpr const char* kDecoder = + "paligemma_with_expert.gemma_expert.model.layers.0"; + +bool load(const SafetensorsFile& file, const std::string& key, + NativeFloatTensor* out) { + const Status st = + flashrt::models::pi05::load_native_float_tensor(file, key, out); + if (!st.ok_status()) std::cerr << st.message << '\n'; + return st.ok_status(); +} + +bool decoder_qkv(const SafetensorsFile& file, NativeFloatTensor* out) { + NativeFloatTensor q; + NativeFloatTensor k; + NativeFloatTensor v; + NativeFloatTensor qr; + NativeFloatTensor kr; + NativeFloatTensor vr; + NativeFloatTensor qi; + NativeFloatTensor ki; + return load(file, std::string(kDecoder) + ".self_attn.q_proj.weight", + &q) && + load(file, std::string(kDecoder) + ".self_attn.k_proj.weight", + &k) && + load(file, std::string(kDecoder) + ".self_attn.v_proj.weight", + &v) && + flashrt::models::pi05::native_round_to_bf16_float(q, &qr) + .ok_status() && + flashrt::models::pi05::native_round_to_bf16_float(k, &kr) + .ok_status() && + flashrt::models::pi05::native_round_to_bf16_float(v, &vr) + .ok_status() && + flashrt::models::pi05::native_interleave_qk_rows(qr, 8, &qi) + .ok_status() && + flashrt::models::pi05::native_interleave_qk_rows(kr, 1, &ki) + .ok_status() && + flashrt::models::pi05::native_concat_rows_transpose( + {&qi, &ki, &vr}, out) + .ok_status(); +} + +std::uint64_t fnv1a(const void* data, std::size_t bytes) { + std::uint64_t hash = 14695981039346656037ull; + const auto* src = static_cast(data); + for (std::size_t i = 0; i < bytes; ++i) { + hash ^= src[i]; + hash *= 1099511628211ull; + } + return hash; +} + +void print_shape(const std::vector& shape) { + for (std::size_t i = 0; i < shape.size(); ++i) { + if (i) std::cout << ','; + std::cout << shape[i]; + } +} + +void print_result(const std::vector& shape, + const void* values, std::size_t value_bytes, + const std::vector& scales) { + std::cout << "shape="; + print_shape(shape); + std::cout << " values_fnv=" << std::hex << std::setw(16) + << std::setfill('0') << fnv1a(values, value_bytes) + << " scale_shape=" << std::dec << scales.size() + << " scales_fnv=" << std::hex << std::setw(16) + << fnv1a(scales.data(), scales.size() * sizeof(float)) << '\n'; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_quant_probe CHECKPOINT OP\n"; + return 2; + } + SafetensorsFile file; + if (!file.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << file.error() << '\n'; + return 2; + } + NativeFloatTensor weight; + if (!decoder_qkv(file, &weight)) return 1; + const std::string op = argv[2]; + if (op == "decoder_qkv0_fp8_kn" || op == "decoder_qkv0_fp8_nk") { + NativeFp8Tensor output; + const bool transpose = op.back() == 'k'; + const Status st = flashrt::models::pi05::native_quantize_fp8_e4m3( + weight, transpose, &output); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + return 1; + } + print_result(output.shape, output.values.data(), output.values.size(), + {output.scale}); + return 0; + } + if (op == "decoder_qkv0_int8") { + NativeInt8Tensor output; + const Status st = + flashrt::models::pi05::native_quantize_int8_per_output( + weight, &output); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + return 1; + } + print_result(output.shape, output.values.data(), output.values.size(), + output.scales); + return 0; + } + std::cerr << "unknown quantization probe operation: " << op << '\n'; + return 2; +} diff --git a/cpp/tests/pi05_native_rope_probe.cpp b/cpp/tests/pi05_native_rope_probe.cpp new file mode 100644 index 00000000..16cb0a6f --- /dev/null +++ b/cpp/tests/pi05_native_rope_probe.cpp @@ -0,0 +1,69 @@ +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#include + +#include +#include +#include +#include +#include + +namespace { + +std::uint64_t fnv1a(const std::vector& values) { + std::uint64_t hash = 14695981039346656037ull; + const auto* bytes = reinterpret_cast(values.data()); + for (std::size_t i = 0; i < values.size() * sizeof(std::uint16_t); ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ull; + } + return hash; +} + +std::vector download( + const flashrt::models::pi05::NativeWorkspaceBuffer& buffer) { + std::vector values( + frt_buffer_bytes(buffer.buffer) / sizeof(std::uint16_t)); + if (cudaMemcpy(values.data(), frt_buffer_dptr(buffer.buffer), + values.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) != cudaSuccess) { + return {}; + } + return values; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 6) { + std::cerr << "usage: pi05_native_rope_probe VIEWS MAX_PROMPT CHUNK " + "POOL PROMPT\n"; + return 2; + } + flashrt::models::pi05::NativeWorkspaceConfig config; + config.num_views = std::stoi(argv[1]); + config.max_prompt_tokens = std::stoi(argv[2]); + config.chunk_size = std::stoi(argv[3]); + config.vision_pool_factor = std::stoi(argv[4]); + const int prompt = std::stoi(argv[5]); + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + flashrt::models::pi05::NativeWorkspace workspace(ctx); + if (!workspace.allocate(config).ok_status() || + !workspace.update_decoder_rope(prompt).ok_status()) { + frt_ctx_destroy(ctx); + return 1; + } + const std::vector encoder = + download(*workspace.find("encoder_rope_weights")); + const std::vector decoder = + download(*workspace.find("decoder_rope_weights")); + std::cout << "encoder_shape=" << workspace.encoder_sequence() << ",256" + << " encoder_fnv=" << std::hex << std::setw(16) + << std::setfill('0') << fnv1a(encoder) + << " decoder_shape=" << std::dec << config.chunk_size << ",256" + << " decoder_fnv=" << std::hex << std::setw(16) + << fnv1a(decoder) << '\n'; + frt_ctx_destroy(ctx); + return 0; +} diff --git a/cpp/tests/pi05_native_style_probe.cpp b/cpp/tests/pi05_native_style_probe.cpp new file mode 100644 index 00000000..0848ebfc --- /dev/null +++ b/cpp/tests/pi05_native_style_probe.cpp @@ -0,0 +1,85 @@ +#include "flashrt/cpp/models/pi05/native_style_precompute.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +#include +#include +#include + +namespace { + +bool write_buffer(const std::string& path, + const flashrt::models::pi05::NativeWorkspaceBuffer& buffer) { + const std::size_t bytes = frt_buffer_bytes(buffer.buffer); + std::vector host(bytes); + if (cudaMemcpy(host.data(), frt_buffer_dptr(buffer.buffer), bytes, + cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + std::ofstream file(path, std::ios::binary | std::ios::trunc); + file.write(reinterpret_cast(host.data()), + static_cast(host.size())); + return file.good(); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_style_probe CHECKPOINT OUTPUT_PREFIX\n"; + return 2; + } + flashrt::loader::SafetensorsFile source; + if (!source.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << source.error() << '\n'; + return 2; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + flashrt::models::pi05::NativeDeviceWeightStore weights(ctx); + flashrt::models::pi05::NativeWeightMaterializer materializer(source, + &weights); + for (int layer = 0; layer < 18; ++layer) { + const flashrt::modalities::Status st = + materializer.materialize_decoder_layer(layer, false); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + } + flashrt::modalities::Status st = + materializer.materialize_decoder_globals(10); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + flashrt::models::pi05::NativeWorkspace workspace(ctx); + flashrt::models::pi05::NativeWorkspaceConfig config; + if (!workspace.allocate(config).ok_status()) { + frt_ctx_destroy(ctx); + return 1; + } + flashrt::models::pi05::NativeKernelDriver driver; + flashrt::models::pi05::NativeStylePrecomputer precomputer(&driver); + st = precomputer.run(weights, &workspace, 0); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + const std::string prefix = argv[2]; + for (const char* name : {"decoder_time_emb", "decoder_style_attn", + "decoder_style_ffn", "decoder_style_final"}) { + const auto* buffer = workspace.find(name); + if (!buffer || !write_buffer(prefix + "." + name + ".bin", *buffer)) { + frt_ctx_destroy(ctx); + return 1; + } + } + std::cout << "PASS native decoder style precompute\n"; + frt_ctx_destroy(ctx); + return 0; +} diff --git a/cpp/tests/pi05_native_vision_probe.cpp b/cpp/tests/pi05_native_vision_probe.cpp new file mode 100644 index 00000000..4ef4771b --- /dev/null +++ b/cpp/tests/pi05_native_vision_probe.cpp @@ -0,0 +1,146 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" + +#include + +#include +#include +#include +#include + +namespace { + +struct CaptureArgs { + const flashrt::models::pi05::NativeBf16Forward* forward = nullptr; + const flashrt::models::pi05::NativeDeviceWeightStore* weights = nullptr; + flashrt::models::pi05::NativeWorkspace* workspace = nullptr; + flashrt::models::pi05::NativeRtxAttentionWorkspace* attention = nullptr; + const flashrt::models::pi05::NativeRtxAttentionDriver* attention_driver = + nullptr; + bool recorded = false; + std::string error; +}; + +void record_vision(void* user, void* stream) { + auto* args = static_cast(user); + const flashrt::modalities::Status st = args->forward->vision( + *args->weights, args->workspace, args->attention, + args->attention_driver, reinterpret_cast(stream)); + args->recorded = st.ok_status(); + args->error = st.message; +} + +bool write_buffer(std::ofstream* file, const void* device, + std::size_t elements) { + std::vector host(elements); + if (cudaMemcpy(host.data(), device, elements * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + file->write(reinterpret_cast(host.data()), + static_cast(host.size() * + sizeof(std::uint16_t))); + return file->good(); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_vision_probe CHECKPOINT OUTPUT\n"; + return 2; + } + using namespace flashrt::models::pi05; + flashrt::loader::SafetensorsFile source; + if (!source.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << source.error() << '\n'; + return 2; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) return 1; + NativeDeviceWeightStore weights(ctx); + NativeWeightMaterializer materializer(source, &weights); + flashrt::modalities::Status st = materializer.materialize_vision_globals(); + for (int layer = 0; layer < 27 && st.ok_status(); ++layer) { + st = materializer.materialize_vision_layer(layer); + } + NativeWorkspace workspace(ctx); + NativeRtxAttentionWorkspace attention(ctx); + if (!st.ok_status() || + !workspace.allocate(NativeWorkspaceConfig{}).ok_status() || + !workspace.expand_vision_position_embedding(weights).ok_status() || + !attention.allocate(NativeRtxAttentionConfig{}).ok_status()) { + std::cerr << st.message << '\n'; + frt_ctx_destroy(ctx); + return 1; + } + const auto* images = workspace.find("observation_images_normalized"); + const auto* vision_x = workspace.find("vision_x"); + const auto* encoder_x = workspace.find("encoder_x"); + std::vector host_images(2 * 224 * 224 * 3); + for (std::size_t i = 0; i < host_images.size(); ++i) { + const float value = static_cast(static_cast(i % 257) - 128) / + 128.0f; + host_images[i] = flashrt::modalities::float_to_bfloat16(value); + } + if (!images || !vision_x || !encoder_x || + cudaMemcpy(frt_buffer_dptr(images->buffer), host_images.data(), + host_images.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess) { + frt_ctx_destroy(ctx); + return 1; + } + NativeKernelDriver driver; + NativeRtxAttentionDriver attention_driver(&attention); + NativeBf16Forward forward(&driver); + frt_graph graph = frt_graph_create(ctx, "native_vision", 512); + cudaStream_t stream = nullptr; + if (!graph || cudaStreamCreate(&stream) != cudaSuccess || + frt_graph_bind(graph, "images", images->buffer) != FRT_OK || + frt_graph_bind(graph, "vision_x", vision_x->buffer) != FRT_OK || + frt_graph_bind(graph, "encoder_x", encoder_x->buffer) != FRT_OK) { + frt_ctx_destroy(ctx); + return 1; + } + CaptureArgs capture{&forward, &weights, &workspace, &attention, + &attention_driver, false, {}}; + const int capture_rc = frt_graph_capture( + graph, 512, record_vision, &capture); + if (capture_rc != FRT_OK || !capture.recorded) { + std::cerr << "vision capture failed: rc=" << capture_rc + << " status=" << capture.error << '\n'; + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + for (int i = 0; i < 100; ++i) { + if (cudaMemcpyAsync(frt_buffer_dptr(images->buffer), host_images.data(), + host_images.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, stream) != cudaSuccess || + frt_graph_replay(graph, 512, stream_id) != FRT_OK) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + } + if (frt_graph_variant_count(graph) != 1 || + cudaStreamSynchronize(stream) != cudaSuccess) { + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + return 1; + } + std::ofstream file(argv[2], std::ios::binary | std::ios::trunc); + const bool ok = file && + write_buffer(&file, frt_buffer_dptr(vision_x->buffer), 512 * 1152) && + write_buffer(&file, frt_buffer_dptr(encoder_x->buffer), 512 * 2048); + frt_graph_destroy(graph); + cudaStreamDestroy(stream); + frt_ctx_destroy(ctx); + if (!ok) return 1; + std::cout << "PASS native vision 27 layers\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_weight_probe.cpp b/cpp/tests/pi05_native_weight_probe.cpp new file mode 100644 index 00000000..a8c4d036 --- /dev/null +++ b/cpp/tests/pi05_native_weight_probe.cpp @@ -0,0 +1,209 @@ +#include "flashrt/cpp/models/pi05/native_weight_ops.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using flashrt::loader::SafetensorsFile; +using flashrt::models::pi05::NativeBf16Tensor; +using flashrt::models::pi05::NativeFloatTensor; +using flashrt::modalities::Status; + +constexpr const char* kVision = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; +constexpr const char* kEncoder = + "paligemma_with_expert.paligemma.model.language_model.layers.0"; +constexpr const char* kDecoder = + "paligemma_with_expert.gemma_expert.model.layers.0"; + +bool load(const SafetensorsFile& file, const std::string& key, + NativeFloatTensor* out) { + const Status st = + flashrt::models::pi05::load_native_float_tensor(file, key, out); + if (!st.ok_status()) std::cerr << st.message << '\n'; + return st.ok_status(); +} + +bool round_bf16(const NativeFloatTensor& input, NativeFloatTensor* out) { + return flashrt::models::pi05::native_round_to_bf16_float(input, out) + .ok_status(); +} + +bool finish(const NativeFloatTensor& input, NativeBf16Tensor* out) { + return flashrt::models::pi05::native_to_bf16(input, out).ok_status(); +} + +bool patch(const SafetensorsFile& file, NativeBf16Tensor* out) { + NativeFloatTensor source; + NativeFloatTensor rounded; + NativeFloatTensor transformed; + return load(file, std::string(kVision) + + ".embeddings.patch_embedding.weight", + &source) && + round_bf16(source, &rounded) && + flashrt::models::pi05::native_patch_oihw_to_hwio( + rounded, &transformed).ok_status() && + finish(transformed, out); +} + +bool qkv(const SafetensorsFile& file, const std::string& prefix, + bool fold_rms, NativeBf16Tensor* out) { + NativeFloatTensor q; + NativeFloatTensor k; + NativeFloatTensor v; + if (!load(file, prefix + ".self_attn.q_proj.weight", &q) || + !load(file, prefix + ".self_attn.k_proj.weight", &k) || + !load(file, prefix + ".self_attn.v_proj.weight", &v)) { + return false; + } + NativeFloatTensor q_input; + NativeFloatTensor k_input; + NativeFloatTensor v_input; + if (fold_rms) { + q_input = std::move(q); + k_input = std::move(k); + v_input = std::move(v); + } else if (!round_bf16(q, &q_input) || !round_bf16(k, &k_input) || + !round_bf16(v, &v_input)) { + return false; + } + + NativeFloatTensor qi; + NativeFloatTensor ki; + if (!flashrt::models::pi05::native_interleave_qk_rows(q_input, 8, &qi) + .ok_status() || + !flashrt::models::pi05::native_interleave_qk_rows(k_input, 1, &ki) + .ok_status()) { + return false; + } + if (fold_rms) { + NativeFloatTensor norm; + NativeFloatTensor qf; + NativeFloatTensor kf; + NativeFloatTensor vf; + if (!load(file, prefix + ".input_layernorm.weight", &norm) || + !flashrt::models::pi05::native_fold_rms_columns(qi, norm, &qf) + .ok_status() || + !flashrt::models::pi05::native_fold_rms_columns(ki, norm, &kf) + .ok_status() || + !flashrt::models::pi05::native_fold_rms_columns(v_input, norm, &vf) + .ok_status()) { + return false; + } + qi = std::move(qf); + ki = std::move(kf); + v_input = std::move(vf); + } + NativeFloatTensor joined; + return flashrt::models::pi05::native_concat_rows_transpose( + {&qi, &ki, &v_input}, &joined).ok_status() && + finish(joined, out); +} + +bool gate_up(const SafetensorsFile& file, NativeBf16Tensor* out) { + NativeFloatTensor gate; + NativeFloatTensor up; + NativeFloatTensor gate_rounded; + NativeFloatTensor up_rounded; + NativeFloatTensor gate_t; + NativeFloatTensor up_t; + NativeFloatTensor joined; + return load(file, std::string(kDecoder) + ".mlp.gate_proj.weight", + &gate) && + load(file, std::string(kDecoder) + ".mlp.up_proj.weight", &up) && + round_bf16(gate, &gate_rounded) && + round_bf16(up, &up_rounded) && + flashrt::models::pi05::native_transpose_2d(gate_rounded, &gate_t) + .ok_status() && + flashrt::models::pi05::native_transpose_2d(up_rounded, &up_t) + .ok_status() && + flashrt::models::pi05::native_concat_columns(gate_t, up_t, &joined) + .ok_status() && + finish(joined, out); +} + +bool action_out(const SafetensorsFile& file, int num_steps, + NativeBf16Tensor* out) { + NativeFloatTensor source; + NativeFloatTensor rounded; + NativeFloatTensor transposed; + NativeFloatTensor scaled; + return load(file, "action_out_proj.weight", &source) && + round_bf16(source, &rounded) && + flashrt::models::pi05::native_transpose_2d(rounded, &transposed) + .ok_status() && + flashrt::models::pi05::native_scale( + transposed, -1.0f / static_cast(num_steps), &scaled) + .ok_status() && + finish(scaled, out); +} + +bool time_embeds(int num_steps, NativeBf16Tensor* out) { + NativeFloatTensor generated; + return flashrt::models::pi05::native_pi05_time_embeddings(num_steps, 1024, + &generated) + .ok_status() && + finish(generated, out); +} + +std::uint64_t fnv1a(const std::vector& values) { + std::uint64_t hash = 14695981039346656037ull; + const auto* bytes = reinterpret_cast(values.data()); + for (std::size_t i = 0; i < values.size() * sizeof(std::uint16_t); ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ull; + } + return hash; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 3) { + std::cerr << "usage: pi05_native_weight_probe CHECKPOINT OP\n"; + return 2; + } + SafetensorsFile file; + if (!file.open(std::string(argv[1]) + "/model.safetensors")) { + std::cerr << file.error() << '\n'; + return 2; + } + NativeBf16Tensor output; + const std::string op = argv[2]; + bool ok = false; + if (op == "patch") { + ok = patch(file, &output); + } else if (op == "encoder_qkv0") { + ok = qkv(file, kEncoder, true, &output); + } else if (op == "decoder_qkv0") { + ok = qkv(file, kDecoder, false, &output); + } else if (op == "decoder_gate_up0") { + ok = gate_up(file, &output); + } else if (op == "action_out10") { + ok = action_out(file, 10, &output); + } else if (op == "action_out5") { + ok = action_out(file, 5, &output); + } else if (op == "time_embeds10") { + ok = time_embeds(10, &output); + } else if (op == "time_embeds5") { + ok = time_embeds(5, &output); + } + if (!ok) { + std::cerr << "weight probe operation failed: " << op << '\n'; + return 1; + } + std::cout << "shape="; + for (std::size_t i = 0; i < output.shape.size(); ++i) { + if (i) std::cout << ','; + std::cout << output.shape[i]; + } + std::cout << " fnv=" << std::hex << std::setw(16) << std::setfill('0') + << fnv1a(output.values) << '\n'; + return 0; +} diff --git a/cpp/tests/pi05_tokenizer_corpus_probe.cpp b/cpp/tests/pi05_tokenizer_corpus_probe.cpp new file mode 100644 index 00000000..d650a919 --- /dev/null +++ b/cpp/tests/pi05_tokenizer_corpus_probe.cpp @@ -0,0 +1,104 @@ +#include "flashrt/cpp/modalities/tokenizer.h" +#include "flashrt/cpp/models/pi05/prompt_format.h" + +#include +#include +#include +#include +#include + +namespace { + +constexpr std::uint32_t kCorpusMagic = 0x50303554u; +constexpr std::uint32_t kOutputMagic = 0x50303549u; + +template +bool read_value(std::ifstream& input, T* value) { + return static_cast(input.read( + reinterpret_cast(value), sizeof(T))); +} + +template +bool write_value(std::ofstream& output, const T& value) { + return static_cast(output.write( + reinterpret_cast(&value), sizeof(T))); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 4) { + std::cerr << "usage: pi05_tokenizer_corpus_probe TOKENIZER CORPUS OUT\n"; + return 2; + } + flashrt::modalities::SentencePieceTokenizer tokenizer; + auto status = tokenizer.load_model(argv[1]); + if (!status.ok_status()) { + std::cerr << "tokenizer load failed: " << status.message << '\n'; + return 1; + } + tokenizer.reserve(200); + std::ifstream input(argv[2], std::ios::binary); + std::ofstream output(argv[3], std::ios::binary | std::ios::trunc); + std::uint32_t magic = 0; + std::uint32_t records = 0; + if (!input || !output || !read_value(input, &magic) || + !read_value(input, &records) || magic != kCorpusMagic || + !write_value(output, kOutputMagic) || !write_value(output, records)) { + std::cerr << "invalid tokenizer corpus header\n"; + return 1; + } + flashrt::modalities::SentencePieceEncodeOptions options; + options.add_bos = true; + options.max_tokens = 200; + std::string task; + std::string formatted; + std::vector state; + std::vector ids; + task.reserve(512); + formatted.reserve(1024); + state.reserve(32); + ids.reserve(200); + for (std::uint32_t record = 0; record < records; ++record) { + std::uint32_t task_bytes = 0; + std::uint32_t state_count = 0; + if (!read_value(input, &task_bytes) || + !read_value(input, &state_count) || task_bytes > 4096 || + state_count > 1024) { + std::cerr << "invalid tokenizer corpus record\n"; + return 1; + } + task.resize(task_bytes); + state.resize(state_count); + if ((task_bytes && !input.read(task.data(), task_bytes)) || + (state_count && !input.read( + reinterpret_cast(state.data()), + static_cast(state_count * sizeof(float))))) { + std::cerr << "truncated tokenizer corpus record\n"; + return 1; + } + flashrt::models::pi05::format_state_prompt_into( + task, state.data(), state.size(), &formatted); + status = tokenizer.encode(formatted, options, &ids); + if (!status.ok_status()) { + std::cerr << "tokenization failed at record " << record << ": " + << status.message << '\n'; + return 1; + } + const std::uint32_t count = static_cast(ids.size()); + if (!write_value(output, count) || + (count && !output.write( + reinterpret_cast(ids.data()), + static_cast(count * sizeof(std::int32_t))))) { + std::cerr << "tokenizer output write failed\n"; + return 1; + } + } + char trailing = 0; + if (input.read(&trailing, 1)) { + std::cerr << "tokenizer corpus has trailing bytes\n"; + return 1; + } + std::cout << "PASS " << records << " tokenized prompt/state records\n"; + return 0; +} diff --git a/cpp/tests/profile_pi05_python_replay.py b/cpp/tests/profile_pi05_python_replay.py new file mode 100644 index 00000000..30e66910 --- /dev/null +++ b/cpp/tests/profile_pi05_python_replay.py @@ -0,0 +1,102 @@ +"""Developer profiling tool for one Pi0.5 Python replay CUDA range. + +This utility produces diagnostic traces; it is not an acceptance test. +""" + +from __future__ import annotations + +import argparse +import ctypes +from pathlib import Path +import sys + +import numpy as np +import torch + + +ROOT = Path(__file__).resolve().parents[2] +for rel in ("", "exec/build-container", "runtime/build-container", + "exec/build", "runtime/build"): + path = str(ROOT / rel) if rel else str(ROOT) + if path not in sys.path: + sys.path.insert(0, path) + +import flash_rt # noqa: E402 + + +def _check_cuda(rc: int, operation: str) -> None: + if rc != 0: + raise RuntimeError(f"{operation} failed with CUDA error {rc}") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--num-views", type=int, default=2) + parser.add_argument("--steps", type=int, default=10) + args = parser.parse_args() + capability = torch.cuda.get_device_capability() + if capability != (12, 0): + raise RuntimeError(f"Pi0.5 native profiling requires SM120, got {capability}") + + rng = np.random.default_rng(7) + images = [ + rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + for _ in range(args.num_views) + ] + state = np.linspace(-0.8, 0.8, 8, dtype=np.float32) + model = flash_rt.load_model( + args.checkpoint, + framework="torch", + config="pi05", + hardware="rtx_sm120", + num_views=args.num_views, + num_steps=args.steps, + cache_frames=1, + use_fp8=False, + state_prompt_mode="fixed", + ) + model.predict(images, prompt="pick up the black bowl", state=state) + + pipe = model._pipe + pipeline = pipe.pipeline + observation = { + "images": images, + "image": images[0], + "state": state, + } + if len(images) >= 2: + observation["wrist_image"] = images[1] + if len(images) >= 3: + observation["wrist_image_right"] = images[2] + + with torch.cuda.stream(pipe._graph_torch_stream): + stream = pipe._graph_torch_stream.cuda_stream + pipe._noise_buf.zero_() + pipe._copy_tensor_to_pipeline_buf_stream( + pipe._noise_buf, pipeline.input_noise_buf, stream) + pipe._fill_img_buf(observation) + pipe._copy_tensor_to_pipeline_buf_stream( + pipe._img_buf, pipeline.input_images_buf, stream) + _check_cuda( + pipe._cudart.cudaStreamSynchronize(ctypes.c_void_p(stream)), + "cudaStreamSynchronize before profiling", + ) + + cudart = ctypes.CDLL("libcudart.so") + cudart.cudaProfilerStart.restype = ctypes.c_int + cudart.cudaProfilerStop.restype = ctypes.c_int + _check_cuda(cudart.cudaProfilerStart(), "cudaProfilerStart") + with torch.cuda.stream(pipe._graph_torch_stream): + pipeline.forward() + _check_cuda( + pipe._cudart.cudaStreamSynchronize(ctypes.c_void_p(stream)), + "cudaStreamSynchronize after replay", + ) + _check_cuda(cudart.cudaProfilerStop(), "cudaProfilerStop") + print("PASS Pi0.5 Python frontend replay profiler range") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index 0c673ebf..f8def582 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -1,9 +1,12 @@ #include "flashrt/cpp/modalities/action.h" +#include "flashrt/cpp/modalities/text.h" #include "flashrt/cpp/modalities/vision.h" #include "flashrt/cpp/models/pi05/spec.h" #include +#include +#include #include #include #include @@ -11,14 +14,19 @@ #include using flashrt::modalities::DType; +using flashrt::modalities::ActionStaging; using flashrt::modalities::Layout; using flashrt::modalities::MemoryPlace; using flashrt::modalities::PixelFormat; using flashrt::modalities::Shape; using flashrt::modalities::TensorView; +using flashrt::modalities::EmbeddingGatherSpec; +using flashrt::modalities::TextEmbeddingStaging; using flashrt::modalities::VisionFrame; using flashrt::modalities::bfloat16_to_float; using flashrt::modalities::float_to_bfloat16; +using flashrt::modalities::gather_token_embeddings; +using flashrt::modalities::gather_token_embeddings_cpu; using flashrt::modalities::postprocess_action; using flashrt::modalities::preprocess_vision_cpu; using flashrt::modalities::preprocess_vision; @@ -37,6 +45,17 @@ bool has_cuda_device() { return n > 0; } +std::uint32_t ordered_bf16(std::uint16_t bits) { + if (bits & 0x8000u) return 0x8000u - (bits & 0x7fffu); + return 0x8000u + bits; +} + +std::uint32_t bf16_ulp_distance(std::uint16_t a, std::uint16_t b) { + const std::uint32_t ao = ordered_bf16(a); + const std::uint32_t bo = ordered_bf16(b); + return ao > bo ? ao - bo : bo - ao; +} + void test_vision_h2d_staging() { const auto spec = flashrt::models::pi05::vision_preprocess_spec(1); const std::uint64_t bytes = required_vision_output_bytes(spec); @@ -110,6 +129,105 @@ void test_vision_h2d_staging() { cudaFree(device); } +void test_vision_resize_matrix() { + struct Case { int width; int height; int padding; }; + const std::array cases{{ + {1, 1, 0}, {3, 2, 5}, {17, 19, 1}, {63, 47, 7}, + {224, 224, 0}, {321, 181, 3}, {181, 321, 9}, + }}; + std::uint64_t max_frame_bytes = 0; + for (const auto& item : cases) { + max_frame_bytes = std::max( + max_frame_bytes, + static_cast(item.width * 3 + item.padding) * + static_cast(item.height)); + } + + const auto spec = flashrt::models::pi05::vision_preprocess_spec(1); + const std::uint64_t output_bytes = required_vision_output_bytes(spec); + void* device = nullptr; + assert(cudaMalloc(&device, output_bytes) == cudaSuccess); + flashrt::modalities::VisionStaging pool; + auto st = flashrt::modalities::vision_staging_create( + &pool, 1, max_frame_bytes); + assert(st.ok_status()); + std::vector actual(output_bytes / 2); + std::vector expected(output_bytes / 2); + std::uint32_t matrix_max_ulp = 0; + float matrix_max_abs = 0.0f; + Case worst_case{}; + std::size_t worst_index = 0; + std::uint16_t worst_actual = 0; + std::uint16_t worst_expected = 0; + + for (const auto& item : cases) { + const int stride = item.width * 3 + item.padding; + std::vector pixels( + static_cast(stride) * item.height, 0xa5); + for (int y = 0; y < item.height; ++y) { + for (int x = 0; x < item.width; ++x) { + for (int c = 0; c < 3; ++c) { + pixels[static_cast(y) * stride + x * 3 + c] = + static_cast( + (x * 13 + y * 17 + c * 71) & 0xff); + } + } + } + VisionFrame frame; + frame.name = "image"; + frame.image = { + pixels.data(), pixels.size(), DType::kUInt8, MemoryPlace::kHost, + Layout::kHWC, + Shape{static_cast(item.height), + static_cast(item.width), 3}}; + frame.format = PixelFormat::kRGB8; + frame.width = item.width; + frame.height = item.height; + frame.stride_bytes = stride; + TensorView device_output{ + device, output_bytes, DType::kBFloat16, MemoryPlace::kDevice, + Layout::kNHWC, Shape{1, 224, 224, 3}}; + st = preprocess_vision(spec, {frame}, device_output, nullptr, &pool); + assert(st.ok_status()); + assert(cudaMemcpy(actual.data(), device, output_bytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + TensorView host_output{ + expected.data(), output_bytes, DType::kBFloat16, + MemoryPlace::kHost, Layout::kNHWC, Shape{1, 224, 224, 3}}; + st = preprocess_vision_cpu(spec, {frame}, host_output); + assert(st.ok_status()); + for (std::size_t i = 0; i < actual.size(); ++i) { + const std::uint32_t ulp = + bf16_ulp_distance(actual[i], expected[i]); + const float absolute = std::fabs( + bfloat16_to_float(actual[i]) - + bfloat16_to_float(expected[i])); + matrix_max_abs = std::max(matrix_max_abs, absolute); + if (ulp > matrix_max_ulp) { + matrix_max_ulp = ulp; + worst_case = item; + worst_index = i; + worst_actual = actual[i]; + worst_expected = expected[i]; + } + } + } + if (matrix_max_ulp > 1) { + std::cerr << "vision resize max_ulp=" << matrix_max_ulp + << " max_abs=" << matrix_max_abs + << " size=" << worst_case.width << 'x' + << worst_case.height << " index=" << worst_index << '\n'; + std::cerr << "vision resize values actual=" + << bfloat16_to_float(worst_actual) << " expected=" + << bfloat16_to_float(worst_expected) << '\n'; + } + std::cout << "vision resize matrix max BF16 ULP: " + << matrix_max_ulp << '\n'; + assert(matrix_max_ulp <= 1); + flashrt::modalities::vision_staging_destroy(&pool); + cudaFree(device); +} + void test_action_d2h_staging() { auto spec = flashrt::models::pi05::action_postprocess_spec( {10.0f, 20.0f, 30.0f}, {2.0f, 3.0f, 4.0f}, @@ -135,9 +253,83 @@ void test_action_d2h_staging() { assert(std::fabs(actions[0] - 12.0f) < 0.01f); assert(std::fabs(actions[1] - 17.0f) < 0.01f); assert(std::fabs(actions[2] - 34.0f) < 0.01f); + ActionStaging staging; + st = flashrt::modalities::action_staging_create(&staging, bytes); + assert(st.ok_status() && staging.host_pinned && staging.bytes == bytes); + const std::size_t action_capacity = actions.capacity(); + for (int round = 0; round < 1000; ++round) { + st = postprocess_action(spec, src, &actions, nullptr, &staging); + assert(st.ok_status()); + assert(actions.capacity() == action_capacity); + } + ActionStaging too_small; + st = flashrt::modalities::action_staging_create(&too_small, bytes - 1); + assert(st.ok_status()); + st = postprocess_action(spec, src, &actions, nullptr, &too_small); + assert(!st.ok_status()); + assert(st.code == flashrt::modalities::StatusCode::kInsufficientStorage); + flashrt::modalities::action_staging_destroy(&too_small); + flashrt::modalities::action_staging_destroy(&staging); + assert(staging.host_pinned == nullptr && staging.bytes == 0); cudaFree(device); } +void test_text_embedding_device_gather() { + const std::vector table = { + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, + 9.0f, 10.0f, 11.0f, 12.0f, + }; + const std::int32_t ids[] = {2, 0}; + std::vector ref(2 * 4, 0.0f); + TensorView host_table{const_cast(table.data()), + static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{3, 4}}; + TensorView host_out{ref.data(), static_cast(ref.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + EmbeddingGatherSpec spec{3, 4, 2.0f}; + auto st = gather_token_embeddings_cpu(spec, ids, 2, host_table, host_out); + assert(st.ok_status()); + + void* d_table = nullptr; + void* d_out = nullptr; + assert(cudaMalloc(&d_table, table.size() * sizeof(float)) == cudaSuccess); + assert(cudaMalloc(&d_out, ref.size() * sizeof(float)) == cudaSuccess); + assert(cudaMemcpy(d_table, table.data(), table.size() * sizeof(float), + cudaMemcpyHostToDevice) == cudaSuccess); + TensorView device_table{d_table, + static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kDevice, + Layout::kFlat, Shape{3, 4}}; + TensorView device_out{d_out, static_cast(ref.size() * 4), + DType::kFloat32, MemoryPlace::kDevice, + Layout::kFlat, Shape{2, 4}}; + + TextEmbeddingStaging staging; + st = flashrt::modalities::text_embedding_staging_create(&staging, 2); + assert(st.ok_status()); + std::vector got(ref.size(), 0.0f); + for (int round = 0; round < 3; ++round) { + assert(cudaMemset(d_out, 0, ref.size() * sizeof(float)) == cudaSuccess); + st = gather_token_embeddings(spec, ids, 2, device_table, device_out, + nullptr, &staging); + assert(st.ok_status()); + assert(cudaMemcpy(got.data(), d_out, got.size() * sizeof(float), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(got == ref); + } + st = gather_token_embeddings(spec, ids, 3, device_table, device_out, + nullptr, &staging); + assert(!st.ok_status()); + assert(st.code == flashrt::modalities::StatusCode::kInsufficientStorage); + flashrt::modalities::text_embedding_staging_destroy(&staging); + assert(staging.device_token_ids == nullptr); + cudaFree(d_out); + cudaFree(d_table); +} + } // namespace int main() { @@ -146,7 +338,9 @@ int main() { return 0; } test_vision_h2d_staging(); + test_vision_resize_matrix(); test_action_d2h_staging(); + test_text_embedding_device_gather(); std::cout << "PASS - CUDA modality kernels/staging\n"; return 0; } diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index 117f11c7..0e264a06 100644 --- a/cpp/tests/test_modalities.cpp +++ b/cpp/tests/test_modalities.cpp @@ -32,6 +32,8 @@ void test_pi05_vision_spec_and_preprocess() { assert(spec.view_order[1] == "wrist_image"); assert(spec.target_width == 224); assert(spec.output_dtype == DType::kBFloat16); + assert(spec.normalize.mode == + flashrt::modalities::NormalizeMode::kDivideShift); const std::uint8_t image_rgb[] = { 0, 127, 255, 255, 127, 0, @@ -70,6 +72,7 @@ void test_pi05_vision_spec_and_preprocess() { assert(std::fabs(first_r - (-1.0f)) < 0.01f); assert(std::fabs(first_g - (127.0f / 127.5f - 1.0f)) < 0.01f); assert(std::fabs(first_b - 1.0f) < 0.01f); + assert(out[1] == float_to_bfloat16(127.0f / 127.5f - 1.0f)); } void test_view_order_guard() { @@ -117,7 +120,7 @@ void test_action_postprocess() { assert(std::fabs(out[1] - 17.0f) < 0.01f); assert(std::fabs(out[2] - 34.0f) < 0.01f); assert(std::fabs(out[3] - 11.0f) < 0.01f); - assert(std::fabs(out[4] - 24.5f) < 0.01f); + assert(std::fabs(out[4] - 23.0f) < 0.01f); assert(std::fabs(out[5] - 26.0f) < 0.01f); } @@ -158,11 +161,39 @@ void test_pi05_runtime_io_adapter() { auto st = io.prepare_vision({image}); assert(st.ok_status()); + VisionFrame invalid = image; + invalid.format = PixelFormat::kBGR8; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.image.dtype = DType::kFloat32; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.image.layout = Layout::kCHW; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.image.shape = Shape{2, 3, 3}; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.stride_bytes = -1; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.stride_bytes = 5; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + st = io.prepare_vision({}); + assert(st.code == StatusCode::kShapeMismatch); + st = io.prepare_vision({image, image}); + assert(st.code == StatusCode::kShapeMismatch); std::vector actions; st = io.read_actions(&actions); assert(st.ok_status()); assert(actions.size() == 3); - assert(std::fabs(actions[0] - 21.0f) < 0.01f); + assert(std::fabs(actions[0] - 11.0f) < 0.01f); assert(std::fabs(actions[1] - 22.0f) < 0.01f); assert(std::fabs(actions[2] - 33.0f) < 0.01f); } diff --git a/cpp/tests/test_pi05_c_api.cpp b/cpp/tests/test_pi05_c_api.cpp index ef3b10ba..e7e29360 100644 --- a/cpp/tests/test_pi05_c_api.cpp +++ b/cpp/tests/test_pi05_c_api.cpp @@ -144,6 +144,33 @@ int main() { frame.pixel_format = FRT_PI05_PIXEL_RGB8; rc = frt_pi05_runtime_prepare_vision(rt, &frame, 1); assert(rc == 0); + std::vector rgb_staged(image_bytes / 2); + assert(cudaMemcpy(rgb_staged.data(), frt_buffer_dptr(image), image_bytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + frt_pi05_vision_frame invalid = frame; + invalid.pixel_format = 999; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == -4); + assert(std::strstr(frt_pi05_runtime_last_error(rt), "pixel format")); + const std::uint8_t bgr[] = { + 255, 127, 0, 0, 127, 255, + 30, 20, 10, 60, 50, 40, + }; + invalid = frame; + invalid.data = bgr; + invalid.bytes = sizeof(bgr); + invalid.pixel_format = FRT_PI05_PIXEL_BGR8; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == 0); + std::vector bgr_staged(image_bytes / 2); + assert(cudaMemcpy(bgr_staged.data(), frt_buffer_dptr(image), image_bytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(bgr_staged == rgb_staged); + invalid = frame; + invalid.stride_bytes = 5; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == -4); + assert(std::strstr(frt_pi05_runtime_last_error(rt), "stride")); float out[3] = {}; uint64_t n_written = 0; diff --git a/cpp/tests/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index e250486e..20555acb 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -11,7 +11,9 @@ #include +#include #include +#include #include #include #include @@ -61,6 +63,13 @@ bool has_cuda_device() { return n > 0; } +int producer_set_input(void*, uint32_t, const void*, uint64_t, int) { + return 0; +} +int producer_get_output(void*, uint32_t, void*, uint64_t, uint64_t*, int) { + return 0; +} + } // namespace int main() { @@ -153,6 +162,12 @@ int main() { m->ports[1].update == FRT_RT_PORT_SWAP && m->ports[1].buffer == action, "noise SWAP port exposes the device window"); + CHECK(std::strcmp(m->ports[2].name, "actions") == 0 && + m->ports[2].dtype == FRT_RT_DTYPE_F32 && + m->ports[2].update == FRT_RT_PORT_STAGED && + m->ports[2].buffer == nullptr && + m->ports[2].bytes == 3 * sizeof(float), + "actions port declares the logical F32 payload"); CHECK(m->stages[0].graph == 0, "stage resolves the infer graph"); /* staged image input lands in the device buffer */ @@ -167,6 +182,29 @@ int main() { view.height = 2; CHECK(m->verbs.set_input(m->self, 0, &view, sizeof(view), -1) == 0, "set_input(images) accepts frt_image_view[]"); + frt_image_view bgr_view = view; + bgr_view.pixel_format = FRT_RT_PIXEL_BGR8; + CHECK(m->verbs.set_input(m->self, 0, &bgr_view, sizeof(bgr_view), -1) + == -4, + "set_input(images) rejects non-RGB image formats"); + frt_image_view invalid_format = view; + invalid_format.pixel_format = 999; + CHECK(m->verbs.set_input(m->self, 0, &invalid_format, + sizeof(invalid_format), -1) == -4, + "set_input(images) rejects unknown pixel formats"); + CHECK(std::strstr(m->verbs.last_error(m->self), "pixel format") != nullptr, + "unknown image format reports a readable error"); + frt_image_view two_views[2] = {view, view}; + CHECK(m->verbs.set_input(m->self, 0, two_views, sizeof(two_views), -1) + == -4, + "set_input(images) rejects the wrong view count"); + frt_image_view bad_stride = view; + bad_stride.stride_bytes = 5; + CHECK(m->verbs.set_input(m->self, 0, &bad_stride, sizeof(bad_stride), -1) + == -4, + "set_input(images) rejects a short row stride"); + CHECK(std::strstr(m->verbs.last_error(m->self), "stride") != nullptr, + "invalid image stride reports a readable error"); std::vector img_host(image_bytes / 2); cudaMemcpy(img_host.data(), frt_buffer_dptr(image), image_bytes, cudaMemcpyDeviceToHost); @@ -242,25 +280,42 @@ int main() { ports[1].bytes = action_bytes; ports[2].name = "actions"; ports[2].modality = FRT_RT_MOD_ACTION; - ports[2].dtype = FRT_RT_DTYPE_BF16; + ports[2].dtype = FRT_RT_DTYPE_F32; ports[2].layout = FRT_RT_LAYOUT_FLAT; ports[2].direction = FRT_RT_PORT_OUT; ports[2].update = FRT_RT_PORT_STAGED; ports[2].shape = action_shape; ports[2].rank = 2; - ports[2].buffer = action; - ports[2].bytes = action_bytes; + ports[2].bytes = 3 * sizeof(float); uint32_t after_action[1] = {0}; frt_runtime_stage_desc stages[2]{}; stages[0].graph = 0; stages[1].graph = 1; stages[1].after = after_action; stages[1].n_after = 1; + frt_model_runtime_verbs producer_verbs{}; + producer_verbs.struct_size = sizeof(producer_verbs); + producer_verbs.set_input = producer_set_input; + producer_verbs.get_output = producer_get_output; frt_model_runtime_v1* producer = frt_model_runtime_wrap( - &exp, ports, 3, stages, 2, nullptr, nullptr, nullptr, nullptr); + &exp, ports, 3, stages, 2, &producer_verbs, nullptr, nullptr, nullptr); CHECK(producer != nullptr, "producer model declaration for create_over"); + frt_runtime_port_desc wrong_action_ports[3] = {}; + for (int i = 0; i < 3; ++i) wrong_action_ports[i] = ports[i]; + wrong_action_ports[2].dtype = FRT_RT_DTYPE_BF16; + frt_model_runtime_v1* wrong_action_producer = frt_model_runtime_wrap( + &exp, wrong_action_ports, 3, stages, 2, &producer_verbs, nullptr, + nullptr, nullptr); + frt_model_runtime_v1* wrong_action_over = nullptr; + CHECK(wrong_action_producer && + frt_pi05_model_runtime_create_over( + wrong_action_producer, &cfg, &wrong_action_over) == -2 && + wrong_action_over == nullptr, + "create_over rejects a non-F32 logical action port"); + wrong_action_producer->release(wrong_action_producer->owner); + frt_model_runtime_v1* over = nullptr; CHECK(frt_pi05_model_runtime_create_over(producer, &cfg, &over) == 0 && over, @@ -272,7 +327,7 @@ int main() { over->stages[1].after[0] == 0, "create_over preserves a producer-declared two-stage DAG"); producer->release(producer->owner); - CHECK(owner.release == retains, + CHECK(owner.release < owner.retain, "producer declaration stays alive through the override"); CHECK(over->verbs.set_input(over->self, 0, &view, sizeof(view), -1) == 0, @@ -290,6 +345,145 @@ int main() { CHECK(owner.release == owner.retain, "create_over releases its native runtime and inherited producer"); + /* A producer may declare prompt only when the native runtime can serve it. */ + const int64_t prompt_shape[1] = {-1}; + frt_runtime_port_desc prompt_ports[4] = {}; + for (int i = 0; i < 3; ++i) prompt_ports[i] = ports[i]; + prompt_ports[3].name = "prompt"; + prompt_ports[3].modality = FRT_RT_MOD_TEXT; + prompt_ports[3].dtype = FRT_RT_DTYPE_U8; + prompt_ports[3].layout = FRT_RT_LAYOUT_FLAT; + prompt_ports[3].direction = FRT_RT_PORT_IN; + prompt_ports[3].update = FRT_RT_PORT_STAGED; + prompt_ports[3].shape = prompt_shape; + prompt_ports[3].rank = 1; + frt_model_runtime_v1* prompt_producer = frt_model_runtime_wrap( + &exp, prompt_ports, 4, stages, 2, &producer_verbs, nullptr, nullptr, + nullptr); + CHECK(prompt_producer != nullptr, + "producer declaration with prompt port"); + frt_model_runtime_v1* prompt_over = nullptr; + CHECK(frt_pi05_model_runtime_create_over(prompt_producer, &cfg, + &prompt_over) == -2 && + prompt_over == nullptr, + "prompt port is refused without prompt staging config"); + + frt_runtime_port_desc state_ports[5] = {}; + for (int i = 0; i < 4; ++i) state_ports[i] = prompt_ports[i]; + const int64_t state_shape[1] = {3}; + state_ports[4].name = "state"; + state_ports[4].modality = FRT_RT_MOD_STATE; + state_ports[4].dtype = FRT_RT_DTYPE_F32; + state_ports[4].layout = FRT_RT_LAYOUT_FLAT; + state_ports[4].direction = FRT_RT_PORT_IN; + state_ports[4].update = FRT_RT_PORT_STAGED; + state_ports[4].required = 1; + state_ports[4].shape = state_shape; + state_ports[4].rank = 1; + frt_model_runtime_v1* state_producer = frt_model_runtime_wrap( + &exp, state_ports, 5, stages, 2, &producer_verbs, nullptr, nullptr, + nullptr); + CHECK(state_producer != nullptr, + "producer declaration with prompt and state ports"); + frt_model_runtime_v1* state_over = nullptr; + CHECK(frt_pi05_model_runtime_create_over(state_producer, &cfg, + &state_over) == -2 && + state_over == nullptr, + "state port is refused without state normalization config"); + +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const char* tokenizer = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + if (tokenizer && tokenizer[0] != '\0') { + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + std::vector prompt_out(max_tokens * hidden, 9.0f); + frt_pi05_runtime_config prompt_cfg = cfg; + prompt_cfg.prompt_tokenizer_model_path = tokenizer; + prompt_cfg.prompt_embedding_table_data = table.data(); + prompt_cfg.prompt_embedding_table_bytes = table.size() * sizeof(float); + prompt_cfg.prompt_embedding_table_dtype = FRT_PI05_DTYPE_FLOAT32; + prompt_cfg.prompt_embedding_vocab_size = vocab; + prompt_cfg.prompt_embedding_hidden_dim = hidden; + prompt_cfg.prompt_embedding_data = prompt_out.data(); + prompt_cfg.prompt_embedding_bytes = + prompt_out.size() * sizeof(float); + prompt_cfg.prompt_embedding_dtype = FRT_PI05_DTYPE_FLOAT32; + prompt_cfg.max_prompt_tokens = max_tokens; + prompt_cfg.prompt_embedding_scale = 0.5f; + const float state_q01[3] = {0.0f, 0.0f, 0.0f}; + const float state_q99[3] = {2.0f, 2.0f, 2.0f}; + prompt_cfg.state_q01 = state_q01; + prompt_cfg.n_state_q01 = 3; + prompt_cfg.state_q99 = state_q99; + prompt_cfg.n_state_q99 = 3; + + CHECK(frt_pi05_model_runtime_create_over(prompt_producer, + &prompt_cfg, + &prompt_over) == 0 && + prompt_over, + "prompt port accepted with prompt staging config"); + const char prompt_text[] = "pick up cube"; + CHECK(prompt_over->verbs.set_input( + prompt_over->self, 3, prompt_text, + sizeof(prompt_text) - 1, -1) == 0, + "set_input(prompt) writes staged embeddings"); + CHECK(std::fabs(prompt_out[0] - 1.0f) < 0.001f && + std::fabs(prompt_out[1] + 1.0f) < 0.001f, + "prompt staging wrote the BOS embedding row"); + prompt_over->release(prompt_over->owner); + + std::fill(prompt_out.begin(), prompt_out.end(), 9.0f); + CHECK(frt_pi05_model_runtime_create_over(state_producer, + &prompt_cfg, + &state_over) == 0 && + state_over, + "state port accepted with prompt staging and norm stats"); + const float raw_state[3] = {1.0f, 2.0f, 0.0f}; + CHECK(state_over->verbs.set_input( + state_over->self, 4, raw_state, sizeof(raw_state), -1) == 0, + "set_input(state) accepts f32 state before prompt"); + CHECK(state_over->verbs.set_input( + state_over->self, 3, prompt_text, + sizeof(prompt_text) - 1, -1) == 0, + "set_input(prompt) renders cached state"); + CHECK(std::fabs(prompt_out[0] - 1.0f) < 0.001f && + std::fabs(prompt_out[1] + 1.0f) < 0.001f, + "state prompt staging wrote embeddings"); + const std::size_t variants_before = frt_graph_variant_count(graph); + for (int tick = 0; tick < 1000; ++tick) { + const float changing_state[3] = { + static_cast(tick % 3), 2.0f, 0.0f}; + CHECK(state_over->verbs.set_input( + state_over->self, 4, changing_state, + sizeof(changing_state), -1) == 0, + "state hot update remains available"); + } + CHECK(frt_graph_variant_count(graph) == variants_before, + "state hot updates do not recapture graph variants"); + const float wrong_state[2] = {0.0f, 0.0f}; + CHECK(state_over->verbs.set_input( + state_over->self, 4, wrong_state, sizeof(wrong_state), -1) == + -4, + "state hot update rejects dimension changes"); + const std::string oversized_prompt(max_tokens * 8 + 1, 'x'); + CHECK(state_over->verbs.set_input( + state_over->self, 3, oversized_prompt.data(), + oversized_prompt.size(), -1) == -4, + "prompt hot update rejects capacity growth"); + state_over->release(state_over->owner); + } else { + std::printf("SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"); + } +#endif + state_producer->release(state_producer->owner); + prompt_producer->release(prompt_producer->owner); + frt_graph_destroy(graph); frt_ctx_destroy(ctx); std::printf(g_fail ? "\n== PI05 MODEL RUNTIME FAILED ==\n" diff --git a/cpp/tests/test_pi05_native_bf16_forward.cpp b/cpp/tests/test_pi05_native_bf16_forward.cpp new file mode 100644 index 00000000..a37e22ce --- /dev/null +++ b/cpp/tests/test_pi05_native_bf16_forward.cpp @@ -0,0 +1,192 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include + +#include +#include +#include +#include + +namespace { + +using flashrt::models::pi05::NativeBf16Forward; +using flashrt::models::pi05::NativeDeviceWeightStore; +using flashrt::models::pi05::NativeKernelDriver; +using flashrt::models::pi05::NativeRtxAttentionWorkspace; +using flashrt::models::pi05::NativeWorkspace; + +std::vector download(const void* device, std::size_t elements) { + std::vector result(elements); + assert(cudaMemcpy(result.data(), device, + result.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + return result; +} + +struct CaptureArgs { + const NativeBf16Forward* forward = nullptr; + const NativeDeviceWeightStore* weights = nullptr; + NativeWorkspace* workspace = nullptr; + NativeRtxAttentionWorkspace* attention = nullptr; + bool recorded = false; +}; + +void record_encoder_qkv(void* user, void* stream) { + auto* args = static_cast(user); + args->recorded = args->forward + ->encoder_qkv(17, *args->weights, args->workspace, args->attention, + reinterpret_cast(stream)) + .ok_status(); +} + +} // namespace + +int main() { + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess || !count) { + cudaGetLastError(); + std::printf("SKIP - no CUDA device\n"); + return 0; + } + using namespace flashrt::models::pi05; + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + NativeWorkspace workspace(ctx); + NativeWorkspaceConfig workspace_config; + workspace_config.num_views = 1; + workspace_config.max_prompt_tokens = 1; + workspace_config.chunk_size = 2; + workspace_config.num_steps = 2; + workspace_config.vision_pool_factor = 4; + assert(workspace.allocate(workspace_config).ok_status()); + assert(workspace.encoder_sequence() == 17); + + NativeRtxAttentionWorkspace attention(ctx); + NativeRtxAttentionConfig attention_config; + attention_config.num_views = 1; + attention_config.encoder_sequence = 17; + attention_config.encoder_vision_sequence = 16; + attention_config.chunk_size = 2; + assert(attention.allocate(attention_config).ok_status()); + + NativeKernelDriver driver; + NativeBf16Forward forward(&driver); + NativeDeviceWeightStore weights(ctx); + assert(!forward.encoder_qkv(17, weights, &workspace, &attention, 0) + .ok_status()); + + NativeBf16Tensor qkv_weight; + qkv_weight.shape = {2048, 2560}; + qkv_weight.values.assign(2048 * 2560, 0); + const std::uint16_t one = + flashrt::modalities::float_to_bfloat16(1.0f); + for (int column = 0; column < 2048; ++column) { + qkv_weight.values[static_cast(column) * 2560 + column] = + one; + } + for (int column = 0; column < 256; ++column) { + qkv_weight.values[static_cast(column) * 2560 + + 2048 + column] = one; + qkv_weight.values[static_cast(256 + column) * 2560 + + 2304 + column] = one; + } + assert(weights.upload("encoder_attn_qkv_w_17", qkv_weight).ok_status()); + + const auto* encoder_x = workspace.find("encoder_x"); + assert(encoder_x); + std::vector host_x(17 * 2048, 0); + for (int row = 0; row < 17; ++row) { + for (int column = 0; column < 512; ++column) { + const float value = float((row + column) % 15 - 7) / 8.0f; + host_x[static_cast(row) * 2048 + column] = + flashrt::modalities::float_to_bfloat16(value); + } + } + assert(cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), + host_x.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); + + cudaStream_t stream = nullptr; + assert(cudaStreamCreate(&stream) == cudaSuccess); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + assert(forward.encoder_qkv(17, weights, &workspace, &attention, + native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + const auto* query_buffer = attention.find("attn_enc_Q"); + assert(query_buffer); + const std::vector expected_q = download( + frt_buffer_dptr(query_buffer->buffer), 17 * 2048); + const std::vector expected_k = + download(attention.encoder_k_layer_dptr(17), 17 * 256); + const std::vector expected_v = + download(attention.encoder_v_layer_dptr(17), 17 * 256); + + assert(cudaMemset(frt_buffer_dptr(query_buffer->buffer), 0, + expected_q.size() * sizeof(std::uint16_t)) == + cudaSuccess); + assert(cudaMemset(attention.encoder_k_layer_dptr(17), 0, + expected_k.size() * sizeof(std::uint16_t)) == + cudaSuccess); + assert(cudaMemset(attention.encoder_v_layer_dptr(17), 0, + expected_v.size() * sizeof(std::uint16_t)) == + cudaSuccess); + const auto* x_norm = workspace.find("encoder_x_norm"); + const auto* qkv = workspace.find("encoder_QKV"); + const auto* rms = workspace.find("encoder_rms_ones"); + const auto* rope = workspace.find("encoder_rope_weights"); + const auto* weight = weights.find("encoder_attn_qkv_w_17"); + assert(x_norm && qkv && rms && rope && weight); + assert(driver.rms_norm_bf16( + frt_buffer_dptr(encoder_x->buffer), + frt_buffer_dptr(rms->buffer), + frt_buffer_dptr(x_norm->buffer), 17, 2048, 1e-6f, + native_stream) + .ok_status()); + assert(driver.bf16_nn( + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(weight->buffer), + frt_buffer_dptr(qkv->buffer), 17, 2560, 2048, + native_stream) + .ok_status()); + assert(driver.qkv_split_rope_bf16( + frt_buffer_dptr(qkv->buffer), frt_buffer_dptr(rope->buffer), + frt_buffer_dptr(query_buffer->buffer), + attention.encoder_k_layer_dptr(17), + attention.encoder_v_layer_dptr(17), 17, 2048, 256, 256, + 256, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + assert(download(frt_buffer_dptr(query_buffer->buffer), 17 * 2048) == + expected_q); + assert(download(attention.encoder_k_layer_dptr(17), 17 * 256) == + expected_k); + assert(download(attention.encoder_v_layer_dptr(17), 17 * 256) == + expected_v); + + frt_graph graph = frt_graph_create(ctx, "native_encoder_qkv", 17); + assert(graph); + assert(frt_graph_bind(graph, "encoder_x", encoder_x->buffer) == FRT_OK); + assert(frt_graph_bind(graph, "encoder_q", query_buffer->buffer) == FRT_OK); + CaptureArgs capture{&forward, &weights, &workspace, &attention, false}; + assert(frt_graph_capture(graph, 17, record_encoder_qkv, &capture) == FRT_OK); + assert(capture.recorded); + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + assert(stream_id >= 0); + for (int i = 0; i < 100; ++i) { + assert(frt_graph_replay(graph, 17, stream_id) == FRT_OK); + } + assert(frt_graph_variant_count(graph) == 1); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + assert(download(attention.encoder_k_layer_dptr(17), 17 * 256) == + expected_k); + + frt_graph_destroy(graph); + assert(cudaStreamDestroy(stream) == cudaSuccess); + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native BF16 encoder QKV\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_device_weights.cpp b/cpp/tests/test_pi05_native_device_weights.cpp new file mode 100644 index 00000000..0c884818 --- /dev/null +++ b/cpp/tests/test_pi05_native_device_weights.cpp @@ -0,0 +1,97 @@ +#include "flashrt/cpp/models/pi05/native_device_weights.h" + +#include + +#include +#include +#include + +namespace { + +bool has_cuda_device() { + int count = 0; + const cudaError_t rc = cudaGetDeviceCount(&count); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return count > 0; +} + +} // namespace + +int main() { + if (!has_cuda_device()) { + std::printf("SKIP - no CUDA device\n"); + return 0; + } + using flashrt::models::pi05::NativeBf16Tensor; + using flashrt::models::pi05::NativeDeviceWeightStore; + using flashrt::models::pi05::NativeWeightDType; + + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + { + NativeDeviceWeightStore store(ctx); + NativeBf16Tensor tensor; + tensor.shape = {2, 3}; + tensor.values = { + flashrt::modalities::float_to_bfloat16(1.0f), + flashrt::modalities::float_to_bfloat16(2.0f), + flashrt::modalities::float_to_bfloat16(3.0f), + flashrt::modalities::float_to_bfloat16(4.0f), + flashrt::modalities::float_to_bfloat16(5.0f), + flashrt::modalities::float_to_bfloat16(6.0f), + }; + assert(store.upload("encoder.layer0.qkv", tensor).ok_status()); + assert(store.size() == 1); + const auto* weight = store.find("encoder.layer0.qkv"); + assert(weight && weight->buffer); + assert(weight->shape == tensor.shape); + assert(weight->dtype == NativeWeightDType::kBf16); + assert(frt_buffer_bytes(weight->buffer) == + tensor.values.size() * sizeof(std::uint16_t)); + assert(std::string(frt_buffer_name(weight->buffer)) == + "encoder.layer0.qkv"); + + std::vector copied(tensor.values.size()); + assert(cudaMemcpy(copied.data(), frt_buffer_dptr(weight->buffer), + copied.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(copied == tensor.values); + + const std::vector int8_values = {-127, -1, 0, 127}; + assert(store.upload_bytes("encoder.layer0.qkv.int8", {2, 2}, + NativeWeightDType::kInt8, + int8_values.data(), int8_values.size()) + .ok_status()); + const auto* int8_weight = store.find("encoder.layer0.qkv.int8"); + assert(int8_weight && int8_weight->dtype == NativeWeightDType::kInt8); + std::vector int8_copied(int8_values.size()); + assert(cudaMemcpy(int8_copied.data(), + frt_buffer_dptr(int8_weight->buffer), + int8_copied.size(), cudaMemcpyDeviceToHost) == + cudaSuccess); + assert(int8_copied == int8_values); + + const std::vector scales = {0.25f, 0.5f}; + assert(store.upload_bytes("encoder.layer0.qkv.scale", {2}, + NativeWeightDType::kFloat32, + scales.data(), + scales.size() * sizeof(float)) + .ok_status()); + assert(store.find("encoder.layer0.qkv.scale")->dtype == + NativeWeightDType::kFloat32); + assert(!store.upload_bytes("bad.bytes", {3}, + NativeWeightDType::kFp8E4M3, + int8_values.data(), int8_values.size()) + .ok_status()); + assert(!store.upload("encoder.layer0.qkv", tensor).ok_status()); + tensor.shape = {3, 3}; + assert(!store.upload("bad", tensor).ok_status()); + } + + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native device weights\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_forward_primitives.cpp b/cpp/tests/test_pi05_native_forward_primitives.cpp new file mode 100644 index 00000000..8648592a --- /dev/null +++ b/cpp/tests/test_pi05_native_forward_primitives.cpp @@ -0,0 +1,348 @@ +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include + +#include +#include +#include +#include +#include + +namespace { + +using flashrt::models::pi05::NativeKernelDriver; + +frt_buffer allocate(frt_ctx ctx, const char* name, std::size_t elements) { + frt_buffer buffer = + frt_buffer_alloc(ctx, name, elements * sizeof(std::uint16_t)); + assert(buffer); + return buffer; +} + +std::vector bf16(const std::vector& values) { + std::vector result(values.size()); + for (std::size_t i = 0; i < values.size(); ++i) { + result[i] = flashrt::modalities::float_to_bfloat16(values[i]); + } + return result; +} + +void upload(frt_buffer buffer, const std::vector& values) { + assert(cudaMemcpy(frt_buffer_dptr(buffer), values.data(), + values.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); +} + +std::vector download(frt_buffer buffer, std::size_t elements) { + std::vector bits(elements); + assert(cudaMemcpy(bits.data(), frt_buffer_dptr(buffer), + bits.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + std::vector result(elements); + for (std::size_t i = 0; i < elements; ++i) { + result[i] = flashrt::modalities::bfloat16_to_float(bits[i]); + } + return result; +} + +void expect_close(const std::vector& actual, + const std::vector& expected, float tolerance) { + assert(actual.size() == expected.size()); + for (std::size_t i = 0; i < actual.size(); ++i) { + assert(std::fabs(actual[i] - expected[i]) <= tolerance); + } +} + +struct CaptureArgs { + const NativeKernelDriver* driver = nullptr; + void* values = nullptr; + const void* weight = nullptr; + void* norm_output = nullptr; + const void* qkv = nullptr; + const void* rope = nullptr; + void* query = nullptr; + void* key = nullptr; + void* value = nullptr; + const void* devpos = nullptr; + bool recorded = false; +}; + +void record_primitives(void* user, void* stream) { + auto* args = static_cast(user); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + args->recorded = + args->driver + ->rms_norm_bf16(args->values, args->weight, args->norm_output, + 2, 4, 1e-6f, native_stream) + .ok_status() && + args->driver + ->qkv_split_rope_devpos_bf16( + args->qkv, args->rope, args->query, args->key, args->value, + args->devpos, 2, 4, 2, 2, 2, native_stream) + .ok_status(); +} + +} // namespace + +int main() { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || !device_count) { + cudaGetLastError(); + std::printf("SKIP - no CUDA device\n"); + return 0; + } + + NativeKernelDriver driver; + assert(driver.status().ok_status()); + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + cudaStream_t stream = nullptr; + assert(cudaStreamCreate(&stream) == cudaSuccess); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + + const std::vector host_x = {1, 2, 3, 4, -1, 0, 1, 2}; + const std::vector host_weight = {1, 1.5f, 0.5f, 2}; + const std::vector host_bias = {0.1f, -0.2f, 0.3f, -0.4f}; + frt_buffer x = allocate(ctx, "primitive_x", 8); + frt_buffer weight = allocate(ctx, "primitive_weight", 4); + frt_buffer bias = allocate(ctx, "primitive_bias", 4); + frt_buffer output = allocate(ctx, "primitive_output", 8); + frt_buffer gate = allocate(ctx, "primitive_gate", 8); + upload(x, bf16(host_x)); + upload(weight, bf16(host_weight)); + upload(bias, bf16(host_bias)); + + assert(driver.rms_norm_bf16( + frt_buffer_dptr(x), frt_buffer_dptr(weight), + frt_buffer_dptr(output), 2, 4, 1e-6f, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector rms_expected(8); + for (int row = 0; row < 2; ++row) { + float sum = 0; + for (int col = 0; col < 4; ++col) { + const float value = host_x[row * 4 + col]; + sum += value * value; + } + const float scale = 1.0f / std::sqrt(sum / 4 + 1e-6f); + for (int col = 0; col < 4; ++col) { + rms_expected[row * 4 + col] = + host_x[row * 4 + col] * scale * host_weight[col]; + } + } + expect_close(download(output, 8), rms_expected, 0.025f); + + assert(driver.layer_norm_bf16( + frt_buffer_dptr(x), frt_buffer_dptr(weight), + frt_buffer_dptr(bias), frt_buffer_dptr(output), 2, 4, + 1e-5f, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector layer_expected(8); + for (int row = 0; row < 2; ++row) { + float mean = 0; + for (int col = 0; col < 4; ++col) mean += host_x[row * 4 + col]; + mean /= 4; + float variance = 0; + for (int col = 0; col < 4; ++col) { + const float delta = host_x[row * 4 + col] - mean; + variance += delta * delta; + } + const float scale = 1.0f / std::sqrt(variance / 4 + 1e-5f); + for (int col = 0; col < 4; ++col) { + layer_expected[row * 4 + col] = + (host_x[row * 4 + col] - mean) * scale * host_weight[col] + + host_bias[col]; + } + } + expect_close(download(output, 8), layer_expected, 0.025f); + + std::vector style(24, 0.0f); + for (int row = 0; row < 2; ++row) { + for (int col = 0; col < 4; ++col) { + style[row * 12 + col] = 0.25f; + style[row * 12 + 4 + col] = -0.5f; + style[row * 12 + 8 + col] = 0.75f; + } + } + frt_buffer style_buffer = allocate(ctx, "primitive_style", 24); + upload(style_buffer, bf16(style)); + assert(driver.ada_rms_norm_style_bf16( + frt_buffer_dptr(x), frt_buffer_dptr(weight), + frt_buffer_dptr(style_buffer), frt_buffer_dptr(output), + frt_buffer_dptr(gate), 2, 4, 1e-6f, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector ada_expected(8); + for (std::size_t i = 0; i < ada_expected.size(); ++i) { + ada_expected[i] = rms_expected[i] * 1.25f - 0.5f; + } + expect_close(download(output, 8), ada_expected, 0.035f); + expect_close(download(gate, 8), std::vector(8, 0.75f), 0.0f); + + frt_buffer residual = allocate(ctx, "primitive_residual", 8); + upload(residual, bf16(std::vector(8, 1.0f))); + upload(output, bf16(std::vector(8, 2.0f))); + upload(gate, bf16(std::vector(8, 0.5f))); + assert(driver.gate_mul_residual_bf16( + frt_buffer_dptr(residual), frt_buffer_dptr(output), + frt_buffer_dptr(gate), 8, native_stream) + .ok_status()); + assert(driver.bias_residual_bf16( + frt_buffer_dptr(residual), frt_buffer_dptr(output), + frt_buffer_dptr(bias), 2, 4, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector residual_expected(8); + for (int i = 0; i < 8; ++i) { + residual_expected[i] = 4.0f + host_bias[i % 4]; + } + expect_close(download(residual, 8), residual_expected, 0.025f); + + upload(residual, bf16(std::vector(8, 1.0f))); + upload(output, bf16(std::vector(8, 2.0f))); + assert(driver.residual_add_bf16( + frt_buffer_dptr(residual), frt_buffer_dptr(output), 8, + native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + expect_close(download(residual, 8), std::vector(8, 3.0f), 0.0f); + + const std::vector activation_input = + {-3, -2, -1, 0, 0.5f, 1, 2, 3}; + upload(gate, bf16(activation_input)); + upload(output, bf16(std::vector(8, 1.5f))); + assert(driver.gate_gelu_bf16( + frt_buffer_dptr(gate), frt_buffer_dptr(output), + frt_buffer_dptr(residual), 8, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector gated_expected(8); + for (int i = 0; i < 8; ++i) { + const float value = activation_input[i]; + const float gelu = value / + (1.0f + std::exp(-1.5957691216057308f * value * + (1.0f + 0.044715f * value * value))); + gated_expected[i] = gelu * 1.5f; + } + expect_close(download(residual, 8), gated_expected, 0.025f); + upload(output, bf16(activation_input)); + assert(driver.gelu_bf16(frt_buffer_dptr(output), 8, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector gelu_expected(8); + for (int i = 0; i < 8; ++i) { + const float value = activation_input[i]; + gelu_expected[i] = value * 0.5f * + (1.0f + std::tanh(0.7978845608f * + (value + 0.044715f * value * value * value))); + } + expect_close(download(output, 8), gelu_expected, 0.025f); + + std::vector pool_input(4 * 4 * 2); + for (int row = 0; row < 4; ++row) { + for (int col = 0; col < 4; ++col) { + pool_input[(row * 4 + col) * 2] = float(row * 4 + col); + pool_input[(row * 4 + col) * 2 + 1] = float(row * 4 + col + 1); + } + } + frt_buffer pool_values = allocate(ctx, "primitive_pool_values", 32); + frt_buffer pool_output = allocate(ctx, "primitive_pool_output", 8); + upload(pool_values, bf16(pool_input)); + assert(driver.avg_pool_vision_bf16( + frt_buffer_dptr(pool_values), frt_buffer_dptr(pool_output), + 1, 4, 4, 2, 2, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + expect_close(download(pool_output, 8), + {2.5f, 3.5f, 4.5f, 5.5f, 10.5f, 11.5f, 12.5f, 13.5f}, + 0.0f); + + const std::vector host_qkv = { + 1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16}; + frt_buffer qkv = allocate(ctx, "primitive_qkv", 16); + frt_buffer rope = allocate(ctx, "primitive_rope", 4); + frt_buffer query = allocate(ctx, "primitive_query", 8); + frt_buffer key = allocate(ctx, "primitive_key", 8); + frt_buffer value = allocate(ctx, "primitive_value", 8); + frt_buffer devpos = frt_buffer_alloc(ctx, "primitive_devpos", sizeof(int)); + assert(devpos); + upload(qkv, bf16(host_qkv)); + upload(rope, bf16({1, 0, 0, 1})); + const int position = 1; + assert(cudaMemcpy(frt_buffer_dptr(devpos), &position, sizeof(position), + cudaMemcpyHostToDevice) == cudaSuccess); + assert(cudaMemset(frt_buffer_dptr(key), 0, 8 * sizeof(std::uint16_t)) == + cudaSuccess); + assert(cudaMemset(frt_buffer_dptr(value), 0, 8 * sizeof(std::uint16_t)) == + cudaSuccess); + assert(driver.qkv_split_rope_devpos_bf16( + frt_buffer_dptr(qkv), frt_buffer_dptr(rope), + frt_buffer_dptr(query), frt_buffer_dptr(key), + frt_buffer_dptr(value), frt_buffer_dptr(devpos), 2, 4, 2, + 2, 2, native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + expect_close(download(query, 8), {1, 2, 3, 4, -10, 9, -12, 11}, 0.0f); + expect_close(download(key, 8), {0, 0, 5, 6, -14, 13, 0, 0}, 0.0f); + expect_close(download(value, 8), {0, 0, 7, 8, 15, 16, 0, 0}, 0.0f); + + const std::size_t image_elements = 224 * 224 * 3; + frt_buffer image = allocate(ctx, "primitive_image", image_elements); + frt_buffer patches = allocate(ctx, "primitive_patches", image_elements); + std::vector image_bits(image_elements); + for (std::size_t i = 0; i < image_bits.size(); ++i) { + image_bits[i] = static_cast(i); + } + upload(image, image_bits); + assert(driver.patch_im2col_16bit( + frt_buffer_dptr(image), frt_buffer_dptr(patches), 1, + native_stream) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + std::vector patch_bits(image_elements); + assert(cudaMemcpy(patch_bits.data(), frt_buffer_dptr(patches), + patch_bits.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + for (int patch = 0; patch < 256; ++patch) { + const int patch_row = patch / 16; + const int patch_col = patch % 16; + for (int feature = 0; feature < 588; ++feature) { + const int pixel_row = feature / 42; + const int pixel_col = (feature % 42) / 3; + const int channel = feature % 3; + const std::size_t source = + static_cast(patch_row * 14 + pixel_row) * 224 * 3 + + (patch_col * 14 + pixel_col) * 3 + channel; + assert(patch_bits[patch * 588 + feature] == image_bits[source]); + } + } + + frt_graph graph = frt_graph_create(ctx, "native_forward_primitives", 7); + assert(graph); + CaptureArgs capture{ + &driver, frt_buffer_dptr(x), frt_buffer_dptr(weight), + frt_buffer_dptr(output), frt_buffer_dptr(qkv), frt_buffer_dptr(rope), + frt_buffer_dptr(query), frt_buffer_dptr(key), frt_buffer_dptr(value), + frt_buffer_dptr(devpos), false}; + assert(frt_graph_capture(graph, 7, record_primitives, &capture) == FRT_OK); + assert(capture.recorded); + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + assert(stream_id >= 0); + for (int i = 0; i < 100; ++i) { + assert(frt_graph_replay(graph, 7, stream_id) == FRT_OK); + } + assert(frt_graph_variant_count(graph) == 1); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + + frt_graph_destroy(graph); + assert(cudaStreamDestroy(stream) == cudaSuccess); + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native forward primitives\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_kernel_driver.cpp b/cpp/tests/test_pi05_native_kernel_driver.cpp new file mode 100644 index 00000000..1bc11324 --- /dev/null +++ b/cpp/tests/test_pi05_native_kernel_driver.cpp @@ -0,0 +1,139 @@ +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include + +#include +#include +#include +#include + +namespace { + +struct CaptureArgs { + flashrt::models::pi05::NativeKernelDriver* driver = nullptr; + void* a = nullptr; + void* b = nullptr; + void* output = nullptr; + const void* bias = nullptr; + bool recorded = false; +}; + +bool has_cuda_device() { + int count = 0; + const cudaError_t rc = cudaGetDeviceCount(&count); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return count > 0; +} + +void record_gemm(void* user, void* stream) { + auto* args = static_cast(user); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + args->recorded = + args->driver + ->bf16_nn(args->a, args->b, args->output, 2, 2, 3, + native_stream) + .ok_status() && + args->driver + ->add_bias_bf16(args->output, args->bias, 2, 2, native_stream) + .ok_status() && + args->driver->silu_bf16(args->output, 4, native_stream).ok_status(); +} + +} // namespace + +int main() { + if (!has_cuda_device()) { + std::printf("SKIP - no CUDA device\n"); + return 0; + } + using flashrt::modalities::bfloat16_to_float; + using flashrt::modalities::float_to_bfloat16; + + flashrt::models::pi05::NativeKernelDriver driver; + assert(driver.status().ok_status()); + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + frt_buffer a = frt_buffer_alloc(ctx, "a", 6 * sizeof(std::uint16_t)); + frt_buffer b = frt_buffer_alloc(ctx, "b", 6 * sizeof(std::uint16_t)); + frt_buffer output = + frt_buffer_alloc(ctx, "output", 4 * sizeof(std::uint16_t)); + frt_buffer bias = frt_buffer_alloc(ctx, "bias", 2 * sizeof(std::uint16_t)); + assert(a && b && output && bias); + const std::vector host_a = { + float_to_bfloat16(1), float_to_bfloat16(2), float_to_bfloat16(3), + float_to_bfloat16(4), float_to_bfloat16(5), float_to_bfloat16(6)}; + const std::vector host_b = { + float_to_bfloat16(1), float_to_bfloat16(2), + float_to_bfloat16(3), float_to_bfloat16(4), + float_to_bfloat16(5), float_to_bfloat16(6)}; + const std::vector host_bias = { + float_to_bfloat16(-24), float_to_bfloat16(-30)}; + assert(cudaMemcpy(frt_buffer_dptr(a), host_a.data(), + host_a.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); + assert(cudaMemcpy(frt_buffer_dptr(b), host_b.data(), + host_b.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); + assert(cudaMemcpy(frt_buffer_dptr(bias), host_bias.data(), + host_bias.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); + + cudaStream_t stream = nullptr; + assert(cudaStreamCreate(&stream) == cudaSuccess); + assert(driver.bf16_nn(frt_buffer_dptr(a), frt_buffer_dptr(b), + frt_buffer_dptr(output), 2, 2, 3, + reinterpret_cast(stream)) + .ok_status()); + assert(driver.add_bias_bf16(frt_buffer_dptr(output), + frt_buffer_dptr(bias), 2, 2, + reinterpret_cast(stream)) + .ok_status()); + assert(driver.silu_bf16(frt_buffer_dptr(output), 4, + reinterpret_cast(stream)) + .ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + + frt_graph graph = frt_graph_create(ctx, "native_bf16_gemm", 1); + assert(graph); + assert(frt_graph_bind(graph, "a", a) == FRT_OK); + assert(frt_graph_bind(graph, "b", b) == FRT_OK); + assert(frt_graph_bind(graph, "output", output) == FRT_OK); + assert(frt_graph_bind(graph, "bias", bias) == FRT_OK); + CaptureArgs capture{&driver, frt_buffer_dptr(a), frt_buffer_dptr(b), + frt_buffer_dptr(output), frt_buffer_dptr(bias), false}; + assert(frt_graph_capture(graph, 1, record_gemm, &capture) == FRT_OK); + assert(capture.recorded); + assert(frt_graph_variant_count(graph) == 1); + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + assert(stream_id >= 0); + for (int i = 0; i < 100; ++i) { + assert(frt_graph_replay(graph, 1, stream_id) == FRT_OK); + } + assert(frt_graph_variant_count(graph) == 1); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + + std::vector host_output(4); + assert(cudaMemcpy(host_output.data(), frt_buffer_dptr(output), + host_output.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + const float expected[] = { + -2.0f / (1.0f + std::exp(2.0f)), + -2.0f / (1.0f + std::exp(2.0f)), + 25.0f / (1.0f + std::exp(-25.0f)), + 34.0f / (1.0f + std::exp(-34.0f))}; + for (std::size_t i = 0; i < host_output.size(); ++i) { + assert(std::fabs(bfloat16_to_float(host_output[i]) - expected[i]) < + 0.02f); + } + frt_graph_destroy(graph); + assert(cudaStreamDestroy(stream) == cudaSuccess); + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native kernel driver capture\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp new file mode 100644 index 00000000..655d94f1 --- /dev/null +++ b/cpp/tests/test_pi05_native_open.cpp @@ -0,0 +1,306 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/models/pi05/native_weights.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +namespace { + +std::string make_temp_dir() { + char tmpl[] = "/tmp/frt_pi05_native_open_XXXXXX"; + char* path = ::mkdtemp(tmpl); + assert(path); + return path; +} + +void write_file(const std::string& path) { + std::ofstream f(path, std::ios::binary); + f << "x"; + assert(f.good()); +} + +void write_raw_safetensors(const std::string& path, + const std::string& header, + const std::string& payload) { + std::ofstream f(path, std::ios::binary); + uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char b = static_cast((n >> (8 * i)) & 0xffu); + f.write(&b, 1); + } + f.write(header.data(), static_cast(header.size())); + f.write(payload.data(), static_cast(payload.size())); + assert(f.good()); +} + +void write_raw_safetensors(const std::string& path, + const std::string& header, + uint64_t payload_bytes) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char b = static_cast((n >> (8 * i)) & 0xffu); + f.write(&b, 1); + } + f.write(header.data(), static_cast(header.size())); + if (payload_bytes) { + f.seekp(static_cast(payload_bytes - 1), std::ios::cur); + const char zero = 0; + f.write(&zero, 1); + } + assert(f.good()); +} + +void append_f32(std::string* out, float value) { + char bytes[sizeof(float)]; + std::memcpy(bytes, &value, sizeof(value)); + out->append(bytes, sizeof(bytes)); +} + +void write_safetensors(const std::string& path, + const std::string& omit = "") { + const auto& entries = + flashrt::models::pi05::native_tensor_requirements(); + std::string header = "{"; + uint64_t offset = 0; + bool first = true; + for (size_t i = 0; i < entries.size(); ++i) { + const auto& e = entries[i]; + if (e.key == omit) continue; + if (!first) header += ","; + first = false; + header += "\""; + header += e.key; + header += "\":{\"dtype\":\"F32\",\"shape\":["; + uint64_t elements = 1; + for (size_t dim = 0; dim < e.shape.size(); ++dim) { + if (dim) header += ","; + header += std::to_string(e.shape[dim]); + elements *= e.shape[dim]; + } + header += "]"; + header += ",\"data_offsets\":["; + header += std::to_string(offset); + header += ","; + offset += elements * sizeof(float); + header += std::to_string(offset); + header += "]}"; + } + header += ",\"__metadata__\":{\"format\":\"pt\"}}"; + write_raw_safetensors(path, header, offset); +} + +void write_bad_safetensors(const std::string& path) { + const uint64_t bytes = 1001ull * 2048ull * 2ull; + write_raw_safetensors( + path, + "{\"paligemma_with_expert.paligemma.lm_head.weight\":{" + "\"dtype\":\"BF16\",\"shape\":[1001,2048]," + "\"data_offsets\":[0," + std::to_string(bytes) + "]}}", + 1024); +} + +void write_lerobot_policy_stats(const std::string& root, bool valid = true) { + std::string state_payload; + for (int i = 0; i < 8; ++i) append_f32(&state_payload, 0.0f); + for (int i = 0; i < 8; ++i) append_f32(&state_payload, valid ? 1.0f : 0.0f); + write_raw_safetensors( + root + "/policy_preprocessor_step_2_normalizer_processor.safetensors", + "{\"observation.state.q01\":{\"dtype\":\"F32\",\"shape\":[8]," + "\"data_offsets\":[0,32]}," + "\"observation.state.q99\":{\"dtype\":\"F32\",\"shape\":[8]," + "\"data_offsets\":[32,64]}}", + state_payload); + std::string action_payload; + for (int i = 0; i < 7; ++i) append_f32(&action_payload, 0.0f); + for (int i = 0; i < 7; ++i) append_f32(&action_payload, valid ? 1.0f : 0.0f); + write_raw_safetensors( + root + "/policy_postprocessor_step_0_unnormalizer_processor.safetensors", + "{\"action.q01\":{\"dtype\":\"F32\",\"shape\":[7]," + "\"data_offsets\":[0,28]}," + "\"action.q99\":{\"dtype\":\"F32\",\"shape\":[7]," + "\"data_offsets\":[28,56]}}", + action_payload); +} + +void write_norm_stats(const std::string& path, bool valid = true) { + std::ofstream f(path); + f << "{" + << "\"norm_stats\":{" + << "\"state\":{\"q01\":[0,0,0,0,0,0,0,0]," + << (valid ? "\"q99\":[1,1,1,1,1,1,1,1]}," + : "\"q99\":[0,0,0,0,0,0,0,0]},") + << "\"actions\":{\"q01\":[0,0,0,0,0,0,0]," + << (valid ? "\"q99\":[1,1,1,1,1,1,1]}" + : "\"q99\":[0,0,0,0,0,0,0]}") + << "}}"; + assert(f.good()); +} + +std::string config(const std::string& ckpt, + const std::string& tokenizer, + const char* extra = "") { + return std::string("{") + + "\"io\":\"native_v2\"," + + "\"checkpoint_path\":\"" + ckpt + "\"," + + "\"tokenizer_model_path\":\"" + tokenizer + "\"," + + "\"state_prompt_mode\":\"fixed\"," + + "\"max_prompt_tokens\":200," + + "\"state_dim\":8," + + "\"num_views\":2," + + "\"chunk\":10" + + extra + "}"; +} + +} // namespace + +int main() { + const auto& inventory = + flashrt::models::pi05::native_tensor_requirements(); + assert(inventory.size() == 812); + const auto has_key = [&inventory](const char* key) { + return std::any_of(inventory.begin(), inventory.end(), + [key](const auto& item) { return item.key == key; }); + }; + assert(has_key( + "paligemma_with_expert.paligemma.model.vision_tower.vision_model." + "encoder.layers.26.mlp.fc2.weight")); + assert(has_key( + "paligemma_with_expert.paligemma.model.language_model.layers.17." + "mlp.down_proj.weight")); + assert(has_key( + "paligemma_with_expert.gemma_expert.model.layers.17." + "post_attention_layernorm.dense.bias")); + + frt_model_runtime_v1* out = reinterpret_cast(0x1); + int rc = frt_model_runtime_open_v1(nullptr, &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "null")); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1("{", &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "JSON")); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + "{\"io\":\"native\",\"checkpoint_path\":\"/tmp\"," + "\"tokenizer_model_path\":\"/tmp/x\"," + "\"state_prompt_mode\":\"fixed\"," + "\"max_prompt_tokens\":200,\"state_dim\":1}", + &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "native_v2")); + + const std::string root = make_temp_dir(); + const std::string tokenizer = root + "/tokenizer.model"; + write_file(tokenizer); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), + "model.safetensors")); + + write_bad_safetensors(root + "/model.safetensors"); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "byte range")); + + const std::string missing_key = + flashrt::models::pi05::native_tensor_requirements().back().key; + write_safetensors(root + "/model.safetensors", missing_key); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), + missing_key.c_str())); + + write_safetensors(root + "/model.safetensors"); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "norm_stats")); + + write_norm_stats(root + "/norm_stats.json", false); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "q01/q99")); + + write_norm_stats(root + "/norm_stats.json"); +#ifndef FLASHRT_CPP_PI05_NATIVE_OPEN_ENABLED + const std::string good = config(root, tokenizer); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(good.c_str(), &out); + assert(rc == -3); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "validated")); +#endif + + const std::string short_prompt = + std::string("{") + + "\"io\":\"native_v2\"," + + "\"checkpoint_path\":\"" + root + "\"," + + "\"tokenizer_model_path\":\"" + tokenizer + "\"," + + "\"state_prompt_mode\":\"fixed\"," + + "\"max_prompt_tokens\":199," + + "\"state_dim\":8}"; + rc = frt_model_runtime_open_v1(short_prompt.c_str(), &out); + assert(rc == -1); + assert(std::strstr(frt_pi05_native_open_last_error(), + "max_prompt_tokens")); + + const std::string lerobot_root = make_temp_dir(); + write_safetensors(lerobot_root + "/model.safetensors"); + write_lerobot_policy_stats(lerobot_root, false); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(lerobot_root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + + write_lerobot_policy_stats(lerobot_root); +#ifndef FLASHRT_CPP_PI05_NATIVE_OPEN_ENABLED + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(lerobot_root, tokenizer).c_str(), &out); + assert(rc == -3); + assert(out == nullptr); +#endif + + ::unlink((lerobot_root + "/model.safetensors").c_str()); + ::unlink((lerobot_root + + "/policy_preprocessor_step_2_normalizer_processor.safetensors") + .c_str()); + ::unlink((lerobot_root + + "/policy_postprocessor_step_0_unnormalizer_processor.safetensors") + .c_str()); + ::rmdir(lerobot_root.c_str()); + + ::unlink(tokenizer.c_str()); + ::unlink((root + "/model.safetensors").c_str()); + ::unlink((root + "/norm_stats.json").c_str()); + ::rmdir(root.c_str()); + std::printf("PASS - Pi05 native open scaffold\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_quantization.cpp b/cpp/tests/test_pi05_native_quantization.cpp new file mode 100644 index 00000000..6ae4dfaa --- /dev/null +++ b/cpp/tests/test_pi05_native_quantization.cpp @@ -0,0 +1,35 @@ +#include "flashrt/cpp/models/pi05/native_quantization.h" + +#include +#include +#include + +int main() { + using namespace flashrt::models::pi05; + + NativeFloatTensor fp8_input{ + {1, 5}, {-448.0f, -1.0f, 0.0f, 1.0f, 448.0f}}; + NativeFp8Tensor fp8; + assert(native_quantize_fp8_e4m3(fp8_input, false, &fp8).ok_status()); + assert(fp8.shape == fp8_input.shape); + assert(fp8.scale == 1.0f); + assert(fp8.values == std::vector( + {0xfe, 0xb8, 0x00, 0x38, 0x7e})); + + NativeFloatTensor int8_input{{2, 3}, {1, 2, 3, 4, 5, 6}}; + NativeInt8Tensor int8; + assert(native_quantize_int8_per_output(int8_input, &int8).ok_status()); + assert(int8.shape == std::vector({3, 2})); + assert(int8.values == + std::vector({32, 127, 51, 127, 64, 127})); + assert(int8.scales.size() == 3); + assert(std::fabs(int8.scales[0] - 4.0f / 127.0f) < 1e-9f); + assert(std::fabs(int8.scales[1] - 5.0f / 127.0f) < 1e-9f); + assert(std::fabs(int8.scales[2] - 6.0f / 127.0f) < 1e-9f); + + NativeFloatTensor invalid{{2}, {1, 2}}; + assert(!native_quantize_fp8_e4m3(invalid, false, &fp8).ok_status()); + assert(!native_quantize_int8_per_output(invalid, &int8).ok_status()); + std::printf("PASS - Pi0.5 native weight quantization\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_rtx_attention.cpp b/cpp/tests/test_pi05_native_rtx_attention.cpp new file mode 100644 index 00000000..e5e8e179 --- /dev/null +++ b/cpp/tests/test_pi05_native_rtx_attention.cpp @@ -0,0 +1,73 @@ +#include "flashrt/cpp/models/pi05/native_rtx_attention.h" + +#include + +#include +#include + +int main() { + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess || count == 0) { + cudaGetLastError(); + std::printf("SKIP - no CUDA device\n"); + return 0; + } + using namespace flashrt::models::pi05; + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + { + NativeRtxAttentionWorkspace attention(ctx); + NativeRtxAttentionConfig bad; + bad.encoder_layers = 17; + assert(!attention.allocate(bad).ok_status()); + NativeRtxAttentionConfig config; + assert(attention.allocate(config).ok_status()); + assert(attention.size() == 22); + assert(attention.allocated_bytes() > 0); + assert(attention.encoder_splits() == 12); + assert(attention.decoder_splits() == 12); + assert(attention.kv_layer_stride_bytes() == 722 * 256 * 2); + assert(attention.find("attn_enc_K")->shape == + std::vector({18, 722, 1, 256})); + assert(attention.find("attn_enc_lse")->shape == + std::vector({1, 8, 768})); + assert(attention.find("attn_dec_lse")->shape == + std::vector({1, 8, 128})); + void* base = frt_buffer_dptr(attention.find("attn_enc_K")->buffer); + assert(attention.encoder_k_layer_dptr(0) == base); + assert(static_cast(attention.encoder_k_layer_dptr(17)) == + static_cast(base) + + 17 * attention.kv_layer_stride_bytes()); + assert(!attention.encoder_k_layer_dptr(18)); + + void* seqused_ptr = + frt_buffer_dptr(attention.find("attn_enc_seqused")->buffer); + const std::size_t bytes = attention.allocated_bytes(); + for (int i = 0; i < 1000; ++i) { + assert(attention.set_fixed_prompt_length(i % 201).ok_status()); + assert(frt_buffer_dptr( + attention.find("attn_enc_seqused")->buffer) == + seqused_ptr); + assert(attention.allocated_bytes() == bytes); + } + std::int32_t enc = 0; + std::int32_t dec = 0; + std::int32_t pos = 0; + assert(cudaMemcpy(&enc, frt_buffer_dptr( + attention.find("attn_enc_seqused")->buffer), + sizeof(enc), cudaMemcpyDeviceToHost) == cudaSuccess); + assert(cudaMemcpy(&dec, frt_buffer_dptr( + attention.find("attn_dec_seqused")->buffer), + sizeof(dec), cudaMemcpyDeviceToHost) == cudaSuccess); + assert(cudaMemcpy(&pos, frt_buffer_dptr( + attention.find("attn_dec_devpos")->buffer), + sizeof(pos), cudaMemcpyDeviceToHost) == cudaSuccess); + assert(enc == 512 + (999 % 201)); + assert(dec == enc + 10); + assert(pos == enc); + assert(!attention.set_fixed_prompt_length(201).ok_status()); + } + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native RTX attention workspace\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_rtx_attention_driver.cpp b/cpp/tests/test_pi05_native_rtx_attention_driver.cpp new file mode 100644 index 00000000..80f01126 --- /dev/null +++ b/cpp/tests/test_pi05_native_rtx_attention_driver.cpp @@ -0,0 +1,157 @@ +#include "flashrt/cpp/models/pi05/native_rtx_attention_driver.h" +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include + +#include +#include +#include +#include +#include + +namespace { + +using flashrt::models::pi05::NativeAttentionBuffer; +using flashrt::models::pi05::NativeRtxAttentionDriver; + +struct CaptureArgs { + const NativeRtxAttentionDriver* driver = nullptr; + bool recorded = false; +}; + +void record_attention(void* user, void* stream) { + auto* args = static_cast(user); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + args->recorded = + args->driver->vision(native_stream).ok_status() && + args->driver->encoder(0, native_stream).ok_status() && + args->driver->decoder(0, native_stream).ok_status(); +} + +std::size_t elements(const NativeAttentionBuffer* buffer) { + assert(buffer); + std::size_t count = 1; + for (std::uint64_t dim : buffer->shape) count *= dim; + return count; +} + +void upload_constant(const NativeAttentionBuffer* buffer, float value) { + std::vector host( + elements(buffer), flashrt::modalities::float_to_bfloat16(value)); + assert(cudaMemcpy(frt_buffer_dptr(buffer->buffer), host.data(), + host.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); +} + +void upload_kv_rows(const NativeAttentionBuffer* buffer, int total_kv) { + std::vector host(elements(buffer), 0); + for (int row = 0; row < total_kv; ++row) { + const std::uint16_t value = + flashrt::modalities::float_to_bfloat16(float(row + 1)); + for (int column = 0; column < 256; ++column) { + host[static_cast(row) * 256 + column] = value; + } + } + assert(cudaMemcpy(frt_buffer_dptr(buffer->buffer), host.data(), + host.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) == cudaSuccess); +} + +void expect_constant(void* device, std::size_t count, float expected) { + std::vector host(count); + assert(cudaMemcpy(host.data(), device, count * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + for (std::uint16_t value : host) { + assert(std::fabs(flashrt::modalities::bfloat16_to_float(value) - + expected) < 0.02f); + } +} + +} // namespace + +int main() { + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess || count == 0) { + cudaGetLastError(); + std::printf("SKIP - no CUDA device\n"); + return 0; + } + cudaDeviceProp properties{}; + assert(cudaGetDeviceProperties(&properties, 0) == cudaSuccess); + if (properties.major < 8) { + std::printf("SKIP - BF16 FA2 needs compute capability 8.0+\n"); + return 0; + } + + using namespace flashrt::models::pi05; + NativeRtxAttentionDriver invalid_driver(nullptr); + assert(!invalid_driver.status().ok_status()); + + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + NativeRtxAttentionWorkspace workspace(ctx); + NativeRtxAttentionConfig config; + config.num_views = 1; + config.encoder_sequence = 128; + config.encoder_vision_sequence = 2; + config.chunk_size = 2; + assert(workspace.allocate(config).ok_status()); + assert(workspace.decoder_splits() == 3); + assert(workspace.set_fixed_prompt_length(1).ok_status()); + + NativeRtxAttentionDriver driver(&workspace); + assert(driver.status().ok_status()); + assert(driver.num_sms() == properties.multiProcessorCount); + + upload_constant(workspace.find("attn_vis_Q"), 0.0f); + upload_constant(workspace.find("attn_vis_K"), 0.0f); + upload_constant(workspace.find("attn_vis_V"), 2.0f); + upload_constant(workspace.find("attn_enc_Q"), 0.0f); + upload_constant(workspace.find("attn_dec_Q"), 0.0f); + upload_kv_rows(workspace.find("attn_enc_K"), 130); + upload_kv_rows(workspace.find("attn_enc_V"), 130); + + cudaStream_t stream = nullptr; + assert(cudaStreamCreate(&stream) == cudaSuccess); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + assert(driver.vision(native_stream).ok_status()); + assert(driver.encoder(0, native_stream).ok_status()); + assert(driver.decoder(0, native_stream).ok_status()); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + expect_constant(driver.vision_output(), 256 * 16 * 72, 2.0f); + expect_constant(driver.encoder_output(), 128 * 8 * 256, 2.0f); + expect_constant(driver.decoder_output(), 2 * 8 * 256, 3.0f); + + frt_graph graph = frt_graph_create(ctx, "native_rtx_attention", 1); + assert(graph); + assert(frt_graph_bind(graph, "vis_q", + workspace.find("attn_vis_Q")->buffer) == FRT_OK); + assert(frt_graph_bind(graph, "enc_q", + workspace.find("attn_enc_Q")->buffer) == FRT_OK); + assert(frt_graph_bind(graph, "dec_q", + workspace.find("attn_dec_Q")->buffer) == FRT_OK); + CaptureArgs capture{&driver, false}; + assert(frt_graph_capture(graph, 1, record_attention, &capture) == FRT_OK); + assert(capture.recorded); + assert(frt_graph_variant_count(graph) == 1); + const int stream_id = frt_ctx_wrap_stream(ctx, stream); + assert(stream_id >= 0); + assert(workspace.set_fixed_prompt_length(0).ok_status()); + for (int i = 0; i < 100; ++i) { + assert(frt_graph_replay(graph, 1, stream_id) == FRT_OK); + } + assert(frt_graph_variant_count(graph) == 1); + assert(cudaStreamSynchronize(stream) == cudaSuccess); + expect_constant(driver.vision_output(), 256 * 16 * 72, 2.0f); + expect_constant(driver.encoder_output(), 128 * 8 * 256, 1.5f); + expect_constant(driver.decoder_output(), 2 * 8 * 256, 2.5f); + + frt_graph_destroy(graph); + assert(cudaStreamDestroy(stream) == cudaSuccess); + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native RTX FA2 driver\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_style_precompute.cpp b/cpp/tests/test_pi05_native_style_precompute.cpp new file mode 100644 index 00000000..0ec5b5dd --- /dev/null +++ b/cpp/tests/test_pi05_native_style_precompute.cpp @@ -0,0 +1,27 @@ +#include "flashrt/cpp/models/pi05/native_style_precompute.h" + +#include + +#include +#include + +int main() { + int count = 0; + if (cudaGetDeviceCount(&count) != cudaSuccess || count == 0) { + cudaGetLastError(); + std::printf("SKIP - no CUDA device\n"); + return 0; + } + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + flashrt::models::pi05::NativeWorkspace workspace(ctx); + flashrt::models::pi05::NativeWorkspaceConfig config; + assert(workspace.allocate(config).ok_status()); + flashrt::models::pi05::NativeDeviceWeightStore weights(ctx); + flashrt::models::pi05::NativeKernelDriver driver; + flashrt::models::pi05::NativeStylePrecomputer precomputer(&driver); + assert(!precomputer.run(weights, &workspace, 0).ok_status()); + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native style precompute validation\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp new file mode 100644 index 00000000..b3f22712 --- /dev/null +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -0,0 +1,340 @@ +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include "flashrt/cpp/models/pi05/native_weight_packer.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct Entry { + std::string key; + std::vector shape; + std::vector values; +}; + +bool has_cuda_device() { + int count = 0; + const cudaError_t rc = cudaGetDeviceCount(&count); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return count > 0; +} + +std::string temp_path() { + char path[] = "/tmp/frt_pi05_materializer_XXXXXX"; + const int fd = ::mkstemp(path); + assert(fd >= 0); + ::close(fd); + return path; +} + +std::vector sequence(std::size_t count, float start) { + std::vector values(count); + for (std::size_t i = 0; i < count; ++i) { + values[i] = start + static_cast(i) * 0.01f; + } + return values; +} + +void write_checkpoint(const std::string& path, + const std::vector& entries) { + std::string header = "{"; + std::uint64_t offset = 0; + for (std::size_t i = 0; i < entries.size(); ++i) { + const Entry& entry = entries[i]; + if (i) header += ','; + header += '"' + entry.key + "\":{\"dtype\":\"F32\",\"shape\":["; + for (std::size_t d = 0; d < entry.shape.size(); ++d) { + if (d) header += ','; + header += std::to_string(entry.shape[d]); + } + header += "],\"data_offsets\":[" + std::to_string(offset) + ','; + offset += entry.values.size() * sizeof(float); + header += std::to_string(offset) + "]}"; + } + header += '}'; + std::ofstream file(path, std::ios::binary | std::ios::trunc); + const std::uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char byte = static_cast((n >> (8 * i)) & 0xffu); + file.write(&byte, 1); + } + file.write(header.data(), static_cast(header.size())); + for (const Entry& entry : entries) { + file.write(reinterpret_cast(entry.values.data()), + static_cast(entry.values.size() * + sizeof(float))); + } + assert(file.good()); +} + +} // namespace + +int main() { + if (!has_cuda_device()) { + std::printf("SKIP - no CUDA device\n"); + return 0; + } + const std::string prefix = + "paligemma_with_expert.paligemma.model.language_model.layers.0"; + const std::string decoder = + "paligemma_with_expert.gemma_expert.model.layers.0"; + const std::string vision = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; + const std::string vision_layer = vision + ".encoder.layers.0"; + const std::vector entries = { + {prefix + ".input_layernorm.weight", {4}, {-0.5f, 0.0f, 0.5f, 1.0f}}, + {prefix + ".self_attn.q_proj.weight", {16, 4}, sequence(64, 0.1f)}, + {prefix + ".self_attn.k_proj.weight", {4, 4}, sequence(16, 1.0f)}, + {prefix + ".self_attn.v_proj.weight", {4, 4}, sequence(16, 2.0f)}, + {prefix + ".self_attn.o_proj.weight", {4, 8}, sequence(32, 3.0f)}, + {prefix + ".post_attention_layernorm.weight", {4}, + {-0.25f, 0.0f, 0.25f, 0.5f}}, + {prefix + ".mlp.gate_proj.weight", {6, 4}, sequence(24, 4.0f)}, + {prefix + ".mlp.up_proj.weight", {6, 4}, sequence(24, 5.0f)}, + {prefix + ".mlp.down_proj.weight", {4, 6}, sequence(24, 6.0f)}, + {decoder + ".self_attn.q_proj.weight", {16, 4}, sequence(64, 7.0f)}, + {decoder + ".self_attn.k_proj.weight", {4, 4}, sequence(16, 8.0f)}, + {decoder + ".self_attn.v_proj.weight", {4, 4}, sequence(16, 9.0f)}, + {decoder + ".self_attn.o_proj.weight", {4, 16}, sequence(64, 10.0f)}, + {decoder + ".mlp.gate_proj.weight", {6, 4}, sequence(24, 11.0f)}, + {decoder + ".mlp.up_proj.weight", {6, 4}, sequence(24, 12.0f)}, + {decoder + ".mlp.down_proj.weight", {4, 6}, sequence(24, 13.0f)}, + {decoder + ".input_layernorm.dense.weight", {12, 4}, + sequence(48, 14.0f)}, + {decoder + ".input_layernorm.dense.bias", {12}, sequence(12, 15.0f)}, + {decoder + ".post_attention_layernorm.dense.weight", {12, 4}, + sequence(48, 16.0f)}, + {decoder + ".post_attention_layernorm.dense.bias", {12}, + sequence(12, 17.0f)}, + {vision + ".embeddings.patch_embedding.weight", {2, 2, 2, 1}, + sequence(8, 18.0f)}, + {vision + ".embeddings.patch_embedding.bias", {2}, + sequence(2, 19.0f)}, + {vision + ".embeddings.position_embedding.weight", {3, 2}, + sequence(6, 20.0f)}, + {vision + ".post_layernorm.weight", {2}, sequence(2, 21.0f)}, + {vision + ".post_layernorm.bias", {2}, sequence(2, 22.0f)}, + {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear." + "weight", {4, 2}, sequence(8, 23.0f)}, + {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear." + "bias", {4}, sequence(4, 24.0f)}, + {vision_layer + ".self_attn.q_proj.weight", {2, 2}, + sequence(4, 25.0f)}, + {vision_layer + ".self_attn.q_proj.bias", {2}, + sequence(2, 26.0f)}, + {vision_layer + ".self_attn.k_proj.weight", {2, 2}, + sequence(4, 27.0f)}, + {vision_layer + ".self_attn.k_proj.bias", {2}, + sequence(2, 28.0f)}, + {vision_layer + ".self_attn.v_proj.weight", {2, 2}, + sequence(4, 29.0f)}, + {vision_layer + ".self_attn.v_proj.bias", {2}, + sequence(2, 30.0f)}, + {vision_layer + ".self_attn.out_proj.weight", {2, 2}, + sequence(4, 31.0f)}, + {vision_layer + ".self_attn.out_proj.bias", {2}, + sequence(2, 32.0f)}, + {vision_layer + ".mlp.fc1.weight", {3, 2}, + sequence(6, 33.0f)}, + {vision_layer + ".mlp.fc1.bias", {3}, sequence(3, 34.0f)}, + {vision_layer + ".mlp.fc2.weight", {2, 3}, + sequence(6, 35.0f)}, + {vision_layer + ".mlp.fc2.bias", {2}, sequence(2, 36.0f)}, + {vision_layer + ".layer_norm1.weight", {2}, sequence(2, 37.0f)}, + {vision_layer + ".layer_norm1.bias", {2}, sequence(2, 38.0f)}, + {vision_layer + ".layer_norm2.weight", {2}, sequence(2, 39.0f)}, + {vision_layer + ".layer_norm2.bias", {2}, sequence(2, 40.0f)}, + {"paligemma_with_expert.gemma_expert.model.norm.dense.weight", + {3, 2}, sequence(6, 41.0f)}, + {"paligemma_with_expert.gemma_expert.model.norm.dense.bias", + {3}, sequence(3, 42.0f)}, + {"time_mlp_in.weight", {2, 2}, sequence(4, 43.0f)}, + {"time_mlp_in.bias", {2}, sequence(2, 44.0f)}, + {"time_mlp_out.weight", {2, 2}, sequence(4, 45.0f)}, + {"time_mlp_out.bias", {2}, sequence(2, 46.0f)}, + {"action_in_proj.weight", {2, 1}, sequence(2, 47.0f)}, + {"action_in_proj.bias", {2}, sequence(2, 48.0f)}, + {"action_out_proj.weight", {1, 2}, sequence(2, 49.0f)}, + {"action_out_proj.bias", {1}, sequence(1, 50.0f)}, + {"paligemma_with_expert.paligemma.lm_head.weight", + {4, 2}, sequence(8, 51.0f)}, + }; + const std::string path = temp_path(); + write_checkpoint(path, entries); + + flashrt::loader::SafetensorsFile source; + assert(source.open(path)); + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + { + flashrt::models::pi05::NativeDeviceWeightStore destination(ctx); + flashrt::models::pi05::NativeWeightMaterializer materializer( + source, &destination); + assert(materializer.materialize_encoder_layer(0).ok_status()); + assert(destination.size() == 5); + const auto* qkv = destination.find("encoder_attn_qkv_w_0"); + assert(qkv && qkv->shape == std::vector({4, 24})); + const auto* gate = destination.find("encoder_ffn_gate_w_0"); + assert(gate && gate->shape == std::vector({4, 6})); + const auto* down = destination.find("encoder_ffn_down_w_0"); + assert(down && down->shape == std::vector({6, 4})); + assert(!materializer.materialize_encoder_layer(0).ok_status()); + assert(!materializer.materialize_encoder_layer(18).ok_status()); + assert(materializer.materialize_decoder_layer(0, true).ok_status()); + assert(destination.size() == 15); + const auto* decoder_qkv = destination.find("decoder_attn_qkv_w_0"); + assert(decoder_qkv && + decoder_qkv->shape == std::vector({4, 24})); + const auto* gate_up = destination.find("decoder_ffn_gate_up_w_0"); + assert(gate_up && + gate_up->shape == std::vector({4, 12})); + const auto* attn_mod = + destination.find("decoder_pre_attn_norm_mod_w_0"); + assert(attn_mod && + attn_mod->shape == std::vector({4, 12})); + assert(!materializer.materialize_decoder_layer(0, true).ok_status()); + assert(!materializer.materialize_decoder_layer(18, true).ok_status()); + assert(materializer.materialize_vision_globals().ok_status()); + assert(materializer.materialize_vision_layer(0).ok_status()); + assert(destination.size() == 34); + const auto* patch = destination.find("vision_patch_embedding_w"); + assert(patch && patch->shape == + std::vector({2, 1, 2, 2})); + const auto* vision_qkv = destination.find("vision_attn_qkv_w_0"); + assert(vision_qkv && + vision_qkv->shape == std::vector({2, 6})); + assert(!materializer.materialize_vision_layer(27).ok_status()); + assert(!materializer.materialize_decoder_globals(0).ok_status()); + assert(materializer.materialize_decoder_globals(10).ok_status()); + assert(destination.size() == 45); + assert(destination.find("decoder_final_norm_mod_w")->shape == + std::vector({2, 3})); + assert(destination.find("decoder_time_embeds")->shape == + std::vector({10, 1024})); + assert(destination.find("decoder_action_out_proj_w")->shape == + std::vector({2, 1})); + assert(materializer.materialize_embedding().ok_status()); + assert(destination.size() == 46); + assert(destination.find("embedding_weight")->shape == + std::vector({4, 2})); + flashrt::models::pi05::NativeWeightPacker packer(&destination); + assert(packer.pack_fp8("decoder_attn_qkv_w_0", false).ok_status()); + assert(packer.pack_int8("decoder_attn_qkv_w_0").ok_status()); + assert(destination.size() == 50); + assert(destination.find("fp8.decoder_attn_qkv_w_0")->shape == + std::vector({4, 24})); + assert(destination.find("int8.decoder_attn_qkv_w_0")->shape == + std::vector({24, 4})); + } + frt_ctx_destroy(ctx); + assert(::unlink(path.c_str()) == 0); + + const char* real_checkpoint = std::getenv("FLASH_RT_PI05_CHECKPOINT"); + if (real_checkpoint && real_checkpoint[0]) { + flashrt::loader::SafetensorsFile real_source; + assert(real_source.open(std::string(real_checkpoint) + + "/model.safetensors")); + frt_ctx real_ctx = frt_ctx_create(); + assert(real_ctx); + { + flashrt::models::pi05::NativeDeviceWeightStore destination( + real_ctx); + flashrt::models::pi05::NativeWeightMaterializer materializer( + real_source, &destination); + assert(materializer.materialize_encoder_layer(0).ok_status()); + assert(destination.size() == 5); + assert(destination.find("encoder_attn_qkv_w_0")->shape == + std::vector({2048, 2560})); + assert(destination.find("encoder_ffn_gate_w_0")->shape == + std::vector({2048, 16384})); + assert(materializer.materialize_decoder_layer(0, true).ok_status()); + assert(destination.size() == 15); + assert(destination.find("decoder_attn_qkv_w_0")->shape == + std::vector({1024, 2560})); + assert(destination.find("decoder_ffn_gate_up_w_0")->shape == + std::vector({1024, 8192})); + assert(materializer.materialize_vision_globals().ok_status()); + assert(materializer.materialize_vision_layer(0).ok_status()); + assert(destination.size() == 34); + assert(destination.find("vision_patch_embedding_w")->shape == + std::vector({14, 14, 3, 1152})); + assert(destination.find("vision_attn_qkv_w_0")->shape == + std::vector({1152, 3456})); + assert(materializer.materialize_decoder_globals(10).ok_status()); + assert(destination.size() == 45); + assert(destination.find("decoder_final_norm_mod_w")->shape == + std::vector({1024, 3072})); + assert(destination.find("decoder_time_embeds")->shape == + std::vector({10, 1024})); + assert(destination.find("decoder_action_in_proj_w")->shape == + std::vector({32, 1024})); + assert(destination.find("decoder_action_out_proj_w")->shape == + std::vector({1024, 32})); + flashrt::models::pi05::NativeWeightPacker packer(&destination); + assert(packer.pack_fp8("decoder_attn_qkv_w_0", false) + .ok_status()); + assert(packer.pack_int8("decoder_attn_qkv_w_0").ok_status()); + assert(destination.size() == 49); + assert(destination.find("fp8.decoder_attn_qkv_w_0")->shape == + std::vector({1024, 2560})); + assert(destination.find("int8.decoder_attn_qkv_w_0")->shape == + std::vector({2560, 1024})); + } + frt_ctx_destroy(real_ctx); + + const char* full = std::getenv("FLASH_RT_PI05_FULL_MATERIALIZE"); + if (full && std::strcmp(full, "1") == 0) { + frt_ctx full_ctx = frt_ctx_create(); + assert(full_ctx); + { + flashrt::models::pi05::NativeDeviceWeightStore destination( + full_ctx); + flashrt::models::pi05::NativeWeightMaterializer materializer( + real_source, &destination); + flashrt::models::pi05::NativeMaterializationOptions options; + assert(materializer.materialize_all(options).ok_status()); + assert(destination.size() == 613); + assert(destination.find("vision_attn_qkv_w_26")->shape == + std::vector({1152, 3456})); + assert(destination.find("encoder_ffn_down_w_17")->shape == + std::vector({16384, 2048})); + assert(destination.find("decoder_ffn_gate_up_w_17")->shape == + std::vector({1024, 8192})); + assert(destination.find("embedding_weight")->shape == + std::vector({257152, 2048})); + const char* pack = + std::getenv("FLASH_RT_PI05_FULL_PACK_FP8"); + if (pack && std::strcmp(pack, "1") == 0) { + flashrt::models::pi05::NativeWeightPacker packer( + &destination); + assert(packer.pack_all_fp8(false).ok_status()); + assert(destination.size() == 1137); + assert(destination.find( + "fp8.vision_attn_qkv_w_26")->dtype == + flashrt::models::pi05::NativeWeightDType:: + kFp8E4M3); + assert(destination.find( + "fp8.encoder_ffn_gate_up_w_17")->shape == + std::vector({2048, 32768})); + assert(destination.find( + "fp8.decoder_ffn_down_w_17.scale")->dtype == + flashrt::models::pi05::NativeWeightDType::kFloat32); + } + } + frt_ctx_destroy(full_ctx); + } + } + std::printf("PASS - Pi0.5 native layer materializer\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp new file mode 100644 index 00000000..d8baf08d --- /dev/null +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -0,0 +1,145 @@ +#include "flashrt/cpp/models/pi05/native_weight_ops.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using flashrt::models::pi05::NativeBf16Tensor; +using flashrt::models::pi05::NativeFloatTensor; + +void expect(const NativeFloatTensor& tensor, + std::initializer_list shape, + std::initializer_list values) { + assert(tensor.shape == std::vector(shape)); + assert(tensor.values.size() == values.size()); + std::size_t i = 0; + for (float value : values) { + assert(std::fabs(tensor.values[i++] - value) < 1e-6f); + } +} + +std::string temp_path() { + char path[] = "/tmp/frt_pi05_weight_ops_XXXXXX"; + const int fd = ::mkstemp(path); + assert(fd >= 0); + ::close(fd); + return path; +} + +void write_tensor_file(const std::string& path) { + const std::string header = + "{\"f32\":{\"dtype\":\"F32\",\"shape\":[2]," + "\"data_offsets\":[0,8]}," + "\"bf16\":{\"dtype\":\"BF16\",\"shape\":[2]," + "\"data_offsets\":[8,12]}," + "\"f16\":{\"dtype\":\"F16\",\"shape\":[2]," + "\"data_offsets\":[12,16]}}"; + const float f32[] = {1.25f, -2.5f}; + const std::uint16_t bf16[] = { + flashrt::modalities::float_to_bfloat16(3.0f), + flashrt::modalities::float_to_bfloat16(-4.0f)}; + const std::uint16_t f16[] = { + flashrt::modalities::float_to_float16(5.0f), + flashrt::modalities::float_to_float16(-6.0f)}; + std::ofstream f(path, std::ios::binary | std::ios::trunc); + const std::uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char byte = static_cast((n >> (8 * i)) & 0xffu); + f.write(&byte, 1); + } + f.write(header.data(), static_cast(header.size())); + f.write(reinterpret_cast(f32), sizeof(f32)); + f.write(reinterpret_cast(bf16), sizeof(bf16)); + f.write(reinterpret_cast(f16), sizeof(f16)); + assert(f.good()); +} + +} // namespace + +int main() { + using namespace flashrt::models::pi05; + + const std::string path = temp_path(); + write_tensor_file(path); + flashrt::loader::SafetensorsFile file; + assert(file.open(path)); + NativeFloatTensor loaded; + assert(load_native_float_tensor(file, "f32", &loaded).ok_status()); + expect(loaded, {2}, {1.25f, -2.5f}); + assert(load_native_float_tensor(file, "bf16", &loaded).ok_status()); + expect(loaded, {2}, {3.0f, -4.0f}); + assert(load_native_float_tensor(file, "f16", &loaded).ok_status()); + expect(loaded, {2}, {5.0f, -6.0f}); + assert(::unlink(path.c_str()) == 0); + + NativeFloatTensor matrix{{2, 3}, {1, 2, 3, 4, 5, 6}}; + NativeFloatTensor result; + assert(native_transpose_2d(matrix, &result).ok_status()); + expect(result, {3, 2}, {1, 4, 2, 5, 3, 6}); + + NativeFloatTensor patch{{2, 2, 2, 1}, {0, 1, 2, 3, 4, 5, 6, 7}}; + assert(native_patch_oihw_to_hwio(patch, &result).ok_status()); + expect(result, {2, 1, 2, 2}, {0, 4, 2, 6, 1, 5, 3, 7}); + + NativeFloatTensor qk{{8, 1}, {0, 1, 2, 3, 4, 5, 6, 7}}; + assert(native_interleave_qk_rows(qk, 2, &result).ok_status()); + expect(result, {8, 1}, {0, 2, 1, 3, 4, 6, 5, 7}); + + NativeFloatTensor norm{{3}, {-0.5f, 0.0f, 1.0f}}; + assert(native_fold_rms_columns(matrix, norm, &result).ok_status()); + expect(result, {2, 3}, {0.5f, 2.0f, 6.0f, 2.0f, 5.0f, 12.0f}); + + NativeFloatTensor q{{2, 2}, {1, 2, 3, 4}}; + NativeFloatTensor k{{1, 2}, {5, 6}}; + NativeFloatTensor v{{1, 2}, {7, 8}}; + assert(native_concat_rows_transpose({&q, &k, &v}, &result).ok_status()); + expect(result, {2, 4}, {1, 3, 5, 7, 2, 4, 6, 8}); + + NativeFloatTensor left{{2, 2}, {1, 2, 3, 4}}; + NativeFloatTensor right{{2, 2}, {5, 6, 7, 8}}; + assert(native_concat_columns(left, right, &result).ok_status()); + expect(result, {2, 4}, {1, 2, 5, 6, 3, 4, 7, 8}); + + NativeFloatTensor a{{2}, {1, 2}}; + NativeFloatTensor b{{3}, {3, 4, 5}}; + assert(native_concat_vectors({&a, &b}, &result).ok_status()); + expect(result, {5}, {1, 2, 3, 4, 5}); + + assert(native_scale(matrix, -0.1f, &result).ok_status()); + expect(result, {2, 3}, {-0.1f, -0.2f, -0.3f, + -0.4f, -0.5f, -0.6f}); + + assert(native_pi05_time_embeddings(2, 4, &result).ok_status()); + assert(result.shape == std::vector({2, 4})); + assert(result.values.size() == 8); + assert(!native_pi05_time_embeddings(0, 4, &result).ok_status()); + assert(!native_pi05_time_embeddings(2, 3, &result).ok_status()); + + NativeFloatTensor unrounded{{2}, {1.003f, -1.003f}}; + assert(native_round_to_bf16_float(unrounded, &result).ok_status()); + assert(result.values[0] == flashrt::modalities::bfloat16_to_float( + flashrt::modalities::float_to_bfloat16( + unrounded.values[0]))); + assert(result.values[1] == flashrt::modalities::bfloat16_to_float( + flashrt::modalities::float_to_bfloat16( + unrounded.values[1]))); + + NativeBf16Tensor converted; + assert(native_to_bf16(matrix, &converted).ok_status()); + assert(converted.shape == matrix.shape); + for (std::size_t i = 0; i < matrix.values.size(); ++i) { + assert(converted.values[i] == + flashrt::modalities::float_to_bfloat16(matrix.values[i])); + } + + assert(!native_interleave_qk_rows(matrix, 2, &result).ok_status()); + assert(!native_concat_columns(matrix, k, &result).ok_status()); + std::printf("PASS - Pi0.5 native weight transforms\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_weight_packer.cpp b/cpp/tests/test_pi05_native_weight_packer.cpp new file mode 100644 index 00000000..83e52180 --- /dev/null +++ b/cpp/tests/test_pi05_native_weight_packer.cpp @@ -0,0 +1,141 @@ +#include "flashrt/cpp/models/pi05/native_weight_packer.h" + +#include + +#include +#include +#include +#include + +namespace { + +bool has_cuda_device() { + int count = 0; + const cudaError_t rc = cudaGetDeviceCount(&count); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return count > 0; +} + +template +std::vector download(const flashrt::models::pi05::NativeDeviceWeight& weight) { + std::vector result(frt_buffer_bytes(weight.buffer) / sizeof(T)); + assert(cudaMemcpy(result.data(), frt_buffer_dptr(weight.buffer), + result.size() * sizeof(T), cudaMemcpyDeviceToHost) == + cudaSuccess); + return result; +} + +void upload(flashrt::models::pi05::NativeDeviceWeightStore* store, + const std::string& name, + const flashrt::models::pi05::NativeBf16Tensor& tensor) { + assert(store->upload(name, tensor).ok_status()); +} + +} // namespace + +int main() { + if (!has_cuda_device()) { + std::printf("SKIP - no CUDA device\n"); + return 0; + } + using namespace flashrt::models::pi05; + NativeFloatTensor source{{2, 3}, {1, 2, 3, 4, 5, 6}}; + NativeBf16Tensor bf16; + assert(native_to_bf16(source, &bf16).ok_status()); + NativeFloatTensor rounded; + assert(native_round_to_bf16_float(source, &rounded).ok_status()); + + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + { + NativeDeviceWeightStore store(ctx); + assert(store.upload("weight", bf16).ok_status()); + NativeBf16Tensor copied; + assert(store.download_bf16("weight", &copied).ok_status()); + assert(copied.values == bf16.values); + + NativeWeightPacker packer(&store); + assert(packer.pack_fp8("weight", false).ok_status()); + assert(packer.pack_fp8("weight", true).ok_status() == false); + assert(packer.pack_int8("weight").ok_status()); + assert(store.size() == 5); + + NativeFp8Tensor expected_fp8; + assert(native_quantize_fp8_e4m3(rounded, false, &expected_fp8) + .ok_status()); + const auto* fp8 = store.find("fp8.weight"); + assert(fp8 && fp8->dtype == NativeWeightDType::kFp8E4M3); + assert(download(*fp8) == expected_fp8.values); + assert(download(*store.find("fp8.weight.scale")) == + std::vector({expected_fp8.scale})); + + NativeInt8Tensor expected_int8; + assert(native_quantize_int8_per_output(rounded, &expected_int8) + .ok_status()); + const auto* int8 = store.find("int8.weight"); + assert(int8 && int8->dtype == NativeWeightDType::kInt8); + assert(download(*int8) == expected_int8.values); + assert(download(*store.find("int8.weight.scale")) == + expected_int8.scales); + assert(!packer.pack_int8("missing").ok_status()); + } + frt_ctx_destroy(ctx); + + ctx = frt_ctx_create(); + assert(ctx); + { + NativeDeviceWeightStore store(ctx); + NativeBf16Tensor tiny; + tiny.shape = {1, 1}; + tiny.values = {flashrt::modalities::float_to_bfloat16(1.0f)}; + for (int layer = 0; layer < 27; ++layer) { + for (const char* stem : { + "vision_attn_qkv_w_", "vision_attn_o_w_", + "vision_ffn_up_w_", "vision_ffn_down_w_"}) { + upload(&store, std::string(stem) + std::to_string(layer), + tiny); + } + } + upload(&store, "encoder_multi_modal_projector_w", tiny); + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + for (const std::string& name : { + "encoder_attn_qkv_w_" + suffix, + "encoder_attn_o_w_" + suffix, + "encoder_ffn_gate_w_" + suffix, + "encoder_ffn_up_w_" + suffix, + "encoder_ffn_down_w_" + suffix, + "decoder_attn_qkv_w_" + suffix, + "decoder_attn_o_w_" + suffix, + "decoder_ffn_gate_w_" + suffix, + "decoder_ffn_up_w_" + suffix, + "decoder_ffn_gate_up_w_" + suffix, + "decoder_ffn_down_w_" + suffix}) { + upload(&store, name, tiny); + } + } + assert(store.size() == 307); + NativeWeightPacker packer(&store); + assert(packer.pack_all_fp8(false).ok_status()); + assert(packer.pack_vision_int8().ok_status()); + assert(packer.pack_encoder_int8().ok_status()); + assert(packer.pack_decoder_int8().ok_status()); + assert(store.size() == 1407); + assert(store.find("fp8.vision_projector_w")->dtype == + NativeWeightDType::kFp8E4M3); + assert(store.find("fp8.encoder_ffn_gate_up_w_17")->shape == + std::vector({1, 2})); + assert(store.find("int8.vision_ffn_down_w_26.scale")->dtype == + NativeWeightDType::kFloat32); + assert(store.find("int8.encoder_ffn_up_w_17")->dtype == + NativeWeightDType::kInt8); + assert(store.find("int8.decoder_ffn_down_w_17")->dtype == + NativeWeightDType::kInt8); + } + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native weight packer\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_workspace.cpp b/cpp/tests/test_pi05_native_workspace.cpp new file mode 100644 index 00000000..78c59586 --- /dev/null +++ b/cpp/tests/test_pi05_native_workspace.cpp @@ -0,0 +1,136 @@ +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#include + +#include +#include +#include + +namespace { + +bool has_cuda_device() { + int count = 0; + const cudaError_t rc = cudaGetDeviceCount(&count); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return count > 0; +} + +void check_ones(const flashrt::models::pi05::NativeWorkspaceBuffer& buffer) { + std::vector values(buffer.shape[0]); + assert(cudaMemcpy(values.data(), frt_buffer_dptr(buffer.buffer), + values.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + for (std::uint16_t value : values) { + assert(value == flashrt::modalities::float_to_bfloat16(1.0f)); + } +} + +} // namespace + +int main() { + if (!has_cuda_device()) { + std::printf("SKIP - no CUDA device\n"); + return 0; + } + using namespace flashrt::models::pi05; + frt_ctx ctx = frt_ctx_create(); + assert(ctx); + { + NativeWorkspace workspace(ctx); + NativeWorkspaceConfig invalid; + invalid.vision_pool_factor = 3; + assert(!workspace.allocate(invalid).ok_status()); + NativeWorkspaceConfig config; + assert(workspace.allocate(config).ok_status()); + assert(workspace.logical_size() == 35); + assert(workspace.allocation_count() == 34); + assert(workspace.allocated_bytes() > 0); + assert(workspace.vision_sequence() == 512); + assert(workspace.encoder_vision_sequence() == 512); + assert(workspace.encoder_sequence() == 712); + assert(workspace.find("prompt_embedding")->shape == + std::vector({200, 2048})); + const auto* vision_x = workspace.find("vision_x"); + const auto* pooled = workspace.find("vision_x_pooled"); + assert(vision_x && pooled && pooled->alias); + assert(vision_x->buffer == pooled->buffer); + assert(workspace.find("decoder_style_attn")->shape == + std::vector({10, 18, 10, 3072})); + assert(workspace.find("rtc_prefix_weights")->dtype == + flashrt::modalities::DType::kFloat32); + check_ones(*workspace.find("encoder_rms_ones")); + check_ones(*workspace.find("decoder_rms_ones")); + assert(workspace.update_decoder_rope(37).ok_status()); + assert(!workspace.update_decoder_rope(201).ok_status()); + void* decoder_rope_ptr = + frt_buffer_dptr(workspace.find("decoder_rope_weights")->buffer); + const std::size_t allocation_count = workspace.allocation_count(); + const std::size_t allocated_bytes = workspace.allocated_bytes(); + for (int i = 0; i < 1000; ++i) { + assert(workspace.update_decoder_rope(i % 201).ok_status()); + assert(frt_buffer_dptr( + workspace.find("decoder_rope_weights")->buffer) == + decoder_rope_ptr); + assert(workspace.allocation_count() == allocation_count); + assert(workspace.allocated_bytes() == allocated_bytes); + } + + NativeDeviceWeightStore weights(ctx); + NativeBf16Tensor position; + position.shape = {256, 1152}; + position.values.resize(256 * 1152); + for (std::size_t i = 0; i < position.values.size(); ++i) { + position.values[i] = flashrt::modalities::float_to_bfloat16( + static_cast(i % 97) / 97.0f); + } + assert(weights.upload("vision_position_embedding", position) + .ok_status()); + assert(workspace.expand_vision_position_embedding(weights) + .ok_status()); + const auto* expanded = workspace.find("vision_pos_embed_expanded"); + std::vector expanded_values(position.values.size() * 2); + assert(cudaMemcpy(expanded_values.data(), + frt_buffer_dptr(expanded->buffer), + expanded_values.size() * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(std::vector(expanded_values.begin(), + expanded_values.begin() + + position.values.size()) == + position.values); + assert(std::vector( + expanded_values.begin() + position.values.size(), + expanded_values.end()) == position.values); + assert(!workspace.allocate(config).ok_status()); + } + frt_ctx_destroy(ctx); + + ctx = frt_ctx_create(); + assert(ctx); + { + NativeWorkspace workspace(ctx); + NativeWorkspaceConfig config; + config.num_views = 3; + config.max_prompt_tokens = 256; + config.chunk_size = 50; + config.num_steps = 5; + config.vision_pool_factor = 2; + assert(workspace.allocate(config).ok_status()); + assert(workspace.logical_size() == 35); + assert(workspace.allocation_count() == 35); + assert(workspace.vision_sequence() == 768); + assert(workspace.encoder_vision_sequence() == 192); + assert(workspace.encoder_sequence() == 448); + const auto* pooled = workspace.find("vision_x_pooled"); + assert(pooled && !pooled->alias); + assert(pooled->shape == std::vector({192, 1152})); + assert(pooled->buffer != workspace.find("vision_x")->buffer); + assert(workspace.find("decoder_time_emb")->shape == + std::vector({5, 50, 1024})); + } + frt_ctx_destroy(ctx); + std::printf("PASS - Pi0.5 native workspace\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_prompt_embed.cpp b/cpp/tests/test_pi05_prompt_embed.cpp new file mode 100644 index 00000000..89c05560 --- /dev/null +++ b/cpp/tests/test_pi05_prompt_embed.cpp @@ -0,0 +1,200 @@ +#include "flashrt/cpp/models/pi05/prompt_embed.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include +#include +#include +#include +#include +#include + +using flashrt::modalities::DType; +using flashrt::modalities::Layout; +using flashrt::modalities::MemoryPlace; +using flashrt::modalities::SentencePieceTokenizer; +using flashrt::modalities::Shape; +using flashrt::modalities::StatusCode; +using flashrt::modalities::TensorView; +using flashrt::modalities::TextEmbeddingStaging; +using flashrt::models::pi05::PromptEmbeddingSpec; +using flashrt::models::pi05::embed_prompt; +using flashrt::models::pi05::embed_prompt_cpu; + +namespace { + +std::string tokenizer_model_path() { + const char* env = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + return env ? std::string(env) : std::string(); +} + +bool has_cuda_device() { +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING + int n = 0; + cudaError_t rc = cudaGetDeviceCount(&n); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return n > 0; +#else + return false; +#endif +} + +void test_requires_loaded_tokenizer() { + SentencePieceTokenizer tokenizer; + std::vector table(8, 0.0f); + std::vector out(8, 0.0f); + std::vector ids; + std::uint64_t prompt_len = 0; + + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + PromptEmbeddingSpec spec{2, 4, 2, 1.0f}; + auto st = embed_prompt_cpu(tokenizer, spec, "pick", nullptr, 0, src, dst, + &ids, &prompt_len); + assert(!st.ok_status()); + assert(st.code == StatusCode::kInvalidArgument); +} + +void test_paligemma_prompt_embedding_when_configured() { +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const std::string path = tokenizer_model_path(); + if (path.empty()) { + std::cout << "SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"; + return; + } + SentencePieceTokenizer tokenizer; + auto st = tokenizer.load_model(path); + assert(st.ok_status()); + + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + std::vector out(max_tokens * hidden, 7.0f); + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{vocab, hidden}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{max_tokens, hidden}}; + + const float state[] = {0.0f, 1.0f, -1.0f}; + PromptEmbeddingSpec spec{vocab, hidden, max_tokens, 0.5f}; + std::vector ids; + ids.reserve(max_tokens + 1); + tokenizer.reserve(max_tokens); + std::string formatted; + formatted.reserve(512); + std::uint64_t prompt_len = 0; + st = embed_prompt(tokenizer, spec, "pick_up_cube", state, 3, src, dst, + &ids, &prompt_len, nullptr, nullptr, &formatted); + assert(st.ok_status()); + const std::vector expected_ids = { + 2, 7071, 235292, 4788, 908, 28660, 235269, 3040, 235292, + 235248, 235274, 235284, 235321, 235248, 235284, 235308, + 235308, 235248, 235276, 235289, 108, 4022, 235292, 235248, + }; + assert(ids == expected_ids); + assert(prompt_len == expected_ids.size()); + for (std::uint64_t i = 0; i < prompt_len; ++i) { + const float id = static_cast(expected_ids[i]); + assert(std::fabs(out[i * hidden + 0] - id * 0.5f) < 0.001f); + assert(std::fabs(out[i * hidden + 1] + id * 0.5f) < 0.001f); + } + for (std::uint64_t i = prompt_len * hidden; i < out.size(); ++i) { + assert(out[i] == 0.0f); + } + const std::size_t id_capacity = ids.capacity(); + const std::size_t formatted_capacity = formatted.capacity(); + const std::uint64_t tokenizer_capacity = tokenizer.workspace_capacity(); + for (int round = 0; round < 1000; ++round) { + st = embed_prompt(tokenizer, spec, "pick_up_cube", state, 3, src, dst, + &ids, &prompt_len, nullptr, nullptr, &formatted); + assert(st.ok_status()); + assert(ids.capacity() == id_capacity); + assert(formatted.capacity() == formatted_capacity); + assert(tokenizer.workspace_capacity() == tokenizer_capacity); + } +#endif +} + +void test_paligemma_prompt_embedding_device_when_configured() { +#if defined(FLASHRT_CPP_HAS_SENTENCEPIECE) && defined(FLASHRT_CPP_WITH_CUDA_STAGING) + const std::string path = tokenizer_model_path(); + if (path.empty() || !has_cuda_device()) { + std::cout << "SKIP - tokenizer or CUDA device not available\n"; + return; + } + SentencePieceTokenizer tokenizer; + auto st = tokenizer.load_model(path); + assert(st.ok_status()); + + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + void* d_table = nullptr; + void* d_out = nullptr; + assert(cudaMalloc(&d_table, table.size() * sizeof(float)) == cudaSuccess); + assert(cudaMalloc(&d_out, max_tokens * hidden * sizeof(float)) == + cudaSuccess); + assert(cudaMemcpy(d_table, table.data(), table.size() * sizeof(float), + cudaMemcpyHostToDevice) == cudaSuccess); + TensorView src{d_table, static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kDevice, Layout::kFlat, + Shape{vocab, hidden}}; + TensorView dst{d_out, + static_cast(max_tokens * hidden * 4), + DType::kFloat32, MemoryPlace::kDevice, Layout::kFlat, + Shape{max_tokens, hidden}}; + TextEmbeddingStaging staging; + st = flashrt::modalities::text_embedding_staging_create(&staging, + max_tokens); + assert(st.ok_status()); + std::vector ids; + std::uint64_t prompt_len = 0; + PromptEmbeddingSpec spec{vocab, hidden, max_tokens, 0.5f}; + st = embed_prompt(tokenizer, spec, "pick up cube", nullptr, 0, src, dst, + &ids, &prompt_len, nullptr, &staging); + assert(st.ok_status()); + std::vector out(max_tokens * hidden, 1.0f); + assert(cudaMemcpy(out.data(), d_out, out.size() * sizeof(float), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(prompt_len == ids.size()); + assert(ids[0] == 2); + assert(std::fabs(out[0] - 1.0f) < 0.001f); + assert(std::fabs(out[1] + 1.0f) < 0.001f); + assert(out[prompt_len * hidden] == 0.0f); + flashrt::modalities::text_embedding_staging_destroy(&staging); + cudaFree(d_out); + cudaFree(d_table); +#endif +} + +} // namespace + +int main() { + test_requires_loaded_tokenizer(); + test_paligemma_prompt_embedding_when_configured(); + test_paligemma_prompt_embedding_device_when_configured(); + std::cout << "PASS - Pi05 prompt embedding\n"; + return 0; +} diff --git a/cpp/tests/test_pi05_prompt_format.cpp b/cpp/tests/test_pi05_prompt_format.cpp new file mode 100644 index 00000000..7184b0ea --- /dev/null +++ b/cpp/tests/test_pi05_prompt_format.cpp @@ -0,0 +1,73 @@ +#include "flashrt/cpp/models/pi05/prompt_format.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +void test_discretize_matches_python_reference() { + const float state[] = {-1.0f, 0.0f, 1.0f, 2.0f, -2.0f}; + const auto bins = flashrt::models::pi05::discretize_state_prompt_bins( + state, 5); + const std::vector expected = {0, 128, 255, 255, -1}; + assert(bins == expected); +} + +void test_prompt_format_matches_python_reference() { + const float state[] = {-1.0f, 0.0f, 1.0f, 2.0f, -2.0f}; + const std::string out = flashrt::models::pi05::format_state_prompt( + "pick_up\nred", state, 5); + assert(out == + "Task: pick up red, State: 0 128 255 255 -1;\nAction: "); + std::string workspace; + workspace.reserve(128); + const std::size_t capacity = workspace.capacity(); + for (int i = 0; i < 1000; ++i) { + flashrt::models::pi05::format_state_prompt_into( + "pick_up\nred", state, 5, &workspace); + assert(workspace == out); + assert(workspace.capacity() == capacity); + } +} + +void test_prompt_without_state_keeps_text_only_format() { + const std::string out = flashrt::models::pi05::format_state_prompt( + " pick_up\nred ", nullptr, 0); + assert(out == "pick up red"); +} + +void test_boundary_values() { + const float eps = 1.0f / 1024.0f; + const float state[] = { + -1.0f - eps, + -1.0f, + -1.0f + eps, + 1.0f - eps, + 1.0f, + std::numeric_limits::quiet_NaN(), + }; + const auto bins = flashrt::models::pi05::discretize_state_prompt_bins( + state, 6); + assert(bins[0] == -1); + assert(bins[1] == 0); + assert(bins[2] == 0); + assert(bins[3] == 255); + assert(bins[4] == 255); + assert(bins[5] == 255); +} + +} // namespace + +int main() { + test_discretize_matches_python_reference(); + test_prompt_format_matches_python_reference(); + test_prompt_without_state_keeps_text_only_format(); + test_boundary_values(); + std::cout << "PASS - Pi05 prompt formatter\n"; + return 0; +} diff --git a/cpp/tests/test_pi05_runtime.cpp b/cpp/tests/test_pi05_runtime.cpp index 607c58dd..fe594dd7 100644 --- a/cpp/tests/test_pi05_runtime.cpp +++ b/cpp/tests/test_pi05_runtime.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -111,6 +112,7 @@ void test_adopted_export_runtime_flow() { assert(runtime.export_runtime() == &exp); assert(runtime.manifest().vision.view_order.size() == 1); assert(runtime.manifest().graphs.infer == "infer"); + assert(runtime.set_prompt("pick up the cube") != 0); const std::uint8_t image_rgb[] = { 0, 127, 255, 255, 127, 0, @@ -127,6 +129,11 @@ void test_adopted_export_runtime_flow() { auto st = runtime.prepare_vision({image}); assert(st.ok_status()); + VisionFrame bgr = image; + bgr.format = PixelFormat::kBGR8; + st = runtime.prepare_vision({bgr}); + assert(!st.ok_status()); + assert(st.code == flashrt::modalities::StatusCode::kShapeMismatch); assert(runtime.replay_tick() == 0); assert(probe.calls == 1); @@ -142,10 +149,76 @@ void test_adopted_export_runtime_flow() { assert(owner.release == 1); } +void test_prompt_staging_when_configured() { +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const char* tokenizer = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + if (!tokenizer || tokenizer[0] == '\0') { + std::cout << "SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"; + return; + } + + Owner owner; + frt_runtime_graph_desc graph{}; + graph.name = "infer"; + graph.handle = reinterpret_cast(0x2000); + graph.default_key = 9; + graph.stream_id = 5; + auto exp = make_export(&owner, &graph); + + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + std::vector prompt(max_tokens * hidden, 3.0f); + std::vector image_input(1 * 224 * 224 * 3); + std::vector action_model(4, 0.0f); + + flashrt::models::pi05::RuntimeConfig cfg; + cfg.num_views = 1; + cfg.chunk = 1; + cfg.model_action_dim = 4; + cfg.robot_action_dim = 3; + cfg.image_input_override = TensorView{ + image_input.data(), static_cast(image_input.size() * 2), + DType::kBFloat16, MemoryPlace::kHost, Layout::kNHWC, + Shape{1, 224, 224, 3}}; + cfg.action_output_override = TensorView{ + action_model.data(), static_cast(action_model.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, Shape{1, 4}}; + cfg.prompt_tokenizer_model_path = tokenizer; + cfg.prompt_vocab_size = vocab; + cfg.prompt_hidden_dim = hidden; + cfg.prompt_max_tokens = max_tokens; + cfg.prompt_embedding_scale = 0.5f; + cfg.prompt_embedding_table = TensorView{ + table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{vocab, hidden}}; + cfg.prompt_embedding_output = TensorView{ + prompt.data(), static_cast(prompt.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{max_tokens, hidden}}; + + flashrt::models::pi05::Runtime runtime(&exp, cfg); + assert(runtime.ok()); + const float state[] = {0.0f, 1.0f, -1.0f}; + assert(runtime.set_prompt_state("pick_up_cube", state, 3) == 0); + assert(runtime.current_prompt_len() == 24); + assert(std::fabs(prompt[0] - 1.0f) < 0.001f); + assert(std::fabs(prompt[1] + 1.0f) < 0.001f); + assert(prompt[24 * hidden] == 0.0f); +#endif +} + } // namespace int main() { test_adopted_export_runtime_flow(); + test_prompt_staging_when_configured(); std::cout << "PASS - Pi05 C++ runtime flow\n"; return 0; } diff --git a/cpp/tests/test_safetensors_loader.cpp b/cpp/tests/test_safetensors_loader.cpp new file mode 100644 index 00000000..a24c83ba --- /dev/null +++ b/cpp/tests/test_safetensors_loader.cpp @@ -0,0 +1,115 @@ +#include "flashrt/cpp/loader/safetensors.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string temp_path() { + char path[] = "/tmp/frt_safetensors_XXXXXX"; + const int fd = ::mkstemp(path); + assert(fd >= 0); + ::close(fd); + return path; +} + +void write_file(const std::string& path, const std::string& header, + const std::string& payload) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + const std::uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char byte = static_cast((n >> (8 * i)) & 0xffu); + f.write(&byte, 1); + } + f.write(header.data(), static_cast(header.size())); + f.write(payload.data(), static_cast(payload.size())); + assert(f.good()); +} + +} // namespace + +int main() { + using flashrt::loader::SafetensorsFile; + + const std::string path = temp_path(); + std::string payload(12, '\0'); + const float expected[] = {1.25f, -2.5f}; + std::memcpy(&payload[4], expected, sizeof(expected)); + write_file( + path, + "{\"__metadata__\":{\"format\":\"pt\"}," + "\"u8\":{\"dtype\":\"U8\",\"shape\":[4]," + "\"data_offsets\":[0,4]}," + "\"values\":{\"shape\":[2],\"data_offsets\":[4,12]," + "\"dtype\":\"F32\"}}", + payload); + + SafetensorsFile file; + assert(file.open(path)); + assert(file.is_open()); + assert(file.tensors().size() == 2); + const auto* values = file.find("values"); + assert(values); + assert(values->dtype == "F32"); + assert(values->shape.size() == 1 && values->shape[0] == 2); + assert(values->bytes == sizeof(expected)); + assert(std::memcmp(file.data(*values), expected, sizeof(expected)) == 0); + + SafetensorsFile moved(std::move(file)); + assert(!file.is_open()); + assert(moved.find("u8")); + assert(::unlink(path.c_str()) == 0); + assert(std::memcmp(moved.data(*moved.find("values")), expected, + sizeof(expected)) == 0); + moved.close(); + assert(!moved.is_open()); + + const std::string invalid = temp_path(); + write_file(invalid, + "{\"x\":{\"dtype\":\"F32\",\"shape\":[2]," + "\"data_offsets\":[0,4]}}", + std::string(4, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("does not match") != std::string::npos); + + write_file(invalid, + "{\"x\":{\"dtype\":\"F4\",\"shape\":[2]," + "\"data_offsets\":[0,1]}}", + std::string(1, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("unsupported") != std::string::npos); + + write_file(invalid, + "{\"x\":{\"dtype\":\"U8\",\"shape\":[2]," + "\"data_offsets\":[0,2]}," + "\"y\":{\"dtype\":\"U8\",\"shape\":[2]," + "\"data_offsets\":[1,3]}}", + std::string(3, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("overlapping") != std::string::npos); + + write_file(invalid, + "{\"x\":{\"dtype\":\"U8\",\"shape\":[1]," + "\"data_offsets\":[0,1]}," + "\"x\":{\"dtype\":\"U8\",\"shape\":[1]," + "\"data_offsets\":[1,2]}}", + std::string(2, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("duplicate") != std::string::npos); + + write_file(invalid, + "{\"x\":{\"dtype\":\"U8\",\"shape\":[4]," + "\"data_offsets\":[0,4]}}", + std::string(2, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("exceeds") != std::string::npos); + + assert(::unlink(invalid.c_str()) == 0); + std::printf("PASS - safetensors mmap loader\n"); + return 0; +} diff --git a/cpp/tests/test_sha256.cpp b/cpp/tests/test_sha256.cpp new file mode 100644 index 00000000..323dca13 --- /dev/null +++ b/cpp/tests/test_sha256.cpp @@ -0,0 +1,30 @@ +#include "flashrt/cpp/loader/sha256.h" + +#include +#include +#include +#include +#include + +int main() { + char path[] = "/tmp/flashrt_sha256_XXXXXX"; + const int fd = ::mkstemp(path); + assert(fd >= 0); + ::close(fd); + { + std::ofstream file(path, std::ios::binary | std::ios::trunc); + file << "abc"; + assert(file.good()); + } + std::string digest; + std::string error; + assert(flashrt::loader::sha256_file(path, &digest, &error)); + assert(digest == + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert(error.empty()); + ::unlink(path); + assert(!flashrt::loader::sha256_file(path, &digest, &error)); + assert(!error.empty()); + std::printf("PASS - SHA-256 file hashing\n"); + return 0; +} diff --git a/cpp/tests/test_text_modalities.cpp b/cpp/tests/test_text_modalities.cpp new file mode 100644 index 00000000..94645572 --- /dev/null +++ b/cpp/tests/test_text_modalities.cpp @@ -0,0 +1,92 @@ +#include "flashrt/cpp/modalities/text.h" + +#include +#include +#include +#include +#include + +using flashrt::modalities::DType; +using flashrt::modalities::EmbeddingGatherSpec; +using flashrt::modalities::Layout; +using flashrt::modalities::MemoryPlace; +using flashrt::modalities::Shape; +using flashrt::modalities::StatusCode; +using flashrt::modalities::TensorView; +using flashrt::modalities::bfloat16_to_float; +using flashrt::modalities::float_to_bfloat16; +using flashrt::modalities::gather_token_embeddings_cpu; + +namespace { + +void test_f32_embedding_gather() { + std::vector table = { + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, + 9.0f, 10.0f, 11.0f, 12.0f, + }; + std::int32_t ids[] = {2, 0}; + std::vector out(2 * 4, 0.0f); + + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{3, 4}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + EmbeddingGatherSpec spec{3, 4, 2.0f}; + auto st = gather_token_embeddings_cpu(spec, ids, 2, src, dst); + assert(st.ok_status()); + const std::vector expected = { + 18.0f, 20.0f, 22.0f, 24.0f, + 2.0f, 4.0f, 6.0f, 8.0f, + }; + assert(out == expected); +} + +void test_bf16_embedding_gather() { + std::vector table = { + float_to_bfloat16(0.5f), float_to_bfloat16(-1.0f), + float_to_bfloat16(2.0f), float_to_bfloat16(3.0f), + }; + std::int32_t ids[] = {1}; + std::vector out(2); + + TensorView src{table.data(), static_cast(table.size() * 2), + DType::kBFloat16, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 2}}; + TensorView dst{out.data(), static_cast(out.size() * 2), + DType::kBFloat16, MemoryPlace::kHost, Layout::kFlat, + Shape{1, 2}}; + EmbeddingGatherSpec spec{2, 2, 1.5f}; + auto st = gather_token_embeddings_cpu(spec, ids, 1, src, dst); + assert(st.ok_status()); + assert(std::fabs(bfloat16_to_float(out[0]) - 3.0f) < 0.01f); + assert(std::fabs(bfloat16_to_float(out[1]) - 4.5f) < 0.01f); +} + +void test_invalid_token_rejected() { + std::vector table(4, 0.0f); + std::vector out(2, 0.0f); + std::int32_t ids[] = {2}; + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 2}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{1, 2}}; + EmbeddingGatherSpec spec{2, 2, 1.0f}; + auto st = gather_token_embeddings_cpu(spec, ids, 1, src, dst); + assert(!st.ok_status()); + assert(st.code == StatusCode::kInvalidArgument); +} + +} // namespace + +int main() { + test_f32_embedding_gather(); + test_bf16_embedding_gather(); + test_invalid_token_rejected(); + std::cout << "PASS - text modality contracts\n"; + return 0; +} diff --git a/cpp/tests/test_text_tokenizer.cpp b/cpp/tests/test_text_tokenizer.cpp new file mode 100644 index 00000000..1f6d7cb1 --- /dev/null +++ b/cpp/tests/test_text_tokenizer.cpp @@ -0,0 +1,81 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +#include +#include +#include +#include +#include +#include + +using flashrt::modalities::SentencePieceEncodeOptions; +using flashrt::modalities::SentencePieceTokenizer; +using flashrt::modalities::StatusCode; + +namespace { + +std::string tokenizer_model_path() { + const char* env = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + return env ? std::string(env) : std::string(); +} + +void test_unavailable_build_reports_unsupported() { + SentencePieceTokenizer tokenizer; +#ifndef FLASHRT_CPP_HAS_SENTENCEPIECE + auto st = tokenizer.load_model("missing.model"); + assert(!st.ok_status()); + assert(st.code == StatusCode::kUnsupported); +#else + (void)tokenizer; +#endif +} + +void test_paligemma_token_exact_when_configured() { +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const std::string path = tokenizer_model_path(); + if (path.empty()) { + std::cout << "SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"; + return; + } + SentencePieceTokenizer tokenizer; + auto st = tokenizer.load_model(path); + assert(st.ok_status()); + assert(tokenizer.loaded()); + assert(tokenizer.vocab_size() == 257152); + assert(tokenizer.bos_id() == 2); + assert(tokenizer.eos_id() == 1); + assert(tokenizer.unk_id() == 3); + assert(tokenizer.pad_id() == 0); + + std::vector ids; + SentencePieceEncodeOptions options; + options.add_bos = true; + st = tokenizer.encode( + "Task: pick up cube, State: 0 128 255;\nAction: ", + options, &ids); + assert(st.ok_status()); + const std::vector expected = { + 2, 7071, 235292, 4788, 908, 28660, 235269, 3040, 235292, + 235248, 235276, 235248, 235274, 235284, 235321, 235248, + 235284, 235308, 235308, 235289, 108, 4022, 235292, 235248, + }; + assert(ids == expected); + + options.max_tokens = expected.size() + 2; + options.pad_to_max_tokens = true; + st = tokenizer.encode( + "Task: pick up cube, State: 0 128 255;\nAction: ", + options, &ids); + assert(st.ok_status()); + assert(ids.size() == expected.size() + 2); + assert(ids[ids.size() - 1] == 0); +#endif +} + +} // namespace + +int main() { + test_unavailable_build_reports_unsupported(); + test_paligemma_token_exact_when_configured(); + std::cout << "PASS - text tokenizer contracts\n"; + return 0; +} diff --git a/csrc/attention/fa2_wrapper.cu b/csrc/attention/fa2_wrapper.cu index dd043327..cc8bfcaa 100644 --- a/csrc/attention/fa2_wrapper.cu +++ b/csrc/attention/fa2_wrapper.cu @@ -24,6 +24,7 @@ #include "flash_attn_2_src/flash_attn/namespace_config.h" #include "flash_attn_2_src/flash_attn/flash.h" +#include "attention/fa2_wrapper.h" namespace FLASH_NAMESPACE { template diff --git a/csrc/attention/fa2_wrapper.h b/csrc/attention/fa2_wrapper.h new file mode 100644 index 00000000..d07de0b0 --- /dev/null +++ b/csrc/attention/fa2_wrapper.h @@ -0,0 +1,73 @@ +#ifndef FLASHRT_ATTENTION_FA2_WRAPPER_H +#define FLASHRT_ATTENTION_FA2_WRAPPER_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +void fvk_attention_fa2_fwd_fp16( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_seqused( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_seqused_splitkv( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_causal( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +#ifdef __cplusplus +} +#endif + +#endif // FLASHRT_ATTENTION_FA2_WRAPPER_H diff --git a/csrc/attention/fa2_wrapper_causal.cu b/csrc/attention/fa2_wrapper_causal.cu index f651b194..fda63187 100644 --- a/csrc/attention/fa2_wrapper_causal.cu +++ b/csrc/attention/fa2_wrapper_causal.cu @@ -25,6 +25,7 @@ #include "flash_attn_2_src/flash_attn/namespace_config.h" #include "flash_attn_2_src/flash_attn/flash.h" +#include "attention/fa2_wrapper.h" namespace FLASH_NAMESPACE { template @@ -180,20 +181,17 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal( int o_batch_stride, int o_row_stride, int o_head_stride, float softmax_scale, int num_sms, cudaStream_t stream) { - if ((head_dim != 128) -#ifdef FA2_HAS_HDIM_256 - && head_dim != 256 + bool supported = false; +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_128) + supported = supported || head_dim == 128; #endif - ) { -#ifdef FA2_HAS_HDIM_256 - fprintf(stderr, - "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. " - "Only head_dim=128 and 256 are currently instantiated.\n", head_dim); -#else +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_256) + supported = supported || head_dim == 256; +#endif + if (!supported) { fprintf(stderr, "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. " - "Only head_dim=128 is currently instantiated.\n", head_dim); -#endif + "Enable its FA2_HDIMS entry and rebuild.\n", head_dim); std::abort(); } @@ -208,26 +206,36 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal( o_batch_stride, o_row_stride, o_head_stride, softmax_scale); - int num_splits = setup_splitkv_causal(params, softmax_lse_accum_ptr, o_accum_ptr, - num_sms, seqlen_q, seqlen_k, - head_dim, batch, num_heads_q); - if (head_dim == 128 && num_splits > 1) { - FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(params, stream); - } else if (head_dim == 128) { - FLASH_NAMESPACE::run_mha_fwd_(params, stream); - } -#ifdef FA2_HAS_HDIM_256 - else if (num_splits > 1) { - FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(params, stream); - } else { - FLASH_NAMESPACE::run_mha_fwd_(params, stream); - } -#else - else { - fprintf(stderr, - "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built " - "(hdim=256 disabled at compile time).\n", head_dim); - std::abort(); - } + int num_splits = setup_splitkv_causal( + params, softmax_lse_accum_ptr, o_accum_ptr, num_sms, seqlen_q, + seqlen_k, head_dim, batch, num_heads_q); + switch (head_dim) { +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_128) + case 128: + if (num_splits > 1) { + FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch< + cutlass::bfloat16_t, 128, true>(params, stream); + } else { + FLASH_NAMESPACE::run_mha_fwd_< + cutlass::bfloat16_t, 128, true>(params, stream); + } + return; +#endif +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_256) + case 256: + if (num_splits > 1) { + FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch< + cutlass::bfloat16_t, 256, true>(params, stream); + } else { + FLASH_NAMESPACE::run_mha_fwd_< + cutlass::bfloat16_t, 256, true>(params, stream); + } + return; #endif + default: + fprintf(stderr, + "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built " + "in this FA2 matrix.\n", head_dim); + std::abort(); + } } diff --git a/csrc/fa2_bindings.cpp b/csrc/fa2_bindings.cpp index 41d13d04..75360388 100644 --- a/csrc/fa2_bindings.cpp +++ b/csrc/fa2_bindings.cpp @@ -22,79 +22,14 @@ #include #include +#include "attention/fa2_wrapper.h" + namespace py = pybind11; static cudaStream_t to_stream(uintptr_t s) { return reinterpret_cast(s); } -// Forward declarations (definitions in csrc/attention/fa2_wrapper.cu). -extern "C" void fvk_attention_fa2_fwd_fp16( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -extern "C" void fvk_attention_fa2_fwd_bf16( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// seqused_k variant — definition in csrc/attention/fa2_wrapper.cu. Reads the -// per-batch K length from device memory so one captured graph serves any pos. -extern "C" void fvk_attention_fa2_fwd_bf16_seqused( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// seqused_k + split-KV variant (experimental; caller pre-inits lse_accum=-inf). -extern "C" void fvk_attention_fa2_fwd_bf16_seqused_splitkv( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// Causal variant — definition in csrc/attention/fa2_wrapper_causal.cu. -// Currently only (bf16, head_dim=128) is built. Used by Qwen3-8B -// prefill (S=N causal self-attention). -extern "C" void fvk_attention_fa2_fwd_bf16_causal( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - - // Shared docstring. pybind::def's doc arg takes a single string; we want the // same text for both fwd_fp16 and fwd_bf16 so deduplicate via static const. static const char* kDocstring = R"(FlashAttention-2 fwd (vendored). GQA-capable cross-attention. diff --git a/docs/cpp_runtime_design.md b/docs/cpp_runtime_design.md index ca3515cc..7f4b86fc 100644 --- a/docs/cpp_runtime_design.md +++ b/docs/cpp_runtime_design.md @@ -53,6 +53,14 @@ binds names and constants, never re-implements a transform. Nothing under The model boundary and the hardware boundary are intentionally different. +The FlashRT ABI is linked as one exec implementation per process. It does not +contain a backend registry and should not gain one. A process that needs +heterogeneous CUDA, CPU, llama.cpp, or future device instances introduces them +at the capsule backend boundary: each adopted `cap_backend` is an instance +vtable. A non-CUDA producer may still represent one invocation as an adopted +graph and expose the same ports/stages/regions. Consumers must not branch on a +backend-kind field; no such frozen field is required. + The **model** is selected by the native overlay/factory that the host loads: `cpp/models/pi05/` exports `frt_pi05_model_runtime_create_over`, a future GROOT runtime would export its own model factory, and so on. That code owns the diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md new file mode 100644 index 00000000..100bba5b --- /dev/null +++ b/docs/mindon_pi05_integration.md @@ -0,0 +1,231 @@ +# Mindon Pi0.5 Integration Guide + +This guide describes how a Mindon C++ host should integrate Pi0.5 through the +existing FlashRT runtime/model-runtime contracts. It is a deployment guide, not +a new ABI. + +## Layer Ownership + +FlashRT owns: + +- checkpoint loading and graph capture; +- ports, stages, streams, buffers, capsule regions, identity, and fingerprint; +- model-specific IO semantics: tokenizer, prompt formatting, image + preprocess, state normalization/discretization, and action postprocess; +- `set_input`, `get_output`, `prepare`, and `step` producer verbs. + +Nexus owns: + +- adoption of `frt_runtime_export_v1` / `frt_model_runtime_v1`; +- capsule snapshot/restore/fork/move over declared regions; +- stage scheduling and interaction modes; +- embedded and transport adapters that map external payloads to declared + ports. + +Mindon owns: + +- the application/control loop; +- camera/state/prompt transport into the adopted ports; +- action publication and deadline policy. + +Nexus should not learn Pi0.5 tokenizer, tensor layout, normalization, or +action schema rules. If a host needs richer state, the producer must export a +richer model-runtime face. + +## Recommended Lanes + +### Lane A: Available Now + +Use a resident Python setup producer, then run the hot loop in C++. + +Flow: + +1. Start a process that embeds CPython or calls a Python setup function. +2. Load the Pi0.5 checkpoint through the FlashRT Python frontend. +3. Capture graphs and call `pipeline.export_model_runtime(io="native", ...)`. +4. Pass the returned `frt_model_runtime_v1*` to the C++ host. +5. Adopt it into Nexus with `flashrt_adopt_model_runtime`. +6. Warm any declared graph variants. +7. Drive `images`, `noise`, and `actions` through the C++ hot loop. + +In Lane A, prompt is setup-time. The current adopted-export face does not +export hot `prompt` or `state` ports. A request may repeat the setup prompt for +bookkeeping, but it cannot change the model prompt dynamically. + +### Lane B: Adopted Prompt/State Staging + +Use the same setup/adopt path as Lane A, but the producer additionally exports +real hot ports: + +- `prompt: TEXT/STAGED` +- `state: STATE/STAGED` + +The C++ host updates these ports with `cap_model_set_input` or the embedded +session equivalent. The producer formats, tokenizes, embeds, and writes the +fixed prompt window. Nexus remains unchanged. + +### Lane C: Native SM120 Producer + +Load a native FlashRT shared object and call: + +```c +int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +``` + +The returned struct must expose the same public model-runtime contract as the +Python setup producer. The host and Nexus adoption code must not change when +switching between Lane A and Lane C. + +The current C++ shared object implements this symbol as a complete SM120 +native-v2 producer when built with CUDA kernels, native FA2, and SentencePiece. +It validates `io`, checkpoint/tokenizer paths, fixed prompt mode, capacities, +the complete 812-tensor inventory, and OpenPI or LeRobot action/state q01/q99 +metadata. It then hashes the model and tokenizer for deployment identity, +materializes context-owned weights/workspace, captures one `infer` variant, +and returns the integrated model runtime. Missing FA2/SentencePiece support or +non-SM120 hardware returns unsupported instead of publishing unusable ports. + +## No-HTTP C++ Host Shape + +For same-process control loops, prefer Nexus embedded/session APIs over HTTP. +The high-level loop is: + +``` +producer setup -> frt_model_runtime_v1 +adopt -> cap_model_runtime +open embedded session +for each control tick: + update declared input ports + tick or fire stages + read declared output ports +optional: + snapshot / restore named capsules +``` + +The C++ loop should discover ports by name and then rely on the declared port +shape, dtype, direction, and update class. It should not hard-code `(10, 7)`, +graph names, or internal buffer names. + +## Port Update Rules + +For `SWAP` ports: + +- write the declared buffer window directly through the capsule/backend copy + mechanism; +- do not call `set_input`; +- verify byte count against `port.bytes`. + +For `STAGED` ports: + +- call the producer verb through `cap_model_set_input` or + `nexus_embedded_set_input`; +- pass bytes in the payload convention declared by `frt_model_runtime_v1`; +- expect shape/status errors for invalid input. + +For `SETUP` ports: + +- never update them inside a control tick. + +## Mapping Existing Mindon Calls + +| Mindon call | Integration point | +|---|---| +| `Prepare` | warm phase, producer `prepare(graph, key)` | +| `Warmup` | host policy: `prepare` plus warm ticks | +| `Infer` | `cap_model_tick`, `nexus_embedded_tick`, or explicit stage firing | +| `Sync` | host/backend stream sync or embedded session synchronization | +| `GetOutput` | `cap_model_get_output` / `nexus_embedded_get_output` | + +Do not introduce a second runtime API beside `frt_model_runtime_v1`. The +existing verbs already carry these phases. + +## Prompt and State + +Pi0.5 state is rendered into the language prompt. It is not an independent +model tensor. The producer path is: + +``` +raw proprioception -> normalize -> 256-bin discretize -> prompt string +-> token ids -> embedding gather -> encoder_x prompt window +``` + +Lane A still requires a setup-time producer refresh for prompt/state changes. +Lanes B and C accept task text through `prompt` and raw proprioception through +`state`. The producer owns all formatting and normalization details. + +## Image Input + +Mindon should pass camera frames as `frt_image_view[]` to the `images` +`IMAGE/STAGED` port, or through the matching Nexus embedded input. Frames are +matched by declared position, not by runtime graph names. + +The current Pi0.5 native producer stages host pixels into the +`observation_images_normalized` device tensor and normalizes to `[-1, 1]`. +Pass `u8` `RGB8` frames in HWC layout. BGR/RGBA/GRAY, CHW, and non-`u8` inputs +are rejected instead of silently reinterpreted. Use the producer documentation +in `docs/pi05_io_contract.md` for accepted formats and shape rules. + +## Action Output + +Read the `actions` port shape to determine chunk length and action dimension. +The output is the host-visible robot action chunk after producer postprocess. +For raw model action state, use `actions_raw` when the producer exports it. In +the Pi0.5 `native_v2` face this is a raw `TENSOR/SWAP` alias of +`diffusion_noise` with shape `(chunk, 32)`. + +## Capsule Boundaries + +Capsules snapshot exactly the producer-declared regions, in declared order. +Mindon should treat capsule contents as opaque bytes. A fingerprint mismatch +on restore is a deployment mismatch and must fail loudly. + +The native-v2 producer currently declares only `rollout_boundary`. Prompt +embedding, attention lengths, RoPE, and CPU prompt/state caches are not a +capsule region because partial restoration would be incorrect. If a future +face makes the entire prompt context restorable, its added ordered regions and +new fingerprint will intentionally reject old capsules. + +## Configuration Sketch + +Lane A setup in a Python producer plugin should export: + +```python +model = pipeline.export_model_runtime( + identity={"deployment": "mindon-pi05"}, + stage_plan="full", + io="native", +) +``` + +A split or RTC deployment may choose another producer-registered stage plan, +but the C++ host still sees only the adopted stage array. + +Lane C opens the native producer with a setup config such as: + +```json +{ + "io": "native_v2", + "checkpoint_path": "/models/pi05", + "tokenizer_model_path": "/models/paligemma/tokenizer.model", + "state_prompt_mode": "fixed", + "max_prompt_tokens": 200, + "state_dim": 8, + "num_views": 2, + "chunk": 10, + "num_steps": 10, + "vision_pool_factor": 1 +} +``` + +## Acceptance Checklist + +- The host discovers ports and shapes from `cap_model_runtime`. +- `images` updates use `STAGED`; `noise` updates use `SWAP`. +- `actions` capacity is computed from the declared output shape and dtype. +- The warm phase finishes before the first control tick. +- The hot loop performs no graph capture, allocation, or graph rebinding. +- Prompt/state/image/action staging capacities are fixed at setup; oversized + payloads fail instead of growing a hot-path workspace. +- Snapshot/restore is tested within one live capture. +- Nexus core code remains unchanged for model-specific semantics. diff --git a/docs/model_runtime_api.md b/docs/model_runtime_api.md index 7730bfc4..2d39b275 100644 --- a/docs/model_runtime_api.md +++ b/docs/model_runtime_api.md @@ -19,6 +19,11 @@ windows are `TENSOR`. A `STAGED` declaration is a promise the port accepts hot updates — a producer that cannot deliver that declares `SETUP` or omits the port, never advertise-and-refuse. +The Python `build_model_runtime()` helper enforces this mechanically: STAGED +inputs require `set_input`, and STAGED outputs require `get_output`. Internal +declaration-only handoffs may bypass that check only while constructing a +native verb overlay; the declaration is not an adoptable runtime. + ## Payload conventions (STAGED `set_input`) | modality | `data` points at | `bytes` | @@ -30,7 +35,7 @@ port, never advertise-and-refuse. ## Descriptors `frt_runtime_port_desc` — one dynamic input/output: -`name`, `modality`, `dtype` (device-side tensor), `layout`, `direction`, +`name`, `modality`, `dtype` (logical payload/target tensor), `layout`, `direction`, `update`, `required`, `shape[rank]` (−1 = bucket-variable), `cadence_hint_hz` (advisory only), and the SWAP window `buffer`/`offset`/ `bytes` (null buffer = staged-only). Strings/arrays are owned by the runtime @@ -58,8 +63,10 @@ frt_model_runtime_v1 { } ``` -**Verbs** (`frt_model_runtime_verbs`; every entry is always callable — absent -producer verbs are filled with unsupported stubs returning `-3`): +**Verbs** (`frt_model_runtime_verbs`; every entry on a successfully constructed +object is callable). Construction rejects a STAGED input without `set_input` +and a STAGED output without `get_output`. Other absent verbs are filled with +unsupported stubs returning `-3`: | verb | phase | semantics | |---|---|---| @@ -129,6 +136,11 @@ The override retains `producer_model`, so inherited port/stage pointers remain valid even if the original producer reference is released first. Deployment identity is unchanged. +The integrated, adapter, and verb-override C paths enforce the same STAGED +verb-presence rule before retaining owners or consuming builders. An internal +intermediate may exist while one factory assembles an override, but it must not +escape that factory or be independently adoptable. + **Native factory (symbol convention)** — a model-runtime `.so` exports `FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL`: `int frt_model_runtime_open_v1(const char* config_json, frt_model_runtime_v1** out)`. @@ -266,3 +278,6 @@ PYTHONPATH=.:./exec/build:./runtime/build \ The consumer side (adoption, hot-input contract, real-model tick) is validated in the FlashRT-Nexus repository. + +For producer implementation and review rules, see +[`native_model_runtime_producer.md`](native_model_runtime_producer.md). diff --git a/docs/native_model_runtime_producer.md b/docs/native_model_runtime_producer.md new file mode 100644 index 00000000..72bfd342 --- /dev/null +++ b/docs/native_model_runtime_producer.md @@ -0,0 +1,96 @@ +# Native model-runtime producer guide + +This guide defines how a native producer joins the stable +`frt_model_runtime_v1` boundary. A model implementation is an example of the +contract, not a reason to specialize the contract. + +## Ownership + +`runtime/` owns opaque handles, descriptors, identity construction, verbs, and +lifetime. `exec/` owns Buffer/Graph/Plan/Event/ShapeKey mechanisms. A producer +under `cpp/models//` owns checkpoint names, model dimensions, tokenizer +and formatter behavior, preprocessing, graph capture, workspace, and output +postprocessing. Nexus and other consumers interpret none of those semantics. + +Do not add `model_kind`, `backend_kind`, model dimensions, checkpoint fields, +or a State object to the frozen ABI. Express the public face with ports, +stages, regions, verbs, and producer identity. + +## Construction + +Use one existing construction path: + +1. `frt_runtime_builder_finish_model`: one producer builds export and model + declarations under one fingerprint. +2. `frt_model_runtime_wrap`: an adapter adds a model face to an existing + export whose identity already covers that face. +3. `frt_model_runtime_override_verbs`: an internal handoff retains an existing + declaration while replacing hot verbs. + +All paths reject STAGED inputs without `set_input` and STAGED outputs without +`get_output`. A factory may use an unpublished intermediate while assembling +an override, but only the final object with real verbs may leave the factory. + +## Identity + +The builder is the only fingerprint implementation. Include actual weights, +tokenizer/configuration, graph/stream placement, port schema and windows, +stage DAG, and ordered restore regions. Query the executing device for hardware +identity; do not copy the requested CMake architecture or a model default. + +Manifest fields are discovery metadata, not a substitute for identity. A +schema or restore change intentionally produces a new fingerprint and rejects +old capsules. + +## Multiple producers and backends + +Python, native CUDA, CPU, llama.cpp, and future producers expose the same +structural boundary but may have different graph counts, internal buffers, +workspace, identities, and synchronization implementations. Validate each +producer's invariants independently. Compare only a deliberately shared +semantic face through checked-in canonical records. + +FlashRT supplies one exec implementation per process. Heterogeneous backend +instances enter above it through capsule backend vtables; do not add a runtime +backend registry or backend-kind ABI field. + +## Hot path + +Setup allocates storage, resolves names, loads weights, captures or adopts +graphs, and prepares variants. A control tick may update SWAP windows, execute +STAGED verbs, fire stages, and read output. It must not allocate device memory, +recapture, rebind graph pointers, or grow capacity. Oversized payloads fail. + +Measure CUDA allocator APIs over the complete service iteration. Host +allocation claims require a host allocation counter scoped to the component; +CUDA traces cannot prove host allocator behavior. + +## Schema workflow + +For a face implemented by more than one producer: + +1. Check in canonical `region:`, `port:`, and `stage:` records. +2. Generate records from every producer independently. +3. Diff each producer against the golden records. +4. Derive expected counts from the records instead of repeating constants. +5. Treat a golden update as a public contract review with fingerprint and + restore-compatibility analysis. + +Do not require implementation-private graph, buffer, manifest, or identity +records to match across backends. + +## Pull request evidence + +- C++ runtime tests in CUDA-off and affected hardware builds. +- Python runtime contract tests when the Python producer is supported. +- Producer-local lifecycle, schema, negative-input, and hot-loop gates. +- Numerical evidence appropriate to the boundary: bit-exact for identical + graph/input bytes; a documented fixed tolerance for genuinely different + backend math. +- Consumer adoption tests when descriptor or enum mapping changes. +- Migration notes for payload, fingerprint, packaging, or compatibility + changes. + +Use placeholders such as `` and `` in public commands. +Do not publish local paths, host/container names, credentials, environment +dumps, internal URLs, or proprietary dataset/checkpoint identifiers. diff --git a/docs/pi05_cpp_runtime_migration.md b/docs/pi05_cpp_runtime_migration.md new file mode 100644 index 00000000..a8f5e728 --- /dev/null +++ b/docs/pi05_cpp_runtime_migration.md @@ -0,0 +1,45 @@ +# PI0.5 C++ runtime migration notes + +This note covers externally visible behavior of the native PI0.5 producer. It +does not add a model-specific ABI: consumers continue to discover and drive +the generic `frt_model_runtime_v1` interface. + +## Image input + +The model-runtime `IMAGE/STAGED` port accepts explicit `RGB8`, `u8`, HWC host +images. It rejects BGR, grayscale, unsupported layouts, invalid dimensions, +and short strides instead of guessing a conversion. + +The legacy `frt_pi05_runtime_prepare_vision` entry remains source-compatible +with its explicit RGB, BGR, RGBA, BGRA, and grayscale formats. Existing OpenCV +BGR callers can remain on that entry or convert to RGB before using the generic +model-runtime face. + +Pixel normalization follows the reference float32 operation order +`value / 127.5 - 1`. Replacing division with a precomputed reciprocal can alter +FP8 quantization boundaries and is not equivalent for this producer. + +## Action output + +The logical `actions` STAGED output is F32 and includes the producer's declared +postprocessing. `actions_raw` is the BF16 SWAP alias for consumers that need +the model-space result. Consumers must select the declared port rather than +infer dtype or normalization from a model name. + +## Runtime adoption + +A published runtime with STAGED input or output ports must install the matching +`set_input` or `get_output` verb. Declaration-only native handoff objects are +internal overlay inputs and are marked as such; they are not independently +adoptable runtimes. + +Port, stage, binding-window, stream-placement, and capsule-region changes alter +the runtime fingerprint. Capsules produced under an older fingerprint must be +regenerated; rejecting their restore is required behavior. + +## Native FA2 dependency + +The Python FA2 adapter and the Python-free `libflashrt_fa2_raw.so` are one +install unit. Native producers link the raw library and Python producers reach +the same symbols through the adapter. Deployment packages must install both in +the same directory so their `$ORIGIN` runtime lookup remains relocatable. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md new file mode 100644 index 00000000..5da5eb10 --- /dev/null +++ b/docs/pi05_io_contract.md @@ -0,0 +1,611 @@ +# Pi0.5 Native Model Runtime IO Contract + +This document is the deployment-facing IO contract for the Pi0.5 native +model-runtime face. It is intentionally limited to the public runtime/model +runtime ABI: + +- `frt_runtime_export_v1` in `runtime/include/flashrt/runtime.h` +- `frt_model_runtime_v1` in `runtime/include/flashrt/model_runtime.h` +- Pi0.5 producer declarations in `flash_rt/models/pi05/runtime_export.py` +- Pi0.5 native verb overlay in `cpp/models/pi05/` + +It does not freeze the C++ implementation classes under `cpp/`. Those classes +may evolve as long as the exported ports, stages, regions, identity, and hot +contract remain valid. + +## Current Native Face + +The current Pi0.5 `io="native"` export declares three host-visible ports. +This is the contract implemented by `frt_pi05_model_runtime_create_over`. + +| port | modality/update | direction | dtype/layout/shape | backing | +|---|---|---|---|---| +| `images` | `IMAGE/STAGED` | input | device tensor dtype, `NHWC`, `(num_views, 224, 224, 3)` | `observation_images_normalized` | +| `noise` | `TENSOR/SWAP` | input | device tensor dtype, `FLAT`, `(chunk_length, 32)` | `diffusion_noise` | +| `actions` | `ACTION/STAGED` | output | host `f32`, `FLAT`, `(chunk_length, robot_action_dim)` | staged-only; no raw window | + +Current source of truth: + +- Export declaration: `flash_rt/models/pi05/runtime_export.py`, + `export_model_runtime(..., io="native")` +- Native verb implementation: `cpp/models/pi05/src/model_runtime.cpp` +- C++ modality binding: `cpp/models/pi05/src/runtime.cpp`, + `cpp/models/pi05/src/io.cpp`, `cpp/models/pi05/src/spec.cpp` + +`io="native"` and `io="native_v2"` are declaration-only handoffs. Their +discovery manifest carries `declaration_only: true`; callers must pass them to +`frt_pi05_model_runtime_create_over` before adoption. Generic Python-produced +model runtimes reject STAGED ports without matching input/output verbs. + +There is deliberately no `prompt` port on the adopted-export path today. The +prompt embedding is prepared by the producer before graph capture/export. A +producer must not declare a `TEXT/STAGED` or `STATE/STAGED` port until the +native verb can really update that input on the hot path. + +## Native V2 Face + +The `io="native_v2"` export adds prompt/state staging and a raw action alias. +Adding these ports changes the model-runtime identity and therefore the +fingerprint. Existing capsules from the old face must refuse restore into this +face. + +| port | modality/update | direction | dtype/layout/shape | backing | +|---|---|---|---|---| +| `prompt` | `TEXT/STAGED` | input | UTF-8 bytes, `FLAT`, variable length | staged by C++ runtime | +| `state` | `STATE/STAGED` | input | host `f32`, `FLAT`, `(state_dim,)` | staged by C++ runtime | +| `images` | `IMAGE/STAGED` | input | device tensor dtype, `NHWC`, `(num_views, 224, 224, 3)` | `observation_images_normalized` | +| `noise` | `TENSOR/SWAP` | input | device tensor dtype, `FLAT`, `(chunk_length, 32)` | `diffusion_noise` | +| `actions` | `ACTION/STAGED` | output | host `f32`, `FLAT`, `(chunk_length, robot_action_dim)` | staged-only; no raw window | +| `actions_raw` | `TENSOR/SWAP` | output | device tensor dtype, `FLAT`, `(chunk_length, 32)` | `diffusion_noise` | + +For Pi0.5, proprioceptive state is not an independent model tensor. It is +normalized, discretized into OpenPI-compatible 256-bin state tokens, rendered +into the prompt text, tokenized, embedded, and written into the language rows +of `encoder_x`. Therefore prompt and state updates are one producer-owned text +staging path. + +Internal model buffers such as `encoder_x`, KV/cache windows, residual +streams, and `diffusion_noise` are not `STATE` ports. They are `TENSOR` ports +when exposed as hot IO, or runtime buffers/capsule regions when they are part +of a restorable boundary. + +## STAGED Payloads + +The payload conventions are inherited from `frt_model_runtime_v1`. + +| modality | `set_input` data | bytes | +|---|---|---| +| `IMAGE` / `DEPTH` | `frt_image_view[]` | `n_frames * sizeof(frt_image_view)` | +| `TEXT` | UTF-8 bytes | byte length | +| `TENSOR` / `STATE` / `ACTION` / `AUDIO` | raw bytes per the port dtype and shape | byte length | + +For `IMAGE`, frames are matched positionally to the producer-declared camera +view order. The Pi0.5 view order is: + +1. `image` +2. `wrist_image` +3. `wrist_image_right` + +Deployments with fewer views export a shorter `num_views` and use the prefix +of that view order. + +## Image Input + +The current native image input accepts host `frt_image_view[]` and stages the +data into the device `observation_images_normalized` buffer before replay. + +Producer-owned preprocessing: + +- view count is checked against the exported `images` port shape; +- frame payloads are host `u8` pixels in `RGB8`/`HWC` layout; +- target tensor is `(num_views, 224, 224, 3)`; +- output layout is `NHWC`; +- output dtype is the exported tensor dtype, normally BF16 for the FP8 path; +- normalization is `x / 127.5 - 1.0`; +- resizing to 224x224 is producer-owned. + +`stride_bytes=0` means tightly packed RGB. A positive stride may include row +padding but must be at least `width * 3`; negative and short strides are shape +errors. The native CUDA and CPU reference paths are bit-exact in BF16 over the +validation matrix (`1x1`, odd dimensions, non-4:3 inputs, wide/tall inputs, +224x224, and padded rows). + +The Pi0.5 native face rejects unsupported input shape, dtype, layout, pixel +format, or view count with a shape/status error. BGR, grayscale, RGBA, CHW, and +non-`u8` frames are not silently converted at the Pi0.5 contract boundary. If a +deployment supports more pixel formats, the supported set must be documented by +the producer and tested against the CPU reference path. + +Compatibility note: the legacy `frt_pi05_runtime_prepare_vision` C API keeps +accepting its explicit `BGR8`, `RGBA8`, `BGRA8`, and `GRAY8` enum values and +converts them through the shared modality implementation. This compatibility +does not expand the `frt_model_runtime_v1` native face: its `IMAGE/STAGED` +deployment contract remains strict `RGB8/u8/HWC`. + +## Noise Input + +`noise` is a `TENSOR/SWAP` port. The host writes its raw bytes directly into +the `diffusion_noise` window, usually by `cap_swap` after Nexus adoption or by +the equivalent runtime/backend copy mechanism. Calling `set_input` on this +port is unsupported by design: SWAP means the device window is the interface. + +Shape is `(chunk_length, 32)`. `chunk_length` is declared by the producer and +must be read from the port shape; host code must not assume `(10, 32)`. + +## Action Output + +`actions` is the host-visible robot action chunk after producer-owned +postprocess. Its declared dtype is F32 because that is the payload returned by +`get_output`; it does not expose the BF16 diffusion window as backing. + +The logical output shape is: + +``` +(chunk_length, robot_action_dim) +``` + +For LIBERO-style Pi0.5 deployments, `robot_action_dim` is typically 7. Other +deployments may export a different fixed robot action dimension. Consumers and +schedulers must read the declared port shape instead of hard-coding `(10, 7)`. + +The internal model output remains `(chunk_length, 32)` in `diffusion_noise`. +The native `actions` STAGED output slices the robot dimensions and applies the +deployment action normalization statistics. With q01/q99 stats, the affine +parameters are: + +``` +mean = (q01 + q99) / 2 +stddev = (q99 - q01) / 2 +``` + +The C++ postprocess path clamps normalized action values to the configured +domain before applying the affine transform. Any raw `(chunk_length, 32)` face +must be exported as a separate `TENSOR/SWAP` output. The Pi0.5 `native_v2` +face declares this as `actions_raw`; RTC stage plans also use the same port +name. Nexus must treat it as a declared raw byte window, not model internals. + +## Lifecycle Mapping + +Mindon-style lifecycle names map to the existing ABI. Do not add a parallel +API family for the same phases. + +| requested name | existing contract | phase | +|---|---|---| +| `Prepare` | `prepare(graph, key)` | warm only | +| `Warmup` | host policy: call `prepare` for needed variants, then run warm ticks | warm | +| `Infer` | `step()` sugar or host-scheduled stage replay | hot | +| `Sync` | host/backend stream synchronization | hot or drain | +| `GetOutput` | `get_output(port, out, capacity, &written, stream)` | hot | + +`prepare` is the only place a shape-bucket miss may capture or materialize a +variant. A hot tick must not recapture, allocate, or rebind graph pointers. +The Pi0.5 C++ face fixes its vision frames, action D2H staging, task/formatted +prompt strings, tokenizer ids, and normalized-state storage during setup. +Payloads that would grow those workspaces return a shape/capacity error; there +is no larger-allocation fallback in the hot path. These workspace changes do +not alter the port schema or deployment fingerprint. + +## Identity and Capsule Regions + +The following changes are deployment identity changes: + +- adding/removing/reordering ports; +- changing a port modality, dtype, layout, direction, update class, required + flag, shape, bound buffer index, offset, or byte window; +- changing graph names or default stream placement; +- changing the stage DAG; +- adding/removing/reordering capsule regions; +- changing a region name, buffer, offset, byte length, or flags. + +The following are not deployment identity changes: + +- editing `manifest_json`; +- changing `cadence_hint_hz`. + +Prompt/state staging does not by itself make prompt context a capsule region. +A restorable prompt context would have to include the embedding, attention +lengths, decoder position/RoPE, and the CPU semantic cache used by later +independent prompt/state updates. The current face can rebuild those values +from its declared inputs, so it does not advertise partial prompt restoration. +Region layout and order are fingerprinted; adding a complete prompt region in +a later face will intentionally invalidate old capsules. + +## Current Integration Lanes + +There are three supported integration lanes: + +- Lane A, current: Python setup/capture/export stays resident in the process; + the hot loop adopts `frt_model_runtime_v1` and runs through C++/Nexus. +- Lane B: an adopted setup producer exposes real hot `prompt`/`state` STAGED + ports and the C++ overlay owns their transforms. +- Lane C, current on RTX SM120: a C++ shared object implements + `frt_model_runtime_open_v1(config_json, &out)` and produces the same public + struct without Python setup. + +The Pi0.5 C++ shared object exports `frt_model_runtime_open_v1` as a complete +native-v2 producer when built with CUDA kernels, native FA2, and SentencePiece. +Execution currently requires RTX SM120. The factory requires `io="native_v2"`, +`checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, +`max_prompt_tokens >= 200`, and a positive `state_dim`; `num_views`, `chunk`, +`num_steps`, and `vision_pool_factor` are optional fixed setup values. It parses +`checkpoint_path/model.safetensors` through the native read-only mmap loader to +verify the complete 812-tensor Pi0.5 inventory: all 27 vision layers, all 18 +language encoder layers, all 18 action-expert layers, embeddings/final norms, +projectors, action projections, and time MLP. It also verifies action/state q01/q99 +dimensions from either openpi `norm_stats.json` or LeRobot policy +normalizer/unnormalizer safetensors. Safetensors tensor byte ranges must match +dtype/shape, and normalization quantiles must be finite ordered pairs. Builds +without native FA2 or SentencePiece validate the config and return unsupported; +they do not advertise a runtime they cannot execute. The mmap and parsed tensor +views are setup-side assets; they never enter the model-runtime ABI or hot path. + +The native setup layer also carries CPU reference transforms matching the +existing PyTorch producer: source BF16 rounding for vision/decoder weights, +`OIHW -> HWIO` patch permutation, Q/K head interleave, QKV and gate/up fusion, +and FP32 encoder RMSNorm fold before the final BF16 rounding. Real-checkpoint +gates compare the resulting BF16 bytes against PyTorch for both bare OpenPI +keys and LeRobot `model.`-prefixed keys. + +Materialized device weights use `frt_buffer` allocations owned by the native +producer's `frt_ctx`. They are internal setup assets, not model ports and not +capsule regions. Upload is complete before capture; duplicate logical names or +typed shape/payload mismatches fail setup. The same store carries BF16, FP8 +E4M3, INT8, and FP32 scale buffers without introducing a model-level state +object. Destroying the context releases the device weights after graphs and +plans, preserving the exec ownership order. + +The composed materializer covers language-encoder, action-expert, and vision +weights. Encoder layers emit the five pipeline groups (`attn_qkv`, `attn_o`, +`ffn_gate`, `ffn_up`, and `ffn_down`) with FP32 RMS folds. Decoder layers emit +those groups plus the four AdaRMS modulation tensors and the optional merged +gate/up buffer used by the FP16 path. Vision setup emits patch/position/final +norm and multimodal-projector globals plus the twelve per-layer attention, +FFN, and normalization buffers. Decoder globals include final AdaRMS +modulation, time MLP, generated time embeddings, and action projections. The +action output projection is pre-scaled by `-1/num_steps` after source BF16 +rounding; 5-step and 10-step schedules are byte-exact with the PyTorch +producer. The prompt embedding table is materialized separately to keep its +approximately 1 GiB allocation explicit. These paths have been exercised +against the two supported real checkpoint layouts. The checkpoint inventory +also validates the language final norm and expert LM head even though the +current Pi0.5 pipeline does not consume them. The native producer materializes +the full BF16 store before capture and keeps it under the graph context lifetime. + +Native setup quantization reproduces the PyTorch producer's per-tensor FP8 +E4M3 weights in either `kn` or `nk` layout and per-output-channel INT8 weights +in `[N,K]` layout. FP8 scalar descales and INT8 channel scales are FP32 device +buffers. Real-checkpoint gates compare both quantized bytes and scale bytes; +the precision choice remains producer setup policy and does not alter ports, +regions, or the exec mechanism. + +The setup packer derives low-precision buffers from the already uploaded BF16 +fallback, so both paths share exactly the same transformed source bytes. It +stores packed weights under `fp8.*` or `int8.*` names and their typed FP32 +scales under the matching `.scale` names in the same context-owned store. + +Full BF16 assembly has one ordered setup path: vision globals and 27 layers, +18 language-encoder layers, 18 action-expert layers, decoder globals, then the +prompt embedding table. With merged decoder gate/up buffers enabled this owns +613 logical device buffers. Assembly options make `num_steps`, merged gate/up, +and the large embedding allocation explicit; they are producer configuration, +not ABI fields. + +Full FP8 packing follows the producer's exact site inventory: four GEMM +weights for each vision layer plus the projector, four for each encoder layer, +and four for each decoder layer. Encoder gate/up columns are merged during +setup. INT8 packing remains independently selectable for vision, encoder, and +decoder and preserves their existing four/five/five weights-per-layer policy. + +The native kernel layer is CPython-independent and links the existing +`GemmRunner` implementation directly. Setup warms required BF16 GEMM shapes, +captures the complete `infer` graph through `frt_graph_capture`, and exports +exactly one shape-key variant (`0`). + +The native core workspace maps every vision, encoder, decoder, style, action, +RTC, and reusable scratch allocation to a context-owned `frt_buffer`. There is +no model-level State object. With vision pooling disabled, `vision_x_pooled` +is an explicit alias of `vision_x` (34 logical names, 33 allocations); pooled +deployments allocate it separately. Buffer shapes are fixed from `num_views`, +`max_prompt_tokens`, `chunk_size`, `num_steps`, and `vision_pool_factor` before +capture, and BF16 RMS-one constants, attention backend storage, and generated +decoder style contents are initialized during setup. + +Native RoPE setup uses the same float64 frequency/phase computation and BF16 +interleaved `[cos, sin]` layout as the Python producer. Encoder and +prompt-relative decoder slices are byte-exact against NumPy/ml_dtypes for +pooled and unpooled configurations. Decoder slice updates reuse one stable +buffer across prompt lengths; vision position embeddings are expanded per view +with setup-side D2D copies from the typed weight store. + +Decoder time/style precompute is also native setup work. It consumes the +generated time embeddings, time-MLP weights, 18 layers of AdaRMS modulation, +and final modulation from the typed store; it reuses existing workspace +buffers as scratch and writes the four persistent style buffers without a +temporary device allocation. The GEMM, explicit BF16 bias round-trip, and +float-SiLU sequence is BF16 bit-exact with the PyTorch producer on both +supported checkpoint layouts. + +The native kernel driver also owns the BF16 forward primitives used around +GEMM and attention: RMS/Layer/AdaRMS normalization, residual and gated +residual updates, GELU/gated GELU, QKV split with fixed or device-position +RoPE, patch im2col, and vision pooling. These are direct typed calls to the +existing CUDA implementations, with CPU-reference and captured-replay gates; +they do not route through pybind or introduce a second kernel implementation. + +The native vision graph composes patch im2col/embedding, all 27 SigLIP layers, +per-view FA2, optional fixed-factor spatial pooling, final LayerNorm, and the +1152-to-2048 multimodal projector. Position embedding expansion remains setup +work. With inputs restored before each of 100 replays, the graph keeps one +variant; final SigLIP and projected encoder tokens reach cosine 0.9999 or +better against the layer-by-layer PyTorch reference on both supported +checkpoint layouts. + +The first composed BF16 forward segment is the encoder QKV path: +RMSNorm, the folded QKV projection, RoPE split, and writes into the selected +layer of the shared K/V cache. Layer 17 is also the complete final encoder +layer behavior because the producer intentionally stops after populating its +cache. Its outputs are bit-exact (`cos=1`, `max=0`) against the PyTorch +checkpoint path for both OpenPI and LeRobot layouts, and the segment captures +and replays with one graph variant. + +Encoder layers 0-16 extend that segment through fixed-shape FA2, output +projection, residual/RMS normalization, the separate gate/up projections, +gated GELU, down projection, and the final residual update. A captured layer 0 +replayed 100 times remains a single variant and reaches cosine 0.999992 versus +the original PyTorch path on both checkpoint layouts. Layer 17 keeps the +intentional cache-only early exit described above. + +The native encoder composes all 18 layers into one captured graph while +preserving that final cache-only behavior. Restoring the input before each of +100 replays produces one graph variant. On both OpenPI and LeRobot checkpoint +layouts, the final encoder state and layer-17 Q/K/V each reach cosine 0.9999 or +better against the layer-by-layer PyTorch reference. This composition owns no +state object: activations and K/V remain context-backed buffers. + +The native decoder composes one BF16 AdaRMS/cross-attention/FFN layer, one +flow-matching update, and the complete 10-step diffusion graph. Decoder K/V is +appended at the device-side fixed-prompt position in the encoder cache; style +and noise remain context-backed buffers. Full 10-step captures replay 100 times +with one variant on both checkpoint layouts. Independent first and final +schedule steps reach cosine 0.9999 or better against PyTorch; the accumulated +endpoint gate remains part of the real-episode end-to-end validation because +synthetic random K/V amplifies SDPA-versus-FA2 rounding across steps. + +The native graph owner now assembles the completed segments into one `infer` +capture: prompt copy, vision, encoder, then diffusion. Prompt embeddings live +in a separate persistent buffer because `encoder_x` is an in-place residual +stream; each replay captures a D2D copy into its language window. Both +checkpoint layouts complete 100 full replays with one variant, bit-identical +outputs for restored inputs, and a constant workspace allocation count. The +persistent prompt source, not the overwritten encoder rows, is the primary +prompt-context capsule candidate. + +RTX attention owns a separate context-backed buffer set rather than borrowing +Torch tensors: SigLIP Q/K/V, encoder Q and 18-layer shared K/V cache, decoder +Q, fixed-shape `seqused/devpos` int32 values, FA2 outputs/LSE, and split-KV +accumulators. Layer K/V pointers are stable offsets into one cache allocation. +Updating a fixed prompt length writes the same three scalar buffers without +allocation or rebinding. The Python-free attention driver calls the vendored +FA2 raw C entries directly for SigLIP, fixed-shape encoder `seqused`, and +decoder `seqused` split-KV. Its graph gate changes the prompt length after +capture, replays 100 times with one variant, and verifies the new device-side +valid length is observed. `flash_rt_fa2` remains a thin Python adapter over the +same `libflashrt_fa2_raw` kernel owner. + +The native builder publishes one `infer` graph/stage and the ordered ports +`prompt`, `state`, `images`, `noise`, `actions`, and `actions_raw`. Identity +includes SM120, model/tokenizer SHA-256 values, prompt mode, fixed shapes, and +schedule parameters. The only capsule region is `rollout_boundary` over the +diffusion/action buffer. Prompt embeddings, encoder/decoder caches, attention +lengths, and RoPE remain context-owned `frt_buffer` workspace that each infer +rebuilds; they are not falsely advertised as independently restorable state. + +The returned verb override retains the builder-produced base model, which +retains the export and graph owner. Releasing the final public model releases +the overlay, export, captured graph, buffers, stream, and context in ownership +order without a second lifecycle owner. + +CUDA graph execs are process-local objects. They are not serialized as a +portable artifact. The native producer therefore loads assets and captures the +graph in the replay process. + +## Validation + +The minimum regression set for this contract is: + +``` +PYTHONPATH=.:./exec/build:./runtime/build python runtime/tests/test_runtime_export.py +PYTHONPATH=.:./exec/build:./runtime/build python runtime/tests/test_model_runtime_py.py +./runtime/build/test_model_runtime +ctest --test-dir cpp/build --output-on-failure +``` + +Real-checkpoint gates: + +``` +python cpp/tests/gate_pi05_native_weight_ops.py \ + --checkpoint \ + --probe cpp/build/pi05_native_weight_probe +``` + +``` +FLASHRT_BUILD_DIR= \ + python cpp/tests/gate_pi05_model_runtime_export.py \ + --lib /libflashrt_cpp_pi05_c.so ... +python cpp/tests/gate_pi05_c_api_export.py ... +``` + +The Python-produced overlay gate uses the FA2-enabled, SentencePiece-off SM120 +build because Python already owns prompt/tokenizer setup. Native-v2 factory +gates use the SentencePiece-enabled SM120 build in a Python-free producer +process. In either lane, pybind modules and the producer library must come from +the same build directory; graph and buffer handles cannot cross exec builds. + +Prompt/state STAGED ports require token-exact, formatter string-exact, +embedding bit-exact, fixed-vs-exact E2E cosine, and hot-contract coverage; a +producer must not retain the declarations if any required verb is unavailable. + +The native factory lifecycle gate is: + +``` +/pi05_native_open_probe \ + +``` + +Run it against both OpenPI and LeRobot checkpoint layouts. It validates the +public schema, one captured variant, prompt/state/image staging, direct SWAP +noise input, finite action output, and retain/release teardown. + +Python and C++ native-v2 producers must also publish identical canonical +port/stage/region records (their producer identity and fingerprints remain +different): + +``` +FLASHRT_BUILD_DIR= \ + python cpp/tests/gate_pi05_native_schema_parity.py \ + --checkpoint \ + --tokenizer \ + --native-probe /pi05_native_open_probe +``` + +Both producers are compared independently against +`cpp/tests/data/pi05_native_v2_schema.records`. Update that golden file only +for an intentional public-face change, and review the resulting fingerprint +and capsule compatibility impact in the same change. + +The native formatter and tokenizer must also remain token-exact over real +prompt/state traffic: + +``` +python cpp/tests/gate_pi05_tokenizer_corpus.py \ + --dataset \ + --checkpoint \ + --tokenizer \ + --probe /pi05_tokenizer_corpus_probe \ + --count 10000 +``` + +This gate normalizes every recorded state with the checkpoint q01/q99 values, +renders the full state prompt through the native formatter, and compares every +valid token ID with OpenPI `PaligemmaTokenizer`. The 10,000-record reference +run used a lightweight oracle whose tokenizer and formatter logic is kept +line-equivalent with upstream OpenPI; it covered 20 token lengths from 43 +through 62 with zero mismatches. This source-level oracle is distinct from the +official-environment end-to-end gate below. + +The real-episode numerical gate compares against the official OpenPI PyTorch +`PI0Pytorch.sample_actions` path, not another native intermediate: + +``` +python cpp/tests/gate_pi05_native_e2e.py \ + --checkpoint \ + --tokenizer \ + --dataset \ + --probe /pi05_native_e2e_probe \ + --episode 0 --frame 0 +``` + +The gate rounds the shared initial noise to the exported BF16 contract before +both runs. Raw and robot action outputs must each reach cosine 0.9999 against +the official FP32 residual path. Separately, the STAGED `actions` bytes must +match q01/q99 postprocess recomputed from the native BF16 `actions_raw` window +at `rtol=atol=1e-6`; this keeps numerical precision and IO semantics as two +independent acceptance checks. Set `OPENPI_BASELINE_SITE_PACKAGES` when the +official OpenPI Transformers replacement is installed in a separate prefix. + +Collect replay-only native and Python BF16 Nsight traces with the same fixed +shape. Setup, graph capture, prompt/image/noise staging, and output copies must +remain outside the CUDA profiler range: + +``` +FLASHRT_PROFILE_RANGE=1 nsys profile --trace=cuda \ + --cuda-graph-trace=node \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + /pi05_native_open_probe \ + + +nsys profile --trace=cuda --cuda-graph-trace=node \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + python cpp/tests/profile_pi05_python_replay.py \ + --checkpoint --num-views 2 --steps 10 + +nsys stats --report cuda_gpu_trace --format csv \ + .nsys-rep > .csv +nsys stats --report cuda_gpu_trace --format csv \ + .nsys-rep > .csv +python cpp/tests/gate_pi05_kernel_sequence.py \ + --native .csv --python .csv +``` + +The comparator rejects unknown kernels and requires equal raw event counts. +Its explicit equivalence list covers only selected GEMM kernel variants, +GEMM workspace-init versus split-K reduction helpers, `add_bias` versus the +equivalent `bias_res` form, and the two negative-infinity fill symbols. On the +reference RTX 5090 SM120 run both traces contained 3,576 raw events and their +3,172 logical-kernel sequences were exactly equal. + +The separate hot allocator gate profiles 1,000 complete service iterations +without tracing individual graph nodes. Each measured iteration updates prompt, +state, image, and noise inputs, launches one graph replay, and reads the logical +action output: + +``` +FLASHRT_PROFILE_REPLAYS=1000 FLASHRT_PROFILE_SERVICE_LOOP=1 \ +nsys profile --trace=cuda \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + /pi05_native_open_probe \ + +nsys stats --report cuda_api_trace --format csv \ + .nsys-rep > .csv +python cpp/tests/gate_pi05_hot_allocator.py \ + --trace .csv --expected-replays 1000 +``` + +The gate requires exactly 1,000 `cudaGraphLaunch` calls and rejects CUDA/driver +device allocation, host registration, mempool creation, virtual-memory map, +and corresponding release APIs. The probe independently requires one graph +variant after the final replay. Omitting `FLASHRT_PROFILE_SERVICE_LOOP` retains +the replay-only diagnostic mode for kernel-sequence comparison. +This trace proves the absence of CUDA/driver allocation APIs across the full +service iteration. Host allocation claims are scoped to components with an +explicit allocation-counter test, such as exec graph-cache LRU maintenance; +the trace does not infer host allocator behavior from CUDA API events. + +The same native probe can gate the complete hot state staging chain +(normalization, formatting, tokenization, embedding gather, and prompt-length +device update): + +``` +FLASHRT_HOT_STATE_UPDATES=1000 FLASHRT_HOT_STATE_P99_US=1000 \ + /pi05_native_open_probe \ + +``` + +The probe varies all eight state dimensions, warms 20 updates, measures 1,000 +updates, and requires the graph variant count to remain one. The reference +RTX 5090 SM120 run measured p50/p99/max of 39.31/41.70/43.44 microseconds, +well below the explicit one-millisecond p99 contract. + +The unload probe has no static dependency on the Pi0.5 producer. It resolves +the factory from the shared object, exercises an extra retain/release pair, +releases the final model reference, and only then unloads the producer: + +``` +/pi05_native_dlopen_probe \ + /libflashrt_cpp_pi05_c.so \ + 1 +``` + +For the ASAN build, instrument the producer, runtime, exec library, and probe, +not only the executable. CUDA needs `protect_shadow_gap=0` in this environment +to avoid an address-space collision with ASAN's default shadow gap: + +``` +ASAN_OPTIONS=detect_leaks=1:halt_on_error=1:protect_shadow_gap=0 \ + /pi05_native_dlopen_probe \ + /libflashrt_cpp_pi05_c.so \ + 1 +``` diff --git a/docs/pr_review_checklist.md b/docs/pr_review_checklist.md index 2a7d78dd..6efc0dd7 100644 --- a/docs/pr_review_checklist.md +++ b/docs/pr_review_checklist.md @@ -688,3 +688,28 @@ Must block: - New model path without routing/import/correctness evidence. - Kernel names, paths, or pybind symbols that hide model/hardware ownership. - `exec/` or common runtime changes that include scenario policy. + +## 24. Native Producer And Public Hygiene Checklist + +Apply this checklist to `runtime/`, `cpp/`, model-runtime adapters, and schema +gates: + +- [ ] No model/backend kind, model dimension, checkpoint field, or scenario + policy was added to the frozen ABI. +- [ ] Integrated, wrap, and override construction reject STAGED inputs without + `set_input` and STAGED outputs without `get_output`. +- [ ] No declaration-only object can escape a factory or reach consumer + adoption. +- [ ] Hardware identity comes from the active device/backend, not a requested + build flag or copied model default. +- [ ] Shared producer faces compare independently with checked-in canonical + records; expected counts are derived rather than repeated. +- [ ] Producer-private graphs, buffers, manifests, and identity pairs are not + incorrectly required to match across backends. +- [ ] CUDA allocator evidence covers the complete claimed service iteration; + host allocation claims use a host-side counter. +- [ ] A heterogeneous backend enters through an instance backend/capsule seam, + not a new backend registry or frozen `backend_kind` field. +- [ ] Public commands use placeholders and the diff contains no absolute local + paths, user/host/container names, tokens, internal URLs, environment + dumps, logs, or proprietary asset identifiers. diff --git a/exec/CMakeLists.txt b/exec/CMakeLists.txt index 2bcc975e..eaa6837a 100644 --- a/exec/CMakeLists.txt +++ b/exec/CMakeLists.txt @@ -33,6 +33,15 @@ target_include_directories(flashrt_exec ${CMAKE_CURRENT_SOURCE_DIR}/backend) target_link_libraries(flashrt_exec PUBLIC CUDA::cudart) +if(BUILD_TESTING) + add_executable(test_graph_lru_alloc tests/test_graph_lru_alloc.cpp) + target_include_directories(test_graph_lru_alloc PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src) + target_link_libraries(test_graph_lru_alloc PRIVATE flashrt_exec) + add_test(NAME graph_lru_alloc COMMAND test_graph_lru_alloc) +endif() + # --- pybind dev module (optional) --- find_package(Python3 COMPONENTS Interpreter Development.Module QUIET) find_package(pybind11 CONFIG QUIET) diff --git a/exec/src/graph.cpp b/exec/src/graph.cpp index 100c67a3..b5a776fc 100644 --- a/exec/src/graph.cpp +++ b/exec/src/graph.cpp @@ -4,7 +4,10 @@ void frt_graph_s::touch(frt_shape_key key) { for (auto it = lru.begin(); it != lru.end(); ++it) { - if (*it == key) { lru.erase(it); break; } + if (*it == key) { + lru.splice(lru.end(), lru, it); + return; + } } lru.push_back(key); // back = most recently used } diff --git a/exec/tests/test_graph_lru_alloc.cpp b/exec/tests/test_graph_lru_alloc.cpp new file mode 100644 index 00000000..f3f572e5 --- /dev/null +++ b/exec/tests/test_graph_lru_alloc.cpp @@ -0,0 +1,46 @@ +#include "internal.h" + +#include +#include +#include +#include + +namespace { + +std::atomic count_allocations{false}; +std::atomic allocation_count{0}; + +} // namespace + +void* operator new(std::size_t bytes) { + if (count_allocations.load(std::memory_order_relaxed)) { + allocation_count.fetch_add(1, std::memory_order_relaxed); + } + if (void* p = std::malloc(bytes)) return p; + throw std::bad_alloc(); +} + +void operator delete(void* p) noexcept { std::free(p); } +void operator delete(void* p, std::size_t) noexcept { std::free(p); } + +int main() { + frt_graph_s graph; + graph.lru = {1, 2, 3}; + + allocation_count.store(0, std::memory_order_relaxed); + count_allocations.store(true, std::memory_order_relaxed); + for (int i = 0; i < 1000; ++i) graph.touch((i % 3) + 1); + count_allocations.store(false, std::memory_order_relaxed); + + assert(allocation_count.load(std::memory_order_relaxed) == 0); + assert(graph.lru.size() == 3); + assert(graph.lru.back() == 1); + + allocation_count.store(0, std::memory_order_relaxed); + count_allocations.store(true, std::memory_order_relaxed); + graph.touch(4); + count_allocations.store(false, std::memory_order_relaxed); + assert(allocation_count.load(std::memory_order_relaxed) > 0); + assert(graph.lru.back() == 4); + return 0; +} diff --git a/flash_rt/frontends/torch/pi05_rtx.py b/flash_rt/frontends/torch/pi05_rtx.py index 5edf211d..6601feb6 100644 --- a/flash_rt/frontends/torch/pi05_rtx.py +++ b/flash_rt/frontends/torch/pi05_rtx.py @@ -1071,6 +1071,7 @@ def _set_prompt_fixed(self, prompt_len: int) -> None: vision_pool_factor=self._vision_pool_factor, vision_num_layers=self._vision_num_layers, fixed_shape=True, + norm_stats=self.norm_stats, **self._pipeline_precision_kwargs()) if self._fixed_pipeline.use_int8_vision_static: self._fixed_pipeline.vis_int8_static_calibrated = False @@ -1121,6 +1122,7 @@ def _set_prompt_per_length(self, state, prompt_len: int) -> None: num_steps=self._num_steps, vision_pool_factor=self._vision_pool_factor, vision_num_layers=self._vision_num_layers, + norm_stats=self.norm_stats, **self._pipeline_precision_kwargs()) self._prompt_pipeline_cache[prompt_len] = self.pipeline # Static INT8 vision scales are per-pipeline-instance. diff --git a/flash_rt/frontends/torch/pi05_rtx_fp16.py b/flash_rt/frontends/torch/pi05_rtx_fp16.py index f0e78266..3d9b6121 100644 --- a/flash_rt/frontends/torch/pi05_rtx_fp16.py +++ b/flash_rt/frontends/torch/pi05_rtx_fp16.py @@ -989,6 +989,7 @@ def set_prompt(self, prompt_text: str, state=None) -> None: num_steps=self._num_steps, vision_pool_factor=self._vision_pool_factor, vision_num_layers=self._vision_num_layers, + norm_stats=self.norm_stats, **self._pipeline_precision_kwargs()) self._prompt_pipeline_cache[prompt_len] = self.pipeline # Static INT8 vision scales are per-pipeline-instance. diff --git a/flash_rt/models/pi05/pipeline_rtx.py b/flash_rt/models/pi05/pipeline_rtx.py index 01be1102..4c43c5ac 100644 --- a/flash_rt/models/pi05/pipeline_rtx.py +++ b/flash_rt/models/pi05/pipeline_rtx.py @@ -137,6 +137,7 @@ class Pi05Pipeline: use_fp8_decoder: Enable FP8 on decoder branch (else BF16). use_int8_decoder: Enable experimental decoder-only INT8 GEMMs. num_steps: Diffusion denoise steps (default 10). + norm_stats: Producer metadata for logical state/action IO contracts. Expected weights dict keys: Vision BF16: @@ -185,11 +186,13 @@ def __init__(self, gemm, fvk, attn_backend, weights, *, vision_pool_factor: int = 1, vision_num_layers: int = VIS_L, num_steps: int = NUM_STEPS_DEFAULT, - fixed_shape: bool = False): + fixed_shape: bool = False, + norm_stats=None): self.gemm = gemm self.fvk = fvk self.attn = attn_backend self.weights = weights + self.norm_stats = norm_stats or {} # Fixed-shape state-prompt mode: one captured graph at the MAX prompt # length serves every length via seqused masking + devpos K/V append. diff --git a/flash_rt/models/pi05/pipeline_rtx_fp16.py b/flash_rt/models/pi05/pipeline_rtx_fp16.py index 09b92584..59065c41 100644 --- a/flash_rt/models/pi05/pipeline_rtx_fp16.py +++ b/flash_rt/models/pi05/pipeline_rtx_fp16.py @@ -186,6 +186,7 @@ class Pi05PipelineFP16: use_fp8_decoder: Enable FP8 on decoder branch (else BF16). use_int8_decoder: Enable experimental decoder-only INT8 GEMMs. num_steps: Diffusion denoise steps (default 10). + norm_stats: Producer metadata for logical state/action IO contracts. Expected weights dict keys: Vision BF16: @@ -233,7 +234,8 @@ def __init__(self, gemm, fvk, attn_backend, weights, *, use_int8_vision_static: bool = False, vision_pool_factor: int = 1, vision_num_layers: int = VIS_L, - num_steps: int = NUM_STEPS_DEFAULT): + num_steps: int = NUM_STEPS_DEFAULT, + norm_stats=None): if use_fp8 or use_fp8_decoder: raise ValueError("Pi05PipelineFP16 supports only use_fp8=False") if use_int8_decoder or use_int8_encoder or use_int8_vision or use_int8_vision_static: @@ -244,6 +246,7 @@ def __init__(self, gemm, fvk, attn_backend, weights, *, self.fvk = _Fp16KernelProxy(fvk) self.attn = attn_backend self.weights = weights + self.norm_stats = norm_stats or {} self.num_views = int(num_views) self.max_prompt_len = int(max_prompt_len) diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index 9ecb55fa..ae2cf05c 100644 --- a/flash_rt/models/pi05/runtime_export.py +++ b/flash_rt/models/pi05/runtime_export.py @@ -9,6 +9,8 @@ from __future__ import annotations +import hashlib + def exec_enable(pl) -> None: """Create the exec ctx/graphs for a captured pipeline and adopt any @@ -84,13 +86,25 @@ def export_model_runtime(pl, identity=None, extra_regions=None, graphs: images/actions are STAGED, noise is SWAP, and the C++ runtime supplies the verbs through ``frt_model_runtime_override_verbs``. + ``io="native_v2"`` extends that face with prompt/state STAGED ports for + fixed state-prompt deployments. The declaration is intended to be consumed + through the C++ verb overlay; port/window/region changes are part of the + export identity and therefore intentionally change the fingerprint. + Prompt staging (text -> embeds) stays with the frontend / the native tokenizer producer. ``stage_plan`` defaults to the full infer graph; an explicit StagePlan or registered plan name may select already-captured graphs from this export. ``stage_plan_kwargs`` are passed only to registered plan factories, for deployment-specific graph cuts. """ - parts = _parts(pl, identity, extra_regions) + identity_for_parts = identity + if io == "native_v2": + _require_native_v2_ready(pl) + identity_for_parts = { + **{str(k): str(v) for k, v in (identity or {}).items()}, + "io": "native_v2", + } + parts = _parts(pl, identity_for_parts, extra_regions) from flash_rt.runtime import export as _rt from flash_rt.subgraphs.pi05 import stage_plans as _pi05_stage_plans # noqa: F401 from flash_rt.subgraphs.stage_plan import resolve_stage_plan @@ -138,7 +152,7 @@ def export_model_runtime(pl, identity=None, extra_regions=None, "in", "swap", shape=(1,), buffer=wrap["rtc_guidance_weight"]), ]) - elif io == "native": + elif io in ("native", "native_v2"): ports = [ _rt.PortSpec("images", "image", tensor_dtype, "nhwc", "in", "staged", required=True, shape=(num_views, 224, 224, 3), @@ -146,10 +160,23 @@ def export_model_runtime(pl, identity=None, extra_regions=None, buffer=wrap["observation_images_normalized"]), _rt.PortSpec("noise", "tensor", tensor_dtype, "flat", "in", "swap", shape=(chunk, 32), buffer=wrap["diffusion_noise"]), - _rt.PortSpec("actions", "action", tensor_dtype, "flat", "out", + _rt.PortSpec("actions", "action", "f32", "flat", "out", "staged", shape=(chunk, robot_action_dim), - buffer=wrap["diffusion_noise"]), + nbytes=chunk * robot_action_dim * 4), ] + if io == "native_v2": + state_dim = _state_dim(pl) + ports = [ + _rt.PortSpec("prompt", "text", "u8", "flat", "in", "staged", + required=True, shape=(-1,)), + _rt.PortSpec("state", "state", "f32", "flat", "in", "staged", + required=True, shape=(state_dim,)), + ] + ports + if not (uses_rtc_prefix or uses_rtc_vjp): + ports.append( + _rt.PortSpec("actions_raw", "tensor", tensor_dtype, + "flat", "out", "swap", shape=(chunk, 32), + buffer=wrap["diffusion_noise"])) if uses_rtc_prefix or uses_rtc_vjp: ports.extend([ _rt.PortSpec("prev_action_chunk", "tensor", tensor_dtype, @@ -193,6 +220,16 @@ def step(): return rc return rc + manifest_extra = {"stage_plan": plan.manifest(), "io": io} + declaration_only = io in ("native", "native_v2") + if declaration_only: + manifest_extra["declaration_only"] = True + if io == "native_v2": + manifest_extra["prompt"] = { + "state_prompt_mode": "fixed", + "max_prompt_len": int(getattr(pl, "max_prompt_len", 0) or 0), + "state_dim": _state_dim(pl), + } return _rt.build_model_runtime( pl._exec_ctx, streams=parts["streams"], @@ -202,9 +239,10 @@ def step(): ports=ports, stages=stages, identity=parts["identity"], - manifest_extra={"stage_plan": plan.manifest(), "io": io}, + manifest_extra=manifest_extra, owner=parts["owner"], step=step, + _allow_incomplete_staged=declaration_only, ) @@ -230,6 +268,37 @@ def _robot_action_dim(pl): return int(LIBERO_ACTION_DIM) +def _state_dim(pl): + """Raw proprioception dimension exposed by native_v2 STATE/STAGED.""" + try: + return int(len(pl.norm_stats["state"]["q01"])) + except Exception as e: + raise ValueError( + "Pi05 native_v2 requires norm_stats['state']['q01']") from e + + +def _tokenizer_sha256() -> str: + from flash_rt.utils.paligemma_tokenizer import ( + resolve_paligemma_tokenizer_path, + ) + h = hashlib.sha256() + with open(resolve_paligemma_tokenizer_path(), "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def _require_native_v2_ready(pl) -> None: + mode = getattr(pl, "_state_prompt_mode", None) + fixed = bool(getattr(pl, "_fixed_shape", False)) + if mode != "fixed" and not fixed: + raise ValueError( + "Pi05 native_v2 requires state_prompt_mode='fixed'") + if int(getattr(pl, "max_prompt_len", 0) or 0) < 200: + raise ValueError("Pi05 native_v2 requires max_prompt_len >= 200") + _state_dim(pl) + + def _parts(pl, identity, extra_regions): """Shared assembly for the plain export and the model-runtime export.""" if getattr(pl, "_graph", None) is None: @@ -296,6 +365,10 @@ def _parts(pl, identity, extra_regions): "robot_action_dim": str(_robot_action_dim(pl)), } ident.update({str(k): str(v) for k, v in (identity or {}).items()}) + if ident.get("io") == "native_v2": + ident["state_prompt_mode"] = "fixed" + ident["state_dim"] = str(_state_dim(pl)) + ident["tokenizer_sha256"] = _tokenizer_sha256() return { "wrap": wrap, diff --git a/flash_rt/runtime/export.py b/flash_rt/runtime/export.py index 05437dc0..8906f7f4 100644 --- a/flash_rt/runtime/export.py +++ b/flash_rt/runtime/export.py @@ -145,7 +145,7 @@ class PortSpec: name: str modality: int | str # MODALITY name or value - dtype: int | str = "bf16" # device-side tensor dtype + dtype: int | str = "bf16" # logical payload/target tensor dtype layout: int | str = "flat" direction: int | str = "in" update: int | str = "swap" @@ -276,6 +276,7 @@ def build_model_runtime( get_output=None, prepare=None, step=None, + _allow_incomplete_staged=False, ) -> ModelRuntime: """Assemble an ``frt_model_runtime_v1``: an export plus the dynamic-IO contract (ports, stage DAG, optional verb callables) under one identity — @@ -289,6 +290,22 @@ def build_model_runtime( them from any thread. SWAP ports need no callable — hosts write the declared buffer window directly. """ + staged_inputs = any( + _enum(UPDATE, port.update) == UPDATE["staged"] and + _enum(DIRECTION, port.direction) == DIRECTION["in"] + for port in ports + ) + staged_outputs = any( + _enum(UPDATE, port.update) == UPDATE["staged"] and + _enum(DIRECTION, port.direction) == DIRECTION["out"] + for port in ports + ) + if not _allow_incomplete_staged: + if staged_inputs and set_input is None: + raise ValueError("STAGED input ports require set_input") + if staged_outputs and get_output is None: + raise ValueError("STAGED output ports require get_output") + b, anchor, manifest_json = _assemble( ctx, streams=streams, graphs=graphs, buffers=buffers, regions=regions, ports=ports, stages=stages, identity=identity, diff --git a/runtime/include/flashrt/model_runtime.h b/runtime/include/flashrt/model_runtime.h index 51dbe689..bf14d025 100644 --- a/runtime/include/flashrt/model_runtime.h +++ b/runtime/include/flashrt/model_runtime.h @@ -131,7 +131,7 @@ typedef struct frt_image_view { typedef struct frt_runtime_port_desc { const char* name; /* "images", "prompt", "state", "actions" */ uint32_t modality; /* frt_rt_modality */ - uint32_t dtype; /* frt_rt_dtype (of the DEVICE-side tensor) */ + uint32_t dtype; /* frt_rt_dtype of logical payload/target */ uint32_t layout; /* frt_rt_layout */ uint32_t direction; /* frt_rt_port_direction */ uint32_t update; /* frt_rt_port_update */ @@ -247,8 +247,10 @@ int frt_runtime_builder_add_stage(frt_runtime_builder, uint32_t graph, /* 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 - * unsupported). Consumes the builder. */ + * is copied; entries may be null only when no matching STAGED declaration + * requires them (other missing verbs report unsupported). A STAGED input + * requires set_input and a STAGED output requires get_output. On validation + * failure the builder is not consumed. Consumes the builder on success. */ frt_model_runtime_v1* frt_runtime_builder_finish_model( frt_runtime_builder, const frt_model_runtime_verbs* verbs, void* verbs_self, @@ -262,7 +264,8 @@ frt_model_runtime_v1* frt_runtime_builder_finish_model( /* producer builds both. Descriptor arrays are copied. The wrapper */ /* takes one export reference and calls `wrapper_release(wrapper_owner)`*/ /* exactly once when its refcount hits zero (use it to destroy the */ -/* producer instance behind `verbs_self`). */ +/* producer instance behind `verbs_self`). STAGED declarations require */ +/* matching input/output verbs, as on construction path 1. */ /* ------------------------------------------------------------------ */ frt_model_runtime_v1* frt_model_runtime_wrap( const frt_runtime_export_v1* exp, @@ -278,7 +281,8 @@ frt_model_runtime_v1* frt_model_runtime_wrap( /* a native runtime owns hot-path transforms. The override retains `in` */ /* so all inherited descriptor pointers stay valid; consumers release */ /* only the returned object. `retain_owner`/`release_owner` manage the */ -/* native verb object, called once at construction/destruction. */ +/* native verb object, called once at construction/destruction. The new */ +/* verbs must satisfy every inherited STAGED input/output declaration. */ /* ------------------------------------------------------------------ */ frt_model_runtime_v1* frt_model_runtime_override_verbs( const frt_model_runtime_v1* in, diff --git a/runtime/src/model_runtime.cpp b/runtime/src/model_runtime.cpp index 9e82e706..e7b284d7 100644 --- a/runtime/src/model_runtime.cpp +++ b/runtime/src/model_runtime.cpp @@ -21,6 +21,21 @@ bool valid_port_args(const char* name, uint32_t direction, uint32_t update, return true; } +bool staged_verbs_present(const frt_runtime_port_desc* ports, uint64_t n_ports, + const frt_model_runtime_verbs* verbs) { + bool needs_input = false; + bool needs_output = false; + for (uint64_t i = 0; i < n_ports; ++i) { + if (ports[i].update != FRT_RT_PORT_STAGED) continue; + needs_input |= ports[i].direction == FRT_RT_PORT_IN; + needs_output |= ports[i].direction == FRT_RT_PORT_OUT; + } + const bool complete = + verbs && verbs->struct_size >= sizeof(frt_model_runtime_verbs); + return (!needs_input || (complete && verbs->set_input)) && + (!needs_output || (complete && verbs->get_output)); +} + /* Default stubs for verbs a producer does not provide: report unsupported * (-3) instead of leaving null function pointers for consumers to crash on. */ int stub_set_input(void*, uint32_t, const void*, uint64_t, int) { return -3; } @@ -108,6 +123,8 @@ extern "C" frt_model_runtime_v1* frt_runtime_builder_finish_model( void (*release_owner)(void*)) { if (!b) return nullptr; Holder* h = b->h; + if (!staged_verbs_present(h->ports.data(), h->ports.size(), verbs)) + return nullptr; frt_rt::finish_export_into(h, b, owner, retain_owner, release_owner); frt_model_runtime_v1& m = h->model; @@ -172,6 +189,7 @@ extern "C" frt_model_runtime_v1* frt_model_runtime_wrap( if (!valid_port_args(ports[i].name, ports[i].direction, ports[i].update, ports[i].shape, ports[i].rank)) return nullptr; + if (!staged_verbs_present(ports, n_ports, verbs)) return nullptr; for (uint64_t i = 0; i < n_stages; ++i) { if (stages[i].graph >= exp->n_graphs) return nullptr; if (stages[i].n_after && !stages[i].after) return nullptr; @@ -255,6 +273,7 @@ extern "C" frt_model_runtime_v1* frt_model_runtime_override_verbs( void* owner, void (*retain_owner)(void*), void (*release_owner)(void*)) { if (!valid_model_runtime(in)) return nullptr; + if (!staged_verbs_present(in->ports, in->n_ports, verbs)) return nullptr; auto* o = new VerbOverride(); o->base = in; diff --git a/runtime/tests/test_model_runtime.cpp b/runtime/tests/test_model_runtime.cpp index 0c9b29a5..a9362b33 100644 --- a/runtime/tests/test_model_runtime.cpp +++ b/runtime/tests/test_model_runtime.cpp @@ -56,16 +56,30 @@ static const int64_t IMG_SHAPE[4] = {3, 224, 224, 3}; static const int64_t ACT_SHAPE[2] = {50, 32}; static void add_ports_and_stages(frt_runtime_builder b, int64_t img_h = 224, - uint64_t act_bytes = 3200) { + uint64_t act_bytes = 3200, + uint32_t cadence_hz = 30, + bool add_prompt_state = false) { const int64_t img[4] = {3, img_h, 224, 3}; frt_runtime_builder_add_port(b, "images", FRT_RT_MOD_IMAGE, FRT_RT_DTYPE_BF16, FRT_RT_LAYOUT_NHWC, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, - img, 4, 30, nullptr, 0, 0); + img, 4, cadence_hz, nullptr, 0, 0); frt_runtime_builder_add_port(b, "actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_BF16, FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, ACT_SHAPE, 2, 0, FAKE_B0, 0, act_bytes); + if (add_prompt_state) { + const int64_t text_shape[1] = {-1}; + const int64_t state_shape[1] = {8}; + 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, + text_shape, 1, 0, nullptr, 0, 0); + 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, + state_shape, 1, 0, nullptr, 0, 0); + } const uint32_t after1[1] = {0}; frt_runtime_builder_add_stage(b, 0, nullptr, 0); frt_runtime_builder_add_stage(b, 1, after1, 1); @@ -87,18 +101,32 @@ 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 */ + CHECK(frt_runtime_builder_finish_model( + b, nullptr, nullptr, nullptr, nullptr, nullptr) == nullptr, + "finish_model rejects STAGED ports without verbs"); + /* The builder survives both refusals; valid verbs consume it. */ + frt_model_runtime_verbs staged_verbs{}; + staged_verbs.struct_size = sizeof(staged_verbs); + staged_verbs.set_input = v_set_input; + staged_verbs.get_output = v_get_output; frt_model_runtime_v1* m = frt_runtime_builder_finish_model( - b, nullptr, nullptr, nullptr, nullptr, nullptr); - CHECK(m != nullptr, "finish_model after refused finish"); - /* absent producer verbs become unsupported stubs, never null */ - CHECK(m->verbs.set_input && m->verbs.step && m->verbs.last_error, - "null verbs are stubbed"); - CHECK(m->verbs.set_input(m->self, 0, nullptr, 0, -1) == -3 && - m->verbs.step(m->self) == -3 && - m->verbs.last_error(m->self)[0] != '\0', - "stubs report unsupported (-3) with an explanation"); + b, &staged_verbs, nullptr, nullptr, nullptr, nullptr); + CHECK(m != nullptr, "finish_model after refused finishes"); m->release(m->owner); + + frt_runtime_builder sb = make_builder(); + const int64_t shape[1] = {16}; + CHECK(frt_runtime_builder_add_port( + sb, "setup", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_SETUP, 0, + shape, 1, 0, nullptr, 0, 0) == 0, + "add non-staged port"); + frt_model_runtime_v1* sm = frt_runtime_builder_finish_model( + sb, nullptr, nullptr, nullptr, nullptr, nullptr); + CHECK(sm && sm->verbs.step(sm->self) == -3 && + sm->verbs.last_error(sm->self)[0] != '\0', + "missing non-staged verbs retain unsupported stubs"); + sm->release(sm->owner); } /* --- integrated build: struct, identity, fingerprint, verbs --- */ @@ -179,6 +207,39 @@ int main() { "graph stream change changes the fingerprint"); m4->release(m4->owner); } + /* prompt/state port additions define a new native face */ + { + frt_runtime_builder b5 = make_builder(); + add_ports_and_stages(b5, 224, 3200, 30, true); + frt_model_runtime_v1* m5 = frt_runtime_builder_finish_model( + b5, &verbs, &vlog, nullptr, nullptr, nullptr); + CHECK(m5 && m5->exp->fingerprint != m->exp->fingerprint, + "adding prompt/state ports changes the fingerprint"); + m5->release(m5->owner); + } + /* cadence is a scheduling hint, not deployment identity */ + { + frt_runtime_builder b6 = make_builder(); + add_ports_and_stages(b6, 224, 3200, 60); + frt_model_runtime_v1* m6 = frt_runtime_builder_finish_model( + b6, &verbs, &vlog, nullptr, nullptr, nullptr); + CHECK(m6 && m6->exp->fingerprint == m->exp->fingerprint, + "cadence hint does not change the fingerprint"); + m6->release(m6->owner); + } + /* manifest is discovery metadata, not deployment identity */ + { + frt_runtime_builder b7 = make_builder(); + add_ports_and_stages(b7); + CHECK(frt_runtime_builder_set_manifest( + b7, "{\"note\":\"discovery-only\"}") == 0, + "set manifest metadata"); + frt_model_runtime_v1* m7 = frt_runtime_builder_finish_model( + b7, &verbs, &vlog, nullptr, nullptr, nullptr); + CHECK(m7 && m7->exp->fingerprint == m->exp->fingerprint, + "manifest does not change the fingerprint"); + m7->release(m7->owner); + } /* verbs plumb through self */ m->verbs.set_input(m->self, 0, nullptr, 0, -1); @@ -224,6 +285,9 @@ int main() { CHECK(frt_model_runtime_wrap(exp, ports, 1, &bad, 1, &verbs, &vlog, nullptr, nullptr) == nullptr, "wrap rejects a stage over a missing graph"); + CHECK(frt_model_runtime_wrap(exp, ports, 1, stages, 1, nullptr, + nullptr, nullptr, nullptr) == nullptr, + "wrap rejects STAGED input without set_input"); frt_model_runtime_v1* wm = frt_model_runtime_wrap( exp, ports, 1, stages, 1, &verbs, &vlog, &wrapper_freed, @@ -262,6 +326,14 @@ int main() { native_verbs.step = v_step; native_verbs.last_error = v_last_error; + frt_model_runtime_verbs incomplete_verbs = native_verbs; + incomplete_verbs.get_output = nullptr; + CHECK(frt_model_runtime_override_verbs( + base, &incomplete_verbs, &native_vlog, &native_owner, + owner_retain, owner_release) == nullptr && + native_owner.retains == 0, + "override rejects missing STAGED output verb without retain"); + frt_model_runtime_v1* over = frt_model_runtime_override_verbs( base, &native_verbs, &native_vlog, &native_owner, owner_retain, owner_release); diff --git a/runtime/tests/test_model_runtime_py.py b/runtime/tests/test_model_runtime_py.py index 39dde2e7..8d982ab0 100644 --- a/runtime/tests/test_model_runtime_py.py +++ b/runtime/tests/test_model_runtime_py.py @@ -95,7 +95,10 @@ def record(stream): def build(setup, img_h=224, verbs=None): ctx, sid, src, dst, g = setup - verbs = verbs or {} + verbs = verbs or { + "set_input": lambda port, payload, stream: 0, + "get_output": lambda port, stream: b"", + } return build_model_runtime( ctx, streams=[StreamSpec("main", sid)], @@ -142,9 +145,48 @@ def build_split(setup): stages=plan.to_stage_specs(export_mod), identity={"model": "trivial", "quant": "none"}, manifest_extra={"stage_plan": plan.manifest()}, + set_input=lambda port, payload, stream: 0, + get_output=lambda port, stream: b"", ) +def check_staged_verb_guards(setup): + ctx, sid, src, dst, g = setup + common = dict( + ctx=ctx, + streams=[StreamSpec("main", sid)], + graphs=[GraphSpec("infer", g, 0, (0,))], + buffers=[BufferSpec("src", src, "input"), + BufferSpec("dst", dst, "output")], + stages=[StageSpec("infer")], + identity={"model": "verb_guard"}, + ) + try: + build_model_runtime( + **common, + ports=[PortSpec("input", "state", "f32", "flat", "in", + "staged", shape=(1,))], + ) + except ValueError as exc: + input_rejected = "require set_input" in str(exc) + else: + input_rejected = False + check("STAGED input without set_input is rejected", input_rejected) + + try: + build_model_runtime( + **common, + ports=[PortSpec("output", "action", "f32", "flat", "out", + "staged", shape=(1,))], + set_input=lambda port, payload, stream: 0, + ) + except ValueError as exc: + output_rejected = "require get_output" in str(exc) + else: + output_rejected = False + check("STAGED output without get_output is rejected", output_rejected) + + def check_stage_plan_registry(): register_stage_plan( "unit_chain", @@ -334,6 +376,9 @@ def py_step(): check_stage_plan_registry() check_vjp_guided_port_lowering(setup) + print("== staged verb declarations ==") + check_staged_verb_guards(setup) + print("== verbs through the C function pointers ==") rc = rt.model_set_input(mr.ptr, 1, b"\xAA\xBB", -1) check("set_input reaches the Python callable",