From 5bcf524989351ed3632064e56d694be1232e9a15 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 13:50:15 -0400 Subject: [PATCH 01/61] docs: add Pi0.5 native IO contract --- docs/mindon_pi05_integration.md | 200 +++++++++++++++++++++++++++++ docs/pi05_io_contract.md | 219 ++++++++++++++++++++++++++++++++ 2 files changed, 419 insertions(+) create mode 100644 docs/mindon_pi05_integration.md create mode 100644 docs/pi05_io_contract.md diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md new file mode 100644 index 00000000..1680e8d2 --- /dev/null +++ b/docs/mindon_pi05_integration.md @@ -0,0 +1,200 @@ +# 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: After 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: Future Native 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. + +## 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 +``` + +Until Lane B lands, prompt and state changes require a setup-time producer +refresh. After Lane B, Mindon should send task text to the `prompt` port and +raw proprioception to the `state` port. 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]`. +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 a producer-declared raw `TENSOR/SWAP` output +such as `actions_raw` when exported by a stage plan. + +## 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. + +If prompt context becomes a region in a future producer face, old capsules +will not restore into that new face. That is expected and protects state +safety. + +## 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. + +## 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. +- Snapshot/restore is tested within one live capture. +- Nexus core code remains unchanged for model-specific semantics. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md new file mode 100644 index 00000000..03dbaa6f --- /dev/null +++ b/docs/pi05_io_contract.md @@ -0,0 +1,219 @@ +# 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-visible robot action chunk, `FLAT`, `(chunk_length, robot_action_dim)` | `diffusion_noise` | + +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` + +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. + +## Target Face After Prompt/State Staging + +After the C++ text path exists, the Pi0.5 native face may add the following +ports. Adding these ports changes the model-runtime identity and therefore the +fingerprint. Existing capsules from the old face must refuse restore into the +new face. + +| port | modality/update | direction | payload | semantics | +|---|---|---|---|---| +| `prompt` | `TEXT/STAGED` | input | UTF-8 bytes, no trailing NUL required | task text only | +| `state` | `STATE/STAGED` | input | `F32`, `(state_dim,)` | raw proprioception, normalized/discretized by the producer | + +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 pixels; +- 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. + +The producer contract should reject unsupported input shape, dtype, layout, or +view count with a shape/status error. It should not silently reinterpret camera +formats in a way that changes model semantics. If a deployment supports more +pixel formats, the supported set must be documented by the producer and tested +against the CPU reference path. + +## 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. + +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, such as the existing RTC +`actions_raw` port, not by changing the meaning of `actions`. + +## 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. + +## 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 should normally add a restorable prompt context region +only after the bytes that define that context are explicit. A valid region +could include the language rows of `encoder_x` plus the fixed-prompt valid +length scalar. Region layout and order are fingerprinted; old capsules should +not restore into the new layout. + +## 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, after prompt/state staging: same as Lane A, plus hot + `prompt`/`state` STAGED ports. +- Lane C, future native producer: a C++ shared object implements + `frt_model_runtime_open_v1(config_json, &out)` and produces the same public + struct without Python setup. + +CUDA graph execs are process-local objects. They are not serialized as a +portable artifact. Removing Python from setup requires a native producer that +loads assets and captures graphs 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_model_runtime_export.py ... +python cpp/tests/gate_pi05_c_api_export.py ... +``` + +For prompt/state staging, add token-exact, formatter string-exact, embedding +bit-exact, fixed-vs-exact E2E cosine, and hot-contract tests before declaring +the new STAGED ports. From a236bae155e419374ec0e8ffa5203186eb23db5a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 13:53:29 -0400 Subject: [PATCH 02/61] feat: add Pi0.5 prompt formatter --- cpp/CMakeLists.txt | 5 ++ .../flashrt/cpp/models/pi05/prompt_format.h | 27 +++++++ cpp/models/pi05/src/prompt_format.cpp | 71 +++++++++++++++++++ cpp/tests/test_pi05_prompt_format.cpp | 64 +++++++++++++++++ 4 files changed, 167 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h create mode 100644 cpp/models/pi05/src/prompt_format.cpp create mode 100644 cpp/tests/test_pi05_prompt_format.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 1b097a85..afba0bfb 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -80,6 +80,7 @@ target_link_libraries(flashrt_cpp_vla add_library(flashrt_cpp_pi05 STATIC models/pi05/src/spec.cpp + models/pi05/src/prompt_format.cpp models/pi05/src/io.cpp models/pi05/src/runtime.cpp) target_include_directories(flashrt_cpp_pi05 @@ -106,6 +107,10 @@ if(BUILD_TESTING) PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) add_test(NAME pi05_runtime COMMAND test_pi05_runtime) + 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) + if(FLASHRT_CPP_WITH_CUDA_STAGING) add_executable(test_device_staging tests/test_device_staging.cpp) target_link_libraries(test_device_staging 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..e3b5cbe6 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h @@ -0,0 +1,27 @@ +#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); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H diff --git a/cpp/models/pi05/src/prompt_format.cpp b/cpp/models/pi05/src/prompt_format.cpp new file mode 100644 index 00000000..6f3f0aad --- /dev/null +++ b/cpp/models/pi05/src/prompt_format.cpp @@ -0,0 +1,71 @@ +#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) { + const std::string cleaned = clean_task_prompt(prompt); + if (!state) return cleaned; + + const auto tokens = discretize_state_prompt_bins(state, n_state); + std::ostringstream oss; + oss << "Task: " << cleaned << ", State: "; + for (std::size_t i = 0; i < tokens.size(); ++i) { + if (i) oss << ' '; + oss << tokens[i]; + } + oss << ";\nAction: "; + return oss.str(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/tests/test_pi05_prompt_format.cpp b/cpp/tests/test_pi05_prompt_format.cpp new file mode 100644 index 00000000..d687873e --- /dev/null +++ b/cpp/tests/test_pi05_prompt_format.cpp @@ -0,0 +1,64 @@ +#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: "); +} + +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; +} From cc5db63e0e899406f3ca3d4073d90afbd9a1e104 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 13:55:38 -0400 Subject: [PATCH 03/61] feat: add text embedding gather --- cpp/CMakeLists.txt | 5 + .../include/flashrt/cpp/modalities/text.h | 26 +++++ cpp/modalities/src/text_cpu.cpp | 99 +++++++++++++++++++ cpp/tests/test_text_modalities.cpp | 92 +++++++++++++++++ 4 files changed, 222 insertions(+) create mode 100644 cpp/modalities/include/flashrt/cpp/modalities/text.h create mode 100644 cpp/modalities/src/text_cpu.cpp create mode 100644 cpp/tests/test_text_modalities.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index afba0bfb..132b5ec0 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -50,6 +50,7 @@ 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_CUDA_KERNELS) list(APPEND FLASHRT_CPP_MODALITY_SRCS modalities/src/vision_cuda.cu) @@ -102,6 +103,10 @@ if(BUILD_TESTING) 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_pi05_runtime tests/test_pi05_runtime.cpp) target_link_libraries(test_pi05_runtime PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) 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..0b1a651f --- /dev/null +++ b/cpp/modalities/include/flashrt/cpp/modalities/text.h @@ -0,0 +1,26 @@ +#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; +}; + +Status gather_token_embeddings_cpu(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output); + +} // namespace modalities +} // namespace flashrt + +#endif // FLASHRT_MODALITIES_TEXT_H diff --git a/cpp/modalities/src/text_cpu.cpp b/cpp/modalities/src/text_cpu.cpp new file mode 100644 index 00000000..4f0ca410 --- /dev/null +++ b/cpp/modalities/src/text_cpu.cpp @@ -0,0 +1,99 @@ +#include "flashrt/cpp/modalities/text.h" + +#include + +namespace flashrt { +namespace modalities { +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 + +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(); +} + +} // namespace modalities +} // namespace flashrt 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; +} From 35040f722c39b4160b62fcebc5f9254111babe98 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:01:43 -0400 Subject: [PATCH 04/61] feat: add native text tokenizer hook --- cpp/CMakeLists.txt | 44 ++++++ .../flashrt/cpp/modalities/tokenizer.h | 53 ++++++++ .../src/sentencepiece_tokenizer.cpp | 127 ++++++++++++++++++ cpp/modalities/src/tokenizer_unavailable.cpp | 46 +++++++ cpp/tests/test_text_tokenizer.cpp | 81 +++++++++++ 5 files changed, 351 insertions(+) create mode 100644 cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h create mode 100644 cpp/modalities/src/sentencepiece_tokenizer.cpp create mode 100644 cpp/modalities/src/tokenizer_unavailable.cpp create mode 100644 cpp/tests/test_text_tokenizer.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 132b5ec0..53329a92 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -22,6 +22,7 @@ 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) if(FLASHRT_CPP_WITH_CUDA_STAGING) find_package(CUDAToolkit REQUIRED) endif() @@ -29,6 +30,30 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) 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) @@ -52,6 +77,13 @@ set(FLASHRT_CPP_MODALITY_SRCS 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) endif() @@ -64,6 +96,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) @@ -107,6 +147,10 @@ if(BUILD_TESTING) 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) 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..0a5f0a4b --- /dev/null +++ b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h @@ -0,0 +1,53 @@ +#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) 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/src/sentencepiece_tokenizer.cpp b/cpp/modalities/src/sentencepiece_tokenizer.cpp new file mode 100644 index 00000000..bc33bf36 --- /dev/null +++ b/cpp/modalities/src/sentencepiece_tokenizer.cpp @@ -0,0 +1,127 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +#include + +#include +#include + +namespace flashrt { +namespace modalities { + +struct SentencePieceTokenizer::Impl { + sentencepiece::SentencePieceProcessor processor; + 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) const { + 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"); + } + + std::vector encoded; + auto status = impl_->processor.Encode(text, &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 (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(encoded.size() + extra); + for (int id : 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 (token_ids->size() > options.max_tokens) { + return Status::error(StatusCode::kShapeMismatch, + "encoded token sequence exceeds 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(); +} + +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/tokenizer_unavailable.cpp b/cpp/modalities/src/tokenizer_unavailable.cpp new file mode 100644 index 00000000..49fa3438 --- /dev/null +++ b/cpp/modalities/src/tokenizer_unavailable.cpp @@ -0,0 +1,46 @@ +#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) const { + (void)text; + (void)options; + if (token_ids) token_ids->clear(); + return Status::error( + StatusCode::kUnsupported, + "native SentencePiece support is not enabled in this build"); +} + +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/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; +} From b21b504735cab16a56fe8ffd1451509e1680e722 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:04:14 -0400 Subject: [PATCH 05/61] feat: add Pi0.5 prompt embedding staging --- cpp/CMakeLists.txt | 5 + .../flashrt/cpp/models/pi05/prompt_embed.h | 39 +++++++ cpp/models/pi05/src/prompt_embed.cpp | 101 +++++++++++++++++ cpp/tests/test_pi05_prompt_embed.cpp | 107 ++++++++++++++++++ 4 files changed, 252 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h create mode 100644 cpp/models/pi05/src/prompt_embed.cpp create mode 100644 cpp/tests/test_pi05_prompt_embed.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 53329a92..c87ca158 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -122,6 +122,7 @@ target_link_libraries(flashrt_cpp_vla add_library(flashrt_cpp_pi05 STATIC models/pi05/src/spec.cpp models/pi05/src/prompt_format.cpp + models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp models/pi05/src/runtime.cpp) target_include_directories(flashrt_cpp_pi05 @@ -160,6 +161,10 @@ if(BUILD_TESTING) 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 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..d3aee69c --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h @@ -0,0 +1,39 @@ +#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_cpu( + const 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/src/prompt_embed.cpp b/cpp/models/pi05/src/prompt_embed.cpp new file mode 100644 index 00000000..4d928eae --- /dev/null +++ b/cpp/models/pi05/src/prompt_embed.cpp @@ -0,0 +1,101 @@ +#include "flashrt/cpp/models/pi05/prompt_embed.h" + +#include "flashrt/cpp/models/pi05/prompt_format.h" + +#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"); + } + auto st = modalities::validate_host_tensor(output, "prompt_embedding"); + if (!st.ok_status()) return st; + 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"); + } + return modalities::Status::ok(); +} + +} // namespace + +modalities::Status embed_prompt_cpu( + const 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) { + 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; + if (state) { + const std::string formatted = format_state_prompt(prompt, state, + n_state); + 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) { + std::memset(output.data, 0, static_cast(output.bytes)); + } + 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_cpu( + gather, token_ids->data(), token_ids->size(), embedding_table, prefix); + if (!st.ok_status()) return st; + *prompt_len = static_cast(token_ids->size()); + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/tests/test_pi05_prompt_embed.cpp b/cpp/tests/test_pi05_prompt_embed.cpp new file mode 100644 index 00000000..0363fe12 --- /dev/null +++ b/cpp/tests/test_pi05_prompt_embed.cpp @@ -0,0 +1,107 @@ +#include "flashrt/cpp/models/pi05/prompt_embed.h" + +#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::models::pi05::PromptEmbeddingSpec; +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(); +} + +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; + std::uint64_t prompt_len = 0; + st = embed_prompt_cpu(tokenizer, spec, "pick_up_cube", state, 3, src, dst, + &ids, &prompt_len); + 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); + } +#endif +} + +} // namespace + +int main() { + test_requires_loaded_tokenizer(); + test_paligemma_prompt_embedding_when_configured(); + std::cout << "PASS - Pi05 prompt embedding\n"; + return 0; +} From cca4ef0084ab3483c47fc8c271d3b3c4e19543bb Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:07:02 -0400 Subject: [PATCH 06/61] feat: wire Pi0.5 prompt staging runtime --- .../include/flashrt/cpp/models/pi05/c_api.h | 19 ++++++ .../include/flashrt/cpp/models/pi05/runtime.h | 28 ++++++++ cpp/models/pi05/src/c_api.cpp | 27 +++++++- cpp/models/pi05/src/config_map.h | 68 +++++++++++++++++++ cpp/models/pi05/src/runtime.cpp | 63 +++++++++++++++-- cpp/tests/test_pi05_runtime.cpp | 68 +++++++++++++++++++ 6 files changed, 267 insertions(+), 6 deletions(-) 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..c3c37293 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,23 @@ 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; } frt_pi05_runtime_config; typedef struct frt_pi05_vision_frame { @@ -75,6 +92,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/runtime.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h index 1f817feb..8af992e3 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -3,6 +3,7 @@ #include "flashrt/cpp/families/vla/runtime.h" #include "flashrt/cpp/models/pi05/io.h" +#include "flashrt/cpp/models/pi05/prompt_embed.h" #include @@ -41,6 +42,17 @@ 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; + ReplayFn replay_fn = nullptr; void* replay_user = nullptr; }; @@ -64,6 +76,13 @@ 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_; + } + 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 +92,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); @@ -83,6 +103,14 @@ class Runtime final : public families::vla::Runtime { modalities::Status status_; modalities::VisionStaging staging_; RuntimeIo io_; + modalities::SentencePieceTokenizer prompt_tokenizer_; + PromptEmbeddingSpec prompt_spec_; + modalities::TensorView prompt_embedding_table_; + modalities::TensorView prompt_embedding_output_; + modalities::Status prompt_status_; + std::vector prompt_token_ids_; + 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..4855e5eb 100644 --- a/cpp/models/pi05/src/c_api.cpp +++ b/cpp/models/pi05/src/c_api.cpp @@ -66,7 +66,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; } diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h index 3c085e28..a58b11b3 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/config_map.h @@ -90,6 +90,74 @@ 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 (cfg.prompt_vocab_size && cfg.prompt_hidden_dim) { + cfg.prompt_embedding_table.place = 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 = 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/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index 37b48b2c..c51f5cde 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -1,5 +1,6 @@ #include "flashrt/cpp/models/pi05/runtime.h" +#include #include namespace flashrt { @@ -170,14 +171,29 @@ modalities::Status Runtime::bind() { 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(); + 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; + } + prompt_status_ = embed_prompt_cpu( + prompt_tokenizer_, prompt_spec_, text, state, n_state, + prompt_embedding_table_, prompt_embedding_output_, &prompt_token_ids_, + ¤t_prompt_len_); + return prompt_status_.ok_status() ? 0 : -1; } modalities::Status Runtime::prepare_vision( @@ -203,6 +219,43 @@ 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)); + prompt_staging_enabled_ = true; + return modalities::Status::ok(); +} + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/tests/test_pi05_runtime.cpp b/cpp/tests/test_pi05_runtime.cpp index 607c58dd..bfe2336f 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, @@ -142,10 +144,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; } From c65a59088a28d2a04f9a2d1154d3376dc6d2e444 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:09:00 -0400 Subject: [PATCH 07/61] feat: support Pi0.5 prompt runtime overlay --- .../include/flashrt/cpp/models/pi05/runtime.h | 1 + cpp/models/pi05/src/model_runtime.cpp | 29 ++++++++ cpp/tests/test_pi05_model_runtime.cpp | 69 +++++++++++++++++++ 3 files changed, 99 insertions(+) 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 8af992e3..b1d0dadf 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -81,6 +81,7 @@ class Runtime final : public families::vla::Runtime { const modalities::Status& prompt_status() const { return prompt_status_; } + bool prompt_staging_enabled() const { return prompt_staging_enabled_; } std::uint64_t current_prompt_len() const { return current_prompt_len_; } modalities::Status prepare_vision( diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index b91988fd..43d31f94 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -31,6 +31,7 @@ struct Adapter { uint32_t images_port = kPortImages; uint32_t noise_port = kPortNoise; uint32_t actions_port = kPortActions; + uint32_t prompt_port = kNoPort; int64_t image_shape[4] = {0, 0, 0, 3}; int64_t noise_shape[2] = {0, 0}; @@ -159,6 +160,24 @@ 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; + } + const char* begin = static_cast(data); + const std::string prompt(begin, begin + bytes); + const int rc = a->runtime->set_prompt(prompt.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 -1; + } + a->last_error.clear(); + return 0; + } a->last_error = "unknown or non-input port"; return -1; } @@ -393,6 +412,7 @@ 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 prompt = find_port_index(model, "prompt"); 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, @@ -404,6 +424,11 @@ extern "C" int frt_pi05_model_runtime_create_over( 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; + } auto cfg = flashrt::models::pi05::cface::make_config(config); if (model->n_stages) { @@ -420,10 +445,14 @@ 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; + } a->source_model = model; a->images_port = images; a->noise_port = noise; a->actions_port = actions; + a->prompt_port = prompt; a->view_order = a->runtime->manifest().vision.view_order; frt_model_runtime_verbs verbs{}; diff --git a/cpp/tests/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index e250486e..eaa56c9c 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -290,6 +291,74 @@ 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, nullptr, 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"); + +#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; + + 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); + } else { + std::printf("SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"); + } +#endif + prompt_producer->release(prompt_producer->owner); + frt_graph_destroy(graph); frt_ctx_destroy(ctx); std::printf(g_fail ? "\n== PI05 MODEL RUNTIME FAILED ==\n" From 436e6feb1a727bf0eac71d2bf14d908d987ced12 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:22:37 -0400 Subject: [PATCH 08/61] feat: add device text embedding staging --- cpp/CMakeLists.txt | 4 +- .../include/flashrt/cpp/modalities/text.h | 18 ++ cpp/modalities/src/text_cpu.cpp | 133 +++++++++++ cpp/modalities/src/text_cuda.cu | 208 ++++++++++++++++++ .../flashrt/cpp/models/pi05/prompt_embed.h | 13 ++ .../include/flashrt/cpp/models/pi05/runtime.h | 1 + cpp/models/pi05/src/prompt_embed.cpp | 80 ++++++- cpp/models/pi05/src/runtime.cpp | 13 +- cpp/tests/test_device_staging.cpp | 62 ++++++ cpp/tests/test_pi05_prompt_embed.cpp | 78 +++++++ 10 files changed, 600 insertions(+), 10 deletions(-) create mode 100644 cpp/modalities/src/text_cuda.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index c87ca158..9e8eeccb 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -85,7 +85,9 @@ else() 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}) diff --git a/cpp/modalities/include/flashrt/cpp/modalities/text.h b/cpp/modalities/include/flashrt/cpp/modalities/text.h index 0b1a651f..b4a3c39c 100644 --- a/cpp/modalities/include/flashrt/cpp/modalities/text.h +++ b/cpp/modalities/include/flashrt/cpp/modalities/text.h @@ -14,12 +14,30 @@ struct EmbeddingGatherSpec { 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 diff --git a/cpp/modalities/src/text_cpu.cpp b/cpp/modalities/src/text_cpu.cpp index 4f0ca410..454b9acb 100644 --- a/cpp/modalities/src/text_cpu.cpp +++ b/cpp/modalities/src/text_cpu.cpp @@ -1,9 +1,25 @@ #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) { @@ -55,6 +71,54 @@ Status validate_matrix(const TensorView& tensor, const char* name, } // 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, @@ -95,5 +159,74 @@ Status gather_token_embeddings_cpu(const EmbeddingGatherSpec& spec, 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/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h index d3aee69c..b9d0824c 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h @@ -21,6 +21,19 @@ struct PromptEmbeddingSpec { bool zero_pad_output = true; }; +modalities::Status embed_prompt( + const 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); + modalities::Status embed_prompt_cpu( const modalities::SentencePieceTokenizer& tokenizer, const PromptEmbeddingSpec& spec, 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 b1d0dadf..40373c6b 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -105,6 +105,7 @@ class Runtime final : public families::vla::Runtime { modalities::VisionStaging 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_; diff --git a/cpp/models/pi05/src/prompt_embed.cpp b/cpp/models/pi05/src/prompt_embed.cpp index 4d928eae..d8261e28 100644 --- a/cpp/models/pi05/src/prompt_embed.cpp +++ b/cpp/models/pi05/src/prompt_embed.cpp @@ -2,6 +2,10 @@ #include "flashrt/cpp/models/pi05/prompt_format.h" +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + #include #include @@ -18,8 +22,18 @@ modalities::Status validate_output_capacity( modalities::StatusCode::kInvalidArgument, "invalid prompt embedding dimensions"); } - auto st = modalities::validate_host_tensor(output, "prompt_embedding"); - if (!st.ok_status()) return st; + 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 || @@ -28,12 +42,46 @@ modalities::Status validate_output_capacity( 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_cpu( +modalities::Status embed_prompt( const modalities::SentencePieceTokenizer& tokenizer, const PromptEmbeddingSpec& spec, const std::string& prompt, @@ -42,7 +90,9 @@ modalities::Status embed_prompt_cpu( modalities::TensorView embedding_table, modalities::TensorView output, std::vector* token_ids, - std::uint64_t* prompt_len) { + std::uint64_t* prompt_len, + void* stream, + modalities::TextEmbeddingStaging* staging) { if (!token_ids || !prompt_len) { return modalities::Status::error( modalities::StatusCode::kInvalidArgument, @@ -78,7 +128,8 @@ modalities::Status embed_prompt_cpu( } if (spec.zero_pad_output) { - std::memset(output.data, 0, static_cast(output.bytes)); + st = zero_prompt_output(output, stream); + if (!st.ok_status()) return st; } modalities::TensorView prefix = output; prefix.shape = modalities::Shape{static_cast( @@ -89,13 +140,28 @@ modalities::Status embed_prompt_cpu( modalities::EmbeddingGatherSpec gather{spec.vocab_size, spec.hidden_dim, spec.scale}; - st = modalities::gather_token_embeddings_cpu( - gather, token_ids->data(), token_ids->size(), embedding_table, prefix); + 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( + const 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/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index c51f5cde..c987f0b3 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -71,6 +71,7 @@ Runtime::Runtime(const frt_runtime_export_v1* exp, RuntimeConfig config) } Runtime::~Runtime() { + modalities::text_embedding_staging_destroy(&prompt_embedding_staging_); modalities::vision_staging_destroy(&staging_); release_export(); } @@ -189,10 +190,13 @@ int Runtime::set_prompt_state(const char* text, const float* state, "prompt text is null"); return -1; } - prompt_status_ = embed_prompt_cpu( + prompt_status_ = embed_prompt( prompt_tokenizer_, prompt_spec_, text, state, n_state, prompt_embedding_table_, prompt_embedding_output_, &prompt_token_ids_, - ¤t_prompt_len_); + ¤t_prompt_len_, find_native_stream(exp_, stream_id_), + prompt_embedding_output_.place == modalities::MemoryPlace::kDevice + ? &prompt_embedding_staging_ + : nullptr); return prompt_status_.ok_status() ? 0 : -1; } @@ -252,6 +256,11 @@ modalities::Status Runtime::bind_prompt_staging() { ? config_.prompt_embedding_scale : std::sqrt(static_cast( config_.prompt_hidden_dim)); + 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(); } diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index 0c673ebf..f772451b 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -1,4 +1,5 @@ #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" @@ -16,9 +17,13 @@ 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; @@ -138,6 +143,62 @@ void test_action_d2h_staging() { 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() { @@ -147,6 +208,7 @@ int main() { } test_vision_h2d_staging(); 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_pi05_prompt_embed.cpp b/cpp/tests/test_pi05_prompt_embed.cpp index 0363fe12..eaabeffd 100644 --- a/cpp/tests/test_pi05_prompt_embed.cpp +++ b/cpp/tests/test_pi05_prompt_embed.cpp @@ -1,5 +1,9 @@ #include "flashrt/cpp/models/pi05/prompt_embed.h" +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + #include #include #include @@ -15,7 +19,9 @@ 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 { @@ -25,6 +31,20 @@ std::string tokenizer_model_path() { 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); @@ -97,11 +117,69 @@ void test_paligemma_prompt_embedding_when_configured() { #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; } From d286fcc09d89610a9c244847ab597e691300f2e8 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:26:01 -0400 Subject: [PATCH 09/61] feat: support Pi0.5 state prompt staging --- .../include/flashrt/cpp/models/pi05/c_api.h | 8 +++ .../include/flashrt/cpp/models/pi05/runtime.h | 8 +++ cpp/models/pi05/src/config_map.h | 14 +++++ cpp/models/pi05/src/model_runtime.cpp | 52 ++++++++++++++++++- cpp/models/pi05/src/runtime.cpp | 19 ++++++- cpp/tests/test_pi05_model_runtime.cpp | 49 +++++++++++++++++ 6 files changed, 147 insertions(+), 3 deletions(-) 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 c3c37293..bce5c057 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 @@ -72,6 +72,14 @@ typedef struct frt_pi05_runtime_config { 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; } frt_pi05_runtime_config; typedef struct frt_pi05_vision_frame { 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 40373c6b..90207e3f 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -6,6 +6,7 @@ #include "flashrt/cpp/models/pi05/prompt_embed.h" #include +#include namespace flashrt { namespace models { @@ -52,6 +53,8 @@ struct RuntimeConfig { 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; @@ -82,6 +85,10 @@ class Runtime final : public families::vla::Runtime { 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( @@ -111,6 +118,7 @@ class Runtime final : public families::vla::Runtime { modalities::TensorView prompt_embedding_output_; modalities::Status prompt_status_; std::vector prompt_token_ids_; + std::vector normalized_state_; std::uint64_t current_prompt_len_ = 0; bool prompt_staging_enabled_ = false; frt_graph graph_ = nullptr; diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h index a58b11b3..1b5222e9 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/config_map.h @@ -146,6 +146,20 @@ inline RuntimeConfig make_config(const frt_pi05_runtime_config* in) { 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 (cfg.prompt_vocab_size && cfg.prompt_hidden_dim) { cfg.prompt_embedding_table.place = modalities::MemoryPlace::kHost; cfg.prompt_embedding_table.layout = modalities::Layout::kFlat; diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index 43d31f94..c2d13b96 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -32,6 +32,11 @@ struct Adapter { 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::string prompt_text; + std::vector state_values; int64_t image_shape[4] = {0, 0, 0, 3}; int64_t noise_shape[2] = {0, 0}; @@ -166,8 +171,14 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, return -1; } const char* begin = static_cast(data); - const std::string prompt(begin, begin + bytes); - const int rc = a->runtime->set_prompt(prompt.c_str()); + a->prompt_text.assign(begin, begin + bytes); + 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() @@ -178,6 +189,31 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, 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 auto* values = static_cast(data); + a->state_values.assign(values, values + bytes / sizeof(float)); + 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 -1; + } + a->last_error.clear(); + return 0; + } a->last_error = "unknown or non-input port"; return -1; } @@ -413,6 +449,7 @@ extern "C" int frt_pi05_model_runtime_create_over( const uint32_t noise = find_port_index(model, "noise"); const uint32_t actions = find_port_index(model, "actions"); 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, @@ -429,6 +466,11 @@ extern "C" int frt_pi05_model_runtime_create_over( 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) { @@ -448,11 +490,17 @@ extern "C" int frt_pi05_model_runtime_create_over( 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; frt_model_runtime_verbs verbs{}; diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index c987f0b3..0ff1254e 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -190,8 +190,25 @@ int Runtime::set_prompt_state(const char* text, const float* state, "prompt text is null"); return -1; } + 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; + } + normalized_state_.resize(n_state); + 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_, text, state, n_state, + prompt_tokenizer_, prompt_spec_, text, 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 diff --git a/cpp/tests/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index eaa56c9c..f3dffcf8 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -11,6 +11,7 @@ #include +#include #include #include #include @@ -313,6 +314,28 @@ int main() { 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, nullptr, 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') { @@ -338,6 +361,12 @@ int main() { 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, @@ -353,10 +382,30 @@ int main() { 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"); + 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); From 76cc7527c15b5dfd1c553aadbd3d9ed5feeb99c1 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:28:18 -0400 Subject: [PATCH 10/61] feat: add Pi0.5 native v2 runtime schema --- flash_rt/models/pi05/runtime_export.py | 78 +++++++++++++++++++++++++- 1 file changed, 75 insertions(+), 3 deletions(-) diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index 9ecb55fa..7013a42c 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), @@ -150,6 +164,14 @@ def export_model_runtime(pl, identity=None, extra_regions=None, "staged", shape=(chunk, robot_action_dim), buffer=wrap["diffusion_noise"]), ] + 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 uses_rtc_prefix or uses_rtc_vjp: ports.extend([ _rt.PortSpec("prev_action_chunk", "tensor", tensor_dtype, @@ -193,6 +215,13 @@ def step(): return rc return rc + manifest_extra = {"stage_plan": plan.manifest(), "io": io} + 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,7 +231,7 @@ 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, ) @@ -230,6 +259,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 +356,18 @@ 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() + prompt_bytes = int(getattr(pl, "max_prompt_len", 0) or 0) * 2048 * 2 + prompt_offset = int(getattr(pl, "num_views", 0) or 0) * 256 * 2048 * 2 + encoder_x = wrap["encoder_x"] + if prompt_bytes <= 0 or prompt_offset + prompt_bytes > encoder_x.nbytes(): + raise ValueError("Pi05 native_v2 prompt_context window is invalid") + regions.append(_rt.RegionSpec( + "prompt_context", encoder_x, offset=prompt_offset, + nbytes=prompt_bytes)) return { "wrap": wrap, From 058a83c176f53a479f17012b3043e458671d1db3 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:29:54 -0400 Subject: [PATCH 11/61] fix: tighten Pi0.5 image input contract --- cpp/models/pi05/src/io.cpp | 29 +++++++++++++++++++++++++++ cpp/tests/test_pi05_model_runtime.cpp | 5 +++++ cpp/tests/test_pi05_runtime.cpp | 5 +++++ docs/mindon_pi05_integration.md | 5 +++-- docs/pi05_io_contract.md | 12 +++++------ 5 files changed, 48 insertions(+), 8 deletions(-) diff --git a/cpp/models/pi05/src/io.cpp b/cpp/models/pi05/src/io.cpp index fbe95220..4072f692 100644 --- a/cpp/models/pi05/src/io.cpp +++ b/cpp/models/pi05/src/io.cpp @@ -3,6 +3,31 @@ namespace flashrt { namespace models { namespace pi05 { +namespace { + +modalities::Status validate_pi05_frame_contract( + const modalities::VisionFrame& frame) { + if (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"); + } + 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, @@ -27,6 +52,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); + if (!st.ok_status()) return st; + } return modalities::preprocess_vision(vision_spec_, frames, image_input_, stream_, staging_); } diff --git a/cpp/tests/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index f3dffcf8..e781bfef 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -169,6 +169,11 @@ 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"); std::vector img_host(image_bytes / 2); cudaMemcpy(img_host.data(), frt_buffer_dptr(image), image_bytes, cudaMemcpyDeviceToHost); diff --git a/cpp/tests/test_pi05_runtime.cpp b/cpp/tests/test_pi05_runtime.cpp index bfe2336f..fe594dd7 100644 --- a/cpp/tests/test_pi05_runtime.cpp +++ b/cpp/tests/test_pi05_runtime.cpp @@ -129,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); diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 1680e8d2..80777ed6 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -154,8 +154,9 @@ 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]`. -Use the producer documentation in `docs/pi05_io_contract.md` for accepted -formats and shape rules. +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 diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 03dbaa6f..2ad50e5e 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -88,18 +88,18 @@ 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 pixels; +- 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. -The producer contract should reject unsupported input shape, dtype, layout, or -view count with a shape/status error. It should not silently reinterpret camera -formats in a way that changes model semantics. If a deployment supports more -pixel formats, the supported set must be documented by the producer and tested -against the CPU reference path. +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. ## Noise Input From 7fe02798e93b6d35ee5c1fc8ad3fd648d8c0c6dd Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:31:55 -0400 Subject: [PATCH 12/61] feat: expose Pi0.5 native v2 raw actions --- docs/mindon_pi05_integration.md | 5 +++-- docs/pi05_io_contract.md | 25 +++++++++++++++---------- flash_rt/models/pi05/runtime_export.py | 5 +++++ 3 files changed, 23 insertions(+), 12 deletions(-) diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 80777ed6..8bd3645e 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -162,8 +162,9 @@ in `docs/pi05_io_contract.md` for accepted formats and shape rules. 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 a producer-declared raw `TENSOR/SWAP` output -such as `actions_raw` when exported by a stage plan. +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 diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 2ad50e5e..71ce20d2 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -37,17 +37,21 @@ 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. -## Target Face After Prompt/State Staging +## Native V2 Face -After the C++ text path exists, the Pi0.5 native face may add the following -ports. Adding these ports changes the model-runtime identity and therefore the -fingerprint. Existing capsules from the old face must refuse restore into the -new 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 | payload | semantics | +| port | modality/update | direction | dtype/layout/shape | backing | |---|---|---|---|---| -| `prompt` | `TEXT/STAGED` | input | UTF-8 bytes, no trailing NUL required | task text only | -| `state` | `STATE/STAGED` | input | `F32`, `(state_dim,)` | raw proprioception, normalized/discretized by the producer | +| `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-visible robot action chunk, `FLAT`, `(chunk_length, robot_action_dim)` | `diffusion_noise` | +| `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 @@ -138,8 +142,9 @@ 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, such as the existing RTC -`actions_raw` port, not by changing the meaning of `actions`. +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 diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index 7013a42c..f11073df 100644 --- a/flash_rt/models/pi05/runtime_export.py +++ b/flash_rt/models/pi05/runtime_export.py @@ -172,6 +172,11 @@ def export_model_runtime(pl, identity=None, extra_regions=None, _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, From 5853428b7b2153a9862f22792d637519cbfcf359 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:36:14 -0400 Subject: [PATCH 13/61] feat: add Pi0.5 native open gate --- cpp/CMakeLists.txt | 8 +- cpp/models/pi05/src/native_open.cpp | 294 ++++++++++++++++++++++++++++ cpp/tests/test_pi05_native_open.cpp | 99 ++++++++++ docs/mindon_pi05_integration.md | 7 + docs/pi05_io_contract.md | 6 + 5 files changed, 413 insertions(+), 1 deletion(-) create mode 100644 cpp/models/pi05/src/native_open.cpp create mode 100644 cpp/tests/test_pi05_native_open.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 9e8eeccb..292323fb 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -134,7 +134,8 @@ target_link_libraries(flashrt_cpp_pi05 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_open.cpp) target_link_libraries(flashrt_cpp_pi05_c PUBLIC flashrt_cpp_pi05 flashrt_runtime) target_include_directories(flashrt_cpp_pi05_c @@ -182,5 +183,10 @@ 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) + add_test(NAME pi05_native_open COMMAND test_pi05_native_open) endif() 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..8609c029 --- /dev/null +++ b/cpp/models/pi05/src/native_open.cpp @@ -0,0 +1,294 @@ +#include "flashrt/model_runtime.h" + +#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); +} + +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) { + 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; + } + + int64_t max_prompt_tokens = 0; + int64_t state_dim = 0; + int64_t num_views = 0; + int64_t chunk = 0; + 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)) { + return -1; + } + if (max_prompt_tokens < 200) { + g_last_error = "max_prompt_tokens must be at least 200"; + return -1; + } + if (state_dim <= 0) { + g_last_error = "state_dim must be positive"; + 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) { + g_last_error = "chunk must be positive"; + return -1; + } + + g_last_error = + "Pi0.5 native open validated config; native graph capture is not " + "implemented yet"; + return -3; +} + +} // 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; + return validate_config(config_json); +} + +extern "C" const char* frt_pi05_native_open_last_error() { + return g_last_error.c_str(); +} diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp new file mode 100644 index 00000000..ce2fc276 --- /dev/null +++ b/cpp/tests/test_pi05_native_open.cpp @@ -0,0 +1,99 @@ +#include "flashrt/model_runtime.h" + +#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()); +} + +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() { + 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); + + 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")); + + 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")); + + ::unlink(tokenizer.c_str()); + ::rmdir(root.c_str()); + std::printf("PASS - Pi05 native open scaffold\n"); + return 0; +} diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 8bd3645e..4b6b52d6 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -77,6 +77,13 @@ 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 exports this symbol as a native-v2 configuration +gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt +mode, prompt capacity, and state dimension, then returns unsupported until the +native checkpoint loader and CUDA graph capture path are implemented. Hosts may +use this to wire dynamic loading and error handling, but must keep using Lane A +or B for execution. + ## No-HTTP C++ Host Shape For same-process control loops, prefer Nexus embedded/session APIs over HTTP. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 71ce20d2..95c22861 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -197,6 +197,12 @@ There are three supported integration lanes: `frt_model_runtime_open_v1(config_json, &out)` and produces the same public struct without Python setup. +The Pi0.5 C++ shared object now exports `frt_model_runtime_open_v1` as a +native-v2 configuration gate. The gate requires `io="native_v2"`, +`checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, +`max_prompt_tokens >= 200`, and a positive `state_dim`; valid configuration +returns unsupported until native asset loading and graph capture are complete. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From 6b712cb9b1a005bd1934cbd91ea1b9136dfccd65 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:46:09 -0400 Subject: [PATCH 14/61] feat: validate Pi0.5 native checkpoint metadata --- cpp/models/pi05/src/native_open.cpp | 244 ++++++++++++++++++++++++++++ cpp/tests/test_pi05_native_open.cpp | 27 +++ docs/mindon_pi05_integration.md | 9 +- docs/pi05_io_contract.md | 6 +- 4 files changed, 280 insertions(+), 6 deletions(-) diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index 8609c029..7331f23c 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -4,9 +4,11 @@ #include #include #include +#include #include #include #include +#include namespace { @@ -20,6 +22,13 @@ struct JsonValue { bool boolean = false; }; +struct TensorMeta { + std::string dtype; + std::vector shape; + uint64_t data_begin = 0; + uint64_t data_end = 0; +}; + class JsonParser { public: explicit JsonParser(const char* src) : cur_(src ? src : "") {} @@ -171,6 +180,238 @@ bool regular_file_exists(const std::string& path) { 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; +} + +bool read_safetensors_header(const std::string& path, std::string* header) { + if (!header) return false; + std::ifstream f(path, std::ios::binary); + if (!f) { + g_last_error = "unable to open safetensors file"; + return false; + } + unsigned char len_bytes[8] = {}; + f.read(reinterpret_cast(len_bytes), sizeof(len_bytes)); + if (f.gcount() != static_cast(sizeof(len_bytes))) { + g_last_error = "safetensors file is too small"; + return false; + } + uint64_t header_len = 0; + for (int i = 7; i >= 0; --i) { + header_len = (header_len << 8) | len_bytes[i]; + } + if (header_len == 0 || header_len > (128ull << 20)) { + g_last_error = "safetensors header length is invalid"; + return false; + } + header->assign(static_cast(header_len), '\0'); + f.read(&(*header)[0], static_cast(header_len)); + if (f.gcount() != static_cast(header_len)) { + g_last_error = "safetensors header is truncated"; + return false; + } + return true; +} + +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_string_property(const std::string& object, + const char* name, + std::string* 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::string value; + while (p < object.size() && object[p] != '"') { + if (object[p] == '\\') return false; + value.push_back(object[p++]); + } + if (p >= object.size()) return false; + if (out) *out = value; + return true; +} + +bool parse_u64_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; + } + if (p >= object.size() || + !std::isdigit(static_cast(object[p]))) { + return false; + } + uint64_t value = 0; + while (p < object.size() && + std::isdigit(static_cast(object[p]))) { + const uint64_t digit = static_cast(object[p] - '0'); + if (value > (UINT64_MAX - digit) / 10ull) return false; + value = value * 10ull + digit; + ++p; + } + values.push_back(value); + 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 tensor_meta(const std::string& header, + const std::string& key, + TensorMeta* meta) { + std::string object; + if (!object_for_key(header, key, &object)) return false; + std::string dtype; + std::vector shape; + std::vector offsets; + if (!parse_string_property(object, "dtype", &dtype) || + !parse_u64_array_property(object, "shape", &shape) || + !parse_u64_array_property(object, "data_offsets", &offsets) || + offsets.size() != 2 || offsets[1] < offsets[0]) { + g_last_error = "safetensors tensor metadata is malformed"; + return false; + } + if (meta) { + meta->dtype = std::move(dtype); + meta->shape = std::move(shape); + meta->data_begin = offsets[0]; + meta->data_end = offsets[1]; + } + return true; +} + +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; + } + std::string header; + if (!read_safetensors_header(path, &header)) return false; + + const char* embedding_keys[] = { + "paligemma_with_expert.paligemma.lm_head.weight", + "model.paligemma_with_expert.paligemma.lm_head.weight", + }; + TensorMeta embedding; + bool found = false; + for (const char* key : embedding_keys) { + if (tensor_meta(header, key, &embedding)) { + found = true; + break; + } + } + if (!found) { + if (g_last_error.empty()) { + g_last_error = "model.safetensors is missing Pi0.5 embedding"; + } + return false; + } + if (embedding.dtype != "BF16" && embedding.dtype != "F16" && + embedding.dtype != "F32") { + g_last_error = "Pi0.5 embedding dtype is unsupported"; + return false; + } + if (embedding.shape.size() != 2 || embedding.shape[1] != 2048 || + embedding.shape[0] < 1000) { + g_last_error = "Pi0.5 embedding shape is invalid"; + return false; + } + g_last_error.clear(); + return true; +} + bool string_field(const std::map& obj, const char* key, std::string* out, @@ -243,6 +484,9 @@ int validate_config(const char* config_json) { 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; diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index ce2fc276..3746a62d 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -27,6 +27,24 @@ void write_file(const std::string& path) { assert(f.good()); } +void write_safetensors(const std::string& path) { + const std::string header = + "{\"paligemma_with_expert.paligemma.lm_head.weight\":{" + "\"dtype\":\"BF16\",\"shape\":[257216,2048]," + "\"data_offsets\":[0,1024]}," + "\"__metadata__\":{\"format\":\"pt\"}}"; + 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())); + std::string payload(1024, '\0'); + f.write(payload.data(), static_cast(payload.size())); + assert(f.good()); +} + std::string config(const std::string& ckpt, const std::string& tokenizer, const char* extra = "") { @@ -72,6 +90,14 @@ int main() { 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_safetensors(root + "/model.safetensors"); const std::string good = config(root, tokenizer); out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1(good.c_str(), &out); @@ -93,6 +119,7 @@ int main() { "max_prompt_tokens")); ::unlink(tokenizer.c_str()); + ::unlink((root + "/model.safetensors").c_str()); ::rmdir(root.c_str()); std::printf("PASS - Pi05 native open scaffold\n"); return 0; diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 4b6b52d6..74229463 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -79,10 +79,11 @@ switching between Lane A and Lane C. The current C++ shared object exports this symbol as a native-v2 configuration gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt -mode, prompt capacity, and state dimension, then returns unsupported until the -native checkpoint loader and CUDA graph capture path are implemented. Hosts may -use this to wire dynamic loading and error handling, but must keep using Lane A -or B for execution. +mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in +`model.safetensors`, then returns unsupported until the native checkpoint +materialization and CUDA graph capture path are implemented. Hosts may use this +to wire dynamic loading and error handling, but must keep using Lane A or B for +execution. ## No-HTTP C++ Host Shape diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 95c22861..5d83aeae 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -200,8 +200,10 @@ There are three supported integration lanes: The Pi0.5 C++ shared object now exports `frt_model_runtime_open_v1` as a native-v2 configuration gate. The gate requires `io="native_v2"`, `checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, -`max_prompt_tokens >= 200`, and a positive `state_dim`; valid configuration -returns unsupported until native asset loading and graph capture are complete. +`max_prompt_tokens >= 200`, and a positive `state_dim`. It also parses +`checkpoint_path/model.safetensors` enough to verify the Pi0.5 prompt +embedding table metadata. Valid configuration returns unsupported until native +asset materialization and graph capture are complete. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 4fa3b3a21fd548963ca77177c30d5b141305d069 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:49:49 -0400 Subject: [PATCH 15/61] feat: validate Pi0.5 native norm assets --- cpp/models/pi05/src/native_open.cpp | 214 ++++++++++++++++++++++++++++ cpp/tests/test_pi05_native_open.cpp | 76 +++++++++- docs/mindon_pi05_integration.md | 4 +- docs/pi05_io_contract.md | 6 +- 4 files changed, 290 insertions(+), 10 deletions(-) diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index 7331f23c..d430983b 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -3,8 +3,11 @@ #include #include #include +#include #include +#include #include +#include #include #include #include @@ -347,6 +350,46 @@ bool parse_u64_array_property(const std::string& object, 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 tensor_meta(const std::string& header, const std::string& key, TensorMeta* meta) { @@ -371,6 +414,174 @@ bool tensor_meta(const std::string& header, 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_dims(const std::string& json, + const char* block_name, + size_t* dims) { + 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) || + q01.empty() || q01.size() != q99.size()) { + return false; + } + if (dims) *dims = q01.size(); + return true; +} + +bool validate_norm_stats_file(const std::string& path, + int64_t state_dim) { + std::string json; + if (!read_text_file(path, &json)) return false; + size_t action_dims = 0; + size_t state_dims = 0; + if (!norm_block_dims(json, "actions", &action_dims) || + !norm_block_dims(json, "state", &state_dims)) { + g_last_error = "norm_stats.json is missing actions/state q01/q99"; + return false; + } + if (action_dims == 0 || action_dims > 32) { + g_last_error = "norm_stats action dimension is invalid"; + return false; + } + if (state_dims != static_cast(state_dim)) { + g_last_error = "norm_stats state dimension does not match config"; + return false; + } + 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 tensor_1d_dim(const std::string& header, + const char* key, + size_t* dims) { + TensorMeta meta; + if (!tensor_meta(header, key, &meta)) return false; + if (meta.dtype != "F32" || meta.shape.size() != 1 || + meta.shape[0] == 0) { + return false; + } + if (dims) *dims = static_cast(meta.shape[0]); + return true; +} + +bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, + int64_t state_dim) { + 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::string pre_header; + std::string post_header; + if (!read_safetensors_header(pre, &pre_header) || + !read_safetensors_header(post, &post_header)) { + return false; + } + size_t state_q01 = 0; + size_t state_q99 = 0; + size_t action_q01 = 0; + size_t action_q99 = 0; + if (!tensor_1d_dim(pre_header, "observation.state.q01", &state_q01) || + !tensor_1d_dim(pre_header, "observation.state.q99", &state_q99) || + !tensor_1d_dim(post_header, "action.q01", &action_q01) || + !tensor_1d_dim(post_header, "action.q99", &action_q99)) { + g_last_error = + "lerobot policy stats are missing action/state q01/q99"; + return false; + } + if (state_q01 != state_q99 || + state_q01 != static_cast(state_dim)) { + g_last_error = + "lerobot policy state dimension does not match config"; + return false; + } + if (action_q01 != action_q99 || action_q01 > 32) { + g_last_error = "lerobot policy action dimension is invalid"; + return false; + } + g_last_error.clear(); + return true; +} + +bool validate_norm_stats(const std::string& checkpoint_path, + int64_t state_dim) { + 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)) return true; + saw_malformed = true; + malformed_error = g_last_error; + } + if (validate_lerobot_policy_norm_stats(checkpoint_path, state_dim)) { + 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)) { @@ -514,6 +725,9 @@ int validate_config(const char* config_json) { g_last_error = "chunk must be positive"; return -1; } + if (!validate_norm_stats(checkpoint_path, state_dim)) { + return -2; + } g_last_error = "Pi0.5 native open validated config; native graph capture is not " diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 3746a62d..84201680 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -27,12 +27,9 @@ void write_file(const std::string& path) { assert(f.good()); } -void write_safetensors(const std::string& path) { - const std::string header = - "{\"paligemma_with_expert.paligemma.lm_head.weight\":{" - "\"dtype\":\"BF16\",\"shape\":[257216,2048]," - "\"data_offsets\":[0,1024]}," - "\"__metadata__\":{\"format\":\"pt\"}}"; +void write_raw_safetensors(const std::string& path, + const std::string& header, + uint64_t payload_bytes) { std::ofstream f(path, std::ios::binary); uint64_t n = header.size(); for (int i = 0; i < 8; ++i) { @@ -40,11 +37,50 @@ void write_safetensors(const std::string& path) { f.write(&b, 1); } f.write(header.data(), static_cast(header.size())); - std::string payload(1024, '\0'); + std::string payload(static_cast(payload_bytes), '\0'); f.write(payload.data(), static_cast(payload.size())); assert(f.good()); } +void write_safetensors(const std::string& path) { + write_raw_safetensors( + path, + "{\"paligemma_with_expert.paligemma.lm_head.weight\":{" + "\"dtype\":\"BF16\",\"shape\":[257216,2048]," + "\"data_offsets\":[0,1024]}," + "\"__metadata__\":{\"format\":\"pt\"}}", + 1024); +} + +void write_lerobot_policy_stats(const std::string& root) { + 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]}}", + 64); + 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]}}", + 56); +} + +void write_norm_stats(const std::string& path) { + std::ofstream f(path); + f << "{" + << "\"norm_stats\":{" + << "\"state\":{\"q01\":[0,0,0,0,0,0,0,0]," + << "\"q99\":[1,1,1,1,1,1,1,1]}," + << "\"actions\":{\"q01\":[0,0,0,0,0,0,0]," + << "\"q99\":[1,1,1,1,1,1,1]}" + << "}}"; + assert(f.good()); +} + std::string config(const std::string& ckpt, const std::string& tokenizer, const char* extra = "") { @@ -98,6 +134,13 @@ int main() { "model.safetensors")); 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"); const std::string good = config(root, tokenizer); out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1(good.c_str(), &out); @@ -118,8 +161,27 @@ int main() { 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); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(lerobot_root, tokenizer).c_str(), &out); + assert(rc == -3); + assert(out == nullptr); + + ::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/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 74229463..9d02f289 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -80,7 +80,9 @@ switching between Lane A and Lane C. The current C++ shared object exports this symbol as a native-v2 configuration gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in -`model.safetensors`, then returns unsupported until the native checkpoint +`model.safetensors`. It also verifies action/state q01/q99 metadata from +openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer +safetensors, then returns unsupported until the native checkpoint materialization and CUDA graph capture path are implemented. Hosts may use this to wire dynamic loading and error handling, but must keep using Lane A or B for execution. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 5d83aeae..b2e1058f 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -202,8 +202,10 @@ native-v2 configuration gate. The gate requires `io="native_v2"`, `checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, `max_prompt_tokens >= 200`, and a positive `state_dim`. It also parses `checkpoint_path/model.safetensors` enough to verify the Pi0.5 prompt -embedding table metadata. Valid configuration returns unsupported until native -asset materialization and graph capture are complete. +embedding table metadata, and verifies action/state q01/q99 dimensions from +either openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer +safetensors. Valid configuration returns unsupported until native asset +materialization and graph capture are complete. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From d965a52ad73d7b252a234ff613e2aa5be70ce5ef Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:56:44 -0400 Subject: [PATCH 16/61] feat: tighten Pi0.5 native asset sanity checks --- cpp/models/pi05/src/native_open.cpp | 175 +++++++++++++++++++++++----- cpp/tests/test_pi05_native_open.cpp | 34 ++++-- docs/mindon_pi05_integration.md | 8 +- docs/pi05_io_contract.md | 6 +- 4 files changed, 183 insertions(+), 40 deletions(-) diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index d430983b..8c22a5f0 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -183,12 +184,22 @@ bool regular_file_exists(const std::string& path) { S_ISREG(st.st_mode); } +uint64_t file_size(const std::string& path) { + struct stat st {}; + if (path.empty() || ::stat(path.c_str(), &st) != 0 || st.st_size < 0) { + return 0; + } + return static_cast(st.st_size); +} + std::string join_path(const std::string& dir, const char* leaf) { if (dir.empty() || dir.back() == '/') return dir + leaf; return dir + "/" + leaf; } -bool read_safetensors_header(const std::string& path, std::string* header) { +bool read_safetensors_header(const std::string& path, std::string* header, + uint64_t* data_start = nullptr, + uint64_t* total_bytes = nullptr) { if (!header) return false; std::ifstream f(path, std::ios::binary); if (!f) { @@ -209,12 +220,20 @@ bool read_safetensors_header(const std::string& path, std::string* header) { g_last_error = "safetensors header length is invalid"; return false; } + const uint64_t start = 8ull + header_len; + const uint64_t size = file_size(path); + if (size < start) { + g_last_error = "safetensors header exceeds file size"; + return false; + } header->assign(static_cast(header_len), '\0'); f.read(&(*header)[0], static_cast(header_len)); if (f.gcount() != static_cast(header_len)) { g_last_error = "safetensors header is truncated"; return false; } + if (data_start) *data_start = start; + if (total_bytes) *total_bytes = size; return true; } @@ -414,6 +433,111 @@ bool tensor_meta(const std::string& header, return true; } +uint64_t dtype_bytes(const std::string& dtype) { + if (dtype == "F32" || dtype == "I32" || dtype == "U32") return 4; + if (dtype == "BF16" || dtype == "F16" || dtype == "I16" || + dtype == "U16") { + return 2; + } + if (dtype == "I64" || dtype == "U64" || dtype == "F64") return 8; + if (dtype == "I8" || dtype == "U8" || dtype == "BOOL") return 1; + return 0; +} + +bool tensor_nbytes(const TensorMeta& meta, uint64_t* out) { + const uint64_t elem = dtype_bytes(meta.dtype); + if (!elem) return false; + uint64_t n = elem; + for (uint64_t dim : meta.shape) { + if (dim == 0 || n > UINT64_MAX / dim) return false; + n *= dim; + } + if (out) *out = n; + return true; +} + +bool tensor_payload_valid(const TensorMeta& meta, + uint64_t data_start, + uint64_t total_bytes) { + uint64_t expected = 0; + if (!tensor_nbytes(meta, &expected)) { + g_last_error = "safetensors tensor dtype/shape is unsupported"; + return false; + } + if (meta.data_end < meta.data_begin || + meta.data_end - meta.data_begin != expected || + data_start > total_bytes || + meta.data_end > total_bytes - data_start) { + g_last_error = "safetensors tensor byte range is invalid"; + return false; + } + return true; +} + +bool read_safetensors_f32_vector(const std::string& path, + const char* key, + std::vector* out) { + if (!out) return false; + std::string header; + uint64_t data_start = 0; + uint64_t total_bytes = 0; + if (!read_safetensors_header(path, &header, &data_start, &total_bytes)) { + return false; + } + TensorMeta meta; + if (!tensor_meta(header, key, &meta)) return false; + if (meta.dtype != "F32" || meta.shape.size() != 1 || + !tensor_payload_valid(meta, data_start, total_bytes)) { + g_last_error = "safetensors F32 vector metadata is invalid"; + return false; + } + const uint64_t n = meta.shape[0]; + if (n > (1ull << 20)) { + g_last_error = "safetensors vector is too large"; + return false; + } + std::ifstream f(path, std::ios::binary); + if (!f) { + g_last_error = "unable to open safetensors file"; + return false; + } + f.seekg(static_cast(data_start + meta.data_begin), + std::ios::beg); + std::vector values(static_cast(n)); + f.read(reinterpret_cast(values.data()), + static_cast(n * sizeof(float))); + if (f.gcount() != static_cast(n * sizeof(float))) { + g_last_error = "safetensors vector payload is truncated"; + return false; + } + *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); @@ -439,7 +563,7 @@ bool norm_block_dims(const std::string& json, std::vector q99; if (!parse_f64_array_property(block, "q01", &q01) || !parse_f64_array_property(block, "q99", &q99) || - q01.empty() || q01.size() != q99.size()) { + !sane_quantile_pair(q01, q99)) { return false; } if (dims) *dims = q01.size(); @@ -496,19 +620,6 @@ std::string find_child(const std::string& dir, return found; } -bool tensor_1d_dim(const std::string& header, - const char* key, - size_t* dims) { - TensorMeta meta; - if (!tensor_meta(header, key, &meta)) return false; - if (meta.dtype != "F32" || meta.shape.size() != 1 || - meta.shape[0] == 0) { - return false; - } - if (dims) *dims = static_cast(meta.shape[0]); - return true; -} - bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, int64_t state_dim) { const std::string pre = find_child( @@ -525,25 +636,28 @@ bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, !read_safetensors_header(post, &post_header)) { return false; } - size_t state_q01 = 0; - size_t state_q99 = 0; - size_t action_q01 = 0; - size_t action_q99 = 0; - if (!tensor_1d_dim(pre_header, "observation.state.q01", &state_q01) || - !tensor_1d_dim(pre_header, "observation.state.q99", &state_q99) || - !tensor_1d_dim(post_header, "action.q01", &action_q01) || - !tensor_1d_dim(post_header, "action.q99", &action_q99)) { + 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 != state_q99 || - state_q01 != static_cast(state_dim)) { + 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 != action_q99 || action_q01 > 32) { + if (action_q01.size() > 32 || + !sane_quantile_pair(action_q01, action_q99)) { g_last_error = "lerobot policy action dimension is invalid"; return false; } @@ -589,7 +703,11 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { return false; } std::string header; - if (!read_safetensors_header(path, &header)) return false; + uint64_t data_start = 0; + uint64_t total_bytes = 0; + if (!read_safetensors_header(path, &header, &data_start, &total_bytes)) { + return false; + } const char* embedding_keys[] = { "paligemma_with_expert.paligemma.lm_head.weight", @@ -619,6 +737,9 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { g_last_error = "Pi0.5 embedding shape is invalid"; return false; } + if (!tensor_payload_valid(embedding, data_start, total_bytes)) { + return false; + } g_last_error.clear(); return true; } diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 84201680..dd888755 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -29,7 +29,7 @@ void write_file(const std::string& path) { void write_raw_safetensors(const std::string& path, const std::string& header, - uint64_t payload_bytes) { + const std::string& payload) { std::ofstream f(path, std::ios::binary); uint64_t n = header.size(); for (int i = 0; i < 8; ++i) { @@ -37,36 +37,56 @@ void write_raw_safetensors(const std::string& path, f.write(&b, 1); } f.write(header.data(), static_cast(header.size())); - std::string payload(static_cast(payload_bytes), '\0'); 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) { + write_raw_safetensors(path, header, + std::string(static_cast(payload_bytes), + '\0')); +} + +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 uint64_t bytes = 1001ull * 2048ull * 2ull; write_raw_safetensors( path, "{\"paligemma_with_expert.paligemma.lm_head.weight\":{" - "\"dtype\":\"BF16\",\"shape\":[257216,2048]," - "\"data_offsets\":[0,1024]}," + "\"dtype\":\"BF16\",\"shape\":[1001,2048]," + "\"data_offsets\":[0," + std::to_string(bytes) + "]}," "\"__metadata__\":{\"format\":\"pt\"}}", - 1024); + bytes); } void write_lerobot_policy_stats(const std::string& root) { + 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, 1.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]}}", - 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, 1.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]}}", - 56); + action_payload); } void write_norm_stats(const std::string& path) { diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 9d02f289..58734c8d 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -82,10 +82,10 @@ gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in `model.safetensors`. It also verifies action/state q01/q99 metadata from openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer -safetensors, then returns unsupported until the native checkpoint -materialization and CUDA graph capture path are implemented. Hosts may use this -to wire dynamic loading and error handling, but must keep using Lane A or B for -execution. +safetensors, including tensor byte ranges and finite ordered quantile payloads, +then returns unsupported until the native checkpoint materialization and CUDA +graph capture path are implemented. Hosts may use this to wire dynamic loading +and error handling, but must keep using Lane A or B for execution. ## No-HTTP C++ Host Shape diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index b2e1058f..254ee452 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -204,8 +204,10 @@ native-v2 configuration gate. The gate requires `io="native_v2"`, `checkpoint_path/model.safetensors` enough to verify the Pi0.5 prompt embedding table metadata, and verifies action/state q01/q99 dimensions from either openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer -safetensors. Valid configuration returns unsupported until native asset -materialization and graph capture are complete. +safetensors. Safetensors tensor byte ranges must match dtype/shape, and +normalization quantiles must be finite ordered pairs. Valid configuration +returns unsupported until native asset materialization and graph capture are +complete. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 295c100f8b32b27b76da5a0bca9fb47868b7a974 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 14:58:16 -0400 Subject: [PATCH 17/61] test: cover Pi0.5 native asset rejection --- cpp/tests/test_pi05_native_open.cpp | 45 +++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 6 deletions(-) diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index dd888755..71d1061d 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -66,10 +66,20 @@ void write_safetensors(const std::string& path) { bytes); } -void write_lerobot_policy_stats(const std::string& root) { +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, 1.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]," @@ -79,7 +89,7 @@ void write_lerobot_policy_stats(const std::string& root) { 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, 1.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]," @@ -89,14 +99,16 @@ void write_lerobot_policy_stats(const std::string& root) { action_payload); } -void write_norm_stats(const std::string& path) { +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]," - << "\"q99\":[1,1,1,1,1,1,1,1]}," + << (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]," - << "\"q99\":[1,1,1,1,1,1,1]}" + << (valid ? "\"q99\":[1,1,1,1,1,1,1]}" + : "\"q99\":[0,0,0,0,0,0,0]}") << "}}"; assert(f.good()); } @@ -153,6 +165,13 @@ int main() { 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")); + write_safetensors(root + "/model.safetensors"); out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); @@ -160,6 +179,13 @@ int main() { 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"); const std::string good = config(root, tokenizer); out = reinterpret_cast(0x1); @@ -183,6 +209,13 @@ int main() { 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); out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1( From fd992d45d29e1f0a110cf8a3bcfa2ecdc4f87430 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 15:06:04 -0400 Subject: [PATCH 18/61] feat: validate Pi0.5 native core weights --- cpp/models/pi05/src/native_open.cpp | 109 +++++++++++++++++++++------- cpp/tests/test_pi05_native_open.cpp | 67 +++++++++++++++-- docs/mindon_pi05_integration.md | 13 ++-- docs/pi05_io_contract.md | 14 ++-- 4 files changed, 156 insertions(+), 47 deletions(-) diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index 8c22a5f0..b3611410 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -33,6 +33,13 @@ struct TensorMeta { uint64_t data_end = 0; }; +struct TensorRequirement { + const char* key; + const char* dtype; + uint64_t rank; + uint64_t dims[4]; +}; + class JsonParser { public: explicit JsonParser(const char* src) : cur_(src ? src : "") {} @@ -709,36 +716,84 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { return false; } - const char* embedding_keys[] = { - "paligemma_with_expert.paligemma.lm_head.weight", - "model.paligemma_with_expert.paligemma.lm_head.weight", + const TensorRequirement requirements[] = { + {"paligemma_with_expert.paligemma.lm_head.weight", + nullptr, 2, {0, 2048, 0, 0}}, + {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".embeddings.patch_embedding.weight", + nullptr, 4, {1152, 3, 14, 14}}, + {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".embeddings.patch_embedding.bias", + nullptr, 1, {1152, 0, 0, 0}}, + {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".embeddings.position_embedding.weight", + nullptr, 2, {256, 1152, 0, 0}}, + {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" + ".weight", + nullptr, 2, {2048, 1152, 0, 0}}, + {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" + ".bias", + nullptr, 1, {2048, 0, 0, 0}}, + {"action_in_proj.weight", nullptr, 2, {1024, 32, 0, 0}}, + {"action_in_proj.bias", nullptr, 1, {1024, 0, 0, 0}}, + {"action_out_proj.weight", nullptr, 2, {32, 1024, 0, 0}}, + {"action_out_proj.bias", nullptr, 1, {32, 0, 0, 0}}, + {"time_mlp_in.weight", nullptr, 2, {1024, 1024, 0, 0}}, + {"time_mlp_in.bias", nullptr, 1, {1024, 0, 0, 0}}, + {"time_mlp_out.weight", nullptr, 2, {1024, 1024, 0, 0}}, + {"time_mlp_out.bias", nullptr, 1, {1024, 0, 0, 0}}, + {"paligemma_with_expert.paligemma.model.language_model.layers.0" + ".self_attn.q_proj.weight", + nullptr, 2, {2048, 2048, 0, 0}}, + {"paligemma_with_expert.gemma_expert.model.layers.0.self_attn" + ".q_proj.weight", + nullptr, 2, {2048, 1024, 0, 0}}, }; - TensorMeta embedding; - bool found = false; - for (const char* key : embedding_keys) { - if (tensor_meta(header, key, &embedding)) { - found = true; - break; + + for (const auto& req : requirements) { + TensorMeta meta; + std::string key = req.key; + if (!tensor_meta(header, key, &meta)) { + key = std::string("model.") + req.key; + if (!tensor_meta(header, key, &meta)) { + g_last_error = std::string("model.safetensors is missing ") + + req.key; + return false; + } } - } - if (!found) { - if (g_last_error.empty()) { - g_last_error = "model.safetensors is missing Pi0.5 embedding"; + if (req.dtype && meta.dtype != req.dtype) { + g_last_error = std::string("Pi0.5 tensor dtype mismatch: ") + + 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.size() != req.rank) { + g_last_error = std::string("Pi0.5 tensor rank mismatch: ") + + req.key; + return false; + } + for (uint64_t i = 0; i < req.rank; ++i) { + if (req.dims[i] && meta.shape[static_cast(i)] != + req.dims[i]) { + g_last_error = std::string("Pi0.5 tensor shape mismatch: ") + + req.key; + return false; + } + } + if (std::string(req.key) == + "paligemma_with_expert.paligemma.lm_head.weight" && + meta.shape[0] < 1000) { + g_last_error = "Pi0.5 embedding vocab size is invalid"; + return false; + } + if (!tensor_payload_valid(meta, data_start, total_bytes)) { + return false; } - return false; - } - if (embedding.dtype != "BF16" && embedding.dtype != "F16" && - embedding.dtype != "F32") { - g_last_error = "Pi0.5 embedding dtype is unsupported"; - return false; - } - if (embedding.shape.size() != 2 || embedding.shape[1] != 2048 || - embedding.shape[0] < 1000) { - g_last_error = "Pi0.5 embedding shape is invalid"; - return false; - } - if (!tensor_payload_valid(embedding, data_start, total_bytes)) { - return false; } g_last_error.clear(); return true; diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 71d1061d..8dbdb1bb 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -56,14 +56,65 @@ void append_f32(std::string* out, float value) { } void write_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) + "]}," - "\"__metadata__\":{\"format\":\"pt\"}}", - bytes); + struct Entry { + const char* key; + const char* dtype; + const char* shape; + uint64_t bytes; + }; + const Entry entries[] = { + {"paligemma_with_expert.paligemma.lm_head.weight", "BF16", + "[1001,2048]", 1001ull * 2048ull * 2ull}, + {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".embeddings.patch_embedding.weight", + "F32", "[1152,3,14,14]", 1152ull * 3ull * 14ull * 14ull * 4ull}, + {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".embeddings.patch_embedding.bias", + "F32", "[1152]", 1152ull * 4ull}, + {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".embeddings.position_embedding.weight", + "F32", "[256,1152]", 256ull * 1152ull * 4ull}, + {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" + ".weight", + "F32", "[2048,1152]", 2048ull * 1152ull * 4ull}, + {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" + ".bias", + "F32", "[2048]", 2048ull * 4ull}, + {"action_in_proj.weight", "F32", "[1024,32]", 1024ull * 32ull * 4ull}, + {"action_in_proj.bias", "F32", "[1024]", 1024ull * 4ull}, + {"action_out_proj.weight", "F32", "[32,1024]", 32ull * 1024ull * 4ull}, + {"action_out_proj.bias", "F32", "[32]", 32ull * 4ull}, + {"time_mlp_in.weight", "F32", "[1024,1024]", 1024ull * 1024ull * 4ull}, + {"time_mlp_in.bias", "F32", "[1024]", 1024ull * 4ull}, + {"time_mlp_out.weight", "F32", "[1024,1024]", 1024ull * 1024ull * 4ull}, + {"time_mlp_out.bias", "F32", "[1024]", 1024ull * 4ull}, + {"paligemma_with_expert.paligemma.model.language_model.layers.0" + ".self_attn.q_proj.weight", + "F32", "[2048,2048]", 2048ull * 2048ull * 4ull}, + {"paligemma_with_expert.gemma_expert.model.layers.0.self_attn" + ".q_proj.weight", + "F32", "[2048,1024]", 2048ull * 1024ull * 4ull}, + }; + std::string header = "{"; + uint64_t offset = 0; + for (size_t i = 0; i < sizeof(entries) / sizeof(entries[0]); ++i) { + const Entry& e = entries[i]; + if (i) header += ","; + header += "\""; + header += e.key; + header += "\":{\"dtype\":\""; + header += e.dtype; + header += "\",\"shape\":"; + header += e.shape; + header += ",\"data_offsets\":["; + header += std::to_string(offset); + header += ","; + offset += e.bytes; + header += std::to_string(offset); + header += "]}"; + } + header += ",\"__metadata__\":{\"format\":\"pt\"}}"; + write_raw_safetensors(path, header, offset); } void write_bad_safetensors(const std::string& path) { diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 58734c8d..16edf7e8 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -80,12 +80,13 @@ switching between Lane A and Lane C. The current C++ shared object exports this symbol as a native-v2 configuration gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in -`model.safetensors`. It also verifies action/state q01/q99 metadata from -openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer -safetensors, including tensor byte ranges and finite ordered quantile payloads, -then returns unsupported until the native checkpoint materialization and CUDA -graph capture path are implemented. Hosts may use this to wire dynamic loading -and error handling, but must keep using Lane A or B for execution. +`model.safetensors`, plus the native builder's minimum required weight set. +It also verifies action/state q01/q99 metadata from openpi `norm_stats.json` or +LeRobot policy normalizer/unnormalizer safetensors, including tensor byte +ranges and finite ordered quantile payloads, then returns unsupported until the +native checkpoint materialization and CUDA graph capture path are implemented. +Hosts may use this to wire dynamic loading and error handling, but must keep +using Lane A or B for execution. ## No-HTTP C++ Host Shape diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 254ee452..14a5a053 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -202,12 +202,14 @@ native-v2 configuration gate. The gate requires `io="native_v2"`, `checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, `max_prompt_tokens >= 200`, and a positive `state_dim`. It also parses `checkpoint_path/model.safetensors` enough to verify the Pi0.5 prompt -embedding table metadata, and 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. Valid configuration -returns unsupported until native asset materialization and graph capture are -complete. +embedding table and the native builder's minimum required weight set +(vision patch/position/projector, action projections, time MLP, and first +encoder/decoder layer sentinels). 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. Valid +configuration returns unsupported until native asset materialization and graph +capture are complete. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 319e7efa4397c7bb1cfdcfd5fc38c82b6b5c87d9 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 15:55:28 -0400 Subject: [PATCH 19/61] feat: add native safetensors mmap loader --- cpp/CMakeLists.txt | 11 +- .../include/flashrt/cpp/loader/safetensors.h | 56 +++ cpp/loader/src/safetensors.cpp | 456 ++++++++++++++++++ cpp/models/pi05/src/native_open.cpp | 256 +--------- cpp/tests/test_safetensors_loader.cpp | 115 +++++ docs/mindon_pi05_integration.md | 4 +- docs/pi05_io_contract.md | 10 +- 7 files changed, 668 insertions(+), 240 deletions(-) create mode 100644 cpp/loader/include/flashrt/cpp/loader/safetensors.h create mode 100644 cpp/loader/src/safetensors.cpp create mode 100644 cpp/tests/test_safetensors_loader.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 292323fb..07b0b787 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -121,6 +121,11 @@ target_include_directories(flashrt_cpp_vla target_link_libraries(flashrt_cpp_vla INTERFACE flashrt_cpp_modalities flashrt_cpp) +add_library(flashrt_cpp_loader STATIC + loader/src/safetensors.cpp) +target_include_directories(flashrt_cpp_loader + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/loader/include) + add_library(flashrt_cpp_pi05 STATIC models/pi05/src/spec.cpp models/pi05/src/prompt_format.cpp @@ -137,11 +142,15 @@ add_library(flashrt_cpp_pi05_c SHARED models/pi05/src/model_runtime.cpp models/pi05/src/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c - PUBLIC flashrt_cpp_pi05 flashrt_runtime) + PUBLIC flashrt_cpp_pi05 flashrt_cpp_loader flashrt_runtime) target_include_directories(flashrt_cpp_pi05_c PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include) 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_cpp_modalities tests/test_modalities.cpp) target_link_libraries(test_cpp_modalities PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) 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/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/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index b3611410..b49cd71b 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -1,4 +1,5 @@ #include "flashrt/model_runtime.h" +#include "flashrt/cpp/loader/safetensors.h" #include #include @@ -26,13 +27,6 @@ struct JsonValue { bool boolean = false; }; -struct TensorMeta { - std::string dtype; - std::vector shape; - uint64_t data_begin = 0; - uint64_t data_end = 0; -}; - struct TensorRequirement { const char* key; const char* dtype; @@ -191,59 +185,11 @@ bool regular_file_exists(const std::string& path) { S_ISREG(st.st_mode); } -uint64_t file_size(const std::string& path) { - struct stat st {}; - if (path.empty() || ::stat(path.c_str(), &st) != 0 || st.st_size < 0) { - return 0; - } - return static_cast(st.st_size); -} - std::string join_path(const std::string& dir, const char* leaf) { if (dir.empty() || dir.back() == '/') return dir + leaf; return dir + "/" + leaf; } -bool read_safetensors_header(const std::string& path, std::string* header, - uint64_t* data_start = nullptr, - uint64_t* total_bytes = nullptr) { - if (!header) return false; - std::ifstream f(path, std::ios::binary); - if (!f) { - g_last_error = "unable to open safetensors file"; - return false; - } - unsigned char len_bytes[8] = {}; - f.read(reinterpret_cast(len_bytes), sizeof(len_bytes)); - if (f.gcount() != static_cast(sizeof(len_bytes))) { - g_last_error = "safetensors file is too small"; - return false; - } - uint64_t header_len = 0; - for (int i = 7; i >= 0; --i) { - header_len = (header_len << 8) | len_bytes[i]; - } - if (header_len == 0 || header_len > (128ull << 20)) { - g_last_error = "safetensors header length is invalid"; - return false; - } - const uint64_t start = 8ull + header_len; - const uint64_t size = file_size(path); - if (size < start) { - g_last_error = "safetensors header exceeds file size"; - return false; - } - header->assign(static_cast(header_len), '\0'); - f.read(&(*header)[0], static_cast(header_len)); - if (f.gcount() != static_cast(header_len)) { - g_last_error = "safetensors header is truncated"; - return false; - } - if (data_start) *data_start = start; - if (total_bytes) *total_bytes = size; - return true; -} - std::string quoted_key(const std::string& key) { std::string out = "\""; for (char c : key) { @@ -306,76 +252,6 @@ bool object_for_key(const std::string& json, return false; } -bool parse_string_property(const std::string& object, - const char* name, - std::string* 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::string value; - while (p < object.size() && object[p] != '"') { - if (object[p] == '\\') return false; - value.push_back(object[p++]); - } - if (p >= object.size()) return false; - if (out) *out = value; - return true; -} - -bool parse_u64_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; - } - if (p >= object.size() || - !std::isdigit(static_cast(object[p]))) { - return false; - } - uint64_t value = 0; - while (p < object.size() && - std::isdigit(static_cast(object[p]))) { - const uint64_t digit = static_cast(object[p] - '0'); - if (value > (UINT64_MAX - digit) / 10ull) return false; - value = value * 10ull + digit; - ++p; - } - values.push_back(value); - 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 parse_f64_array_property(const std::string& object, const char* name, std::vector* out) { @@ -416,107 +292,28 @@ bool parse_f64_array_property(const std::string& object, return false; } -bool tensor_meta(const std::string& header, - const std::string& key, - TensorMeta* meta) { - std::string object; - if (!object_for_key(header, key, &object)) return false; - std::string dtype; - std::vector shape; - std::vector offsets; - if (!parse_string_property(object, "dtype", &dtype) || - !parse_u64_array_property(object, "shape", &shape) || - !parse_u64_array_property(object, "data_offsets", &offsets) || - offsets.size() != 2 || offsets[1] < offsets[0]) { - g_last_error = "safetensors tensor metadata is malformed"; - return false; - } - if (meta) { - meta->dtype = std::move(dtype); - meta->shape = std::move(shape); - meta->data_begin = offsets[0]; - meta->data_end = offsets[1]; - } - return true; -} - -uint64_t dtype_bytes(const std::string& dtype) { - if (dtype == "F32" || dtype == "I32" || dtype == "U32") return 4; - if (dtype == "BF16" || dtype == "F16" || dtype == "I16" || - dtype == "U16") { - return 2; - } - if (dtype == "I64" || dtype == "U64" || dtype == "F64") return 8; - if (dtype == "I8" || dtype == "U8" || dtype == "BOOL") return 1; - return 0; -} - -bool tensor_nbytes(const TensorMeta& meta, uint64_t* out) { - const uint64_t elem = dtype_bytes(meta.dtype); - if (!elem) return false; - uint64_t n = elem; - for (uint64_t dim : meta.shape) { - if (dim == 0 || n > UINT64_MAX / dim) return false; - n *= dim; - } - if (out) *out = n; - return true; -} - -bool tensor_payload_valid(const TensorMeta& meta, - uint64_t data_start, - uint64_t total_bytes) { - uint64_t expected = 0; - if (!tensor_nbytes(meta, &expected)) { - g_last_error = "safetensors tensor dtype/shape is unsupported"; - return false; - } - if (meta.data_end < meta.data_begin || - meta.data_end - meta.data_begin != expected || - data_start > total_bytes || - meta.data_end > total_bytes - data_start) { - g_last_error = "safetensors tensor byte range is invalid"; - return false; - } - return true; -} - bool read_safetensors_f32_vector(const std::string& path, const char* key, std::vector* out) { if (!out) return false; - std::string header; - uint64_t data_start = 0; - uint64_t total_bytes = 0; - if (!read_safetensors_header(path, &header, &data_start, &total_bytes)) { + flashrt::loader::SafetensorsFile file; + if (!file.open(path)) { + g_last_error = file.error(); return false; } - TensorMeta meta; - if (!tensor_meta(header, key, &meta)) return false; - if (meta.dtype != "F32" || meta.shape.size() != 1 || - !tensor_payload_valid(meta, data_start, total_bytes)) { + 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 = meta.shape[0]; + const uint64_t n = tensor->shape[0]; if (n > (1ull << 20)) { g_last_error = "safetensors vector is too large"; return false; } - std::ifstream f(path, std::ios::binary); - if (!f) { - g_last_error = "unable to open safetensors file"; - return false; - } - f.seekg(static_cast(data_start + meta.data_begin), - std::ios::beg); std::vector values(static_cast(n)); - f.read(reinterpret_cast(values.data()), - static_cast(n * sizeof(float))); - if (f.gcount() != static_cast(n * sizeof(float))) { - g_last_error = "safetensors vector payload is truncated"; - return false; - } + std::memcpy(values.data(), file.data(*tensor), + static_cast(tensor->bytes)); *out = std::move(values); return true; } @@ -637,12 +434,6 @@ bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, "_unnormalizer_processor.safetensors"); if (pre.empty() || post.empty()) return false; - std::string pre_header; - std::string post_header; - if (!read_safetensors_header(pre, &pre_header) || - !read_safetensors_header(post, &post_header)) { - return false; - } std::vector state_q01; std::vector state_q99; std::vector action_q01; @@ -709,10 +500,9 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { g_last_error = "checkpoint_path must contain model.safetensors"; return false; } - std::string header; - uint64_t data_start = 0; - uint64_t total_bytes = 0; - if (!read_safetensors_header(path, &header, &data_start, &total_bytes)) { + flashrt::loader::SafetensorsFile file; + if (!file.open(path)) { + g_last_error = file.error(); return false; } @@ -751,34 +541,35 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { }; for (const auto& req : requirements) { - TensorMeta meta; std::string key = req.key; - if (!tensor_meta(header, key, &meta)) { + const flashrt::loader::SafetensorInfo* meta = file.find(key); + if (!meta) { key = std::string("model.") + req.key; - if (!tensor_meta(header, key, &meta)) { + meta = file.find(key); + if (!meta) { g_last_error = std::string("model.safetensors is missing ") + req.key; return false; } } - if (req.dtype && meta.dtype != req.dtype) { + if (req.dtype && meta->dtype != req.dtype) { g_last_error = std::string("Pi0.5 tensor dtype mismatch: ") + req.key; return false; } - if (meta.dtype != "BF16" && meta.dtype != "F16" && - meta.dtype != "F32") { + 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.size() != req.rank) { + if (meta->shape.size() != req.rank) { g_last_error = std::string("Pi0.5 tensor rank mismatch: ") + req.key; return false; } for (uint64_t i = 0; i < req.rank; ++i) { - if (req.dims[i] && meta.shape[static_cast(i)] != + if (req.dims[i] && meta->shape[static_cast(i)] != req.dims[i]) { g_last_error = std::string("Pi0.5 tensor shape mismatch: ") + req.key; @@ -787,13 +578,10 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { } if (std::string(req.key) == "paligemma_with_expert.paligemma.lm_head.weight" && - meta.shape[0] < 1000) { + meta->shape[0] < 1000) { g_last_error = "Pi0.5 embedding vocab size is invalid"; return false; } - if (!tensor_payload_valid(meta, data_start, total_bytes)) { - return false; - } } g_last_error.clear(); return true; 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/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 16edf7e8..f4d92a14 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -80,7 +80,9 @@ switching between Lane A and Lane C. The current C++ shared object exports this symbol as a native-v2 configuration gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in -`model.safetensors`, plus the native builder's minimum required weight set. +`model.safetensors`, plus the native builder's minimum required weight set. A +read-only mmap loader owns the checkpoint mapping and exposes bounded tensor +views for the later device materialization step. It also verifies action/state q01/q99 metadata from openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer safetensors, including tensor byte ranges and finite ordered quantile payloads, then returns unsupported until the diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 14a5a053..a2819141 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -201,15 +201,17 @@ The Pi0.5 C++ shared object now exports `frt_model_runtime_open_v1` as a native-v2 configuration gate. The gate requires `io="native_v2"`, `checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, `max_prompt_tokens >= 200`, and a positive `state_dim`. It also parses -`checkpoint_path/model.safetensors` enough to verify the Pi0.5 prompt -embedding table and the native builder's minimum required weight set +`checkpoint_path/model.safetensors` through the native read-only mmap loader to +verify the Pi0.5 prompt embedding table and the native builder's minimum +required weight set (vision patch/position/projector, action projections, time MLP, and first encoder/decoder layer sentinels). 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. Valid -configuration returns unsupported until native asset materialization and graph -capture are complete. +configuration returns unsupported until device weight materialization and graph +capture are complete. The mmap and parsed tensor views are setup-side assets; +they never enter the model-runtime ABI or the hot path. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From bd43c4fe700d7aef0f6d1bfcf97041c62c1048ef Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:05:00 -0400 Subject: [PATCH 20/61] fix: preallocate Pi0.5 hot staging --- .../include/flashrt/cpp/modalities/action.h | 11 +++- .../flashrt/cpp/modalities/tokenizer.h | 4 +- cpp/modalities/src/action_cpu.cpp | 55 ++++++++++++++-- .../src/sentencepiece_tokenizer.cpp | 30 ++++++--- cpp/modalities/src/tokenizer_unavailable.cpp | 8 ++- .../pi05/include/flashrt/cpp/models/pi05/io.h | 4 +- .../flashrt/cpp/models/pi05/prompt_embed.h | 7 +- .../flashrt/cpp/models/pi05/prompt_format.h | 5 ++ .../include/flashrt/cpp/models/pi05/runtime.h | 3 + cpp/models/pi05/src/c_api.cpp | 49 ++++++++++---- cpp/models/pi05/src/io.cpp | 7 +- cpp/models/pi05/src/model_runtime.cpp | 65 ++++++++++++++----- cpp/models/pi05/src/prompt_embed.cpp | 16 +++-- cpp/models/pi05/src/prompt_format.cpp | 46 +++++++++---- cpp/models/pi05/src/runtime.cpp | 54 +++++++++++++-- cpp/tests/test_device_staging.cpp | 19 ++++++ cpp/tests/test_modalities.cpp | 4 +- cpp/tests/test_pi05_model_runtime.cpp | 21 ++++++ cpp/tests/test_pi05_prompt_embed.cpp | 19 +++++- cpp/tests/test_pi05_prompt_format.cpp | 9 +++ docs/mindon_pi05_integration.md | 2 + docs/pi05_io_contract.md | 5 ++ 22 files changed, 366 insertions(+), 77 deletions(-) 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/tokenizer.h b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h index 0a5f0a4b..2a4667c0 100644 --- a/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h +++ b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h @@ -33,7 +33,9 @@ class SentencePieceTokenizer final { Status load_model(const std::string& model_path); Status encode(const std::string& text, const SentencePieceEncodeOptions& options, - std::vector* token_ids) const; + 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; 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 index bc33bf36..363e61e7 100644 --- a/cpp/modalities/src/sentencepiece_tokenizer.cpp +++ b/cpp/modalities/src/sentencepiece_tokenizer.cpp @@ -10,6 +10,7 @@ namespace modalities { struct SentencePieceTokenizer::Impl { sentencepiece::SentencePieceProcessor processor; + std::vector encoded; bool loaded = false; }; @@ -37,7 +38,7 @@ Status SentencePieceTokenizer::load_model(const std::string& model_path) { Status SentencePieceTokenizer::encode( const std::string& text, const SentencePieceEncodeOptions& options, - std::vector* token_ids) const { + std::vector* token_ids) { if (!token_ids) { return Status::error(StatusCode::kInvalidArgument, "token_ids output is null"); @@ -48,14 +49,19 @@ Status SentencePieceTokenizer::encode( "SentencePiece model is not loaded"); } - std::vector encoded; - auto status = impl_->processor.Encode(text, &encoded); + 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 (encoded.size() + extra > + 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"); @@ -69,8 +75,8 @@ Status SentencePieceTokenizer::encode( } token_ids->push_back(static_cast(bos)); } - token_ids->reserve(encoded.size() + extra); - for (int id : encoded) { + token_ids->reserve(impl_->encoded.size() + extra); + for (int id : impl_->encoded) { token_ids->push_back(static_cast(id)); } if (options.add_eos) { @@ -83,10 +89,6 @@ Status SentencePieceTokenizer::encode( } if (options.max_tokens) { - if (token_ids->size() > options.max_tokens) { - return Status::error(StatusCode::kShapeMismatch, - "encoded token sequence exceeds max_tokens"); - } if (options.pad_to_max_tokens) { token_ids->resize(options.max_tokens, options.pad_id); } @@ -97,6 +99,14 @@ Status SentencePieceTokenizer::encode( 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; } diff --git a/cpp/modalities/src/tokenizer_unavailable.cpp b/cpp/modalities/src/tokenizer_unavailable.cpp index 49fa3438..7cb70158 100644 --- a/cpp/modalities/src/tokenizer_unavailable.cpp +++ b/cpp/modalities/src/tokenizer_unavailable.cpp @@ -26,7 +26,7 @@ Status SentencePieceTokenizer::load_model(const std::string& model_path) { Status SentencePieceTokenizer::encode( const std::string& text, const SentencePieceEncodeOptions& options, - std::vector* token_ids) const { + std::vector* token_ids) { (void)text; (void)options; if (token_ids) token_ids->clear(); @@ -35,6 +35,12 @@ Status SentencePieceTokenizer::encode( "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; } 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..65bcad97 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,8 @@ 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); modalities::Status prepare_vision( const std::vector& frames) const; @@ -42,6 +43,7 @@ class RuntimeIo { modalities::TensorView action_output_; void* stream_ = nullptr; modalities::VisionStaging* staging_ = nullptr; /* borrowed */ + modalities::ActionStaging* action_staging_ = nullptr; /* borrowed */ modalities::VisionPreprocessSpec vision_spec_; modalities::ActionPostprocessSpec action_spec_; }; 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 index b9d0824c..f5bd0d36 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h @@ -22,7 +22,7 @@ struct PromptEmbeddingSpec { }; modalities::Status embed_prompt( - const modalities::SentencePieceTokenizer& tokenizer, + modalities::SentencePieceTokenizer& tokenizer, const PromptEmbeddingSpec& spec, const std::string& prompt, const float* state, @@ -32,10 +32,11 @@ modalities::Status embed_prompt( std::vector* token_ids, std::uint64_t* prompt_len, void* stream = nullptr, - modalities::TextEmbeddingStaging* staging = nullptr); + modalities::TextEmbeddingStaging* staging = nullptr, + std::string* formatted_workspace = nullptr); modalities::Status embed_prompt_cpu( - const modalities::SentencePieceTokenizer& tokenizer, + modalities::SentencePieceTokenizer& tokenizer, const PromptEmbeddingSpec& spec, const std::string& prompt, const float* state, 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 index e3b5cbe6..a5ecc9d7 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h @@ -20,6 +20,11 @@ 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 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 90207e3f..cb2e2b65 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -110,6 +110,7 @@ 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_; @@ -119,6 +120,8 @@ class Runtime final : public families::vla::Runtime { modalities::Status prompt_status_; std::vector prompt_token_ids_; std::vector normalized_state_; + std::string task_prompt_workspace_; + std::string formatted_prompt_workspace_; std::uint64_t current_prompt_len_ = 0; bool prompt_staging_enabled_ = false; frt_graph graph_ = nullptr; diff --git a/cpp/models/pi05/src/c_api.cpp b/cpp/models/pi05/src/c_api.cpp index 4855e5eb..45d90327 100644 --- a/cpp/models/pi05/src/c_api.cpp +++ b/cpp/models/pi05/src/c_api.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -14,6 +15,9 @@ 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 { @@ -54,6 +58,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; } @@ -100,8 +112,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) || @@ -109,8 +124,19 @@ 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; + 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; @@ -125,9 +151,8 @@ extern "C" int frt_pi05_runtime_prepare_vision( 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); @@ -148,19 +173,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/io.cpp b/cpp/models/pi05/src/io.cpp index 4072f692..c2e109f4 100644 --- a/cpp/models/pi05/src/io.cpp +++ b/cpp/models/pi05/src/io.cpp @@ -39,11 +39,13 @@ 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) : image_input_(image_input), action_output_(action_output), stream_(stream), staging_(staging), + action_staging_(action_staging), vision_spec_(vision_preprocess_spec(num_views)), action_spec_(action_postprocess_spec(action_mean, action_stddev, chunk, model_action_dim, robot_action_dim)) { @@ -63,7 +65,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 c2d13b96..99fe7d48 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -37,6 +37,8 @@ struct Adapter { bool has_state = false; 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}; @@ -125,18 +127,17 @@ 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); + 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; @@ -150,9 +151,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); @@ -170,8 +170,16 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, a->last_error = "prompt payload is null"; return -1; } + if (bytes > a->prompt_text.capacity()) { + a->last_error = "prompt payload exceeds the hot-path capacity"; + return -4; + } const char* begin = static_cast(data); - a->prompt_text.assign(begin, begin + bytes); + 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( @@ -184,7 +192,7 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, a->last_error = st.message.empty() ? "prompt staging is not configured" : st.message; - return -1; + return status_code(st); } a->last_error.clear(); return 0; @@ -194,8 +202,14 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, 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); - a->state_values.assign(values, values + bytes / sizeof(float)); + std::memcpy(a->state_values.data(), values, + static_cast(bytes)); a->has_state = true; if (!a->has_prompt_text) { a->last_error.clear(); @@ -209,7 +223,7 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, a->last_error = st.message.empty() ? "state staging failed" : st.message; - return -1; + return status_code(st); } a->last_error.clear(); return 0; @@ -227,19 +241,18 @@ 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; } @@ -361,6 +374,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; @@ -502,6 +521,20 @@ extern "C" int frt_pi05_model_runtime_create_over( 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.reserve(static_cast( + cfg.prompt_max_tokens * 8ull)); + } + 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/prompt_embed.cpp b/cpp/models/pi05/src/prompt_embed.cpp index d8261e28..50a9c046 100644 --- a/cpp/models/pi05/src/prompt_embed.cpp +++ b/cpp/models/pi05/src/prompt_embed.cpp @@ -82,7 +82,7 @@ modalities::Status zero_prompt_output(const modalities::TensorView& output, } // namespace modalities::Status embed_prompt( - const modalities::SentencePieceTokenizer& tokenizer, + modalities::SentencePieceTokenizer& tokenizer, const PromptEmbeddingSpec& spec, const std::string& prompt, const float* state, @@ -92,7 +92,8 @@ modalities::Status embed_prompt( std::vector* token_ids, std::uint64_t* prompt_len, void* stream, - modalities::TextEmbeddingStaging* staging) { + modalities::TextEmbeddingStaging* staging, + std::string* formatted_workspace) { if (!token_ids || !prompt_len) { return modalities::Status::error( modalities::StatusCode::kInvalidArgument, @@ -110,10 +111,13 @@ modalities::Status embed_prompt( modalities::SentencePieceEncodeOptions options; options.add_bos = true; + options.max_tokens = spec.max_tokens; if (state) { - const std::string formatted = format_state_prompt(prompt, state, - n_state); - st = tokenizer.encode(formatted, options, token_ids); + 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) { @@ -149,7 +153,7 @@ modalities::Status embed_prompt( } modalities::Status embed_prompt_cpu( - const modalities::SentencePieceTokenizer& tokenizer, + modalities::SentencePieceTokenizer& tokenizer, const PromptEmbeddingSpec& spec, const std::string& prompt, const float* state, diff --git a/cpp/models/pi05/src/prompt_format.cpp b/cpp/models/pi05/src/prompt_format.cpp index 6f3f0aad..5e1c36a1 100644 --- a/cpp/models/pi05/src/prompt_format.cpp +++ b/cpp/models/pi05/src/prompt_format.cpp @@ -2,7 +2,7 @@ #include #include -#include +#include namespace flashrt { namespace models { @@ -52,18 +52,42 @@ 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) { - const std::string cleaned = clean_task_prompt(prompt); - if (!state) return cleaned; + 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; - const auto tokens = discretize_state_prompt_bins(state, n_state); - std::ostringstream oss; - oss << "Task: " << cleaned << ", State: "; - for (std::size_t i = 0; i < tokens.size(); ++i) { - if (i) oss << ' '; - oss << tokens[i]; + 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); } - oss << ";\nAction: "; - return oss.str(); + out->append(";\nAction: "); } } // namespace pi05 diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index 0ff1254e..6fd94c20 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -2,6 +2,7 @@ #include #include +#include namespace flashrt { namespace models { @@ -72,6 +73,7 @@ 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(); } @@ -168,10 +170,24 @@ 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); + config_.robot_action_dim, config_.image_dtype, staging, + action_staging); return bind_prompt_staging(); } @@ -190,6 +206,14 @@ int Runtime::set_prompt_state(const char* text, const float* state, "prompt text is null"); return -1; } + const std::size_t text_bytes = std::strlen(text); + if (text_bytes > task_prompt_workspace_.capacity()) { + 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()) { @@ -198,7 +222,6 @@ int Runtime::set_prompt_state(const char* text, const float* state, "state dimension does not match norm stats"); return -1; } - normalized_state_.resize(n_state); for (std::uint64_t i = 0; i < n_state; ++i) { const float lo = config_.state_q01[i]; const float hi = config_.state_q99[i]; @@ -208,12 +231,14 @@ int Runtime::set_prompt_state(const char* text, const float* state, state_for_prompt = normalized_state_.data(); } prompt_status_ = embed_prompt( - prompt_tokenizer_, prompt_spec_, text, state_for_prompt, n_state, + 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); + : nullptr, + &formatted_prompt_workspace_); return prompt_status_.ok_status() ? 0 : -1; } @@ -273,6 +298,27 @@ modalities::Status Runtime::bind_prompt_staging() { ? 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); + 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); diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index f772451b..9b3329ed 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -12,6 +12,7 @@ #include using flashrt::modalities::DType; +using flashrt::modalities::ActionStaging; using flashrt::modalities::Layout; using flashrt::modalities::MemoryPlace; using flashrt::modalities::PixelFormat; @@ -140,6 +141,24 @@ 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); } diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index 117f11c7..3636c0ec 100644 --- a/cpp/tests/test_modalities.cpp +++ b/cpp/tests/test_modalities.cpp @@ -117,7 +117,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); } @@ -162,7 +162,7 @@ void test_pi05_runtime_io_adapter() { 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_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index e781bfef..895f3162 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -405,6 +405,27 @@ int main() { 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"); diff --git a/cpp/tests/test_pi05_prompt_embed.cpp b/cpp/tests/test_pi05_prompt_embed.cpp index eaabeffd..89c05560 100644 --- a/cpp/tests/test_pi05_prompt_embed.cpp +++ b/cpp/tests/test_pi05_prompt_embed.cpp @@ -95,9 +95,13 @@ void test_paligemma_prompt_embedding_when_configured() { 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_cpu(tokenizer, spec, "pick_up_cube", state, 3, src, dst, - &ids, &prompt_len); + 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, @@ -114,6 +118,17 @@ void test_paligemma_prompt_embedding_when_configured() { 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 } diff --git a/cpp/tests/test_pi05_prompt_format.cpp b/cpp/tests/test_pi05_prompt_format.cpp index d687873e..7184b0ea 100644 --- a/cpp/tests/test_pi05_prompt_format.cpp +++ b/cpp/tests/test_pi05_prompt_format.cpp @@ -24,6 +24,15 @@ void test_prompt_format_matches_python_reference() { "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() { diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index f4d92a14..f12724b9 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -211,5 +211,7 @@ but the C++ host still sees only the adopted stage array. - `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/pi05_io_contract.md b/docs/pi05_io_contract.md index a2819141..1e1a92b7 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -161,6 +161,11 @@ API family for the same phases. `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 From bd645075d9cdf80ef16965563f23c82a5f98f192 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:09:02 -0400 Subject: [PATCH 21/61] feat: validate complete Pi0.5 weight inventory --- cpp/CMakeLists.txt | 1 + .../flashrt/cpp/models/pi05/native_weights.h | 23 ++++ cpp/models/pi05/src/native_open.cpp | 68 +---------- cpp/models/pi05/src/native_weights.cpp | 115 ++++++++++++++++++ cpp/tests/test_pi05_native_open.cpp | 112 +++++++++-------- docs/mindon_pi05_integration.md | 3 +- docs/pi05_io_contract.md | 7 +- 7 files changed, 210 insertions(+), 119 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weights.h create mode 100644 cpp/models/pi05/src/native_weights.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 07b0b787..4feeadfe 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -128,6 +128,7 @@ target_include_directories(flashrt_cpp_loader add_library(flashrt_cpp_pi05 STATIC models/pi05/src/spec.cpp + models/pi05/src/native_weights.cpp models/pi05/src/prompt_format.cpp models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp 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/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index b49cd71b..2b56da38 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -1,5 +1,6 @@ #include "flashrt/model_runtime.h" #include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/models/pi05/native_weights.h" #include #include @@ -27,13 +28,6 @@ struct JsonValue { bool boolean = false; }; -struct TensorRequirement { - const char* key; - const char* dtype; - uint64_t rank; - uint64_t dims[4]; -}; - class JsonParser { public: explicit JsonParser(const char* src) : cur_(src ? src : "") {} @@ -506,41 +500,8 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { return false; } - const TensorRequirement requirements[] = { - {"paligemma_with_expert.paligemma.lm_head.weight", - nullptr, 2, {0, 2048, 0, 0}}, - {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".embeddings.patch_embedding.weight", - nullptr, 4, {1152, 3, 14, 14}}, - {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".embeddings.patch_embedding.bias", - nullptr, 1, {1152, 0, 0, 0}}, - {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".embeddings.position_embedding.weight", - nullptr, 2, {256, 1152, 0, 0}}, - {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" - ".weight", - nullptr, 2, {2048, 1152, 0, 0}}, - {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" - ".bias", - nullptr, 1, {2048, 0, 0, 0}}, - {"action_in_proj.weight", nullptr, 2, {1024, 32, 0, 0}}, - {"action_in_proj.bias", nullptr, 1, {1024, 0, 0, 0}}, - {"action_out_proj.weight", nullptr, 2, {32, 1024, 0, 0}}, - {"action_out_proj.bias", nullptr, 1, {32, 0, 0, 0}}, - {"time_mlp_in.weight", nullptr, 2, {1024, 1024, 0, 0}}, - {"time_mlp_in.bias", nullptr, 1, {1024, 0, 0, 0}}, - {"time_mlp_out.weight", nullptr, 2, {1024, 1024, 0, 0}}, - {"time_mlp_out.bias", nullptr, 1, {1024, 0, 0, 0}}, - {"paligemma_with_expert.paligemma.model.language_model.layers.0" - ".self_attn.q_proj.weight", - nullptr, 2, {2048, 2048, 0, 0}}, - {"paligemma_with_expert.gemma_expert.model.layers.0.self_attn" - ".q_proj.weight", - nullptr, 2, {2048, 1024, 0, 0}}, - }; - - for (const auto& req : requirements) { + 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) { @@ -552,36 +513,17 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { return false; } } - if (req.dtype && meta->dtype != req.dtype) { - g_last_error = std::string("Pi0.5 tensor dtype mismatch: ") + - 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.size() != req.rank) { - g_last_error = std::string("Pi0.5 tensor rank mismatch: ") + + if (meta->shape != req.shape) { + g_last_error = std::string("Pi0.5 tensor shape mismatch: ") + req.key; return false; } - for (uint64_t i = 0; i < req.rank; ++i) { - if (req.dims[i] && meta->shape[static_cast(i)] != - req.dims[i]) { - g_last_error = std::string("Pi0.5 tensor shape mismatch: ") + - req.key; - return false; - } - } - if (std::string(req.key) == - "paligemma_with_expert.paligemma.lm_head.weight" && - meta->shape[0] < 1000) { - g_last_error = "Pi0.5 embedding vocab size is invalid"; - return false; - } } g_last_error.clear(); return true; 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/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 8dbdb1bb..e27f1399 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -1,5 +1,7 @@ #include "flashrt/model_runtime.h" +#include "flashrt/cpp/models/pi05/native_weights.h" +#include #include #include #include @@ -44,9 +46,19 @@ void write_raw_safetensors(const std::string& path, void write_raw_safetensors(const std::string& path, const std::string& header, uint64_t payload_bytes) { - write_raw_safetensors(path, header, - std::string(static_cast(payload_bytes), - '\0')); + 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) { @@ -55,61 +67,32 @@ void append_f32(std::string* out, float value) { out->append(bytes, sizeof(bytes)); } -void write_safetensors(const std::string& path) { - struct Entry { - const char* key; - const char* dtype; - const char* shape; - uint64_t bytes; - }; - const Entry entries[] = { - {"paligemma_with_expert.paligemma.lm_head.weight", "BF16", - "[1001,2048]", 1001ull * 2048ull * 2ull}, - {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".embeddings.patch_embedding.weight", - "F32", "[1152,3,14,14]", 1152ull * 3ull * 14ull * 14ull * 4ull}, - {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".embeddings.patch_embedding.bias", - "F32", "[1152]", 1152ull * 4ull}, - {"paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".embeddings.position_embedding.weight", - "F32", "[256,1152]", 256ull * 1152ull * 4ull}, - {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" - ".weight", - "F32", "[2048,1152]", 2048ull * 1152ull * 4ull}, - {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear" - ".bias", - "F32", "[2048]", 2048ull * 4ull}, - {"action_in_proj.weight", "F32", "[1024,32]", 1024ull * 32ull * 4ull}, - {"action_in_proj.bias", "F32", "[1024]", 1024ull * 4ull}, - {"action_out_proj.weight", "F32", "[32,1024]", 32ull * 1024ull * 4ull}, - {"action_out_proj.bias", "F32", "[32]", 32ull * 4ull}, - {"time_mlp_in.weight", "F32", "[1024,1024]", 1024ull * 1024ull * 4ull}, - {"time_mlp_in.bias", "F32", "[1024]", 1024ull * 4ull}, - {"time_mlp_out.weight", "F32", "[1024,1024]", 1024ull * 1024ull * 4ull}, - {"time_mlp_out.bias", "F32", "[1024]", 1024ull * 4ull}, - {"paligemma_with_expert.paligemma.model.language_model.layers.0" - ".self_attn.q_proj.weight", - "F32", "[2048,2048]", 2048ull * 2048ull * 4ull}, - {"paligemma_with_expert.gemma_expert.model.layers.0.self_attn" - ".q_proj.weight", - "F32", "[2048,1024]", 2048ull * 1024ull * 4ull}, - }; +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; - for (size_t i = 0; i < sizeof(entries) / sizeof(entries[0]); ++i) { - const Entry& e = entries[i]; - if (i) header += ","; + 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\":\""; - header += e.dtype; - header += "\",\"shape\":"; - header += e.shape; + 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 += e.bytes; + offset += elements * sizeof(float); header += std::to_string(offset); header += "]}"; } @@ -182,6 +165,23 @@ std::string config(const std::string& ckpt, } // 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); @@ -223,6 +223,16 @@ int main() { 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); diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index f12724b9..d756bf4e 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -80,7 +80,8 @@ switching between Lane A and Lane C. The current C++ shared object exports this symbol as a native-v2 configuration gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in -`model.safetensors`, plus the native builder's minimum required weight set. A +`model.safetensors`, plus the native builder's complete 812-tensor weight +inventory across vision, language encoder, and action expert. A read-only mmap loader owns the checkpoint mapping and exposes bounded tensor views for the later device materialization step. It also verifies action/state q01/q99 metadata from openpi `norm_stats.json` or diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 1e1a92b7..a78b2ba2 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -207,10 +207,9 @@ native-v2 configuration gate. The gate requires `io="native_v2"`, `checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, `max_prompt_tokens >= 200`, and a positive `state_dim`. It also parses `checkpoint_path/model.safetensors` through the native read-only mmap loader to -verify the Pi0.5 prompt embedding table and the native builder's minimum -required weight set -(vision patch/position/projector, action projections, time MLP, and first -encoder/decoder layer sentinels). It also verifies action/state q01/q99 +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. Valid From 7354d64e342c31353862acf3fd4adb7f414f611b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:17:05 -0400 Subject: [PATCH 22/61] feat: add Pi0.5 native weight transforms --- cpp/CMakeLists.txt | 14 +- .../cpp/models/pi05/native_weight_ops.h | 71 ++++ cpp/models/pi05/src/native_weight_ops.cpp | 310 ++++++++++++++++++ cpp/tests/gate_pi05_native_weight_ops.py | 99 ++++++ cpp/tests/pi05_native_weight_probe.cpp | 177 ++++++++++ cpp/tests/test_pi05_native_weight_ops.cpp | 134 ++++++++ docs/pi05_io_contract.md | 13 + 7 files changed, 816 insertions(+), 2 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h create mode 100644 cpp/models/pi05/src/native_weight_ops.cpp create mode 100644 cpp/tests/gate_pi05_native_weight_ops.py create mode 100644 cpp/tests/pi05_native_weight_probe.cpp create mode 100644 cpp/tests/test_pi05_native_weight_ops.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4feeadfe..f5db83c9 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -129,6 +129,7 @@ target_include_directories(flashrt_cpp_loader add_library(flashrt_cpp_pi05 STATIC models/pi05/src/spec.cpp models/pi05/src/native_weights.cpp + models/pi05/src/native_weight_ops.cpp models/pi05/src/prompt_format.cpp models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp @@ -136,14 +137,14 @@ add_library(flashrt_cpp_pi05 STATIC 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) add_library(flashrt_cpp_pi05_c SHARED models/pi05/src/c_api.cpp models/pi05/src/model_runtime.cpp models/pi05/src/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c - PUBLIC flashrt_cpp_pi05 flashrt_cpp_loader flashrt_runtime) + PUBLIC flashrt_cpp_pi05 flashrt_runtime) target_include_directories(flashrt_cpp_pi05_c PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include) @@ -170,6 +171,15 @@ if(BUILD_TESTING) 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) + 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) 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..cf40be14 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h @@ -0,0 +1,71 @@ +#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_scale(const NativeFloatTensor& input, + float scale, + NativeFloatTensor* out); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_OPS_H 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..4bc10e84 --- /dev/null +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -0,0 +1,310 @@ +#include "flashrt/cpp/models/pi05/native_weight_ops.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 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_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(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt 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..dea4df75 --- /dev/null +++ b/cpp/tests/gate_pi05_native_weight_ops.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +import argparse +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() + + 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/pi05_native_weight_probe.cpp b/cpp/tests/pi05_native_weight_probe.cpp new file mode 100644 index 00000000..ebdc8ad5 --- /dev/null +++ b/cpp/tests/pi05_native_weight_probe.cpp @@ -0,0 +1,177 @@ +#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); +} + +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); + } + 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/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp new file mode 100644 index 00000000..477eb206 --- /dev/null +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -0,0 +1,134 @@ +#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}); + + 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}); + + 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index a78b2ba2..3581ad85 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -217,6 +217,13 @@ configuration returns unsupported until device weight materialization and graph capture are complete. The mmap and parsed tensor views are setup-side assets; they never enter the model-runtime ABI or the 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. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. @@ -234,6 +241,12 @@ 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 +``` + ``` python cpp/tests/gate_pi05_model_runtime_export.py ... python cpp/tests/gate_pi05_c_api_export.py ... From 96e8e2dcef14f062d180d2ab230c07f34874b7c4 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:19:08 -0400 Subject: [PATCH 23/61] feat: add Pi0.5 device weight store --- cpp/CMakeLists.txt | 8 ++ .../cpp/models/pi05/native_device_weights.h | 42 +++++++++ cpp/models/pi05/src/native_device_weights.cpp | 86 +++++++++++++++++++ cpp/tests/test_pi05_native_device_weights.cpp | 68 +++++++++++++++ docs/pi05_io_contract.md | 6 ++ 5 files changed, 210 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h create mode 100644 cpp/models/pi05/src/native_device_weights.cpp create mode 100644 cpp/tests/test_pi05_native_device_weights.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index f5db83c9..fbad4b47 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -130,6 +130,7 @@ add_library(flashrt_cpp_pi05 STATIC 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/prompt_format.cpp models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp @@ -194,6 +195,13 @@ if(BUILD_TESTING) 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_c_api tests/test_pi05_c_api.cpp) target_link_libraries(test_pi05_c_api PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) 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..d09291d5 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h @@ -0,0 +1,42 @@ +#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 { + +struct NativeDeviceWeight { + frt_buffer buffer = nullptr; + std::vector shape; +}; + +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); + 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/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp new file mode 100644 index 00000000..8dd6da72 --- /dev/null +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -0,0 +1,86 @@ +#include "flashrt/cpp/models/pi05/native_device_weights.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#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; +} + +} // namespace + +modalities::Status NativeDeviceWeightStore::upload( + const std::string& name, + const NativeBf16Tensor& tensor) { + 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; + if (!element_count(tensor.shape, &elements) || + elements != tensor.values.size() || + elements > std::numeric_limits::max() / + sizeof(std::uint16_t)) { + return invalid("device weight shape does not match BF16 payload"); + } + const std::size_t bytes = elements * sizeof(std::uint16_t); + if (!bytes) return invalid("device weight payload is empty"); + +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)tensor; + 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) { + return modalities::Status::error(modalities::StatusCode::kBackend, + "device weight allocation failed"); + } + const cudaError_t rc = cudaMemcpy(frt_buffer_dptr(buffer), + tensor.values.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, tensor.shape}); + 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; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt 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..0b8cc14a --- /dev/null +++ b/cpp/tests/test_pi05_native_device_weights.cpp @@ -0,0 +1,68 @@ +#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; + + 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(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); + 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 3581ad85..eab2bcb9 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -224,6 +224,12 @@ 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 +shape/payload mismatches fail setup. Destroying the context releases the device +weights after graphs and plans, preserving the exec ownership order. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From bc3d5657fd4a45ff23d632c7b010410d049056dd Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:22:07 -0400 Subject: [PATCH 24/61] feat: materialize Pi0.5 encoder weights --- cpp/CMakeLists.txt | 8 + .../models/pi05/native_weight_materializer.h | 39 +++++ .../pi05/src/native_weight_materializer.cpp | 135 ++++++++++++++++ .../test_pi05_native_weight_materializer.cpp | 149 ++++++++++++++++++ docs/pi05_io_contract.md | 8 + 5 files changed, 339 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h create mode 100644 cpp/models/pi05/src/native_weight_materializer.cpp create mode 100644 cpp/tests/test_pi05_native_weight_materializer.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index fbad4b47..ee8ee8bd 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -131,6 +131,7 @@ add_library(flashrt_cpp_pi05 STATIC 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_materializer.cpp models/pi05/src/prompt_format.cpp models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp @@ -202,6 +203,13 @@ if(BUILD_TESTING) 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) 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..d97666ef --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h @@ -0,0 +1,39 @@ +#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 { + +class NativeWeightMaterializer { +public: + NativeWeightMaterializer(const loader::SafetensorsFile& source, + NativeDeviceWeightStore* destination) + : source_(source), destination_(destination) {} + + modalities::Status materialize_encoder_layer(int layer); + +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_folded_transpose( + const std::string& source_key, + const NativeFloatTensor& norm, + const std::string& destination_name); + + 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/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp new file mode 100644 index 00000000..e382e296 --- /dev/null +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -0,0 +1,135 @@ +#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 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_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::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)); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt 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..c4f2d8ed --- /dev/null +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -0,0 +1,149 @@ +#include "flashrt/cpp/models/pi05/native_weight_materializer.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::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)}, + }; + 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()); + } + 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})); + } + frt_ctx_destroy(real_ctx); + } + std::printf("PASS - Pi0.5 encoder weight materializer\n"); + return 0; +} diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index eab2bcb9..3241293c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -230,6 +230,14 @@ capsule regions. Upload is complete before capture; duplicate logical names or shape/payload mismatches fail setup. Destroying the context releases the device weights after graphs and plans, preserving the exec ownership order. +The first composed materializer covers one language-encoder layer and emits +the five names consumed by the existing pipeline (`attn_qkv`, `attn_o`, +`ffn_gate`, `ffn_up`, and `ffn_down`). It performs FP32 RMS folds before BF16 +upload and has been exercised against both supported real checkpoint layouts. +Vision, decoder, global weights, and precision-specific packing remain setup +work; `open_v1` stays unsupported until that inventory and graph capture are +complete. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From 53de52c09bb288ea5fcf0cb71219bacb5a810b38 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:26:35 -0400 Subject: [PATCH 25/61] feat: materialize Pi0.5 decoder weights --- .../models/pi05/native_weight_materializer.h | 5 + .../pi05/src/native_weight_materializer.cpp | 110 ++++++++++++++++++ .../test_pi05_native_weight_materializer.cpp | 38 +++++- docs/pi05_io_contract.md | 14 +-- 4 files changed, 159 insertions(+), 8 deletions(-) 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 index d97666ef..27072df4 100644 --- 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 @@ -15,6 +15,8 @@ class NativeWeightMaterializer { : source_(source), destination_(destination) {} modalities::Status materialize_encoder_layer(int layer); + modalities::Status materialize_decoder_layer(int layer, + bool merge_gate_up); private: modalities::Status load(const std::string& key, NativeFloatTensor* out); @@ -23,6 +25,9 @@ class NativeWeightMaterializer { 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, diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp index e382e296..348ffc9c 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -17,6 +17,11 @@ std::string encoder_prefix(int layer) { std::to_string(layer); } +std::string decoder_prefix(int layer) { + return "paligemma_with_expert.gemma_expert.model.layers." + + std::to_string(layer); +} + std::string layer_name(const char* stem, int layer) { return std::string(stem) + std::to_string(layer); } @@ -54,6 +59,18 @@ modalities::Status NativeWeightMaterializer::upload_rounded_transpose( 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, @@ -130,6 +147,99 @@ modalities::Status NativeWeightMaterializer::materialize_encoder_layer( 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)); +} + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index c4f2d8ed..55a26af7 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -86,6 +86,8 @@ int main() { } 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::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)}, @@ -97,6 +99,20 @@ int main() { {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)}, }; const std::string path = temp_path(); write_checkpoint(path, entries); @@ -119,6 +135,20 @@ int main() { 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()); } frt_ctx_destroy(ctx); assert(::unlink(path.c_str()) == 0); @@ -141,9 +171,15 @@ int main() { 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})); } frt_ctx_destroy(real_ctx); } - std::printf("PASS - Pi0.5 encoder weight materializer\n"); + std::printf("PASS - Pi0.5 native layer materializer\n"); return 0; } diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 3241293c..04051c7c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -230,13 +230,13 @@ capsule regions. Upload is complete before capture; duplicate logical names or shape/payload mismatches fail setup. Destroying the context releases the device weights after graphs and plans, preserving the exec ownership order. -The first composed materializer covers one language-encoder layer and emits -the five names consumed by the existing pipeline (`attn_qkv`, `attn_o`, -`ffn_gate`, `ffn_up`, and `ffn_down`). It performs FP32 RMS folds before BF16 -upload and has been exercised against both supported real checkpoint layouts. -Vision, decoder, global weights, and precision-specific packing remain setup -work; `open_v1` stays unsupported until that inventory and graph capture are -complete. +The composed layer materializer covers language-encoder and action-expert +layers. 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. Both have been exercised against the two +supported real checkpoint layouts. Vision, global weights, precision-specific +packing, and graph capture remain incomplete, so `open_v1` stays unsupported. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 413f762ef3b483a14e6a76fdb341ccccb41f294b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:30:17 -0400 Subject: [PATCH 26/61] feat: materialize Pi0.5 vision weights --- .../models/pi05/native_weight_materializer.h | 2 + .../cpp/models/pi05/native_weight_ops.h | 4 + .../pi05/src/native_weight_materializer.cpp | 120 ++++++++++++++++++ cpp/models/pi05/src/native_weight_ops.cpp | 24 ++++ .../test_pi05_native_weight_materializer.cpp | 58 +++++++++ cpp/tests/test_pi05_native_weight_ops.cpp | 5 + docs/pi05_io_contract.md | 13 +- 7 files changed, 221 insertions(+), 5 deletions(-) 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 index 27072df4..634ca458 100644 --- 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 @@ -17,6 +17,8 @@ class NativeWeightMaterializer { 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(); private: modalities::Status load(const std::string& key, NativeFloatTensor* out); 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 index cf40be14..fa0bc543 100644 --- 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 @@ -60,6 +60,10 @@ modalities::Status native_concat_columns( 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); diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp index 348ffc9c..120ec81c 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -22,6 +22,12 @@ std::string decoder_prefix(int layer) { 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); } @@ -240,6 +246,120 @@ modalities::Status NativeWeightMaterializer::materialize_decoder_layer( 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"); +} + } // 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 index 4bc10e84..6316da12 100644 --- a/cpp/models/pi05/src/native_weight_ops.cpp +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -295,6 +295,30 @@ modalities::Status native_concat_columns( 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) { diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index 55a26af7..029be97d 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -88,6 +88,9 @@ int main() { "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)}, @@ -113,6 +116,44 @@ int main() { 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)}, }; const std::string path = temp_path(); write_checkpoint(path, entries); @@ -149,6 +190,16 @@ int main() { 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()); } frt_ctx_destroy(ctx); assert(::unlink(path.c_str()) == 0); @@ -177,6 +228,13 @@ int main() { 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})); } frt_ctx_destroy(real_ctx); } diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp index 477eb206..e57b1856 100644 --- a/cpp/tests/test_pi05_native_weight_ops.cpp +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -106,6 +106,11 @@ int main() { 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}); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 04051c7c..0604aca2 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -230,13 +230,16 @@ capsule regions. Upload is complete before capture; duplicate logical names or shape/payload mismatches fail setup. Destroying the context releases the device weights after graphs and plans, preserving the exec ownership order. -The composed layer materializer covers language-encoder and action-expert -layers. Encoder layers emit the five pipeline groups (`attn_qkv`, `attn_o`, +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. Both have been exercised against the two -supported real checkpoint layouts. Vision, global weights, precision-specific -packing, and graph capture remain incomplete, so `open_v1` stays unsupported. +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. These paths have been exercised against the +two supported real checkpoint layouts. Remaining global language/action/time +weights, precision-specific packing, and graph capture are incomplete, so +`open_v1` stays unsupported. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 10bd2babc802d47cfe0b3d2ba4eaa18c41b25cae Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:37:24 -0400 Subject: [PATCH 27/61] feat: materialize Pi0.5 global weights --- .../models/pi05/native_weight_materializer.h | 7 ++ .../cpp/models/pi05/native_weight_ops.h | 5 ++ .../pi05/src/native_weight_materializer.cpp | 77 +++++++++++++++++++ cpp/models/pi05/src/native_weight_ops.cpp | 38 +++++++++ cpp/tests/gate_pi05_native_weight_ops.py | 28 +++++++ cpp/tests/pi05_native_weight_probe.cpp | 32 ++++++++ .../test_pi05_native_weight_materializer.cpp | 37 +++++++++ cpp/tests/test_pi05_native_weight_ops.cpp | 6 ++ docs/pi05_io_contract.md | 14 +++- 9 files changed, 240 insertions(+), 4 deletions(-) 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 index 634ca458..fea7902b 100644 --- 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 @@ -19,6 +19,8 @@ class NativeWeightMaterializer { 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(); private: modalities::Status load(const std::string& key, NativeFloatTensor* out); @@ -34,6 +36,11 @@ class NativeWeightMaterializer { 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; 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 index fa0bc543..6e278640 100644 --- 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 @@ -68,6 +68,11 @@ 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 diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp index 120ec81c..c9a87ae4 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -93,6 +93,30 @@ modalities::Status NativeWeightMaterializer::upload_folded_transpose( 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_) { @@ -360,6 +384,59 @@ modalities::Status NativeWeightMaterializer::materialize_vision_globals() { "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"); +} + } // 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 index 6316da12..d3dd5134 100644 --- a/cpp/models/pi05/src/native_weight_ops.cpp +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -1,5 +1,6 @@ #include "flashrt/cpp/models/pi05/native_weight_ops.h" +#include #include #include #include @@ -329,6 +330,43 @@ modalities::Status native_scale(const NativeFloatTensor& input, 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/tests/gate_pi05_native_weight_ops.py b/cpp/tests/gate_pi05_native_weight_ops.py index dea4df75..ad75773a 100644 --- a/cpp/tests/gate_pi05_native_weight_ops.py +++ b/cpp/tests/gate_pi05_native_weight_ops.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import argparse +import math import subprocess import torch @@ -82,6 +83,33 @@ def bf16(key: str) -> torch.Tensor: 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 diff --git a/cpp/tests/pi05_native_weight_probe.cpp b/cpp/tests/pi05_native_weight_probe.cpp index ebdc8ad5..a8c4d036 100644 --- a/cpp/tests/pi05_native_weight_probe.cpp +++ b/cpp/tests/pi05_native_weight_probe.cpp @@ -128,6 +128,30 @@ bool gate_up(const SafetensorsFile& file, NativeBf16Tensor* out) { 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()); @@ -161,6 +185,14 @@ int main(int argc, char** argv) { 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'; diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index 029be97d..f369f7eb 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -154,6 +154,20 @@ int main() { {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); @@ -200,6 +214,19 @@ int main() { 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})); } frt_ctx_destroy(ctx); assert(::unlink(path.c_str()) == 0); @@ -235,6 +262,16 @@ int main() { 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})); } frt_ctx_destroy(real_ctx); } diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp index e57b1856..d8baf08d 100644 --- a/cpp/tests/test_pi05_native_weight_ops.cpp +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -115,6 +115,12 @@ int main() { 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( diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 0604aca2..64914d9d 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -236,10 +236,16 @@ weights. Encoder layers emit the five pipeline groups (`attn_qkv`, `attn_o`, 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. These paths have been exercised against the -two supported real checkpoint layouts. Remaining global language/action/time -weights, precision-specific packing, and graph capture are incomplete, so -`open_v1` stays unsupported. +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. Precision-specific packing and +graph capture remain incomplete, so `open_v1` stays unsupported. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 47a3eeed55820ea4bd4de731e56574574315d42b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:39:50 -0400 Subject: [PATCH 28/61] feat: support typed Pi0.5 device weights --- .../cpp/models/pi05/native_device_weights.h | 14 +++++++ cpp/models/pi05/src/native_device_weights.cpp | 38 ++++++++++++++----- cpp/tests/test_pi05_native_device_weights.cpp | 29 ++++++++++++++ docs/pi05_io_contract.md | 6 ++- 4 files changed, 76 insertions(+), 11 deletions(-) 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 index d09291d5..d4963cb1 100644 --- 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 @@ -13,9 +13,17 @@ 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 { @@ -27,6 +35,12 @@ class NativeDeviceWeightStore { 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); const NativeDeviceWeight* find(const std::string& name) const; std::size_t size() const { return weights_.size(); } diff --git a/cpp/models/pi05/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp index 8dd6da72..fdafe1d2 100644 --- a/cpp/models/pi05/src/native_device_weights.cpp +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -31,27 +31,47 @@ bool element_count(const std::vector& shape, 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; - if (!element_count(tensor.shape, &elements) || - elements != tensor.values.size() || - elements > std::numeric_limits::max() / - sizeof(std::uint16_t)) { - return invalid("device weight shape does not match BF16 payload"); + 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"); } - const std::size_t bytes = elements * sizeof(std::uint16_t); if (!bytes) return invalid("device weight payload is empty"); #ifndef FLASHRT_CPP_WITH_CUDA_STAGING - (void)tensor; + (void)data; return modalities::Status::error( modalities::StatusCode::kUnsupported, "device weight upload requires the CUDA build"); @@ -62,7 +82,7 @@ modalities::Status NativeDeviceWeightStore::upload( "device weight allocation failed"); } const cudaError_t rc = cudaMemcpy(frt_buffer_dptr(buffer), - tensor.values.data(), bytes, + data, bytes, cudaMemcpyHostToDevice); if (rc != cudaSuccess) { return modalities::Status::error( @@ -70,7 +90,7 @@ modalities::Status NativeDeviceWeightStore::upload( std::string("device weight upload failed: ") + cudaGetErrorString(rc)); } - weights_.emplace(name, NativeDeviceWeight{buffer, tensor.shape}); + weights_.emplace(name, NativeDeviceWeight{buffer, shape, dtype}); return modalities::Status::ok(); #endif } diff --git a/cpp/tests/test_pi05_native_device_weights.cpp b/cpp/tests/test_pi05_native_device_weights.cpp index 0b8cc14a..0c884818 100644 --- a/cpp/tests/test_pi05_native_device_weights.cpp +++ b/cpp/tests/test_pi05_native_device_weights.cpp @@ -27,6 +27,7 @@ int main() { } using flashrt::models::pi05::NativeBf16Tensor; using flashrt::models::pi05::NativeDeviceWeightStore; + using flashrt::models::pi05::NativeWeightDType; frt_ctx ctx = frt_ctx_create(); assert(ctx); @@ -47,6 +48,7 @@ int main() { 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)) == @@ -57,6 +59,33 @@ int main() { 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()); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 64914d9d..56cfd5a6 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -227,8 +227,10 @@ 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 -shape/payload mismatches fail setup. Destroying the context releases the device -weights after graphs and plans, preserving the exec ownership order. +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`, From be97bd248d727ccf0b79c0b061ccc5e4db5dc898 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:47:09 -0400 Subject: [PATCH 29/61] feat: add Pi0.5 native weight quantization --- cpp/CMakeLists.txt | 23 +++- .../cpp/models/pi05/native_quantization.h | 38 +++++ cpp/models/pi05/src/native_quantization.cu | 111 +++++++++++++++ .../src/native_quantization_unavailable.cpp | 31 +++++ cpp/tests/gate_pi05_native_quantization.py | 112 +++++++++++++++ cpp/tests/pi05_native_quant_probe.cpp | 130 ++++++++++++++++++ cpp/tests/test_pi05_native_quantization.cpp | 35 +++++ docs/pi05_io_contract.md | 11 +- 8 files changed, 488 insertions(+), 3 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h create mode 100644 cpp/models/pi05/src/native_quantization.cu create mode 100644 cpp/models/pi05/src/native_quantization_unavailable.cpp create mode 100644 cpp/tests/gate_pi05_native_quantization.py create mode 100644 cpp/tests/pi05_native_quant_probe.cpp create mode 100644 cpp/tests/test_pi05_native_quantization.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ee8ee8bd..ce39eafe 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -126,7 +126,7 @@ add_library(flashrt_cpp_loader STATIC target_include_directories(flashrt_cpp_loader PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/loader/include) -add_library(flashrt_cpp_pi05 STATIC +set(FLASHRT_CPP_PI05_SRCS models/pi05/src/spec.cpp models/pi05/src/native_weights.cpp models/pi05/src/native_weight_ops.cpp @@ -136,6 +136,15 @@ add_library(flashrt_cpp_pi05 STATIC 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 @@ -182,6 +191,18 @@ if(BUILD_TESTING) 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) + 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) 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/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/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/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/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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 56cfd5a6..17f4ecdb 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -246,8 +246,15 @@ 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. Precision-specific packing and -graph capture remain incomplete, so `open_v1` stays unsupported. +current Pi0.5 pipeline does not consume them. Full-model precision-store +assembly and graph capture remain incomplete, so `open_v1` stays unsupported. + +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. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 93d3dabe64a40fd6f0a2ceb560819c0327e0b7a6 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:50:34 -0400 Subject: [PATCH 30/61] feat: pack Pi0.5 low-precision weights --- cpp/CMakeLists.txt | 8 ++ .../cpp/models/pi05/native_device_weights.h | 3 + .../cpp/models/pi05/native_weight_packer.h | 30 +++++++ cpp/models/pi05/src/native_device_weights.cpp | 36 +++++++++ cpp/models/pi05/src/native_weight_packer.cpp | 74 +++++++++++++++++ .../test_pi05_native_weight_materializer.cpp | 18 +++++ cpp/tests/test_pi05_native_weight_packer.cpp | 81 +++++++++++++++++++ docs/pi05_io_contract.md | 5 ++ 8 files changed, 255 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h create mode 100644 cpp/models/pi05/src/native_weight_packer.cpp create mode 100644 cpp/tests/test_pi05_native_weight_packer.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ce39eafe..e36666c8 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -131,6 +131,7 @@ set(FLASHRT_CPP_PI05_SRCS 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/prompt_format.cpp models/pi05/src/prompt_embed.cpp @@ -201,6 +202,13 @@ if(BUILD_TESTING) 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) endif() add_executable(test_pi05_prompt_format tests/test_pi05_prompt_format.cpp) 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 index d4963cb1..2cee6402 100644 --- 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 @@ -41,6 +41,9 @@ class NativeDeviceWeightStore { 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(); } 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..6ae21a8a --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h @@ -0,0 +1,30 @@ +#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_int8(const std::string& name); + +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/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp index fdafe1d2..c310b488 100644 --- a/cpp/models/pi05/src/native_device_weights.cpp +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -5,6 +5,7 @@ #endif #include +#include namespace flashrt { namespace models { @@ -101,6 +102,41 @@ const NativeDeviceWeight* NativeDeviceWeightStore::find( 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_weight_packer.cpp b/cpp/models/pi05/src/native_weight_packer.cpp new file mode 100644 index 00000000..4a70c01e --- /dev/null +++ b/cpp/models/pi05/src/native_weight_packer.cpp @@ -0,0 +1,74 @@ +#include "flashrt/cpp/models/pi05/native_weight_packer.h" + +#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) { + NativeFloatTensor source; + modalities::Status st = load_bf16(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." + 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)); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index f369f7eb..4e2060ba 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -1,4 +1,5 @@ #include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include "flashrt/cpp/models/pi05/native_weight_packer.h" #include @@ -227,6 +228,14 @@ int main() { 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); @@ -272,6 +281,15 @@ int main() { 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); } 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..4165a9a6 --- /dev/null +++ b/cpp/tests/test_pi05_native_weight_packer.cpp @@ -0,0 +1,81 @@ +#include "flashrt/cpp/models/pi05/native_weight_packer.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; +} + +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; +} + +} // 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); + std::printf("PASS - Pi0.5 native weight packer\n"); + return 0; +} diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 17f4ecdb..4969dc60 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -256,6 +256,11 @@ 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. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From c4afdbb0d3ec16aacb122eb75a7de0dbb8e9884d Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 16:54:54 -0400 Subject: [PATCH 31/61] feat: assemble complete Pi0.5 BF16 weights --- .../models/pi05/native_weight_materializer.h | 8 ++++++ .../pi05/src/native_weight_materializer.cpp | 25 +++++++++++++++++++ .../test_pi05_native_weight_materializer.cpp | 24 ++++++++++++++++++ docs/pi05_io_contract.md | 7 ++++++ 4 files changed, 64 insertions(+) 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 index fea7902b..ece1efb3 100644 --- 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 @@ -8,6 +8,12 @@ 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, @@ -21,6 +27,8 @@ class NativeWeightMaterializer { 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); diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp index c9a87ae4..251f7420 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -437,6 +437,31 @@ modalities::Status NativeWeightMaterializer::materialize_embedding() { "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/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index 4e2060ba..6bfb8888 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -292,6 +292,30 @@ int main() { 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})); + } + frt_ctx_destroy(full_ctx); + } } std::printf("PASS - Pi0.5 native layer materializer\n"); return 0; diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 4969dc60..64286c5c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -261,6 +261,13 @@ 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. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From e259a2ee6d91a456038be2d913153b612a118164 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:00:26 -0400 Subject: [PATCH 32/61] feat: assemble Pi0.5 precision stores --- .../cpp/models/pi05/native_weight_packer.h | 10 ++ cpp/models/pi05/src/native_weight_packer.cpp | 124 +++++++++++++++++- .../test_pi05_native_weight_materializer.cpp | 18 +++ cpp/tests/test_pi05_native_weight_packer.cpp | 60 +++++++++ docs/pi05_io_contract.md | 6 + 5 files changed, 216 insertions(+), 2 deletions(-) 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 index 6ae21a8a..281885d7 100644 --- 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 @@ -14,7 +14,17 @@ class NativeWeightPacker { : 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, diff --git a/cpp/models/pi05/src/native_weight_packer.cpp b/cpp/models/pi05/src/native_weight_packer.cpp index 4a70c01e..54e3d9b3 100644 --- a/cpp/models/pi05/src/native_weight_packer.cpp +++ b/cpp/models/pi05/src/native_weight_packer.cpp @@ -1,5 +1,6 @@ #include "flashrt/cpp/models/pi05/native_weight_packer.h" +#include #include namespace flashrt { @@ -35,13 +36,20 @@ modalities::Status NativeWeightPacker::load_bf16( 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(name, &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." + name; + const std::string prefix = "fp8." + packed_name; st = weights_->upload_bytes(prefix, packed.shape, NativeWeightDType::kFp8E4M3, packed.values.data(), packed.values.size()); @@ -69,6 +77,118 @@ modalities::Status NativeWeightPacker::pack_int8(const std::string& name) { 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/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index 6bfb8888..b3f22712 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -313,6 +313,24 @@ int main() { 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); } diff --git a/cpp/tests/test_pi05_native_weight_packer.cpp b/cpp/tests/test_pi05_native_weight_packer.cpp index 4165a9a6..83e52180 100644 --- a/cpp/tests/test_pi05_native_weight_packer.cpp +++ b/cpp/tests/test_pi05_native_weight_packer.cpp @@ -4,6 +4,7 @@ #include #include +#include #include namespace { @@ -27,6 +28,12 @@ std::vector download(const flashrt::models::pi05::NativeDeviceWeight& weight) 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() { @@ -76,6 +83,59 @@ int main() { 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 64286c5c..4495b2ba 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -268,6 +268,12 @@ prompt embedding table. With merged decoder gate/up buffers enabled this owns 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. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From fd18aaeaab93a76d98cc0c640ad449f298efeac6 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:04:15 -0400 Subject: [PATCH 33/61] feat: add Pi0.5 native kernel capture driver --- cpp/CMakeLists.txt | 29 +++++ .../cpp/models/pi05/native_kernel_driver.h | 37 ++++++ cpp/models/pi05/src/native_kernel_driver.cu | 72 +++++++++++ cpp/tests/test_pi05_native_kernel_driver.cpp | 114 ++++++++++++++++++ docs/pi05_io_contract.md | 8 ++ 5 files changed, 260 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h create mode 100644 cpp/models/pi05/src/native_kernel_driver.cu create mode 100644 cpp/tests/test_pi05_native_kernel_driver.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index e36666c8..8e5fec9f 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -151,6 +151,24 @@ target_include_directories(flashrt_cpp_pi05 target_link_libraries(flashrt_cpp_pi05 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_kernel_driver.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu) + target_include_directories(flashrt_cpp_pi05_kernels + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm) + target_link_libraries(flashrt_cpp_pi05_kernels + PUBLIC flashrt_cpp_modalities 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>) +endif() + add_library(flashrt_cpp_pi05_c SHARED models/pi05/src/c_api.cpp models/pi05/src/model_runtime.cpp @@ -159,6 +177,10 @@ 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) @@ -209,6 +231,13 @@ if(BUILD_TESTING) 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) endif() add_executable(test_pi05_prompt_format tests/test_pi05_prompt_format.cpp) 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..02131c0d --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h @@ -0,0 +1,37 @@ +#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 + +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; + +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/src/native_kernel_driver.cu b/cpp/models/pi05/src/native_kernel_driver.cu new file mode 100644 index 00000000..11576a1a --- /dev/null +++ b/cpp/models/pi05/src/native_kernel_driver.cu @@ -0,0 +1,72 @@ +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" + +#include "gemm_runner.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); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +} // 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"); + } +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt 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..f8c26c6a --- /dev/null +++ b/cpp/tests/test_pi05_native_kernel_driver.cpp @@ -0,0 +1,114 @@ +#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; + 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); + args->recorded = args->driver + ->bf16_nn(args->a, args->b, args->output, + 2, 2, 3, + reinterpret_cast(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)); + assert(a && b && output); + 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)}; + 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); + + 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(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); + CaptureArgs capture{&driver, frt_buffer_dptr(a), frt_buffer_dptr(b), + frt_buffer_dptr(output), 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[] = {22, 28, 49, 64}; + for (std::size_t i = 0; i < host_output.size(); ++i) { + assert(std::fabs(bfloat16_to_float(host_output[i]) - expected[i]) < + 0.01f); + } + 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 4495b2ba..c39775fb 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -274,6 +274,14 @@ 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. A capture gate warms a BF16 GEMM shape, +records it through `frt_graph_capture`, binds context-owned input/output +buffers, and replays the owned graph exec 100 times without adding variants. +This establishes the 4b kernel/capture ownership path; it does not yet +constitute the complete Pi0.5 forward graph, so the native open gate remains +unsupported. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From dcf83f5103f3336cd7fc2bbee5869f5d466ab24f Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:11:01 -0400 Subject: [PATCH 34/61] feat: add Pi0.5 native workspace --- cpp/CMakeLists.txt | 8 + .../cpp/models/pi05/native_workspace.h | 71 ++++++ cpp/models/pi05/src/native_workspace.cpp | 208 ++++++++++++++++++ cpp/tests/test_pi05_native_workspace.cpp | 94 ++++++++ docs/pi05_io_contract.md | 10 + 5 files changed, 391 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h create mode 100644 cpp/models/pi05/src/native_workspace.cpp create mode 100644 cpp/tests/test_pi05_native_workspace.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 8e5fec9f..c929ea67 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -133,6 +133,7 @@ set(FLASHRT_CPP_PI05_SRCS 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/prompt_format.cpp models/pi05/src/prompt_embed.cpp models/pi05/src/io.cpp @@ -238,6 +239,13 @@ if(BUILD_TESTING) 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_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) endif() add_executable(test_pi05_prompt_format tests/test_pi05_prompt_format.cpp) 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..94458594 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h @@ -0,0 +1,71 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H + +#include "flashrt/cpp/modalities/types.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); + 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(); + + 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; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H diff --git a/cpp/models/pi05/src/native_workspace.cpp b/cpp/models/pi05/src/native_workspace.cpp new file mode 100644 index 00000000..2d1b8688 --- /dev/null +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -0,0 +1,208 @@ +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#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::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; + 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("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 + return initialize_rms_ones(); +} + +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/tests/test_pi05_native_workspace.cpp b/cpp/tests/test_pi05_native_workspace.cpp new file mode 100644 index 00000000..bebd3c7d --- /dev/null +++ b/cpp/tests/test_pi05_native_workspace.cpp @@ -0,0 +1,94 @@ +#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() == 34); + assert(workspace.allocation_count() == 33); + assert(workspace.allocated_bytes() > 0); + assert(workspace.vision_sequence() == 512); + assert(workspace.encoder_vision_sequence() == 512); + assert(workspace.encoder_sequence() == 712); + 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.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() == 34); + assert(workspace.allocation_count() == 34); + 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index c39775fb..5d1a7da5 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -282,6 +282,16 @@ This establishes the 4b kernel/capture ownership path; it does not yet constitute the complete Pi0.5 forward graph, so the native open gate remains unsupported. +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 are initialized during setup. Attention +backend buffers and generated RoPE/style contents are the remaining workspace +subsystems before the complete forward can be captured. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From f9284bc9d3cea271a6c9cd851b13030e0b5ddb53 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:14:03 -0400 Subject: [PATCH 35/61] feat: initialize Pi0.5 native workspace --- cpp/CMakeLists.txt | 5 + .../cpp/models/pi05/native_workspace.h | 9 ++ cpp/models/pi05/src/native_workspace.cpp | 111 +++++++++++++++++- cpp/tests/gate_pi05_native_rope.py | 70 +++++++++++ cpp/tests/pi05_native_rope_probe.cpp | 69 +++++++++++ cpp/tests/test_pi05_native_workspace.cpp | 40 +++++++ docs/pi05_io_contract.md | 11 +- 7 files changed, 312 insertions(+), 3 deletions(-) create mode 100644 cpp/tests/gate_pi05_native_rope.py create mode 100644 cpp/tests/pi05_native_rope_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index c929ea67..8c4b519f 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -246,6 +246,11 @@ if(BUILD_TESTING) PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) add_test(NAME pi05_native_workspace COMMAND test_pi05_native_workspace) + + 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) 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 index 94458594..efd6c67c 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h @@ -2,6 +2,7 @@ #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 @@ -37,6 +38,9 @@ class NativeWorkspace { 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(); } @@ -54,6 +58,7 @@ class NativeWorkspace { 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_; @@ -62,6 +67,10 @@ class NativeWorkspace { 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; + std::vector rope_table_; }; } // namespace pi05 diff --git a/cpp/models/pi05/src/native_workspace.cpp b/cpp/models/pi05/src/native_workspace.cpp index 2d1b8688..753ba714 100644 --- a/cpp/models/pi05/src/native_workspace.cpp +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -5,6 +5,7 @@ #endif #include +#include namespace flashrt { namespace models { @@ -108,6 +109,109 @@ modalities::Status NativeWorkspace::initialize_rms_ones() { #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 + const NativeWorkspaceBuffer* decoder = find("decoder_rope_weights"); + if (!decoder) 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->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 || @@ -121,6 +225,9 @@ modalities::Status NativeWorkspace::allocate( } 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_ = @@ -194,7 +301,9 @@ modalities::Status NativeWorkspace::allocate( FRT_ADD("gate_buf", {ds, 1024}, modalities::DType::kBFloat16); FRT_ADD("decoder_rms_ones", {1024}, modalities::DType::kBFloat16); #undef FRT_ADD - return initialize_rms_ones(); + st = initialize_rms_ones(); + if (!st.ok_status()) return st; + return initialize_rope(); } const NativeWorkspaceBuffer* NativeWorkspace::find( 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/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/test_pi05_native_workspace.cpp b/cpp/tests/test_pi05_native_workspace.cpp index bebd3c7d..28b13323 100644 --- a/cpp/tests/test_pi05_native_workspace.cpp +++ b/cpp/tests/test_pi05_native_workspace.cpp @@ -61,6 +61,46 @@ int main() { 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); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 5d1a7da5..7ddaac1c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -289,8 +289,15 @@ 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 are initialized during setup. Attention -backend buffers and generated RoPE/style contents are the remaining workspace -subsystems before the complete forward can be captured. +backend buffers and generated decoder style contents are the remaining +workspace subsystems before the complete forward can be captured. + +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. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From d71923d3d536f909863a3a15da5c0148d349e406 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:22:23 -0400 Subject: [PATCH 36/61] feat: precompute Pi0.5 decoder styles --- cpp/CMakeLists.txt | 21 +- .../cpp/models/pi05/native_kernel_driver.h | 6 + .../cpp/models/pi05/native_style_precompute.h | 28 +++ cpp/models/pi05/src/native_kernel_driver.cu | 52 ++++- .../pi05/src/native_style_precompute.cu | 197 ++++++++++++++++++ cpp/tests/gate_pi05_native_style.py | 132 ++++++++++++ cpp/tests/pi05_native_style_probe.cpp | 85 ++++++++ cpp/tests/test_pi05_native_kernel_driver.cpp | 43 +++- .../test_pi05_native_style_precompute.cpp | 27 +++ docs/pi05_io_contract.md | 8 + 10 files changed, 586 insertions(+), 13 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_style_precompute.h create mode 100644 cpp/models/pi05/src/native_style_precompute.cu create mode 100644 cpp/tests/gate_pi05_native_style.py create mode 100644 cpp/tests/pi05_native_style_probe.cpp create mode 100644 cpp/tests/test_pi05_native_style_precompute.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 8c4b519f..6c4ba005 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -155,12 +155,15 @@ target_link_libraries(flashrt_cpp_pi05 if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_library(flashrt_cpp_pi05_kernels STATIC models/pi05/src/native_kernel_driver.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu) + models/pi05/src/native_style_precompute.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu) target_include_directories(flashrt_cpp_pi05_kernels PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm) + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels) target_link_libraries(flashrt_cpp_pi05_kernels - PUBLIC flashrt_cpp_modalities CUDA::cublas CUDA::cublasLt CUDA::cudart) + 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) @@ -240,6 +243,18 @@ if(BUILD_TESTING) add_test(NAME pi05_native_kernel_driver COMMAND test_pi05_native_kernel_driver) + 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 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 index 02131c0d..d2daf230 100644 --- 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 @@ -3,6 +3,7 @@ #include "flashrt/cpp/modalities/types.h" +#include #include #include #include @@ -23,6 +24,11 @@ class NativeKernelDriver { 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; private: struct Impl; 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/src/native_kernel_driver.cu b/cpp/models/pi05/src/native_kernel_driver.cu index 11576a1a..c2779782 100644 --- a/cpp/models/pi05/src/native_kernel_driver.cu +++ b/cpp/models/pi05/src/native_kernel_driver.cu @@ -3,15 +3,29 @@ #include "gemm_runner.h" #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); @@ -67,6 +81,42 @@ modalities::Status NativeKernelDriver::bf16_nn( } } +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)); + const cudaError_t rc = cudaGetLastError(); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +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); + const cudaError_t rc = cudaGetLastError(); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + } // 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/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/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/test_pi05_native_kernel_driver.cpp b/cpp/tests/test_pi05_native_kernel_driver.cpp index f8c26c6a..1bc11324 100644 --- a/cpp/tests/test_pi05_native_kernel_driver.cpp +++ b/cpp/tests/test_pi05_native_kernel_driver.cpp @@ -16,6 +16,7 @@ struct CaptureArgs { void* a = nullptr; void* b = nullptr; void* output = nullptr; + const void* bias = nullptr; bool recorded = false; }; @@ -31,11 +32,17 @@ bool has_cuda_device() { void record_gemm(void* user, void* stream) { auto* args = static_cast(user); - args->recorded = args->driver - ->bf16_nn(args->a, args->b, args->output, - 2, 2, 3, - reinterpret_cast(stream)) - .ok_status(); + 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 @@ -56,7 +63,8 @@ int main() { 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)); - assert(a && b && output); + 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)}; @@ -64,12 +72,17 @@ int main() { 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); @@ -77,6 +90,13 @@ int main() { 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); @@ -84,8 +104,9 @@ int main() { 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), false}; + 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); @@ -101,10 +122,14 @@ int main() { assert(cudaMemcpy(host_output.data(), frt_buffer_dptr(output), host_output.size() * sizeof(std::uint16_t), cudaMemcpyDeviceToHost) == cudaSuccess); - const float expected[] = {22, 28, 49, 64}; + 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.01f); + 0.02f); } frt_graph_destroy(graph); assert(cudaStreamDestroy(stream) == cudaSuccess); 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 7ddaac1c..02def20e 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -299,6 +299,14 @@ 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. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From ed0d905639355f6cd21ffd83f94c6782f978db99 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:26:46 -0400 Subject: [PATCH 37/61] feat: own Pi0.5 RTX attention buffers --- cpp/CMakeLists.txt | 8 + .../cpp/models/pi05/native_rtx_attention.h | 75 +++++++ cpp/models/pi05/src/native_rtx_attention.cpp | 201 ++++++++++++++++++ cpp/tests/test_pi05_native_rtx_attention.cpp | 73 +++++++ docs/pi05_io_contract.md | 8 + 5 files changed, 365 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h create mode 100644 cpp/models/pi05/src/native_rtx_attention.cpp create mode 100644 cpp/tests/test_pi05_native_rtx_attention.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 6c4ba005..205d0c5f 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -134,6 +134,7 @@ set(FLASHRT_CPP_PI05_SRCS 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 @@ -262,6 +263,13 @@ if(BUILD_TESTING) 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) + add_executable(pi05_native_rope_probe tests/pi05_native_rope_probe.cpp) target_link_libraries(pi05_native_rope_probe 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..15908119 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h @@ -0,0 +1,75 @@ +#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; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H 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..657a11a2 --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_attention.cpp @@ -0,0 +1,201 @@ +#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 + 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}; + const char* names[] = {"attn_enc_seqused", "attn_dec_seqused", + "attn_dec_devpos"}; + for (int i = 0; i < 3; ++i) { + const NativeAttentionBuffer* target = find(names[i]); + if (!target || + cudaMemcpy(frt_buffer_dptr(target->buffer), &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/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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 02def20e..0efa6a03 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -307,6 +307,14 @@ 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. +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 remaining attention task is wiring these buffers +to the vendored FA2 C++ entry points and validating captured outputs. + CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that loads assets and captures graphs in the replay process. From ad1c7984777c9396509f6b0b416751e54ac842bf Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:43:21 -0400 Subject: [PATCH 38/61] feat: add native Pi0.5 FA2 driver --- CMakeLists.txt | 39 +++- cpp/CMakeLists.txt | 42 ++++ .../models/pi05/native_rtx_attention_driver.h | 44 ++++ .../pi05/src/native_rtx_attention_driver.cu | 208 ++++++++++++++++++ .../test_pi05_native_rtx_attention_driver.cpp | 157 +++++++++++++ csrc/attention/fa2_wrapper.cu | 1 + csrc/attention/fa2_wrapper.h | 73 ++++++ csrc/attention/fa2_wrapper_causal.cu | 72 +++--- csrc/fa2_bindings.cpp | 69 +----- docs/pi05_io_contract.md | 9 +- 10 files changed, 603 insertions(+), 111 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h create mode 100644 cpp/models/pi05/src/native_rtx_attention_driver.cu create mode 100644 cpp/tests/test_pi05_native_rtx_attention_driver.cpp create mode 100644 csrc/attention/fa2_wrapper.h 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/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 205d0c5f..10d936e7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -23,6 +23,9 @@ 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() @@ -30,6 +33,26 @@ 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) @@ -172,6 +195,16 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) $<$: --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_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 @@ -270,6 +303,15 @@ if(BUILD_TESTING) add_test(NAME pi05_native_rtx_attention COMMAND test_pi05_native_rtx_attention) + if(FLASHRT_CPP_WITH_FA2) + 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() + add_executable(pi05_native_rope_probe tests/pi05_native_rope_probe.cpp) target_link_libraries(pi05_native_rope_probe 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/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/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/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/pi05_io_contract.md b/docs/pi05_io_contract.md index 0efa6a03..2f4b3a88 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -312,8 +312,13 @@ 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 remaining attention task is wiring these buffers -to the vendored FA2 C++ entry points and validating captured outputs. +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 remaining native producer task is +assembling these primitives into the complete model forward and capture. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 7ae1fedd092cf5c8cca38b960e53648ecfd12089 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:51:41 -0400 Subject: [PATCH 39/61] feat: expose Pi0.5 native forward primitives --- cpp/CMakeLists.txt | 14 +- .../cpp/models/pi05/native_kernel_driver.h | 43 +++ cpp/models/pi05/src/native_kernel_driver.cu | 239 +++++++++++- .../test_pi05_native_forward_primitives.cpp | 348 ++++++++++++++++++ docs/pi05_io_contract.md | 7 + 5 files changed, 642 insertions(+), 9 deletions(-) create mode 100644 cpp/tests/test_pi05_native_forward_primitives.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 10d936e7..cdd0c099 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -181,7 +181,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) 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/dit_bf16.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 @@ -277,6 +282,13 @@ if(BUILD_TESTING) 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_style_precompute tests/test_pi05_native_style_precompute.cpp) target_link_libraries(test_pi05_native_style_precompute 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 index d2daf230..c6678413 100644 --- 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 @@ -29,6 +29,49 @@ class NativeKernelDriver { 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; diff --git a/cpp/models/pi05/src/native_kernel_driver.cu b/cpp/models/pi05/src/native_kernel_driver.cu index c2779782..cff13040 100644 --- a/cpp/models/pi05/src/native_kernel_driver.cu +++ b/cpp/models/pi05/src/native_kernel_driver.cu @@ -1,10 +1,16 @@ #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, @@ -36,6 +42,12 @@ modalities::Status backend(const std::string& message) { message); } +modalities::Status launch_status() { + const cudaError_t rc = cudaGetLastError(); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + } // namespace struct NativeKernelDriver::Impl { @@ -94,10 +106,7 @@ modalities::Status NativeKernelDriver::add_bias_bf16( ::add_bias_bf16(static_cast<__nv_bfloat16*>(values), static_cast(bias), rows, columns, reinterpret_cast(stream)); - const cudaError_t rc = cudaGetLastError(); - return rc == cudaSuccess - ? modalities::Status::ok() - : backend(cudaGetErrorString(rc)); + return launch_status(); } modalities::Status NativeKernelDriver::silu_bf16( @@ -111,10 +120,224 @@ modalities::Status NativeKernelDriver::silu_bf16( native_silu_bf16_kernel<<<(elements + 255) / 256, 256, 0, reinterpret_cast(stream)>>>( static_cast<__nv_bfloat16*>(values), elements); - const cudaError_t rc = cudaGetLastError(); - return rc == cudaSuccess - ? modalities::Status::ok() - : backend(cudaGetErrorString(rc)); + 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 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 2f4b3a88..f60dbbe0 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -307,6 +307,13 @@ 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. + 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 From 66211074024d8b8ffb818d666a0c555bbab2e996 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 17:59:28 -0400 Subject: [PATCH 40/61] feat: compose Pi0.5 native encoder QKV --- cpp/CMakeLists.txt | 13 ++ .../cpp/models/pi05/native_bf16_forward.h | 30 +++ cpp/models/pi05/src/native_bf16_forward.cpp | 91 +++++++++ cpp/tests/gate_pi05_native_encoder_qkv.py | 100 +++++++++ cpp/tests/pi05_native_encoder_qkv_probe.cpp | 99 +++++++++ cpp/tests/test_pi05_native_bf16_forward.cpp | 192 ++++++++++++++++++ docs/pi05_io_contract.md | 8 + 7 files changed, 533 insertions(+) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h create mode 100644 cpp/models/pi05/src/native_bf16_forward.cpp create mode 100644 cpp/tests/gate_pi05_native_encoder_qkv.py create mode 100644 cpp/tests/pi05_native_encoder_qkv_probe.cpp create mode 100644 cpp/tests/test_pi05_native_bf16_forward.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index cdd0c099..62944804 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -178,6 +178,7 @@ target_link_libraries(flashrt_cpp_pi05 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 @@ -289,6 +290,18 @@ if(BUILD_TESTING) 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 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..1f895934 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h @@ -0,0 +1,30 @@ +#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_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; + +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/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp new file mode 100644 index 00000000..2fc44b49 --- /dev/null +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -0,0 +1,91 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" + +#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); +} + +} // 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); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt 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/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/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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index f60dbbe0..c6eee3c2 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -314,6 +314,14 @@ 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 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. + 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 From ca60085211da6ec304dfd6e670b3e60dff42be95 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 18:06:43 -0400 Subject: [PATCH 41/61] feat: compose Pi0.5 native encoder layer --- cpp/CMakeLists.txt | 5 + .../cpp/models/pi05/native_bf16_forward.h | 8 ++ cpp/models/pi05/src/native_bf16_forward.cpp | 86 +++++++++++ cpp/tests/gate_pi05_native_encoder_layer.py | 125 ++++++++++++++++ cpp/tests/pi05_native_encoder_layer_probe.cpp | 133 ++++++++++++++++++ docs/pi05_io_contract.md | 7 + 6 files changed, 364 insertions(+) create mode 100644 cpp/tests/gate_pi05_native_encoder_layer.py create mode 100644 cpp/tests/pi05_native_encoder_layer_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 62944804..1720cef4 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -329,6 +329,11 @@ if(BUILD_TESTING) 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(test_pi05_native_rtx_attention_driver tests/test_pi05_native_rtx_attention_driver.cpp) target_link_libraries(test_pi05_native_rtx_attention_driver 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 index 1f895934..f50ab65f 100644 --- 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 @@ -3,6 +3,7 @@ #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 { @@ -18,6 +19,13 @@ class NativeBf16Forward { int layer, const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, std::uintptr_t stream) const; +#ifdef FLASHRT_CPP_WITH_FA2 + modalities::Status encoder_layer( + int layer, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; +#endif private: const NativeKernelDriver* driver_ = nullptr; diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp index 2fc44b49..202c79db 100644 --- a/cpp/models/pi05/src/native_bf16_forward.cpp +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -25,6 +25,14 @@ bool shape_is(const NativeAttentionBuffer* buffer, 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( @@ -86,6 +94,84 @@ modalities::Status NativeBf16Forward::encoder_qkv( 256, stream); } +#ifdef FLASHRT_CPP_WITH_FA2 +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); +} +#endif + } // namespace pi05 } // namespace models } // namespace flashrt 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/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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index c6eee3c2..50c42066 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -322,6 +322,13 @@ 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. + 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 From 4cdc69a380e16dae7767fd72aee8b2222733a8ec Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 18:23:25 -0400 Subject: [PATCH 42/61] feat: compose Pi0.5 native encoder --- cpp/CMakeLists.txt | 5 + .../cpp/models/pi05/native_bf16_forward.h | 5 + cpp/models/pi05/src/native_bf16_forward.cpp | 14 ++ cpp/tests/gate_pi05_native_encoder.py | 184 ++++++++++++++++++ cpp/tests/pi05_native_encoder_probe.cpp | 143 ++++++++++++++ docs/pi05_io_contract.md | 9 +- 6 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 cpp/tests/gate_pi05_native_encoder.py create mode 100644 cpp/tests/pi05_native_encoder_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 1720cef4..5ae0859e 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -334,6 +334,11 @@ if(BUILD_TESTING) 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(test_pi05_native_rtx_attention_driver tests/test_pi05_native_rtx_attention_driver.cpp) target_link_libraries(test_pi05_native_rtx_attention_driver 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 index f50ab65f..3be95a3c 100644 --- 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 @@ -25,6 +25,11 @@ class NativeBf16Forward { 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; #endif private: diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp index 202c79db..fe66dde7 100644 --- a/cpp/models/pi05/src/native_bf16_forward.cpp +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -170,6 +170,20 @@ modalities::Status NativeBf16Forward::encoder_layer( 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(); +} #endif } // namespace pi05 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/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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 50c42066..18a8bce5 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -329,6 +329,13 @@ 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. + 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 @@ -340,7 +347,7 @@ 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 remaining native producer task is -assembling these primitives into the complete model forward and capture. +assembling vision, decoder, and diffusion around the completed encoder graph. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 8f4b4c276c2b58aa3379276505a834be1947958e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 18:30:12 -0400 Subject: [PATCH 43/61] feat: compose Pi0.5 native vision --- cpp/CMakeLists.txt | 5 + .../cpp/models/pi05/native_bf16_forward.h | 10 + cpp/models/pi05/src/native_bf16_forward.cpp | 237 ++++++++++++++++++ cpp/tests/gate_pi05_native_vision.py | 190 ++++++++++++++ cpp/tests/pi05_native_vision_probe.cpp | 146 +++++++++++ docs/pi05_io_contract.md | 11 +- 6 files changed, 598 insertions(+), 1 deletion(-) create mode 100644 cpp/tests/gate_pi05_native_vision.py create mode 100644 cpp/tests/pi05_native_vision_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5ae0859e..7fd3eb7c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -339,6 +339,11 @@ if(BUILD_TESTING) 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(test_pi05_native_rtx_attention_driver tests/test_pi05_native_rtx_attention_driver.cpp) target_link_libraries(test_pi05_native_rtx_attention_driver 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 index 3be95a3c..8b7368fd 100644 --- 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 @@ -20,6 +20,16 @@ class NativeBf16Forward { 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, diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp index fe66dde7..1f1969e4 100644 --- a/cpp/models/pi05/src/native_bf16_forward.cpp +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -95,6 +95,243 @@ modalities::Status NativeBf16Forward::encoder_qkv( } #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, 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/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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 18a8bce5..e9272947 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -314,6 +314,14 @@ 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 @@ -347,7 +355,8 @@ 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 remaining native producer task is -assembling vision, decoder, and diffusion around the completed encoder graph. +assembling decoder and diffusion around the completed vision and encoder +graphs. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 85b76add7698b74769e982739187907bcfe32623 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 18:43:25 -0400 Subject: [PATCH 44/61] feat: compose Pi0.5 native diffusion --- cpp/CMakeLists.txt | 5 + .../cpp/models/pi05/native_bf16_forward.h | 15 + cpp/models/pi05/src/native_bf16_forward.cpp | 247 ++++++++++++++++ cpp/tests/gate_pi05_native_diffusion.py | 273 ++++++++++++++++++ cpp/tests/pi05_native_diffusion_probe.cpp | 195 +++++++++++++ docs/pi05_io_contract.md | 13 +- 6 files changed, 746 insertions(+), 2 deletions(-) create mode 100644 cpp/tests/gate_pi05_native_diffusion.py create mode 100644 cpp/tests/pi05_native_diffusion_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7fd3eb7c..fb29bf45 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -344,6 +344,11 @@ if(BUILD_TESTING) 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(test_pi05_native_rtx_attention_driver tests/test_pi05_native_rtx_attention_driver.cpp) target_link_libraries(test_pi05_native_rtx_attention_driver 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 index 8b7368fd..16f4b610 100644 --- 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 @@ -40,6 +40,21 @@ class NativeBf16Forward { 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: diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp index 1f1969e4..10d0e9a1 100644 --- a/cpp/models/pi05/src/native_bf16_forward.cpp +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -1,5 +1,6 @@ #include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include #include #include @@ -421,6 +422,252 @@ modalities::Status NativeBf16Forward::encoder( } 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 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/pi05_native_diffusion_probe.cpp b/cpp/tests/pi05_native_diffusion_probe.cpp new file mode 100644 index 00000000..02556842 --- /dev/null +++ b/cpp/tests/pi05_native_diffusion_probe.cpp @@ -0,0 +1,195 @@ +#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_style_precompute.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; + 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index e9272947..e368e0de 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -344,6 +344,15 @@ 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. + 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 @@ -355,8 +364,8 @@ 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 remaining native producer task is -assembling decoder and diffusion around the completed vision and encoder -graphs. +combining the completed vision, encoder, and diffusion graphs with the native +builder and producer lifetime. CUDA graph execs are process-local objects. They are not serialized as a portable artifact. Removing Python from setup requires a native producer that From 6fbd85eab34f03a5398a24e4d5fc2311cfe542dc Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 18:53:14 -0400 Subject: [PATCH 45/61] feat: capture Pi0.5 native full graph --- cpp/CMakeLists.txt | 6 + .../cpp/models/pi05/native_graph_owner.h | 72 ++++++ cpp/models/pi05/src/native_graph_owner.cpp | 239 ++++++++++++++++++ cpp/models/pi05/src/native_workspace.cpp | 3 + cpp/tests/pi05_native_graph_probe.cpp | 103 ++++++++ cpp/tests/test_pi05_native_workspace.cpp | 10 +- docs/pi05_io_contract.md | 9 + 7 files changed, 438 insertions(+), 4 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h create mode 100644 cpp/models/pi05/src/native_graph_owner.cpp create mode 100644 cpp/tests/pi05_native_graph_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index fb29bf45..a4fe77b5 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -203,6 +203,7 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) --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) @@ -349,6 +350,11 @@ if(BUILD_TESTING) 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) + 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 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/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_workspace.cpp b/cpp/models/pi05/src/native_workspace.cpp index 753ba714..8638e56f 100644 --- a/cpp/models/pi05/src/native_workspace.cpp +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -264,6 +264,9 @@ modalities::Status NativeWorkspace::allocate( 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); 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/test_pi05_native_workspace.cpp b/cpp/tests/test_pi05_native_workspace.cpp index 28b13323..78c59586 100644 --- a/cpp/tests/test_pi05_native_workspace.cpp +++ b/cpp/tests/test_pi05_native_workspace.cpp @@ -45,12 +45,14 @@ int main() { assert(!workspace.allocate(invalid).ok_status()); NativeWorkspaceConfig config; assert(workspace.allocate(config).ok_status()); - assert(workspace.logical_size() == 34); - assert(workspace.allocation_count() == 33); + 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); @@ -116,8 +118,8 @@ int main() { config.num_steps = 5; config.vision_pool_factor = 2; assert(workspace.allocate(config).ok_status()); - assert(workspace.logical_size() == 34); - assert(workspace.allocation_count() == 34); + 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); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index e368e0de..9019d87c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -353,6 +353,15 @@ 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 From a37307e2c35d140369788039681b9ab7d95b8463 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 19:27:54 -0400 Subject: [PATCH 46/61] feat: open Pi0.5 native runtime --- cpp/CMakeLists.txt | 20 +- .../include/flashrt/cpp/loader/sha256.h | 15 + cpp/loader/src/sha256.cpp | 175 ++++++++++ .../include/flashrt/cpp/models/pi05/c_api.h | 6 + .../include/flashrt/cpp/models/pi05/runtime.h | 3 + cpp/models/pi05/src/config_map.h | 23 +- cpp/models/pi05/src/model_runtime.cpp | 20 +- cpp/models/pi05/src/native_device_weights.cpp | 10 +- cpp/models/pi05/src/native_model_runtime.cpp | 326 ++++++++++++++++++ cpp/models/pi05/src/native_open.cpp | 115 ++++-- cpp/models/pi05/src/native_open_internal.h | 36 ++ cpp/models/pi05/src/runtime.cpp | 9 + cpp/tests/pi05_native_open_probe.cpp | 149 ++++++++ cpp/tests/test_pi05_native_open.cpp | 4 + cpp/tests/test_sha256.cpp | 30 ++ docs/mindon_pi05_integration.md | 57 +-- docs/pi05_io_contract.md | 90 +++-- flash_rt/models/pi05/runtime_export.py | 8 - 18 files changed, 995 insertions(+), 101 deletions(-) create mode 100644 cpp/loader/include/flashrt/cpp/loader/sha256.h create mode 100644 cpp/loader/src/sha256.cpp create mode 100644 cpp/models/pi05/src/native_model_runtime.cpp create mode 100644 cpp/models/pi05/src/native_open_internal.h create mode 100644 cpp/tests/pi05_native_open_probe.cpp create mode 100644 cpp/tests/test_sha256.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a4fe77b5..5558a2a7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -145,7 +145,9 @@ target_link_libraries(flashrt_cpp_vla INTERFACE flashrt_cpp_modalities flashrt_cpp) add_library(flashrt_cpp_loader STATIC - loader/src/safetensors.cpp) + 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) @@ -217,6 +219,7 @@ endif() add_library(flashrt_cpp_pi05_c SHARED models/pi05/src/c_api.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) @@ -232,6 +235,10 @@ if(BUILD_TESTING) 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) @@ -355,6 +362,13 @@ if(BUILD_TESTING) 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) + 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 @@ -410,6 +424,10 @@ if(BUILD_TESTING) 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/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/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/models/pi05/include/flashrt/cpp/models/pi05/c_api.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/c_api.h index bce5c057..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 @@ -80,6 +80,12 @@ typedef struct frt_pi05_runtime_config { 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 { 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 cb2e2b65..61f391ce 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -14,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; @@ -58,6 +59,8 @@ struct RuntimeConfig { 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 { diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h index 1b5222e9..4e933777 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/config_map.h @@ -160,14 +160,33 @@ inline RuntimeConfig make_config(const frt_pi05_runtime_config* in) { 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 = modalities::MemoryPlace::kHost; + 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 = modalities::MemoryPlace::kHost; + 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}; diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index 99fe7d48..bd3e0063 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -258,11 +258,23 @@ int get_output(void* self, uint32_t port, void* out, uint64_t capacity, } 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) { diff --git a/cpp/models/pi05/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp index c310b488..a6289b55 100644 --- a/cpp/models/pi05/src/native_device_weights.cpp +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -5,6 +5,7 @@ #endif #include +#include #include namespace flashrt { @@ -79,8 +80,15 @@ modalities::Status NativeDeviceWeightStore::upload_bytes( #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, - "device weight allocation failed"); + message.str()); } const cudaError_t rc = cudaMemcpy(frt_buffer_dptr(buffer), data, bytes, 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..508629bf --- /dev/null +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -0,0 +1,326 @@ +#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 fail_builder(frt_runtime_builder builder, std::string* error, + const char* message) { + frt_model_runtime_v1* discarded = frt_runtime_builder_finish_model( + builder, nullptr, 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; + } + + 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", "sm120") && + 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\":\"sm120\",\"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())}; + 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_BF16, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + action_shape, 2, 0, noise->buffer, 0, + frt_buffer_bytes(noise->buffer)) == 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(); + frt_model_runtime_v1* base = frt_runtime_builder_finish_model( + builder, nullptr, 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_mean[i] = (config.action_q01[i] + config.action_q99[i]) * 0.5f; + action_stddev[i] = + (config.action_q99[i] - config.action_q01[i]) * 0.5f; + } + 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 index 2b56da38..fb078cf2 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -1,9 +1,11 @@ #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 @@ -352,9 +354,10 @@ std::string dirname(const std::string& path) { return path.substr(0, p); } -bool norm_block_dims(const std::string& json, - const char* block_name, - size_t* dims) { +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; @@ -364,29 +367,39 @@ bool norm_block_dims(const std::string& json, !sane_quantile_pair(q01, q99)) { return false; } - if (dims) *dims = q01.size(); + 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) { + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* config) { std::string json; if (!read_text_file(path, &json)) return false; - size_t action_dims = 0; - size_t state_dims = 0; - if (!norm_block_dims(json, "actions", &action_dims) || - !norm_block_dims(json, "state", &state_dims)) { + 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_dims == 0 || action_dims > 32) { + if (action_q01.empty() || action_q01.size() > 32) { g_last_error = "norm_stats action dimension is invalid"; return false; } - if (state_dims != static_cast(state_dim)) { + 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; } @@ -419,7 +432,9 @@ std::string find_child(const std::string& dir, } bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, - int64_t state_dim) { + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* + config) { const std::string pre = find_child( checkpoint_path, "policy_preprocessor_step_", "_normalizer_processor.safetensors"); @@ -453,12 +468,19 @@ bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, 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) { + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* config) { const std::string parent = dirname(checkpoint_path); const std::string candidates[] = { join_path(checkpoint_path, @@ -475,11 +497,12 @@ bool validate_norm_stats(const std::string& checkpoint_path, std::string malformed_error; for (const std::string& path : candidates) { if (!regular_file_exists(path)) continue; - if (validate_norm_stats_file(path, state_dim)) return true; + 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)) { + if (validate_lerobot_policy_norm_stats(checkpoint_path, state_dim, + config)) { return true; } g_last_error = saw_malformed @@ -561,7 +584,9 @@ bool integer_field(const std::map& obj, return true; } -int validate_config(const char* config_json) { +int validate_config( + const char* config_json, + flashrt::models::pi05::NativeOpenConfig* parsed) { if (!config_json) { g_last_error = "config_json is null"; return -1; @@ -609,36 +634,57 @@ int validate_config(const char* config_json) { 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, "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) { - g_last_error = "max_prompt_tokens must be at least 200"; + 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) { - g_last_error = "state_dim must be positive"; + 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) { - g_last_error = "chunk must be positive"; + 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 (!validate_norm_stats(checkpoint_path, state_dim)) { + 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; } - g_last_error = - "Pi0.5 native open validated config; native graph capture is not " - "implemented yet"; - return -3; + if (parsed) *parsed = std::move(config); + g_last_error.clear(); + return 0; } } // namespace @@ -650,7 +696,18 @@ extern "C" int frt_model_runtime_open_v1(const char* config_json, return -1; } *out = nullptr; - return validate_config(config_json); + 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() { 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/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index 6fd94c20..2bb04f82 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -239,6 +239,15 @@ int Runtime::set_prompt_state(const char* text, const float* state, ? &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; } diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp new file mode 100644 index 00000000..0f86a655 --- /dev/null +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -0,0 +1,149 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/modalities/types.h" + +#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; + } + 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; + 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") && + 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[5].update == FRT_RT_PORT_SWAP; + if (!ok) { + std::cerr << "native schema validation failed\n"; + model->release(model->owner); + return 1; + } + 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"; + const 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; + } + 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); + } + 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 || + model->verbs.step(model->self) != 0) { + std::cerr << "native step failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + float actions[10 * 7]{}; + std::uint64_t written = 0; + 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/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index e27f1399..655d94f1 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -248,12 +248,14 @@ int main() { 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("{") + @@ -278,11 +280,13 @@ int main() { 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 + 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/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index d756bf4e..100bba5b 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -52,7 +52,7 @@ 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: After Prompt/State Staging +### Lane B: Adopted Prompt/State Staging Use the same setup/adopt path as Lane A, but the producer additionally exports real hot ports: @@ -64,7 +64,7 @@ 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: Future Native Producer +### Lane C: Native SM120 Producer Load a native FlashRT shared object and call: @@ -77,19 +77,14 @@ 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 exports this symbol as a native-v2 configuration -gate. It validates `io`, checkpoint path, tokenizer model path, fixed prompt -mode, prompt capacity, state dimension, and the Pi0.5 embedding metadata in -`model.safetensors`, plus the native builder's complete 812-tensor weight -inventory across vision, language encoder, and action expert. A -read-only mmap loader owns the checkpoint mapping and exposes bounded tensor -views for the later device materialization step. -It also verifies action/state q01/q99 metadata from openpi `norm_stats.json` or -LeRobot policy normalizer/unnormalizer safetensors, including tensor byte -ranges and finite ordered quantile payloads, then returns unsupported until the -native checkpoint materialization and CUDA graph capture path are implemented. -Hosts may use this to wire dynamic loading and error handling, but must keep -using Lane A or B for execution. +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 @@ -155,10 +150,9 @@ raw proprioception -> normalize -> 256-bin discretize -> prompt string -> token ids -> embedding gather -> encoder_x prompt window ``` -Until Lane B lands, prompt and state changes require a setup-time producer -refresh. After Lane B, Mindon should send task text to the `prompt` port and -raw proprioception to the `state` port. The producer owns all formatting and -normalization details. +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 @@ -186,9 +180,11 @@ 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. -If prompt context becomes a region in a future producer face, old capsules -will not restore into that new face. That is expected and protects state -safety. +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 @@ -205,6 +201,23 @@ model = pipeline.export_model_runtime( 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`. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 9019d87c..bd86fabc 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -184,11 +184,13 @@ The following are not deployment identity changes: - editing `manifest_json`; - changing `cadence_hint_hz`. -Prompt/state staging should normally add a restorable prompt context region -only after the bytes that define that context are explicit. A valid region -could include the language rows of `encoder_x` plus the fixed-prompt valid -length scalar. Region layout and order are fingerprinted; old capsules should -not restore into the new layout. +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 @@ -196,26 +198,28 @@ 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, after prompt/state staging: same as Lane A, plus hot - `prompt`/`state` STAGED ports. -- Lane C, future native producer: a C++ shared object implements +- 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 now exports `frt_model_runtime_open_v1` as a -native-v2 configuration gate. The gate requires `io="native_v2"`, +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`. It also parses +`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. Valid -configuration returns unsupported until device weight materialization and graph -capture are complete. The mmap and parsed tensor views are setup-side assets; -they never enter the model-runtime ABI or the hot path. +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, @@ -246,8 +250,8 @@ 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. Full-model precision-store -assembly and graph capture remain incomplete, so `open_v1` stays unsupported. +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 @@ -275,12 +279,9 @@ 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. A capture gate warms a BF16 GEMM shape, -records it through `frt_graph_capture`, binds context-owned input/output -buffers, and replays the owned graph exec 100 times without adding variants. -This establishes the 4b kernel/capture ownership path; it does not yet -constitute the complete Pi0.5 forward graph, so the native open gate remains -unsupported. +`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 @@ -288,9 +289,8 @@ 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 are initialized during setup. Attention -backend buffers and generated decoder style contents are the remaining -workspace subsystems before the complete forward can be captured. +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 @@ -372,13 +372,24 @@ 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 remaining native producer task is -combining the completed vision, encoder, and diffusion graphs with the native -builder and producer lifetime. +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. Removing Python from setup requires a native producer that -loads assets and captures graphs in the replay process. +portable artifact. The native producer therefore loads assets and captures the +graph in the replay process. ## Validation @@ -404,6 +415,17 @@ python cpp/tests/gate_pi05_model_runtime_export.py ... python cpp/tests/gate_pi05_c_api_export.py ... ``` -For prompt/state staging, add token-exact, formatter string-exact, embedding -bit-exact, fixed-vs-exact E2E cosine, and hot-contract tests before declaring -the new STAGED ports. +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: + +``` +cpp/build-sm120-spm-debug/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. diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index f11073df..18c33e54 100644 --- a/flash_rt/models/pi05/runtime_export.py +++ b/flash_rt/models/pi05/runtime_export.py @@ -365,14 +365,6 @@ def _parts(pl, identity, extra_regions): ident["state_prompt_mode"] = "fixed" ident["state_dim"] = str(_state_dim(pl)) ident["tokenizer_sha256"] = _tokenizer_sha256() - prompt_bytes = int(getattr(pl, "max_prompt_len", 0) or 0) * 2048 * 2 - prompt_offset = int(getattr(pl, "num_views", 0) or 0) * 256 * 2048 * 2 - encoder_x = wrap["encoder_x"] - if prompt_bytes <= 0 or prompt_offset + prompt_bytes > encoder_x.nbytes(): - raise ValueError("Pi05 native_v2 prompt_context window is invalid") - regions.append(_rt.RegionSpec( - "prompt_context", encoder_x, offset=prompt_offset, - nbytes=prompt_bytes)) return { "wrap": wrap, From 5e6b9cdb5e9a27a10a9014d0522bbd1698e59400 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 19:50:39 -0400 Subject: [PATCH 47/61] test: gate Pi0.5 native E2E --- cpp/CMakeLists.txt | 5 + cpp/models/pi05/src/native_model_runtime.cpp | 4 +- cpp/tests/gate_pi05_native_e2e.py | 275 +++++++++++++++++++ cpp/tests/pi05_native_e2e_probe.cpp | 168 +++++++++++ docs/pi05_io_contract.md | 20 ++ 5 files changed, 470 insertions(+), 2 deletions(-) create mode 100644 cpp/tests/gate_pi05_native_e2e.py create mode 100644 cpp/tests/pi05_native_e2e_probe.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5558a2a7..3c2afe11 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -367,6 +367,11 @@ if(BUILD_TESTING) 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) endif() add_executable(test_pi05_native_rtx_attention_driver diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 508629bf..7d5ba080 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -247,9 +247,9 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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_mean[i] = (config.action_q01[i] + config.action_q99[i]) * 0.5f; action_stddev[i] = - (config.action_q99[i] - config.action_q01[i]) * 0.5f; + (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); 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/pi05_native_e2e_probe.cpp b/cpp/tests/pi05_native_e2e_probe.cpp new file mode 100644 index 00000000..bfa99038 --- /dev/null +++ b/cpp/tests/pi05_native_e2e_probe.cpp @@ -0,0 +1,168 @@ +#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"); + } + } + + 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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index bd86fabc..8e1743f8 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -429,3 +429,23 @@ cpp/build-sm120-spm-debug/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. + +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 cpp/build-sm120-spm-debug/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. From 1d02f9ca4d13ebfd154b69a029cf268e9145f5aa Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 20:40:44 -0400 Subject: [PATCH 48/61] test: close Pi0.5 native validation gates --- cpp/CMakeLists.txt | 15 +++ cpp/tests/gate_pi05_kernel_sequence.py | 119 ++++++++++++++++++ cpp/tests/gate_pi05_tokenizer_corpus.py | 140 ++++++++++++++++++++++ cpp/tests/pi05_native_dlopen_probe.cpp | 82 +++++++++++++ cpp/tests/pi05_native_open_probe.cpp | 15 ++- cpp/tests/pi05_tokenizer_corpus_probe.cpp | 104 ++++++++++++++++ cpp/tests/profile_pi05_python_replay.py | 99 +++++++++++++++ docs/pi05_io_contract.md | 71 +++++++++++ 8 files changed, 644 insertions(+), 1 deletion(-) create mode 100644 cpp/tests/gate_pi05_kernel_sequence.py create mode 100644 cpp/tests/gate_pi05_tokenizer_corpus.py create mode 100644 cpp/tests/pi05_native_dlopen_probe.cpp create mode 100644 cpp/tests/pi05_tokenizer_corpus_probe.cpp create mode 100644 cpp/tests/profile_pi05_python_replay.py diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 3c2afe11..c57f00e0 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -372,6 +372,14 @@ if(BUILD_TESTING) 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 @@ -382,6 +390,13 @@ if(BUILD_TESTING) 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 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_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 { + +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_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index 0f86a655..f92b36ed 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -1,9 +1,11 @@ #include "flashrt/model_runtime.h" #include "flashrt/cpp/modalities/types.h" +#include #include #include +#include #include #include #include @@ -113,12 +115,23 @@ int main(int argc, char** argv) { host_noise[i] = flashrt::modalities::float_to_bfloat16( static_cast(static_cast(i % 23) - 11) / 12.0f); } + const bool profile_range = std::getenv("FLASHRT_PROFILE_RANGE") != nullptr; 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 || - model->verbs.step(model->self) != 0) { + (profile_range && cudaProfilerStart() != cudaSuccess)) { + std::cerr << "native step failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + const int step_rc = model->verbs.step(model->self); + const cudaError_t sync_rc = cudaDeviceSynchronize(); + const cudaError_t profiler_rc = + profile_range ? cudaProfilerStop() : cudaSuccess; + if (step_rc != 0 || sync_rc != cudaSuccess || profiler_rc != cudaSuccess) { std::cerr << "native step failed: " << model->verbs.last_error(model->self) << '\n'; model->release(model->owner); 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..4a6eedce --- /dev/null +++ b/cpp/tests/profile_pi05_python_replay.py @@ -0,0 +1,99 @@ +"""Emit one replay-only CUDA profiler range for the Pi0.5 Python frontend.""" + +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/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 8e1743f8..d518b3ff 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -430,6 +430,23 @@ 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. +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 cpp/build-sm120-spm-debug/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 reference SM120 run covered +10,000 records and 20 token lengths from 43 through 62 with zero mismatches. + The real-episode numerical gate compares against the official OpenPI PyTorch `PI0Pytorch.sample_actions` path, not another native intermediate: @@ -449,3 +466,57 @@ 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 \ + cpp/build-sm120-spm-debug/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 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: + +``` +cpp/build-sm120-spm-debug/pi05_native_dlopen_probe \ + cpp/build-sm120-spm-debug/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 \ + cpp/build-sm120-spm-asan/pi05_native_dlopen_probe \ + cpp/build-sm120-spm-asan/libflashrt_cpp_pi05_c.so \ + 1 +``` From e913a2d282814c2ed553f1330ffd680d7414b47a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 21:00:37 -0400 Subject: [PATCH 49/61] fix: enforce Pi0.5 hot IO contracts --- cpp/modalities/src/vision_cpu.cpp | 8 ++ cpp/modalities/src/vision_cuda.cu | 43 ++++++++-- cpp/models/pi05/src/c_api.cpp | 5 ++ cpp/models/pi05/src/config_map.h | 4 + cpp/models/pi05/src/io.cpp | 9 ++ cpp/models/pi05/src/model_runtime.cpp | 4 + cpp/tests/gate_pi05_hot_allocator.py | 86 ++++++++++++++++++++ cpp/tests/pi05_native_open_probe.cpp | 103 +++++++++++++++++++++-- cpp/tests/test_device_staging.cpp | 113 ++++++++++++++++++++++++++ cpp/tests/test_modalities.cpp | 28 +++++++ cpp/tests/test_pi05_c_api.cpp | 14 ++++ cpp/tests/test_pi05_model_runtime.cpp | 18 ++++ docs/pi05_io_contract.md | 42 ++++++++++ 13 files changed, 463 insertions(+), 14 deletions(-) create mode 100644 cpp/tests/gate_pi05_hot_allocator.py diff --git a/cpp/modalities/src/vision_cpu.cpp b/cpp/modalities/src/vision_cpu.cpp index e6c58e49..8bea9747 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 @@ -142,6 +143,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..1ce624dd 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) * @@ -111,8 +119,11 @@ __device__ __forceinline__ float normalize_value(float raw, 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); + } + return __fmul_rn( + __fsub_rn(__fdiv_rn(raw, 255.0f), mean[c]), inv_std[c]); } __device__ __forceinline__ void store_value(void* out, @@ -152,10 +163,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,9 +191,15 @@ __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 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, shift, mean, inv_std); const std::uint64_t out_idx = diff --git a/cpp/models/pi05/src/c_api.cpp b/cpp/models/pi05/src/c_api.cpp index 45d90327..e603490d 100644 --- a/cpp/models/pi05/src/c_api.cpp +++ b/cpp/models/pi05/src/c_api.cpp @@ -25,6 +25,7 @@ namespace { using flashrt::models::pi05::cface::make_config; using flashrt::models::pi05::cface::pixel_format; using flashrt::models::pi05::cface::status_code; +using flashrt::models::pi05::cface::valid_pixel_format; } // namespace @@ -124,6 +125,10 @@ extern "C" int frt_pi05_runtime_prepare_vision( h->last_error = "invalid Pi05 vision frame"; return -1; } + 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) { diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h index 4e933777..b18e0098 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/config_map.h @@ -40,6 +40,10 @@ 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 modalities::DType dtype(int value) { using modalities::DType; switch (value) { diff --git a/cpp/models/pi05/src/io.cpp b/cpp/models/pi05/src/io.cpp index c2e109f4..621f6f2c 100644 --- a/cpp/models/pi05/src/io.cpp +++ b/cpp/models/pi05/src/io.cpp @@ -18,6 +18,15 @@ modalities::Status validate_pi05_frame_contract( modalities::StatusCode::kShapeMismatch, "Pi05 image input must be u8 HWC"); } + 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] != 3) { + 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( diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index bd3e0063..39fccee0 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -137,6 +137,10 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, a->last_error = "invalid image view"; return -1; } + 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; 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/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index f92b36ed..8a39485d 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -4,6 +4,8 @@ #include #include +#include +#include #include #include #include @@ -21,6 +23,39 @@ int main(int argc, char** argv) { 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); + } + 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] @@ -76,8 +111,8 @@ int main(int argc, char** argv) { } const std::string prompt = "pick up the black bowl"; - const float state[8] = {0.1f, -0.2f, 0.3f, -0.4f, - 0.5f, -0.6f, 0.7f, -0.8f}; + 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) { @@ -86,6 +121,49 @@ int main(int argc, char** argv) { 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) { @@ -115,7 +193,8 @@ int main(int argc, char** argv) { host_noise[i] = flashrt::modalities::float_to_bfloat16( static_cast(static_cast(i % 23) - 11) / 12.0f); } - const bool profile_range = std::getenv("FLASHRT_PROFILE_RANGE") != nullptr; + const bool profile_range = replay_env || + std::getenv("FLASHRT_PROFILE_RANGE") != nullptr; 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(), @@ -127,11 +206,25 @@ int main(int argc, char** argv) { model->release(model->owner); return 1; } - const int step_rc = model->verbs.step(model->self); + int step_rc = 0; + cudaError_t upload_rc = cudaSuccess; + for (int replay = 0; replay < replay_count; ++replay) { + if (replay != 0) { + 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; + } const cudaError_t sync_rc = cudaDeviceSynchronize(); const cudaError_t profiler_rc = profile_range ? cudaProfilerStop() : cudaSuccess; - if (step_rc != 0 || sync_rc != cudaSuccess || profiler_rc != 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); diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index 9b3329ed..f8def582 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -5,6 +5,8 @@ #include +#include +#include #include #include #include @@ -43,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); @@ -116,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}, @@ -226,6 +338,7 @@ 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"; diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index 3636c0ec..90218aae 100644 --- a/cpp/tests/test_modalities.cpp +++ b/cpp/tests/test_modalities.cpp @@ -158,6 +158,34 @@ 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()); diff --git a/cpp/tests/test_pi05_c_api.cpp b/cpp/tests/test_pi05_c_api.cpp index ef3b10ba..82778a55 100644 --- a/cpp/tests/test_pi05_c_api.cpp +++ b/cpp/tests/test_pi05_c_api.cpp @@ -144,6 +144,20 @@ int main() { frame.pixel_format = FRT_PI05_PIXEL_RGB8; rc = frt_pi05_runtime_prepare_vision(rt, &frame, 1); assert(rc == 0); + 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")); + invalid = frame; + invalid.pixel_format = FRT_PI05_PIXEL_BGR8; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == -4); + 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 895f3162..2c92187f 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -174,6 +174,24 @@ int main() { 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); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index d518b3ff..47d456fb 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -99,6 +99,12 @@ Producer-owned preprocessing: - 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 @@ -500,6 +506,42 @@ 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 graph replays without tracing +individual graph nodes: + +``` +FLASHRT_PROFILE_REPLAYS=1000 nsys profile --trace=cuda \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + cpp/build-sm120-spm-debug/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. The reference SM120 run observed zero allocator +calls across 2,001 CUDA API calls. + +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 \ + cpp/build-sm120-spm-debug/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: From ac4757740af39148032a0f1d7d713594242d663a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 21:15:04 -0400 Subject: [PATCH 50/61] test: pin model runtime identity rules --- runtime/tests/test_model_runtime.cpp | 51 ++++++++++++++++++++++++++-- 1 file changed, 49 insertions(+), 2 deletions(-) diff --git a/runtime/tests/test_model_runtime.cpp b/runtime/tests/test_model_runtime.cpp index 0c9b29a5..38fc0d81 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); @@ -179,6 +193,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); From 10791264f7225c8cca3e6cbb976587f53f6e7cf5 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 9 Jul 2026 21:15:16 -0400 Subject: [PATCH 51/61] fix: align Pi0.5 action port payloads --- cpp/models/pi05/src/model_runtime.cpp | 9 ++++++++- cpp/models/pi05/src/native_model_runtime.cpp | 8 +++++--- cpp/tests/pi05_native_e2e_probe.cpp | 9 +++++++++ cpp/tests/pi05_native_open_probe.cpp | 10 +++++++++- 4 files changed, 31 insertions(+), 5 deletions(-) diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index 39fccee0..599d1fe9 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -483,6 +483,7 @@ 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, @@ -496,6 +497,11 @@ 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)) { @@ -514,7 +520,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; diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 7d5ba080..2a4b0694 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -204,6 +204,9 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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, @@ -223,10 +226,9 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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_BF16, + 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, noise->buffer, 0, - frt_buffer_bytes(noise->buffer)) == 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, diff --git a/cpp/tests/pi05_native_e2e_probe.cpp b/cpp/tests/pi05_native_e2e_probe.cpp index bfa99038..0b3e5249 100644 --- a/cpp/tests/pi05_native_e2e_probe.cpp +++ b/cpp/tests/pi05_native_e2e_probe.cpp @@ -97,6 +97,15 @@ int main(int argc, char** argv) { 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; diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index 8a39485d..aa99b7fc 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -97,7 +97,15 @@ int main(int argc, char** argv) { 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[5].update == FRT_RT_PORT_SWAP; + 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); From 43282a7f85a80cf219aa3a6ca5ee5b0d52a90059 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 10 Jul 2026 04:24:34 -0400 Subject: [PATCH 52/61] fix: align staged action payload contracts --- cpp/models/pi05/src/model_runtime.cpp | 10 +-- cpp/tests/gate_pi05_model_runtime_export.py | 75 ++++++++++++++++++--- cpp/tests/test_pi05_model_runtime.cpp | 27 ++++++-- docs/model_runtime_api.md | 2 +- docs/pi05_io_contract.md | 17 +++-- flash_rt/models/pi05/runtime_export.py | 4 +- flash_rt/runtime/export.py | 2 +- runtime/include/flashrt/model_runtime.h | 2 +- 8 files changed, 112 insertions(+), 27 deletions(-) diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index 599d1fe9..89cb492f 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -425,11 +425,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"; @@ -489,7 +490,8 @@ extern "C" int frt_pi05_model_runtime_create_over( 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 && diff --git a/cpp/tests/gate_pi05_model_runtime_export.py b/cpp/tests/gate_pi05_model_runtime_export.py index 7af2a09c..e3b609aa 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 [ @@ -395,6 +413,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 +491,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) @@ -523,15 +566,18 @@ def main() -> None: _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 (args.fp8 or 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}") @@ -542,19 +588,28 @@ def main() -> None: 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(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 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/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index 2c92187f..d1fd12fe 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -155,6 +155,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 */ @@ -267,14 +273,13 @@ 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; @@ -286,6 +291,20 @@ int main() { &exp, ports, 3, stages, 2, nullptr, 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, nullptr, 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, @@ -297,7 +316,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, diff --git a/docs/model_runtime_api.md b/docs/model_runtime_api.md index 7730bfc4..ef8af850 100644 --- a/docs/model_runtime_api.md +++ b/docs/model_runtime_api.md @@ -30,7 +30,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 diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 47d456fb..5060ce57 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -22,7 +22,7 @@ This is the contract implemented by `frt_pi05_model_runtime_create_over`. |---|---|---|---|---| | `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-visible robot action chunk, `FLAT`, `(chunk_length, robot_action_dim)` | `diffusion_noise` | +| `actions` | `ACTION/STAGED` | output | host `f32`, `FLAT`, `(chunk_length, robot_action_dim)` | staged-only; no raw window | Current source of truth: @@ -50,7 +50,7 @@ face. | `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-visible robot action chunk, `FLAT`, `(chunk_length, robot_action_dim)` | `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 @@ -124,7 +124,8 @@ 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. +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: @@ -417,10 +418,18 @@ python cpp/tests/gate_pi05_native_weight_ops.py \ ``` ``` -python cpp/tests/gate_pi05_model_runtime_export.py ... +FLASHRT_BUILD_DIR=cpp/build-sm120-debug \ + python cpp/tests/gate_pi05_model_runtime_export.py \ + --lib cpp/build-sm120-debug/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. diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index 18c33e54..6c192c8d 100644 --- a/flash_rt/models/pi05/runtime_export.py +++ b/flash_rt/models/pi05/runtime_export.py @@ -160,9 +160,9 @@ 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) diff --git a/flash_rt/runtime/export.py b/flash_rt/runtime/export.py index 05437dc0..8e505394 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" diff --git a/runtime/include/flashrt/model_runtime.h b/runtime/include/flashrt/model_runtime.h index 51dbe689..d8db836c 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 */ From 38fd6065c0b328f73d9f15e25b1908b1f4f44152 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 10 Jul 2026 04:32:51 -0400 Subject: [PATCH 53/61] fix: align native v2 producer schemas --- cpp/tests/gate_pi05_native_schema_parity.py | 106 ++++++++++++++++++++ cpp/tests/pi05_native_open_probe.cpp | 24 +++++ docs/pi05_io_contract.md | 12 +++ flash_rt/frontends/torch/pi05_rtx.py | 2 + flash_rt/frontends/torch/pi05_rtx_fp16.py | 1 + flash_rt/models/pi05/pipeline_rtx.py | 5 +- flash_rt/models/pi05/pipeline_rtx_fp16.py | 5 +- 7 files changed, 153 insertions(+), 2 deletions(-) create mode 100644 cpp/tests/gate_pi05_native_schema_parity.py 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..274dd892 --- /dev/null +++ b/cpp/tests/gate_pi05_native_schema_parity.py @@ -0,0 +1,106 @@ +#!/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] +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 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) + + 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()) != 6 or len(producer.stages()) != 1 or \ + counts.get("capsule_regions") != 1: + 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) + 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() + + if python_records != native_records: + diff = "\n".join(difflib.unified_diff( + python_records, native_records, + fromfile="python-native-v2", tofile="cpp-native-v2", + lineterm="", + )) + raise AssertionError(f"native-v2 schema mismatch:\n{diff}") + + print("\n===== PI0.5 NATIVE-V2 SCHEMA PARITY =====") + print(f"records : {len(python_records)}") + print(f"ports/stage : 6 / 1") + print(f"regions : 1 (rollout_boundary)") + print("PASS - Python and C++ native-v2 schemas are canonical-record exact") + return 0 + finally: + producer.release() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index aa99b7fc..fc73838b 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -111,6 +112,29 @@ int main(int argc, char** argv) { 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"; diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 5060ce57..6723bfbd 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -445,6 +445,18 @@ 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=cpp/build-sm120-debug \ + python cpp/tests/gate_pi05_native_schema_parity.py \ + --checkpoint \ + --tokenizer \ + --native-probe cpp/build-sm120-spm-debug/pi05_native_open_probe +``` + The native formatter and tokenizer must also remain token-exact over real prompt/state traffic: 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) From f97fd8444601d65ea9434cdae5a57e638bfc1b5d Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 07:35:37 -0400 Subject: [PATCH 54/61] fix(pi05): preserve reference image normalization --- .../include/flashrt/cpp/modalities/vision.h | 2 + cpp/modalities/src/vision_cpu.cpp | 3 + cpp/modalities/src/vision_cuda.cu | 16 +++-- cpp/models/pi05/src/spec.cpp | 4 +- cpp/tests/gate_pi05_model_runtime_export.py | 60 ++++++++++++++++++- cpp/tests/test_modalities.cpp | 3 + 6 files changed, 80 insertions(+), 8 deletions(-) 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/vision_cpu.cpp b/cpp/modalities/src/vision_cpu.cpp index 8bea9747..cc217d4e 100644 --- a/cpp/modalities/src/vision_cpu.cpp +++ b/cpp/modalities/src/vision_cpu.cpp @@ -112,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]; } diff --git a/cpp/modalities/src/vision_cuda.cu b/cpp/modalities/src/vision_cuda.cu index 1ce624dd..50ee7e73 100644 --- a/cpp/modalities/src/vision_cuda.cu +++ b/cpp/modalities/src/vision_cuda.cu @@ -116,12 +116,16 @@ __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 __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]); } @@ -152,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, @@ -200,8 +205,8 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, 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, shift, - mean, inv_std); + 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); @@ -330,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/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/gate_pi05_model_runtime_export.py b/cpp/tests/gate_pi05_model_runtime_export.py index e3b609aa..d0fd8335 100644 --- a/cpp/tests/gate_pi05_model_runtime_export.py +++ b/cpp/tests/gate_pi05_model_runtime_export.py @@ -349,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) @@ -527,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) @@ -560,6 +579,23 @@ 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( @@ -570,7 +606,7 @@ def main() -> None: act_close = bool(np.allclose( py_actions, full_actions, rtol=1e-4, atol=1e-3 )) - act_ok = act_cos >= 0.999 and (args.fp8 or act_close) + 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( @@ -587,6 +623,15 @@ 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 : accepted={act_ok} " @@ -600,7 +645,18 @@ def main() -> None: 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 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}" ) diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index 90218aae..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() { From c67aef577326630104b0af1a5555fd36801bbfe5 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 07:40:47 -0400 Subject: [PATCH 55/61] fix(runtime): enforce staged verb declarations --- docs/model_runtime_api.md | 5 +++ docs/pi05_io_contract.md | 5 +++ flash_rt/models/pi05/runtime_export.py | 4 +++ flash_rt/runtime/export.py | 17 ++++++++++ runtime/tests/test_model_runtime_py.py | 47 +++++++++++++++++++++++++- 5 files changed, 77 insertions(+), 1 deletion(-) diff --git a/docs/model_runtime_api.md b/docs/model_runtime_api.md index ef8af850..5b6554d9 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` | diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 6723bfbd..2b248bd8 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -32,6 +32,11 @@ Current source of truth: - 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 diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index 6c192c8d..ae2cf05c 100644 --- a/flash_rt/models/pi05/runtime_export.py +++ b/flash_rt/models/pi05/runtime_export.py @@ -221,6 +221,9 @@ def step(): 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", @@ -239,6 +242,7 @@ def step(): manifest_extra=manifest_extra, owner=parts["owner"], step=step, + _allow_incomplete_staged=declaration_only, ) diff --git a/flash_rt/runtime/export.py b/flash_rt/runtime/export.py index 8e505394..8906f7f4 100644 --- a/flash_rt/runtime/export.py +++ b/flash_rt/runtime/export.py @@ -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/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", From b915efb4623b6ea0d3ec24b25950a646b60e2bba Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 07:43:44 -0400 Subject: [PATCH 56/61] fix(runtime): remove host hot-path allocations --- .../cpp/models/pi05/native_rtx_attention.h | 1 + .../cpp/models/pi05/native_workspace.h | 1 + .../include/flashrt/cpp/models/pi05/runtime.h | 1 + cpp/models/pi05/src/model_runtime.cpp | 8 ++-- cpp/models/pi05/src/native_rtx_attention.cpp | 14 ++++-- cpp/models/pi05/src/native_workspace.cpp | 9 ++-- cpp/models/pi05/src/runtime.cpp | 3 +- exec/CMakeLists.txt | 9 ++++ exec/src/graph.cpp | 5 +- exec/tests/test_graph_lru_alloc.cpp | 46 +++++++++++++++++++ 10 files changed, 84 insertions(+), 13 deletions(-) create mode 100644 exec/tests/test_graph_lru_alloc.cpp 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 index 15908119..cec18243 100644 --- 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 @@ -66,6 +66,7 @@ class NativeRtxAttentionWorkspace { int encoder_layers_ = 0; int encoder_splits_ = 0; int decoder_splits_ = 0; + frt_buffer prompt_length_buffers_[3] = {nullptr, nullptr, nullptr}; }; } // namespace pi05 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 index efd6c67c..1098bb24 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h @@ -70,6 +70,7 @@ class NativeWorkspace { int num_views_ = 0; int max_prompt_tokens_ = 0; int chunk_size_ = 0; + frt_buffer decoder_rope_buffer_ = nullptr; std::vector rope_table_; }; 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 61f391ce..5c3ffc92 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -125,6 +125,7 @@ class Runtime final : public families::vla::Runtime { 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; diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/model_runtime.cpp index 89cb492f..d2730124 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/model_runtime.cpp @@ -35,6 +35,7 @@ struct Adapter { 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; @@ -174,7 +175,7 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, a->last_error = "prompt payload is null"; return -1; } - if (bytes > a->prompt_text.capacity()) { + if (bytes > a->prompt_text_limit) { a->last_error = "prompt payload exceeds the hot-path capacity"; return -4; } @@ -554,8 +555,9 @@ extern "C" int frt_pi05_model_runtime_create_over( a->action_values.resize(static_cast(action.chunk * action.robot_dim)); if (cfg.prompt_max_tokens) { - a->prompt_text.reserve(static_cast( - cfg.prompt_max_tokens * 8ull)); + 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()); diff --git a/cpp/models/pi05/src/native_rtx_attention.cpp b/cpp/models/pi05/src/native_rtx_attention.cpp index 657a11a2..f77a6b91 100644 --- a/cpp/models/pi05/src/native_rtx_attention.cpp +++ b/cpp/models/pi05/src/native_rtx_attention.cpp @@ -145,6 +145,13 @@ modalities::Status NativeRtxAttentionWorkspace::allocate( {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); } @@ -161,12 +168,9 @@ modalities::Status NativeRtxAttentionWorkspace::set_fixed_prompt_length( #else const std::int32_t valid = encoder_vision_sequence_ + prompt_tokens; const std::int32_t values[] = {valid, valid + chunk_size_, valid}; - const char* names[] = {"attn_enc_seqused", "attn_dec_seqused", - "attn_dec_devpos"}; for (int i = 0; i < 3; ++i) { - const NativeAttentionBuffer* target = find(names[i]); - if (!target || - cudaMemcpy(frt_buffer_dptr(target->buffer), &values[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"); diff --git a/cpp/models/pi05/src/native_workspace.cpp b/cpp/models/pi05/src/native_workspace.cpp index 8638e56f..81a0004b 100644 --- a/cpp/models/pi05/src/native_workspace.cpp +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -156,8 +156,8 @@ modalities::Status NativeWorkspace::update_decoder_rope(int prompt_tokens) { modalities::StatusCode::kUnsupported, "decoder RoPE update requires the CUDA build"); #else - const NativeWorkspaceBuffer* decoder = find("decoder_rope_weights"); - if (!decoder) return invalid("decoder RoPE buffer was not allocated"); + 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; @@ -168,7 +168,7 @@ modalities::Status NativeWorkspace::update_decoder_rope(int prompt_tokens) { return invalid("decoder RoPE slice exceeds the generated table"); } const cudaError_t rc = cudaMemcpy( - frt_buffer_dptr(decoder->buffer), rope_table_.data() + start, + frt_buffer_dptr(decoder_rope_buffer_), rope_table_.data() + start, elements * sizeof(std::uint16_t), cudaMemcpyHostToDevice); return rc == cudaSuccess ? modalities::Status::ok() @@ -304,6 +304,9 @@ modalities::Status NativeWorkspace::allocate( 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(); diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index 2bb04f82..cd3b8af6 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -207,7 +207,7 @@ int Runtime::set_prompt_state(const char* text, const float* state, return -1; } const std::size_t text_bytes = std::strlen(text); - if (text_bytes > task_prompt_workspace_.capacity()) { + if (text_bytes > max_task_prompt_bytes_) { prompt_status_ = modalities::Status::error( modalities::StatusCode::kShapeMismatch, "prompt text exceeds the configured hot-path capacity"); @@ -322,6 +322,7 @@ modalities::Status Runtime::bind_prompt_staging() { } 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( 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; +} From 431c1da498b8fef4c7137bca66a5eb0ee6312f4e Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 07:45:55 -0400 Subject: [PATCH 57/61] fix(pi05): preserve legacy image format support --- .../pi05/include/flashrt/cpp/models/pi05/io.h | 4 +++- .../include/flashrt/cpp/models/pi05/runtime.h | 1 + cpp/models/pi05/src/c_api.cpp | 8 ++++++-- cpp/models/pi05/src/config_map.h | 9 +++++++++ cpp/models/pi05/src/io.cpp | 19 ++++++++++++++----- cpp/models/pi05/src/runtime.cpp | 2 +- cpp/tests/test_pi05_c_api.cpp | 15 ++++++++++++++- docs/pi05_io_contract.md | 6 ++++++ 8 files changed, 54 insertions(+), 10 deletions(-) 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 65bcad97..66ef4673 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h @@ -24,7 +24,8 @@ class RuntimeIo { int robot_action_dim = kLiberoActionDim, modalities::DType image_dtype = modalities::DType::kBFloat16, modalities::VisionStaging* staging = nullptr, - modalities::ActionStaging* action_staging = nullptr); + modalities::ActionStaging* action_staging = nullptr, + bool strict_rgb8 = true); modalities::Status prepare_vision( const std::vector& frames) const; @@ -44,6 +45,7 @@ class RuntimeIo { 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/runtime.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h index 5c3ffc92..2c65967d 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h @@ -36,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 diff --git a/cpp/models/pi05/src/c_api.cpp b/cpp/models/pi05/src/c_api.cpp index e603490d..ea696b2b 100644 --- a/cpp/models/pi05/src/c_api.cpp +++ b/cpp/models/pi05/src/c_api.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include struct frt_pi05_runtime_s { @@ -23,6 +24,7 @@ struct frt_pi05_runtime_s { 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; @@ -43,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; @@ -150,7 +154,7 @@ 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; diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h index b18e0098..a0bc1d21 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/config_map.h @@ -44,6 +44,15 @@ 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) { diff --git a/cpp/models/pi05/src/io.cpp b/cpp/models/pi05/src/io.cpp index 621f6f2c..27eb5c6e 100644 --- a/cpp/models/pi05/src/io.cpp +++ b/cpp/models/pi05/src/io.cpp @@ -6,8 +6,8 @@ namespace pi05 { namespace { modalities::Status validate_pi05_frame_contract( - const modalities::VisionFrame& frame) { - if (frame.format != modalities::PixelFormat::kRGB8) { + 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"); @@ -18,11 +18,18 @@ modalities::Status validate_pi05_frame_contract( 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] != 3) { + frame.image.shape.dims[2] != channels) { return modalities::Status::error( modalities::StatusCode::kShapeMismatch, "Pi05 image shape must match HWC dimensions"); @@ -49,12 +56,14 @@ RuntimeIo::RuntimeIo(int num_views, int robot_action_dim, modalities::DType image_dtype, modalities::VisionStaging* staging, - modalities::ActionStaging* action_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)) { @@ -64,7 +73,7 @@ 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); + auto st = validate_pi05_frame_contract(frame, strict_rgb8_); if (!st.ok_status()) return st; } return modalities::preprocess_vision(vision_spec_, frames, image_input_, diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/runtime.cpp index cd3b8af6..ecdbae61 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/runtime.cpp @@ -187,7 +187,7 @@ modalities::Status Runtime::bind() { config_.action_stddev, find_native_stream(exp_, stream_id_), config_.chunk, config_.model_action_dim, config_.robot_action_dim, config_.image_dtype, staging, - action_staging); + action_staging, config_.strict_rgb8); return bind_prompt_staging(); } diff --git a/cpp/tests/test_pi05_c_api.cpp b/cpp/tests/test_pi05_c_api.cpp index 82778a55..e7e29360 100644 --- a/cpp/tests/test_pi05_c_api.cpp +++ b/cpp/tests/test_pi05_c_api.cpp @@ -144,15 +144,28 @@ 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 == -4); + 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); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 2b248bd8..a5b9d568 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -116,6 +116,12 @@ 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 From 86a93b350faf2a8e322b6634aa078e55152c0c12 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 08:16:14 -0400 Subject: [PATCH 58/61] docs(pi05): clarify native runtime migration --- cpp/tests/profile_pi05_python_replay.py | 5 ++- docs/pi05_cpp_runtime_migration.md | 45 +++++++++++++++++++++++++ docs/pi05_io_contract.md | 35 ++++++++++--------- 3 files changed, 68 insertions(+), 17 deletions(-) create mode 100644 docs/pi05_cpp_runtime_migration.md diff --git a/cpp/tests/profile_pi05_python_replay.py b/cpp/tests/profile_pi05_python_replay.py index 4a6eedce..30e66910 100644 --- a/cpp/tests/profile_pi05_python_replay.py +++ b/cpp/tests/profile_pi05_python_replay.py @@ -1,4 +1,7 @@ -"""Emit one replay-only CUDA profiler range for the Pi0.5 Python frontend.""" +"""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 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 index a5b9d568..1aee5a32 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -429,9 +429,9 @@ python cpp/tests/gate_pi05_native_weight_ops.py \ ``` ``` -FLASHRT_BUILD_DIR=cpp/build-sm120-debug \ +FLASHRT_BUILD_DIR= \ python cpp/tests/gate_pi05_model_runtime_export.py \ - --lib cpp/build-sm120-debug/libflashrt_cpp_pi05_c.so ... + --lib /libflashrt_cpp_pi05_c.so ... python cpp/tests/gate_pi05_c_api_export.py ... ``` @@ -448,7 +448,7 @@ producer must not retain the declarations if any required verb is unavailable. The native factory lifecycle gate is: ``` -cpp/build-sm120-spm-debug/pi05_native_open_probe \ +/pi05_native_open_probe \ ``` @@ -461,11 +461,11 @@ port/stage/region records (their producer identity and fingerprints remain different): ``` -FLASHRT_BUILD_DIR=cpp/build-sm120-debug \ +FLASHRT_BUILD_DIR= \ python cpp/tests/gate_pi05_native_schema_parity.py \ --checkpoint \ --tokenizer \ - --native-probe cpp/build-sm120-spm-debug/pi05_native_open_probe + --native-probe /pi05_native_open_probe ``` The native formatter and tokenizer must also remain token-exact over real @@ -476,14 +476,17 @@ python cpp/tests/gate_pi05_tokenizer_corpus.py \ --dataset \ --checkpoint \ --tokenizer \ - --probe cpp/build-sm120-spm-debug/pi05_tokenizer_corpus_probe \ + --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 reference SM120 run covered -10,000 records and 20 token lengths from 43 through 62 with zero mismatches. +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: @@ -493,7 +496,7 @@ python cpp/tests/gate_pi05_native_e2e.py \ --checkpoint \ --tokenizer \ --dataset \ - --probe cpp/build-sm120-spm-debug/pi05_native_e2e_probe \ + --probe /pi05_native_e2e_probe \ --episode 0 --frame 0 ``` @@ -514,7 +517,7 @@ FLASHRT_PROFILE_RANGE=1 nsys profile --trace=cuda \ --cuda-graph-trace=node \ --capture-range=cudaProfilerApi --capture-range-end=stop \ -o \ - cpp/build-sm120-spm-debug/pi05_native_open_probe \ + /pi05_native_open_probe \ nsys profile --trace=cuda --cuda-graph-trace=node \ @@ -545,7 +548,7 @@ individual graph nodes: FLASHRT_PROFILE_REPLAYS=1000 nsys profile --trace=cuda \ --capture-range=cudaProfilerApi --capture-range-end=stop \ -o \ - cpp/build-sm120-spm-debug/pi05_native_open_probe \ + /pi05_native_open_probe \ nsys stats --report cuda_api_trace --format csv \ .nsys-rep > .csv @@ -565,7 +568,7 @@ device update): ``` FLASHRT_HOT_STATE_UPDATES=1000 FLASHRT_HOT_STATE_P99_US=1000 \ - cpp/build-sm120-spm-debug/pi05_native_open_probe \ + /pi05_native_open_probe \ ``` @@ -579,8 +582,8 @@ the factory from the shared object, exercises an extra retain/release pair, releases the final model reference, and only then unloads the producer: ``` -cpp/build-sm120-spm-debug/pi05_native_dlopen_probe \ - cpp/build-sm120-spm-debug/libflashrt_cpp_pi05_c.so \ +/pi05_native_dlopen_probe \ + /libflashrt_cpp_pi05_c.so \ 1 ``` @@ -590,7 +593,7 @@ 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 \ - cpp/build-sm120-spm-asan/pi05_native_dlopen_probe \ - cpp/build-sm120-spm-asan/libflashrt_cpp_pi05_c.so \ + /pi05_native_dlopen_probe \ + /libflashrt_cpp_pi05_c.so \ 1 ``` From 89a955fd7caa14276d66087c105afc47af39feff Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 08:32:57 -0400 Subject: [PATCH 59/61] test(pi05): profile the complete hot service loop --- cpp/tests/pi05_native_open_probe.cpp | 39 +++++++++++++++++++++++++--- docs/pi05_io_contract.md | 17 ++++++++---- 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index fc73838b..78bba03c 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -35,6 +35,13 @@ int main(int argc, char** argv) { } 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) { @@ -227,6 +234,8 @@ int main(int argc, char** argv) { } 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(), @@ -241,7 +250,26 @@ int main(int argc, char** argv) { int step_rc = 0; cudaError_t upload_rc = cudaSuccess; for (int replay = 0; replay < replay_count; ++replay) { - if (replay != 0) { + 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), @@ -250,6 +278,13 @@ int main(int argc, char** argv) { } 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 = @@ -262,8 +297,6 @@ int main(int argc, char** argv) { model->release(model->owner); return 1; } - float actions[10 * 7]{}; - std::uint64_t written = 0; if (model->verbs.get_output(model->self, 4, actions, sizeof(actions), &written, -1) != 0 || written != sizeof(actions)) { diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 1aee5a32..d63f9df8 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -541,11 +541,14 @@ 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 graph replays without tracing -individual graph nodes: +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 nsys profile --trace=cuda \ +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 \ @@ -559,8 +562,12 @@ python cpp/tests/gate_pi05_hot_allocator.py \ 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. The reference SM120 run observed zero allocator -calls across 2,001 CUDA API calls. +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 From 924c7c625a33c146552cdaeba926e0371937a72b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 11:12:22 -0400 Subject: [PATCH 60/61] fix(runtime): harden native producer contracts --- cpp/models/pi05/src/native_model_runtime.cpp | 31 ++++++++++++-- cpp/tests/data/pi05_native_v2_schema.records | 8 ++++ cpp/tests/gate_pi05_native_schema_parity.py | 34 +++++++++------ cpp/tests/pi05_native_open_probe.cpp | 17 ++++++++ cpp/tests/test_pi05_model_runtime.cpp | 23 +++++++--- runtime/include/flashrt/model_runtime.h | 12 ++++-- runtime/src/model_runtime.cpp | 19 +++++++++ runtime/tests/test_model_runtime.cpp | 45 +++++++++++++++----- 8 files changed, 154 insertions(+), 35 deletions(-) create mode 100644 cpp/tests/data/pi05_native_v2_schema.records diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 2a4b0694..a334f32d 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -41,10 +41,26 @@ bool add_identity(frt_runtime_builder builder, const char* key, 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, nullptr, nullptr, nullptr, nullptr, nullptr); + builder, &discard_verbs, nullptr, nullptr, nullptr, nullptr); if (discarded) discarded->release(discarded->owner); if (error) *error = message; return -6; @@ -71,6 +87,8 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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; @@ -169,7 +187,7 @@ int build_native_model_runtime(const NativeOpenConfig& config, ok = add_identity(builder, "model", "pi05") && add_identity(builder, "producer", "native") && add_identity(builder, "pipeline", "NativeBf16") && - add_identity(builder, "hardware", "sm120") && + 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) && @@ -190,7 +208,8 @@ int build_native_model_runtime(const NativeOpenConfig& config, std::ostringstream manifest; manifest << "{\"model\":\"pi05\",\"producer\":\"native\"," - << "\"hardware\":\"sm120\",\"io\":\"native_v2\"," + << "\"hardware\":\"" << hardware_id + << "\",\"io\":\"native_v2\"," << "\"stage_plan\":{\"name\":\"full\"," << "\"stages\":[{\"name\":\"infer\"," << "\"graph\":\"infer\",\"after\":[]}]}}"; @@ -238,8 +257,12 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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, nullptr, nullptr, raw_graph, nullptr, release_graph_owner); + builder, &base_verbs, nullptr, raw_graph, nullptr, + release_graph_owner); if (!base) { delete raw_graph; if (error) *error = "native integrated runtime finish failed"; 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_native_schema_parity.py b/cpp/tests/gate_pi05_native_schema_parity.py index 274dd892..845cf698 100644 --- a/cpp/tests/gate_pi05_native_schema_parity.py +++ b/cpp/tests/gate_pi05_native_schema_parity.py @@ -15,6 +15,7 @@ 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: @@ -32,6 +33,15 @@ def canonical_records(identity: str) -> list[str]: 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) @@ -43,6 +53,10 @@ def main() -> int: 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 = [ @@ -64,14 +78,16 @@ def main() -> int: ) try: counts = dict(runtime_abi.export_counts(producer.export_ptr)) - if len(producer.ports()) != 6 or len(producer.stages()) != 1 or \ - counts.get("capsule_regions") != 1: + 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) @@ -84,19 +100,13 @@ def main() -> int: ) native_records = native_path.read_text().splitlines() - if python_records != native_records: - diff = "\n".join(difflib.unified_diff( - python_records, native_records, - fromfile="python-native-v2", tofile="cpp-native-v2", - lineterm="", - )) - raise AssertionError(f"native-v2 schema mismatch:\n{diff}") + 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 : 6 / 1") - print(f"regions : 1 (rollout_boundary)") - print("PASS - Python and C++ native-v2 schemas are canonical-record exact") + 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() diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index 78bba03c..edb0043d 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -81,6 +81,19 @@ int main(int argc, char** argv) { 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 && @@ -90,6 +103,10 @@ int main(int argc, char** argv) { 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 && diff --git a/cpp/tests/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index d1fd12fe..20555acb 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -63,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() { @@ -286,17 +293,21 @@ int main() { 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, nullptr, nullptr, nullptr, - nullptr); + &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( @@ -347,7 +358,8 @@ int main() { 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, nullptr, nullptr, nullptr, nullptr); + &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; @@ -369,7 +381,8 @@ int main() { 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, nullptr, nullptr, nullptr, nullptr); + &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; diff --git a/runtime/include/flashrt/model_runtime.h b/runtime/include/flashrt/model_runtime.h index d8db836c..bf14d025 100644 --- a/runtime/include/flashrt/model_runtime.h +++ b/runtime/include/flashrt/model_runtime.h @@ -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 38fc0d81..a9362b33 100644 --- a/runtime/tests/test_model_runtime.cpp +++ b/runtime/tests/test_model_runtime.cpp @@ -101,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 --- */ @@ -271,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, @@ -309,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); From 62ee5504351e1cf3ec612d8c8630c4a76edd21a5 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sat, 11 Jul 2026 11:12:35 -0400 Subject: [PATCH 61/61] docs(runtime): define native producer review standards --- .github/pull_request_template.md | 25 +++++++ CONTRIBUTING.md | 18 ++++- docs/cpp_runtime_design.md | 8 +++ docs/model_runtime_api.md | 14 +++- docs/native_model_runtime_producer.md | 96 +++++++++++++++++++++++++++ docs/pi05_io_contract.md | 5 ++ docs/pr_review_checklist.md | 25 +++++++ 7 files changed, 187 insertions(+), 4 deletions(-) create mode 100644 .github/pull_request_template.md create mode 100644 docs/native_model_runtime_producer.md 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/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/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/model_runtime_api.md b/docs/model_runtime_api.md index 5b6554d9..2d39b275 100644 --- a/docs/model_runtime_api.md +++ b/docs/model_runtime_api.md @@ -63,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 | |---|---|---| @@ -134,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)`. @@ -271,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_io_contract.md b/docs/pi05_io_contract.md index d63f9df8..5da5eb10 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -468,6 +468,11 @@ FLASHRT_BUILD_DIR= \ --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: 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.