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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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/83] 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. From 3a107849169d4fc61af45af2364c512e3371dec4 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 14 Jul 2026 04:37:08 -0400 Subject: [PATCH 62/83] perf(pi05): reduce native checkpoint startup --- cpp/CMakeLists.txt | 6 + cpp/loader/src/sha256.cpp | 39 ++- .../cpp/models/pi05/native_weight_ops.h | 20 ++ cpp/models/pi05/src/native_model_runtime.cpp | 27 +- .../pi05/src/native_weight_materializer.cpp | 62 +++++ cpp/models/pi05/src/native_weight_ops.cpp | 256 +++++++++++++++--- cpp/tests/test_pi05_native_weight_ops.cpp | 33 +++ docs/mindon_pi05_integration.md | 9 + docs/pi05_io_contract.md | 23 ++ 9 files changed, 433 insertions(+), 42 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index c57f00e0..964db7b0 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -150,6 +150,12 @@ add_library(flashrt_cpp_loader STATIC 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) +find_package(OpenSSL QUIET) +if(OpenSSL_FOUND) + target_compile_definitions(flashrt_cpp_loader + PRIVATE FLASHRT_CPP_HAS_OPENSSL=1) + target_link_libraries(flashrt_cpp_loader PRIVATE OpenSSL::Crypto) +endif() set(FLASHRT_CPP_PI05_SRCS models/pi05/src/spec.cpp diff --git a/cpp/loader/src/sha256.cpp b/cpp/loader/src/sha256.cpp index 780a4873..63acb488 100644 --- a/cpp/loader/src/sha256.cpp +++ b/cpp/loader/src/sha256.cpp @@ -8,6 +8,10 @@ #include #include +#ifdef FLASHRT_CPP_HAS_OPENSSL +#include +#endif + namespace flashrt { namespace loader { namespace { @@ -149,20 +153,53 @@ bool sha256_file(const std::string& path, std::string* hex, if (error) *error = "unable to open file for SHA-256: " + path; return false; } +#ifdef FLASHRT_CPP_HAS_OPENSSL + EVP_MD_CTX* context = EVP_MD_CTX_new(); + if (!context || EVP_DigestInit_ex(context, EVP_sha256(), nullptr) != 1) { + if (context) EVP_MD_CTX_free(context); + if (error) *error = "unable to initialize SHA-256"; + return false; + } +#else Sha256 hash; - std::vector buffer(1 << 20); +#endif + std::vector buffer(4 << 20); while (file) { file.read(reinterpret_cast(buffer.data()), buffer.size()); const std::streamsize count = file.gcount(); if (count > 0) { +#ifdef FLASHRT_CPP_HAS_OPENSSL + if (EVP_DigestUpdate(context, buffer.data(), + static_cast(count)) != 1) { + EVP_MD_CTX_free(context); + if (error) *error = "failed while hashing file: " + path; + return false; + } +#else hash.update(buffer.data(), static_cast(count)); +#endif } } if (!file.eof()) { +#ifdef FLASHRT_CPP_HAS_OPENSSL + EVP_MD_CTX_free(context); +#endif if (error) *error = "failed while reading file for SHA-256: " + path; return false; } +#ifdef FLASHRT_CPP_HAS_OPENSSL + std::array digest{}; + unsigned int digest_bytes = 0; + if (EVP_DigestFinal_ex(context, digest.data(), &digest_bytes) != 1 || + digest_bytes != digest.size()) { + EVP_MD_CTX_free(context); + if (error) *error = "unable to finalize SHA-256"; + return false; + } + EVP_MD_CTX_free(context); +#else const auto digest = hash.finish(); +#endif std::ostringstream output; output << std::hex << std::setfill('0'); for (unsigned char byte : digest) output << std::setw(2) << unsigned(byte); 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 6e278640..b68634fd 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 @@ -30,6 +30,26 @@ modalities::Status load_native_float_tensor( modalities::Status native_to_bf16(const NativeFloatTensor& input, NativeBf16Tensor* out); +modalities::Status native_f32_to_bf16( + const float* input, + const std::vector& shape, + bool transpose, + NativeBf16Tensor* out); + +modalities::Status native_f32_fold_rms_columns_transpose( + const float* weight, + std::uint64_t rows, + std::uint64_t cols, + const NativeFloatTensor& norm, + NativeBf16Tensor* out); + +modalities::Status native_f32_round_scale_to_bf16( + const float* input, + const std::vector& shape, + float scale, + bool transpose, + NativeBf16Tensor* out); + modalities::Status native_round_to_bf16_float( const NativeFloatTensor& input, NativeFloatTensor* out); diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index a334f32d..4da75e70 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -90,12 +91,23 @@ int build_native_model_runtime(const NativeOpenConfig& config, const std::string hardware_id = "sm" + std::to_string(properties.major * 10 + properties.minor); - std::string weights_sha256; + struct HashResult { + bool ok = false; + std::string digest; + std::string error; + }; + const std::string weights_path = + config.checkpoint_path + "/model.safetensors"; + std::future weights_hash = std::async( + std::launch::async, [weights_path] { + HashResult result; + result.ok = loader::sha256_file( + weights_path, &result.digest, &result.error); + return result; + }); 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, + if (!loader::sha256_file(config.tokenizer_model_path, &tokenizer_sha256, &hash_error)) { if (error) *error = hash_error; return -2; @@ -114,6 +126,11 @@ int build_native_model_runtime(const NativeOpenConfig& config, if (error) *error = st.message; return cface::status_code(st); } + HashResult weights_sha256 = weights_hash.get(); + if (!weights_sha256.ok) { + if (error) *error = weights_sha256.error; + return -2; + } const NativeWorkspaceBuffer* images = workspace_buffer(*graph, "observation_images_normalized"); @@ -189,7 +206,7 @@ int build_native_model_runtime(const NativeOpenConfig& config, add_identity(builder, "pipeline", "NativeBf16") && add_identity(builder, "hardware", hardware_id) && add_identity(builder, "tensor_dtype", "bf16") && - add_identity(builder, "weights_sha256", weights_sha256) && + add_identity(builder, "weights_sha256", weights_sha256.digest) && add_identity(builder, "tokenizer_sha256", tokenizer_sha256) && add_identity(builder, "io", "native_v2") && add_identity(builder, "state_prompt_mode", "fixed") && diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp index 251f7420..62dfc74f 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -28,6 +28,24 @@ const std::string& vision_prefix() { return prefix; } +const loader::SafetensorInfo* source_tensor( + const loader::SafetensorsFile& source, + const std::string& key) { + const loader::SafetensorInfo* tensor = source.find(key); + if (!tensor) tensor = source.find(std::string("model.") + key); + return tensor; +} + +const float* f32_data(const loader::SafetensorsFile& source, + const loader::SafetensorInfo* tensor) { + if (!tensor || tensor->dtype != "F32") return nullptr; + const void* data = source.data(*tensor); + if (reinterpret_cast(data) % alignof(float) != 0) { + return nullptr; + } + return static_cast(data); +} + std::string layer_name(const char* stem, int layer) { return std::string(stem) + std::to_string(layer); } @@ -53,6 +71,17 @@ modalities::Status NativeWeightMaterializer::upload( modalities::Status NativeWeightMaterializer::upload_rounded_transpose( const std::string& source_key, const std::string& destination_name) { + const loader::SafetensorInfo* info = source_tensor(source_, source_key); + if (destination_ && info && info->shape.size() == 2) { + const float* data = f32_data(source_, info); + if (data) { + NativeBf16Tensor converted; + modalities::Status st = native_f32_to_bf16( + data, info->shape, true, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); + } + } NativeFloatTensor source; NativeFloatTensor rounded; NativeFloatTensor transposed; @@ -68,6 +97,17 @@ modalities::Status NativeWeightMaterializer::upload_rounded_transpose( modalities::Status NativeWeightMaterializer::upload_rounded_copy( const std::string& source_key, const std::string& destination_name) { + const loader::SafetensorInfo* info = source_tensor(source_, source_key); + if (destination_ && info) { + const float* data = f32_data(source_, info); + if (data) { + NativeBf16Tensor converted; + modalities::Status st = native_f32_to_bf16( + data, info->shape, false, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); + } + } NativeFloatTensor source; NativeFloatTensor rounded; modalities::Status st = load(source_key, &source); @@ -81,6 +121,17 @@ modalities::Status NativeWeightMaterializer::upload_folded_transpose( const std::string& source_key, const NativeFloatTensor& norm, const std::string& destination_name) { + const loader::SafetensorInfo* info = source_tensor(source_, source_key); + if (destination_ && info && info->shape.size() == 2) { + const float* data = f32_data(source_, info); + if (data) { + NativeBf16Tensor converted; + modalities::Status st = native_f32_fold_rms_columns_transpose( + data, info->shape[0], info->shape[1], norm, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); + } + } NativeFloatTensor source; NativeFloatTensor folded; NativeFloatTensor transposed; @@ -98,6 +149,17 @@ modalities::Status NativeWeightMaterializer::upload_rounded_scaled( const std::string& destination_name, float scale, bool transpose) { + const loader::SafetensorInfo* info = source_tensor(source_, source_key); + if (destination_ && info) { + const float* data = f32_data(source_, info); + if (data) { + NativeBf16Tensor converted; + modalities::Status st = native_f32_round_scale_to_bf16( + data, info->shape, scale, transpose, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); + } + } NativeFloatTensor source; NativeFloatTensor rounded; NativeFloatTensor arranged; diff --git a/cpp/models/pi05/src/native_weight_ops.cpp b/cpp/models/pi05/src/native_weight_ops.cpp index d3dd5134..ba9716f8 100644 --- a/cpp/models/pi05/src/native_weight_ops.cpp +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -1,8 +1,10 @@ #include "flashrt/cpp/models/pi05/native_weight_ops.h" +#include #include #include #include +#include #include namespace flashrt { @@ -36,6 +38,31 @@ bool valid_tensor(const NativeFloatTensor& tensor) { expected == tensor.values.size(); } +template +void parallel_ranges(std::size_t count, + std::size_t minimum_items, + const Fn& fn) { + if (!count) return; + const unsigned available = std::thread::hardware_concurrency(); + const std::size_t wanted = + (count + minimum_items - 1) / minimum_items; + const std::size_t workers = std::max( + 1, std::min({16, available ? available : 1, wanted})); + if (workers == 1) { + fn(0, count); + return; + } + std::vector threads; + threads.reserve(workers - 1); + for (std::size_t worker = 1; worker < workers; ++worker) { + const std::size_t begin = count * worker / workers; + const std::size_t end = count * (worker + 1) / workers; + threads.emplace_back([begin, end, &fn] { fn(begin, end); }); + } + fn(0, count / workers); + for (std::thread& thread : threads) thread.join(); +} + const loader::SafetensorInfo* find_source_tensor( const loader::SafetensorsFile& file, const std::string& key) { @@ -70,14 +97,20 @@ modalities::Status load_native_float_tensor( 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]); - } + parallel_ranges(count, 1 << 18, [&](std::size_t begin, + std::size_t end) { + for (std::size_t i = begin; i < end; ++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]); - } + parallel_ranges(count, 1 << 18, [&](std::size_t begin, + std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + loaded.values[i] = modalities::float16_to_float(src[i]); + } + }); } else { return modalities::Status::error( modalities::StatusCode::kUnsupported, @@ -94,13 +127,134 @@ modalities::Status native_to_bf16(const NativeFloatTensor& 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]); + parallel_ranges(input.values.size(), 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + converted.values[i] = + modalities::float_to_bfloat16(input.values[i]); + } + }); + *out = std::move(converted); + return modalities::Status::ok(); +} + +modalities::Status native_f32_to_bf16( + const float* input, + const std::vector& shape, + bool transpose, + NativeBf16Tensor* out) { + std::size_t count = 0; + if (!input || !out || !element_count(shape, &count) || + (transpose && shape.size() != 2)) { + return invalid("invalid direct F32 to BF16 input"); + } + NativeBf16Tensor converted; + converted.shape = transpose + ? std::vector{shape[1], shape[0]} + : shape; + converted.values.resize(count); + if (!transpose) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + converted.values[i] = modalities::float_to_bfloat16(input[i]); + } + }); + } else { + const std::size_t rows = static_cast(shape[0]); + const std::size_t cols = static_cast(shape[1]); + parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++row) { + for (std::size_t col = 0; col < cols; ++col) { + converted.values[col * rows + row] = + modalities::float_to_bfloat16( + input[row * cols + col]); + } + } + }); } *out = std::move(converted); return modalities::Status::ok(); } +modalities::Status native_f32_fold_rms_columns_transpose( + const float* weight, + std::uint64_t rows_u64, + std::uint64_t cols_u64, + const NativeFloatTensor& norm, + NativeBf16Tensor* out) { + if (!weight || !out || !valid_tensor(norm) || norm.shape.size() != 1 || + norm.shape[0] != cols_u64 || + rows_u64 > std::numeric_limits::max() || + cols_u64 > std::numeric_limits::max()) { + return invalid("invalid direct F32 RMS fold input"); + } + const std::size_t rows = static_cast(rows_u64); + const std::size_t cols = static_cast(cols_u64); + if (rows && cols > std::numeric_limits::max() / rows) { + return invalid("direct F32 RMS fold shape overflows size_t"); + } + NativeBf16Tensor folded; + folded.shape = {cols_u64, rows_u64}; + folded.values.resize(rows * cols); + parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++row) { + for (std::size_t col = 0; col < cols; ++col) { + folded.values[col * rows + row] = + modalities::float_to_bfloat16( + weight[row * cols + col] * + (1.0f + norm.values[col])); + } + } + }); + *out = std::move(folded); + return modalities::Status::ok(); +} + +modalities::Status native_f32_round_scale_to_bf16( + const float* input, + const std::vector& shape, + float scale, + bool transpose, + NativeBf16Tensor* out) { + std::size_t count = 0; + if (!input || !out || !element_count(shape, &count) || + (transpose && shape.size() != 2)) { + return invalid("invalid direct F32 scale input"); + } + NativeBf16Tensor scaled; + scaled.shape = transpose + ? std::vector{shape[1], shape[0]} + : shape; + scaled.values.resize(count); + const auto convert = [scale](float value) { + const float rounded = modalities::bfloat16_to_float( + modalities::float_to_bfloat16(value)); + return modalities::float_to_bfloat16(rounded * scale); + }; + if (!transpose) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + scaled.values[i] = convert(input[i]); + } + }); + } else { + const std::size_t rows = static_cast(shape[0]); + const std::size_t cols = static_cast(shape[1]); + parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++row) { + for (std::size_t col = 0; col < cols; ++col) { + scaled.values[col * rows + row] = + convert(input[row * cols + col]); + } + } + }); + } + *out = std::move(scaled); + return modalities::Status::ok(); +} + modalities::Status native_round_to_bf16_float( const NativeFloatTensor& input, NativeFloatTensor* out) { @@ -108,10 +262,13 @@ modalities::Status native_round_to_bf16_float( 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)); - } + parallel_ranges(rounded.values.size(), 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + rounded.values[i] = modalities::bfloat16_to_float( + modalities::float_to_bfloat16(rounded.values[i])); + } + }); *out = std::move(rounded); return modalities::Status::ok(); } @@ -126,12 +283,26 @@ modalities::Status native_transpose_2d(const NativeFloatTensor& input, 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]; + constexpr std::size_t kTile = 32; + const std::size_t row_tiles = (rows + kTile - 1) / kTile; + parallel_ranges(row_tiles, 2, [&](std::size_t first, + std::size_t last) { + for (std::size_t tile = first; tile < last; ++tile) { + const std::size_t row_begin = tile * kTile; + const std::size_t row_end = std::min(rows, row_begin + kTile); + for (std::size_t col_begin = 0; col_begin < cols; + col_begin += kTile) { + const std::size_t col_end = + std::min(cols, col_begin + kTile); + for (std::size_t col = col_begin; col < col_end; ++col) { + for (std::size_t row = row_begin; row < row_end; ++row) { + transposed.values[col * rows + row] = + input.values[row * cols + col]; + } + } + } } - } + }); *out = std::move(transposed); return modalities::Status::ok(); } @@ -212,11 +383,13 @@ modalities::Status native_fold_rms_columns( 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]; + parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++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(); } @@ -248,13 +421,19 @@ modalities::Status native_concat_rows_transpose( 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)]; + const std::uint64_t input_rows = input->shape[0]; + const std::uint64_t output_offset = row_offset; + parallel_ranges(static_cast(input_rows), 32, + [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++row) { + for (std::uint64_t col = 0; col < cols; ++col) { + joined.values[static_cast( + col * total_rows + output_offset + row)] = + input->values[static_cast( + row * cols + col)]; + } } - } + }); row_offset += input->shape[0]; } *out = std::move(joined); @@ -284,14 +463,16 @@ modalities::Status native_concat_columns( 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)); - } + parallel_ranges(rows, 32, [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++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(); } @@ -325,7 +506,10 @@ modalities::Status native_scale(const NativeFloatTensor& input, NativeFloatTensor* out) { if (!out || !valid_tensor(input)) return invalid("invalid scale input"); NativeFloatTensor scaled = input; - for (float& value : scaled.values) value *= scale; + parallel_ranges(scaled.values.size(), 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) scaled.values[i] *= scale; + }); *out = std::move(scaled); return modalities::Status::ok(); } diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp index d8baf08d..545040ac 100644 --- a/cpp/tests/test_pi05_native_weight_ops.cpp +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -138,6 +138,39 @@ int main() { flashrt::modalities::float_to_bfloat16(matrix.values[i])); } + NativeBf16Tensor direct; + assert(native_f32_to_bf16(matrix.values.data(), matrix.shape, false, + &direct).ok_status()); + assert(direct.shape == converted.shape && + direct.values == converted.values); + assert(native_f32_to_bf16(matrix.values.data(), matrix.shape, true, + &direct).ok_status()); + NativeFloatTensor transposed; + NativeBf16Tensor transposed_bf16; + assert(native_transpose_2d(matrix, &transposed).ok_status()); + assert(native_to_bf16(transposed, &transposed_bf16).ok_status()); + assert(direct.shape == transposed_bf16.shape && + direct.values == transposed_bf16.values); + + assert(native_f32_fold_rms_columns_transpose( + matrix.values.data(), 2, 3, norm, &direct).ok_status()); + NativeFloatTensor folded; + assert(native_fold_rms_columns(matrix, norm, &folded).ok_status()); + assert(native_transpose_2d(folded, &transposed).ok_status()); + assert(native_to_bf16(transposed, &transposed_bf16).ok_status()); + assert(direct.shape == transposed_bf16.shape && + direct.values == transposed_bf16.values); + + constexpr float kScale = -0.1f; + assert(native_f32_round_scale_to_bf16( + unrounded.values.data(), unrounded.shape, kScale, false, + &direct).ok_status()); + NativeFloatTensor rounded_scaled; + assert(native_scale(result, kScale, &rounded_scaled).ok_status()); + assert(native_to_bf16(rounded_scaled, &converted).ok_status()); + assert(direct.shape == converted.shape && + direct.values == converted.values); + 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"); diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 100bba5b..7788a62b 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -86,6 +86,15 @@ 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. +Use a Release build for MindOn deployment. Native startup includes full-content +checkpoint hashing and CUDA graph capture in addition to safetensors parsing +and H2D weight upload, so time the complete `frt_model_runtime_open_v1` call and +do not compare it directly with a Python weight-load-only timer. OpenSSL is an +optional configure-time acceleration for SHA-256; builds without it retain the +portable implementation with identical identity bytes. The model hash and +weight materialization execute concurrently, but the factory publishes no +runtime until both have succeeded. + ## 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 5da5eb10..c865a134 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -246,6 +246,29 @@ 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. +Build native producers with `CMAKE_BUILD_TYPE=Release` for deployment. The +Pi0.5 checkpoint stores F32 tensors, so setup converts and lays out independent +weight ranges in parallel and writes the final BF16 payload directly instead +of creating intermediate rounded and transposed F32 copies. These direct +transforms are byte-compared with the reference multi-step transforms in the +weight-op tests. Debug builds intentionally retain unoptimized setup code and +are not a startup-latency reference. + +Checkpoint identity remains a full-file SHA-256, not a path, timestamp, or +partial-content surrogate. When OpenSSL is available at configure time the +loader uses its accelerated EVP implementation; the portable in-tree SHA-256 +remains the build fallback. Model hashing runs concurrently with independent +weight materialization, and both must complete successfully before the runtime +descriptor is published. + +For a 14.47 GB F32 Pi0.5 safetensors checkpoint on RTX 5090 SM120, the Release +`pi05_native_open_probe` full lifecycle measured 9.47 seconds with a warm file +cache and 10.91 seconds after evicting that file from the page cache. This +scope includes identity hashing, weight conversion/upload, workspace setup, +CUDA graph capture, one inference, and teardown. Compare producer startup +using this complete scope; a Python `load_model` timer that excludes graph +capture and identity construction is not the same metric. + 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 From 04b95000a63622c40d20adcdacf335a32f807d97 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Tue, 14 Jul 2026 05:12:28 -0400 Subject: [PATCH 63/83] perf(pi05): fuse native weight materialization --- USAGE.md | 16 +- .../cpp/models/pi05/native_weight_ops.h | 65 +- cpp/models/pi05/src/native_graph_owner.cpp | 27 + .../pi05/src/native_weight_materializer.cpp | 326 ++++----- cpp/models/pi05/src/native_weight_ops.cpp | 632 +++++++++++++----- cpp/tests/gate_pi05_native_weight_ops.py | 14 + cpp/tests/pi05_native_weight_probe.cpp | 177 +++-- .../test_pi05_native_weight_materializer.cpp | 7 +- cpp/tests/test_pi05_native_weight_ops.cpp | 91 ++- docs/cpp_runtime_modalities.md | 6 +- docs/mindon_pi05_integration.md | 8 + docs/pi05_io_contract.md | 32 +- docs/pr_review_checklist.md | 14 + 13 files changed, 899 insertions(+), 516 deletions(-) diff --git a/USAGE.md b/USAGE.md index b493e8f7..55192325 100644 --- a/USAGE.md +++ b/USAGE.md @@ -212,7 +212,7 @@ model = flash_rt.load_model( | `num_views` | `int` | `2` | Number of camera views. LIBERO uses 2 (base + wrist). | | `autotune` | `int\|bool` | `3` | CUDA Graph autotune intensity. See [Autotune](#autotune). | | `recalibrate` | `bool` | `False` | Force fresh FP8 calibration (and weight cache for JAX), ignoring cache. See [Calibration](#calibration). | -| `weight_cache` | `bool` | `True` | Cache FP8-quantized weights to disk. **JAX only** — reduces cold start from ~42s to ~6s. Torch loads in ~3s and ignores this. See [Weight Cache](#weight-cache-jax-only). | +| `weight_cache` | `bool` | `True` | Cache FP8-quantized weights to disk. **JAX only** — reduces cold start from ~42s to ~6s. Torch loads in ~3s and ignores this. See [Python JAX Weight Cache](#python-jax-weight-cache). | | `config` | `str` | `"pi05"` | Model architecture config: `"pi05"`, `"pi0"`, `"groot"`, `"groot_n17"`, `"pi0fast"`, `"motus"`, `"wan22_ti2v_5b"`, `"cosmos3_video"`. `"cosmos3_video"` is a non-VLA text2video denoise model — drive it with `set_prompt(ref=...)` + `infer(...)`, not `predict()`. | | `decode_cuda_graph` | `bool` | `False` | **Pi0-FAST only.** Capture action-phase decode as CUDA Graph. Trades startup time for per-token speed. See [Pi0-FAST](#pi0-fast). | | `decode_graph_steps` | `int` | `80` | **Pi0-FAST only.** Number of action tokens to capture in the decode graph. Should cover your longest expected action sequence. | @@ -1155,7 +1155,7 @@ actions = model.predict(images=[img3, img4], prompt="task B") # fresh calibrati --- -## Weight Cache (JAX only) +## Python JAX Weight Cache JAX (Orbax) checkpoint loading takes ~42s due to OCDBT deserialization + weight transform + FP8 quantization. The weight cache saves the final FP8-quantized engine weights to disk after first load, so subsequent loads skip all three steps. @@ -1163,10 +1163,18 @@ JAX (Orbax) checkpoint loading takes ~42s due to OCDBT deserialization + weight | Framework | Cold start | Bottleneck | |-----------|-----------|------------| -| **Torch** (safetensors) | ~3s | mmap load — already fast | +| **Python Torch** (safetensors) | ~3s | mmap-backed conversion and GPU upload | | **JAX** (Orbax) | ~42s → **~6s with cache** | OCDBT deserialize + transform + FP8 quant | -Torch uses safetensors which is essentially a flat binary mmap — there's nothing to cache. JAX's Orbax format requires complex deserialization that the weight cache eliminates. +This cache belongs to the Python JAX producer. Python Torch uses a flat, +mmap-backed safetensors source and does not write this cache. That does not mean +that mmap alone completes model startup: layout transforms, dtype conversion, +GPU upload, identity construction, workspace setup, and graph capture still +belong to the selected producer. The native C++ producer fuses its safetensors +conversion and layout transforms and reports complete setup phases with +`FLASHRT_PROFILE_NATIVE_SETUP=1`; it does not consume the JAX cache format. +JAX's Orbax format requires the additional deserialization work that this cache +eliminates. ### How it works 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 b68634fd..0f7dc4b5 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 @@ -22,34 +22,69 @@ struct NativeBf16Tensor { std::vector values; }; -modalities::Status load_native_float_tensor( +enum class NativeSourceDType { + kF32, + kBf16, + kF16, +}; + +struct NativeSourceTensorView { + const void* data = nullptr; + std::vector shape; + NativeSourceDType dtype = NativeSourceDType::kF32; +}; + +modalities::Status load_native_source_tensor( const loader::SafetensorsFile& file, const std::string& key, - NativeFloatTensor* out); + NativeSourceTensorView* out); -modalities::Status native_to_bf16(const NativeFloatTensor& input, - NativeBf16Tensor* out); - -modalities::Status native_f32_to_bf16( - const float* input, - const std::vector& shape, +modalities::Status native_source_to_bf16( + const NativeSourceTensorView& input, bool transpose, NativeBf16Tensor* out); -modalities::Status native_f32_fold_rms_columns_transpose( - const float* weight, - std::uint64_t rows, - std::uint64_t cols, +modalities::Status native_source_fold_rms_columns_transpose( + const NativeSourceTensorView& weight, const NativeFloatTensor& norm, NativeBf16Tensor* out); -modalities::Status native_f32_round_scale_to_bf16( - const float* input, - const std::vector& shape, +modalities::Status native_source_round_scale_to_bf16( + const NativeSourceTensorView& input, float scale, bool transpose, NativeBf16Tensor* out); +modalities::Status native_source_qkv_to_bf16( + const NativeSourceTensorView& q, + const NativeSourceTensorView& k, + const NativeSourceTensorView& v, + std::uint64_t q_heads, + std::uint64_t k_heads, + const NativeFloatTensor* norm, + NativeBf16Tensor* out); + +modalities::Status native_source_concat_vectors_to_bf16( + const std::vector& inputs, + NativeBf16Tensor* out); + +modalities::Status native_source_patch_oihw_to_hwio_bf16( + const NativeSourceTensorView& input, + NativeBf16Tensor* out); + +modalities::Status native_source_pair_transpose_concat_bf16( + const NativeSourceTensorView& left, + const NativeSourceTensorView& right, + NativeBf16Tensor* out); + +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); diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index ad449eb5..da678126 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -5,6 +5,9 @@ #include +#include +#include +#include #include namespace flashrt { @@ -82,11 +85,24 @@ std::unique_ptr NativeGraphOwner::create( modalities::Status NativeGraphOwner::initialize( const std::string& checkpoint_path) { + const bool profile_setup = std::getenv("FLASHRT_PROFILE_NATIVE_SETUP"); + const auto setup_begin = std::chrono::steady_clock::now(); + auto checkpoint = setup_begin; + const auto report = [&](const char* phase) { + const auto now = std::chrono::steady_clock::now(); + if (profile_setup) { + std::fprintf(stderr, "native_setup %s_ms=%.3f\n", phase, + std::chrono::duration( + now - checkpoint).count()); + } + checkpoint = now; + }; loader::SafetensorsFile source; if (!source.open(checkpoint_path + "/model.safetensors")) { return modalities::Status::error(modalities::StatusCode::kNotFound, source.error()); } + report("header"); NativeWeightMaterializer materializer(source, &weights_); NativeMaterializationOptions options; options.num_steps = config_.num_steps; @@ -94,6 +110,7 @@ modalities::Status NativeGraphOwner::initialize( options.include_embedding = true; modalities::Status st = materializer.materialize_all(options); if (!st.ok_status()) return st; + report("materialize"); NativeWorkspaceConfig workspace_config; workspace_config.num_views = config_.num_views; @@ -127,6 +144,7 @@ modalities::Status NativeGraphOwner::initialize( } st = attention_driver_->status(); if (!st.ok_status()) return st; + report("workspace_style"); for (const char* name : {"observation_images_normalized", "prompt_embedding", "diffusion_noise"}) { @@ -140,6 +158,7 @@ modalities::Status NativeGraphOwner::initialize( if (cudaDeviceSynchronize() != cudaSuccess) { return backend("native graph setup synchronization failed"); } + report("input_init"); infer_graph_ = frt_graph_create(ctx_, "infer", 1); const NativeWorkspaceBuffer* images = @@ -166,6 +185,7 @@ modalities::Status NativeGraphOwner::initialize( ? backend("native full graph variant is missing") : capture_status_; } + report("capture"); cudaStream_t stream = nullptr; if (cudaStreamCreate(&stream) != cudaSuccess) { @@ -174,6 +194,13 @@ modalities::Status NativeGraphOwner::initialize( replay_stream_ = stream; stream_id_ = frt_ctx_wrap_stream(ctx_, replay_stream_); if (stream_id_ < 0) return backend("native replay stream wrapping failed"); + report("stream"); + if (profile_setup) { + const auto now = std::chrono::steady_clock::now(); + std::fprintf(stderr, "native_setup total_ms=%.3f\n", + std::chrono::duration( + now - setup_begin).count()); + } return modalities::Status::ok(); } diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/native_weight_materializer.cpp index 62dfc74f..7531dbbb 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/native_weight_materializer.cpp @@ -1,5 +1,8 @@ #include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include +#include +#include #include namespace flashrt { @@ -28,24 +31,6 @@ const std::string& vision_prefix() { return prefix; } -const loader::SafetensorInfo* source_tensor( - const loader::SafetensorsFile& source, - const std::string& key) { - const loader::SafetensorInfo* tensor = source.find(key); - if (!tensor) tensor = source.find(std::string("model.") + key); - return tensor; -} - -const float* f32_data(const loader::SafetensorsFile& source, - const loader::SafetensorInfo* tensor) { - if (!tensor || tensor->dtype != "F32") return nullptr; - const void* data = source.data(*tensor); - if (reinterpret_cast(data) % alignof(float) != 0) { - return nullptr; - } - return static_cast(data); -} - std::string layer_name(const char* stem, int layer) { return std::string(stem) + std::to_string(layer); } @@ -71,77 +56,44 @@ modalities::Status NativeWeightMaterializer::upload( modalities::Status NativeWeightMaterializer::upload_rounded_transpose( const std::string& source_key, const std::string& destination_name) { - const loader::SafetensorInfo* info = source_tensor(source_, source_key); - if (destination_ && info && info->shape.size() == 2) { - const float* data = f32_data(source_, info); - if (data) { - NativeBf16Tensor converted; - modalities::Status st = native_f32_to_bf16( - data, info->shape, true, &converted); - if (!st.ok_status()) return st; - return destination_->upload(destination_name, converted); - } - } - 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 (!destination_) return invalid("native weight destination is null"); + NativeSourceTensorView source; + NativeBf16Tensor converted; + modalities::Status st = + load_native_source_tensor(source_, source_key, &source); if (!st.ok_status()) return st; - st = native_transpose_2d(rounded, &transposed); + st = native_source_to_bf16(source, true, &converted); if (!st.ok_status()) return st; - return upload(destination_name, transposed); + return destination_->upload(destination_name, converted); } modalities::Status NativeWeightMaterializer::upload_rounded_copy( const std::string& source_key, const std::string& destination_name) { - const loader::SafetensorInfo* info = source_tensor(source_, source_key); - if (destination_ && info) { - const float* data = f32_data(source_, info); - if (data) { - NativeBf16Tensor converted; - modalities::Status st = native_f32_to_bf16( - data, info->shape, false, &converted); - if (!st.ok_status()) return st; - return destination_->upload(destination_name, converted); - } - } - NativeFloatTensor source; - NativeFloatTensor rounded; - modalities::Status st = load(source_key, &source); + if (!destination_) return invalid("native weight destination is null"); + NativeSourceTensorView source; + NativeBf16Tensor converted; + modalities::Status st = + load_native_source_tensor(source_, source_key, &source); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(source, &rounded); + st = native_source_to_bf16(source, false, &converted); if (!st.ok_status()) return st; - return upload(destination_name, rounded); + return destination_->upload(destination_name, converted); } modalities::Status NativeWeightMaterializer::upload_folded_transpose( const std::string& source_key, const NativeFloatTensor& norm, const std::string& destination_name) { - const loader::SafetensorInfo* info = source_tensor(source_, source_key); - if (destination_ && info && info->shape.size() == 2) { - const float* data = f32_data(source_, info); - if (data) { - NativeBf16Tensor converted; - modalities::Status st = native_f32_fold_rms_columns_transpose( - data, info->shape[0], info->shape[1], norm, &converted); - if (!st.ok_status()) return st; - return destination_->upload(destination_name, converted); - } - } - 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 (!destination_) return invalid("native weight destination is null"); + NativeSourceTensorView source; + NativeBf16Tensor converted; + modalities::Status st = + load_native_source_tensor(source_, source_key, &source); if (!st.ok_status()) return st; - st = native_transpose_2d(folded, &transposed); + st = native_source_fold_rms_columns_transpose(source, norm, &converted); if (!st.ok_status()) return st; - return upload(destination_name, transposed); + return destination_->upload(destination_name, converted); } modalities::Status NativeWeightMaterializer::upload_rounded_scaled( @@ -149,34 +101,16 @@ modalities::Status NativeWeightMaterializer::upload_rounded_scaled( const std::string& destination_name, float scale, bool transpose) { - const loader::SafetensorInfo* info = source_tensor(source_, source_key); - if (destination_ && info) { - const float* data = f32_data(source_, info); - if (data) { - NativeBf16Tensor converted; - modalities::Status st = native_f32_round_scale_to_bf16( - data, info->shape, scale, transpose, &converted); - if (!st.ok_status()) return st; - return destination_->upload(destination_name, converted); - } - } - NativeFloatTensor source; - NativeFloatTensor rounded; - NativeFloatTensor arranged; - NativeFloatTensor scaled; - modalities::Status st = load(source_key, &source); + if (!destination_) return invalid("native weight destination is null"); + NativeSourceTensorView source; + NativeBf16Tensor converted; + modalities::Status st = + load_native_source_tensor(source_, source_key, &source); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(source, &rounded); + st = native_source_round_scale_to_bf16( + source, scale, transpose, &converted); 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); + return destination_->upload(destination_name, converted); } modalities::Status NativeWeightMaterializer::materialize_encoder_layer( @@ -189,34 +123,22 @@ modalities::Status NativeWeightMaterializer::materialize_encoder_layer( 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); + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + NativeBf16Tensor qkv; + st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.weight", &q); if (!st.ok_status()) return st; - st = native_interleave_qk_rows(q, 8, &qi); + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.weight", &k); if (!st.ok_status()) return st; - st = native_interleave_qk_rows(k, 1, &ki); + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.weight", &v); if (!st.ok_status()) return st; - st = native_fold_rms_columns(qi, norm, &qf); + st = native_source_qkv_to_bf16(q, k, v, 8, 1, &norm, &qkv); 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); + st = destination_->upload(layer_name("encoder_attn_qkv_w_", layer), qkv); if (!st.ok_status()) return st; st = upload_rounded_transpose( @@ -246,34 +168,22 @@ modalities::Status NativeWeightMaterializer::materialize_decoder_layer( 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); + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + NativeBf16Tensor qkv; + modalities::Status st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.weight", &q); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(v, &vr); + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.weight", &k); if (!st.ok_status()) return st; - st = native_interleave_qk_rows(qr, 8, &qi); + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.weight", &v); if (!st.ok_status()) return st; - st = native_interleave_qk_rows(kr, 1, &ki); + st = native_source_qkv_to_bf16(q, k, v, 8, 1, nullptr, &qkv); 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); + st = destination_->upload(layer_name("decoder_attn_qkv_w_", layer), qkv); if (!st.ok_status()) return st; st = upload_rounded_transpose( @@ -281,33 +191,28 @@ modalities::Status NativeWeightMaterializer::materialize_decoder_layer( 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); + st = upload_rounded_transpose( + prefix + ".mlp.gate_proj.weight", + layer_name("decoder_ffn_gate_w_", layer)); if (!st.ok_status()) return st; - st = upload(layer_name("decoder_ffn_up_w_", layer), up_t); + st = upload_rounded_transpose( + prefix + ".mlp.up_proj.weight", + layer_name("decoder_ffn_up_w_", layer)); if (!st.ok_status()) return st; if (merge_gate_up) { - NativeFloatTensor gate_up; - st = native_concat_columns(gate_t, up_t, &gate_up); + NativeSourceTensorView gate; + NativeSourceTensorView up; + NativeBf16Tensor gate_up; + st = load_native_source_tensor( + source_, prefix + ".mlp.gate_proj.weight", &gate); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".mlp.up_proj.weight", &up); + if (!st.ok_status()) return st; + st = native_source_pair_transpose_concat_bf16(gate, up, &gate_up); if (!st.ok_status()) return st; - st = upload(layer_name("decoder_ffn_gate_up_w_", layer), gate_up); + st = destination_->upload( + layer_name("decoder_ffn_gate_up_w_", layer), gate_up); if (!st.ok_status()) return st; } st = upload_rounded_transpose( @@ -339,45 +244,36 @@ modalities::Status NativeWeightMaterializer::materialize_vision_layer( } 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); + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + NativeBf16Tensor qkv; + modalities::Status st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.weight", &q); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(k, &kr); + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.weight", &k); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(v, &vr); + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.weight", &v); if (!st.ok_status()) return st; - st = native_concat_rows_transpose({&qr, &kr, &vr}, &qkv); + st = native_source_qkv_to_bf16(q, k, v, 0, 0, nullptr, &qkv); if (!st.ok_status()) return st; - st = upload(layer_name("vision_attn_qkv_w_", layer), qkv); + st = destination_->upload(layer_name("vision_attn_qkv_w_", layer), qkv); if (!st.ok_status()) return st; - st = load(prefix + ".self_attn.q_proj.bias", &q); + st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.bias", &q); if (!st.ok_status()) return st; - st = load(prefix + ".self_attn.k_proj.bias", &k); + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.bias", &k); if (!st.ok_status()) return st; - st = load(prefix + ".self_attn.v_proj.bias", &v); + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.bias", &v); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(q, &qr); + st = native_source_concat_vectors_to_bf16({&q, &k, &v}, &qkv); 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); + st = destination_->upload(layer_name("vision_attn_qkv_b_", layer), qkv); if (!st.ok_status()) return st; const struct { @@ -412,17 +308,14 @@ modalities::Status NativeWeightMaterializer::materialize_vision_layer( 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); + NativeSourceTensorView patch; + NativeBf16Tensor permuted; + modalities::Status st = load_native_source_tensor( + source_, prefix + ".embeddings.patch_embedding.weight", &patch); if (!st.ok_status()) return st; - st = native_round_to_bf16_float(patch, &rounded); + st = native_source_patch_oihw_to_hwio_bf16(patch, &permuted); 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); + st = destination_->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"); @@ -504,24 +397,41 @@ modalities::Status NativeWeightMaterializer::materialize_all( if (!destination_ || options.num_steps <= 0) { return invalid("Pi0.5 materialization options are invalid"); } + const bool profile = std::getenv("FLASHRT_PROFILE_NATIVE_SETUP"); + auto checkpoint = std::chrono::steady_clock::now(); + const auto report = [&](const char* phase) { + const auto now = std::chrono::steady_clock::now(); + if (profile) { + std::fprintf(stderr, "native_weights %s_ms=%.3f\n", phase, + std::chrono::duration( + now - checkpoint).count()); + } + checkpoint = now; + }; 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; } + report("vision"); for (int layer = 0; layer < 18; ++layer) { st = materialize_encoder_layer(layer); if (!st.ok_status()) return st; } + report("encoder"); for (int layer = 0; layer < 18; ++layer) { st = materialize_decoder_layer( layer, options.merge_decoder_gate_up); if (!st.ok_status()) return st; } + report("decoder"); st = materialize_decoder_globals(options.num_steps); if (!st.ok_status() || !options.include_embedding) return st; - return materialize_embedding(); + report("globals"); + st = materialize_embedding(); + report("embedding"); + return st; } } // namespace pi05 diff --git a/cpp/models/pi05/src/native_weight_ops.cpp b/cpp/models/pi05/src/native_weight_ops.cpp index ba9716f8..11f672b6 100644 --- a/cpp/models/pi05/src/native_weight_ops.cpp +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -47,7 +47,7 @@ void parallel_ranges(std::size_t count, const std::size_t wanted = (count + minimum_items - 1) / minimum_items; const std::size_t workers = std::max( - 1, std::min({16, available ? available : 1, wanted})); + 1, std::min({32, available ? available : 1, wanted})); if (workers == 1) { fn(0, count); return; @@ -63,6 +63,45 @@ void parallel_ranges(std::size_t count, for (std::thread& thread : threads) thread.join(); } +template +void tiled_transform_transpose(std::size_t rows, + std::size_t cols, + std::size_t output_rows, + std::size_t output_row_offset, + T* output, + const Fn& transform) { + constexpr std::size_t kTile = 32; + const std::size_t row_tiles = (rows + kTile - 1) / kTile; + parallel_ranges(row_tiles, 2, [&](std::size_t first, + std::size_t last) { + for (std::size_t tile = first; tile < last; ++tile) { + const std::size_t row_begin = tile * kTile; + const std::size_t row_end = std::min(rows, row_begin + kTile); + T scratch[kTile * kTile]; + for (std::size_t col_begin = 0; col_begin < cols; + col_begin += kTile) { + const std::size_t col_end = + std::min(cols, col_begin + kTile); + for (std::size_t row = row_begin; row < row_end; ++row) { + for (std::size_t col = col_begin; col < col_end; ++col) { + scratch[(row - row_begin) * kTile + + (col - col_begin)] = transform(row, col); + } + } + for (std::size_t col = col_begin; col < col_end; ++col) { + T* destination = output + col * output_rows + + output_row_offset + row_begin; + for (std::size_t row = row_begin; row < row_end; ++row) { + destination[row - row_begin] = + scratch[(row - row_begin) * kTile + + (col - col_begin)]; + } + } + } + } + }); +} + const loader::SafetensorInfo* find_source_tensor( const loader::SafetensorsFile& file, const std::string& key) { @@ -71,8 +110,437 @@ const loader::SafetensorInfo* find_source_tensor( return tensor; } +struct F32Reader { + const float* data; + float operator[](std::size_t index) const { return data[index]; } + std::uint16_t bf16(std::size_t index) const { + return modalities::float_to_bfloat16(data[index]); + } +}; + +struct Bf16Reader { + const std::uint16_t* data; + float operator[](std::size_t index) const { + return modalities::bfloat16_to_float(data[index]); + } + std::uint16_t bf16(std::size_t index) const { return data[index]; } +}; + +struct F16Reader { + const std::uint16_t* data; + float operator[](std::size_t index) const { + return modalities::float16_to_float(data[index]); + } + std::uint16_t bf16(std::size_t index) const { + return modalities::float_to_bfloat16( + modalities::float16_to_float(data[index])); + } +}; + +struct UnalignedF32Reader { + const unsigned char* data; + float operator[](std::size_t index) const { + float value = 0.0f; + std::memcpy(&value, data + index * sizeof(value), sizeof(value)); + return value; + } + std::uint16_t bf16(std::size_t index) const { + return modalities::float_to_bfloat16((*this)[index]); + } +}; + +template +struct UnalignedU16Reader { + const unsigned char* data; + std::uint16_t bits(std::size_t index) const { + std::uint16_t value = 0; + std::memcpy(&value, data + index * sizeof(value), sizeof(value)); + return value; + } + float operator[](std::size_t index) const { + return IsBf16 ? modalities::bfloat16_to_float(bits(index)) + : modalities::float16_to_float(bits(index)); + } + std::uint16_t bf16(std::size_t index) const { + return IsBf16 ? bits(index) + : modalities::float_to_bfloat16((*this)[index]); + } +}; + +template +modalities::Status dispatch_source(const NativeSourceTensorView& source, + const Fn& fn) { + if (!source.data) return invalid("native source tensor has no payload"); + switch (source.dtype) { + case NativeSourceDType::kF32: + if (reinterpret_cast(source.data) % + alignof(float) != 0) { + return fn(UnalignedF32Reader{ + static_cast(source.data)}); + } + return fn(F32Reader{static_cast(source.data)}); + case NativeSourceDType::kBf16: + if (reinterpret_cast(source.data) % + alignof(std::uint16_t) != 0) { + return fn(UnalignedU16Reader{ + static_cast(source.data)}); + } + return fn(Bf16Reader{ + static_cast(source.data)}); + case NativeSourceDType::kF16: + if (reinterpret_cast(source.data) % + alignof(std::uint16_t) != 0) { + return fn(UnalignedU16Reader{ + static_cast(source.data)}); + } + return fn(F16Reader{ + static_cast(source.data)}); + } + return invalid("native source tensor dtype is invalid"); +} + +std::size_t interleaved_source_row(std::size_t output_row, + std::size_t rows, + std::size_t heads) { + const std::size_t head_dim = rows / heads; + const std::size_t head = output_row / head_dim; + const std::size_t within = output_row % head_dim; + const std::size_t pair = within / 2; + const std::size_t half = within % 2; + return head * head_dim + half * (head_dim / 2) + pair; +} + } // namespace +modalities::Status load_native_source_tensor( + const loader::SafetensorsFile& file, + const std::string& key, + NativeSourceTensorView* out) { + if (!file.is_open() || !out) return invalid("invalid native tensor view"); + const loader::SafetensorInfo* tensor = find_source_tensor(file, key); + if (!tensor) { + return modalities::Status::error(modalities::StatusCode::kNotFound, + "native tensor not found: " + key); + } + NativeSourceDType dtype; + if (tensor->dtype == "F32") { + dtype = NativeSourceDType::kF32; + } else if (tensor->dtype == "BF16") { + dtype = NativeSourceDType::kBf16; + } else if (tensor->dtype == "F16") { + dtype = NativeSourceDType::kF16; + } else { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "native tensor dtype is not a floating-point weight: " + + tensor->dtype); + } + 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 payload is missing"); + *out = NativeSourceTensorView{data, tensor->shape, dtype}; + return modalities::Status::ok(); +} + +modalities::Status native_source_to_bf16( + const NativeSourceTensorView& input, + bool transpose, + NativeBf16Tensor* out) { + std::size_t count = 0; + if (!out || !element_count(input.shape, &count) || + (transpose && input.shape.size() != 2)) { + return invalid("invalid direct source to BF16 input"); + } + NativeBf16Tensor converted; + converted.shape = transpose + ? std::vector{input.shape[1], input.shape[0]} + : input.shape; + converted.values.resize(count); + modalities::Status st = dispatch_source(input, [&](const auto& reader) { + if (!transpose) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + converted.values[i] = reader.bf16(i); + } + }); + } else { + const std::size_t rows = static_cast(input.shape[0]); + const std::size_t cols = static_cast(input.shape[1]); + tiled_transform_transpose( + rows, cols, rows, 0, converted.values.data(), + [&](std::size_t row, std::size_t col) { + return reader.bf16(row * cols + col); + }); + } + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *out = std::move(converted); + return modalities::Status::ok(); +} + +modalities::Status native_source_fold_rms_columns_transpose( + const NativeSourceTensorView& weight, + const NativeFloatTensor& norm, + NativeBf16Tensor* out) { + std::size_t count = 0; + if (!out || weight.shape.size() != 2 || + !element_count(weight.shape, &count) || !valid_tensor(norm) || + norm.shape.size() != 1 || weight.shape[1] != norm.shape[0]) { + return invalid("invalid direct source RMS fold input"); + } + const std::size_t rows = static_cast(weight.shape[0]); + const std::size_t cols = static_cast(weight.shape[1]); + NativeBf16Tensor folded; + folded.shape = {weight.shape[1], weight.shape[0]}; + folded.values.resize(count); + modalities::Status st = dispatch_source(weight, [&](const auto& reader) { + tiled_transform_transpose( + rows, cols, rows, 0, folded.values.data(), + [&](std::size_t row, std::size_t col) { + return modalities::float_to_bfloat16( + reader[row * cols + col] * (1.0f + norm.values[col])); + }); + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *out = std::move(folded); + return modalities::Status::ok(); +} + +modalities::Status native_source_round_scale_to_bf16( + const NativeSourceTensorView& input, + float scale, + bool transpose, + NativeBf16Tensor* out) { + std::size_t count = 0; + if (!out || !element_count(input.shape, &count) || + (transpose && input.shape.size() != 2)) { + return invalid("invalid direct source scale input"); + } + NativeBf16Tensor scaled; + scaled.shape = transpose + ? std::vector{input.shape[1], input.shape[0]} + : input.shape; + scaled.values.resize(count); + modalities::Status st = dispatch_source(input, [&](const auto& reader) { + const auto convert = [&](std::size_t index) { + const float rounded = modalities::bfloat16_to_float( + modalities::float_to_bfloat16(reader[index])); + return modalities::float_to_bfloat16(rounded * scale); + }; + if (!transpose) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + scaled.values[i] = convert(i); + } + }); + } else { + const std::size_t rows = static_cast(input.shape[0]); + const std::size_t cols = static_cast(input.shape[1]); + tiled_transform_transpose( + rows, cols, rows, 0, scaled.values.data(), + [&](std::size_t row, std::size_t col) { + return convert(row * cols + col); + }); + } + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *out = std::move(scaled); + return modalities::Status::ok(); +} + +modalities::Status native_source_qkv_to_bf16( + const NativeSourceTensorView& q, + const NativeSourceTensorView& k, + const NativeSourceTensorView& v, + std::uint64_t q_heads, + std::uint64_t k_heads, + const NativeFloatTensor* norm, + NativeBf16Tensor* out) { + std::size_t q_count = 0; + std::size_t k_count = 0; + std::size_t v_count = 0; + if (!out || q.shape.size() != 2 || k.shape.size() != 2 || + v.shape.size() != 2 || q.shape[1] != k.shape[1] || + q.shape[1] != v.shape[1] || + !element_count(q.shape, &q_count) || + !element_count(k.shape, &k_count) || + !element_count(v.shape, &v_count) || + q.shape[0] > std::numeric_limits::max() - k.shape[0] || + q.shape[0] + k.shape[0] > + std::numeric_limits::max() - v.shape[0] || + (q_heads && (q.shape[0] % q_heads || + (q.shape[0] / q_heads) % 2)) || + (k_heads && (k.shape[0] % k_heads || + (k.shape[0] / k_heads) % 2)) || + (norm && (!valid_tensor(*norm) || norm->shape.size() != 1 || + norm->shape[0] != q.shape[1]))) { + return invalid("QKV source tensors have incompatible shapes"); + } + const std::size_t cols = static_cast(q.shape[1]); + const std::uint64_t total_rows_u64 = + q.shape[0] + k.shape[0] + v.shape[0]; + if (total_rows_u64 > std::numeric_limits::max() || + q.shape[1] > std::numeric_limits::max()) { + return invalid("QKV output shape overflows size_t"); + } + const std::size_t total_rows = static_cast(total_rows_u64); + NativeBf16Tensor joined; + joined.shape = {q.shape[1], total_rows_u64}; + std::size_t joined_count = 0; + if (!element_count(joined.shape, &joined_count)) { + return invalid("QKV output shape overflows size_t"); + } + joined.values.resize(joined_count); + const auto write = [&](const NativeSourceTensorView& source, + std::size_t heads, bool interleave, + std::size_t offset) { + return dispatch_source(source, [&](const auto& reader) { + const std::size_t rows = + static_cast(source.shape[0]); + tiled_transform_transpose( + rows, cols, total_rows, offset, joined.values.data(), + [&](std::size_t output_row, std::size_t col) { + const std::size_t source_row = interleave + ? interleaved_source_row(output_row, rows, heads) + : output_row; + const std::size_t index = source_row * cols + col; + if (!norm) return reader.bf16(index); + return modalities::float_to_bfloat16( + reader[index] * (1.0f + norm->values[col])); + }); + return modalities::Status::ok(); + }); + }; + modalities::Status st = write(q, q_heads, q_heads != 0, 0); + if (!st.ok_status()) return st; + st = write(k, k_heads, k_heads != 0, + static_cast(q.shape[0])); + if (!st.ok_status()) return st; + st = write(v, 1, false, + static_cast(q.shape[0] + k.shape[0])); + if (!st.ok_status()) return st; + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_source_concat_vectors_to_bf16( + const std::vector& inputs, + NativeBf16Tensor* out) { + if (!out || inputs.empty()) return invalid("vector concat has no inputs"); + std::size_t total = 0; + for (const NativeSourceTensorView* input : inputs) { + if (!input || input->shape.size() != 1 || + input->shape[0] > std::numeric_limits::max() || + input->shape[0] > std::numeric_limits::max() - total) { + return invalid("vector concat tensors have incompatible shapes"); + } + total += static_cast(input->shape[0]); + } + NativeBf16Tensor joined; + joined.shape = {static_cast(total)}; + joined.values.resize(total); + std::size_t offset = 0; + for (const NativeSourceTensorView* input : inputs) { + const std::size_t count = static_cast(input->shape[0]); + modalities::Status st = dispatch_source(*input, [&](const auto& reader) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + joined.values[offset + i] = reader.bf16(i); + } + }); + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + offset += count; + } + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_source_patch_oihw_to_hwio_bf16( + const NativeSourceTensorView& input, + NativeBf16Tensor* out) { + std::size_t count = 0; + if (!out || input.shape.size() != 4 || + !element_count(input.shape, &count)) { + return invalid("patch permutation requires a 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]); + NativeBf16Tensor result; + result.shape = {input.shape[2], input.shape[3], input.shape[1], + input.shape[0]}; + result.values.resize(count); + modalities::Status st = dispatch_source(input, [&](const auto& reader) { + parallel_ranges(height * width * channels, 32, + [&](std::size_t begin, std::size_t end) { + for (std::size_t hwc = begin; hwc < end; ++hwc) { + const std::size_t c = hwc % channels; + const std::size_t hw = hwc / channels; + const std::size_t h = hw / width; + const std::size_t w = hw % width; + for (std::size_t o = 0; o < outputs; ++o) { + const std::size_t src = + ((o * channels + c) * height + h) * width + w; + result.values[hwc * outputs + o] = reader.bf16(src); + } + } + }); + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *out = std::move(result); + return modalities::Status::ok(); +} + +modalities::Status native_source_pair_transpose_concat_bf16( + const NativeSourceTensorView& left, + const NativeSourceTensorView& right, + NativeBf16Tensor* out) { + std::size_t source_count = 0; + if (!out || left.shape.size() != 2 || right.shape.size() != 2 || + left.shape != right.shape || + !element_count(left.shape, &source_count) || + left.shape[0] > std::numeric_limits::max() / 2 || + source_count > std::numeric_limits::max() / 2) { + return invalid("paired transpose tensors have incompatible shapes"); + } + const std::size_t source_rows = static_cast(left.shape[0]); + const std::size_t output_rows = static_cast(left.shape[1]); + NativeBf16Tensor joined; + joined.shape = {left.shape[1], left.shape[0] + right.shape[0]}; + joined.values.resize(source_count * 2); + const auto write = [&](const NativeSourceTensorView& source, + std::size_t offset) { + return dispatch_source(source, [&](const auto& reader) { + tiled_transform_transpose( + source_rows, output_rows, source_rows * 2, offset, + joined.values.data(), + [&](std::size_t row, std::size_t col) { + return reader.bf16(row * output_rows + col); + }); + return modalities::Status::ok(); + }); + }; + modalities::Status st = write(left, 0); + if (!st.ok_status()) return st; + st = write(right, source_rows); + if (!st.ok_status()) return st; + *out = std::move(joined); + return modalities::Status::ok(); +} + modalities::Status load_native_float_tensor( const loader::SafetensorsFile& file, const std::string& key, @@ -138,123 +606,6 @@ modalities::Status native_to_bf16(const NativeFloatTensor& input, return modalities::Status::ok(); } -modalities::Status native_f32_to_bf16( - const float* input, - const std::vector& shape, - bool transpose, - NativeBf16Tensor* out) { - std::size_t count = 0; - if (!input || !out || !element_count(shape, &count) || - (transpose && shape.size() != 2)) { - return invalid("invalid direct F32 to BF16 input"); - } - NativeBf16Tensor converted; - converted.shape = transpose - ? std::vector{shape[1], shape[0]} - : shape; - converted.values.resize(count); - if (!transpose) { - parallel_ranges(count, 1 << 18, - [&](std::size_t begin, std::size_t end) { - for (std::size_t i = begin; i < end; ++i) { - converted.values[i] = modalities::float_to_bfloat16(input[i]); - } - }); - } else { - const std::size_t rows = static_cast(shape[0]); - const std::size_t cols = static_cast(shape[1]); - parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { - for (std::size_t row = begin; row < end; ++row) { - for (std::size_t col = 0; col < cols; ++col) { - converted.values[col * rows + row] = - modalities::float_to_bfloat16( - input[row * cols + col]); - } - } - }); - } - *out = std::move(converted); - return modalities::Status::ok(); -} - -modalities::Status native_f32_fold_rms_columns_transpose( - const float* weight, - std::uint64_t rows_u64, - std::uint64_t cols_u64, - const NativeFloatTensor& norm, - NativeBf16Tensor* out) { - if (!weight || !out || !valid_tensor(norm) || norm.shape.size() != 1 || - norm.shape[0] != cols_u64 || - rows_u64 > std::numeric_limits::max() || - cols_u64 > std::numeric_limits::max()) { - return invalid("invalid direct F32 RMS fold input"); - } - const std::size_t rows = static_cast(rows_u64); - const std::size_t cols = static_cast(cols_u64); - if (rows && cols > std::numeric_limits::max() / rows) { - return invalid("direct F32 RMS fold shape overflows size_t"); - } - NativeBf16Tensor folded; - folded.shape = {cols_u64, rows_u64}; - folded.values.resize(rows * cols); - parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { - for (std::size_t row = begin; row < end; ++row) { - for (std::size_t col = 0; col < cols; ++col) { - folded.values[col * rows + row] = - modalities::float_to_bfloat16( - weight[row * cols + col] * - (1.0f + norm.values[col])); - } - } - }); - *out = std::move(folded); - return modalities::Status::ok(); -} - -modalities::Status native_f32_round_scale_to_bf16( - const float* input, - const std::vector& shape, - float scale, - bool transpose, - NativeBf16Tensor* out) { - std::size_t count = 0; - if (!input || !out || !element_count(shape, &count) || - (transpose && shape.size() != 2)) { - return invalid("invalid direct F32 scale input"); - } - NativeBf16Tensor scaled; - scaled.shape = transpose - ? std::vector{shape[1], shape[0]} - : shape; - scaled.values.resize(count); - const auto convert = [scale](float value) { - const float rounded = modalities::bfloat16_to_float( - modalities::float_to_bfloat16(value)); - return modalities::float_to_bfloat16(rounded * scale); - }; - if (!transpose) { - parallel_ranges(count, 1 << 18, - [&](std::size_t begin, std::size_t end) { - for (std::size_t i = begin; i < end; ++i) { - scaled.values[i] = convert(input[i]); - } - }); - } else { - const std::size_t rows = static_cast(shape[0]); - const std::size_t cols = static_cast(shape[1]); - parallel_ranges(rows, 16, [&](std::size_t begin, std::size_t end) { - for (std::size_t row = begin; row < end; ++row) { - for (std::size_t col = 0; col < cols; ++col) { - scaled.values[col * rows + row] = - convert(input[row * cols + col]); - } - } - }); - } - *out = std::move(scaled); - return modalities::Status::ok(); -} - modalities::Status native_round_to_bf16_float( const NativeFloatTensor& input, NativeFloatTensor* out) { @@ -283,26 +634,11 @@ modalities::Status native_transpose_2d(const NativeFloatTensor& input, NativeFloatTensor transposed; transposed.shape = {input.shape[1], input.shape[0]}; transposed.values.resize(input.values.size()); - constexpr std::size_t kTile = 32; - const std::size_t row_tiles = (rows + kTile - 1) / kTile; - parallel_ranges(row_tiles, 2, [&](std::size_t first, - std::size_t last) { - for (std::size_t tile = first; tile < last; ++tile) { - const std::size_t row_begin = tile * kTile; - const std::size_t row_end = std::min(rows, row_begin + kTile); - for (std::size_t col_begin = 0; col_begin < cols; - col_begin += kTile) { - const std::size_t col_end = - std::min(cols, col_begin + kTile); - for (std::size_t col = col_begin; col < col_end; ++col) { - for (std::size_t row = row_begin; row < row_end; ++row) { - transposed.values[col * rows + row] = - input.values[row * cols + col]; - } - } - } - } - }); + tiled_transform_transpose( + rows, cols, rows, 0, transposed.values.data(), + [&](std::size_t row, std::size_t col) { + return input.values[row * cols + col]; + }); *out = std::move(transposed); return modalities::Status::ok(); } @@ -423,17 +759,15 @@ modalities::Status native_concat_rows_transpose( for (const NativeFloatTensor* input : inputs) { const std::uint64_t input_rows = input->shape[0]; const std::uint64_t output_offset = row_offset; - parallel_ranges(static_cast(input_rows), 32, - [&](std::size_t begin, std::size_t end) { - for (std::size_t row = begin; row < end; ++row) { - for (std::uint64_t col = 0; col < cols; ++col) { - joined.values[static_cast( - col * total_rows + output_offset + row)] = - input->values[static_cast( - row * cols + col)]; - } - } - }); + tiled_transform_transpose( + static_cast(input_rows), + static_cast(cols), + static_cast(total_rows), + static_cast(output_offset), joined.values.data(), + [&](std::size_t row, std::size_t col) { + return input->values[ + row * static_cast(cols) + col]; + }); row_offset += input->shape[0]; } *out = std::move(joined); diff --git a/cpp/tests/gate_pi05_native_weight_ops.py b/cpp/tests/gate_pi05_native_weight_ops.py index ad75773a..2276606d 100644 --- a/cpp/tests/gate_pi05_native_weight_ops.py +++ b/cpp/tests/gate_pi05_native_weight_ops.py @@ -83,6 +83,20 @@ 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() + expected["encoder_o0_fast"] = bf16( + f"{ENCODER}.self_attn.o_proj.weight" + ).t().contiguous() + ffn_norm = 1.0 + raw( + f"{ENCODER}.post_attention_layernorm.weight" + ).float() + expected["encoder_gate0_fast"] = ( + raw(f"{ENCODER}.mlp.gate_proj.weight").float() + * ffn_norm.unsqueeze(0) + ).t().to(torch.bfloat16).contiguous() + expected["decoder_mod_bias0_fast"] = bf16( + f"{DECODER}.input_layernorm.dense.bias" + ).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 diff --git a/cpp/tests/pi05_native_weight_probe.cpp b/cpp/tests/pi05_native_weight_probe.cpp index a8c4d036..872a041d 100644 --- a/cpp/tests/pi05_native_weight_probe.cpp +++ b/cpp/tests/pi05_native_weight_probe.cpp @@ -5,7 +5,6 @@ #include #include #include -#include #include namespace { @@ -13,6 +12,7 @@ namespace { using flashrt::loader::SafetensorsFile; using flashrt::models::pi05::NativeBf16Tensor; using flashrt::models::pi05::NativeFloatTensor; +using flashrt::models::pi05::NativeSourceTensorView; using flashrt::modalities::Status; constexpr const char* kVision = @@ -22,6 +22,14 @@ constexpr const char* kEncoder = constexpr const char* kDecoder = "paligemma_with_expert.gemma_expert.model.layers.0"; +bool source_view(const SafetensorsFile& file, const std::string& key, + NativeSourceTensorView* out) { + const Status st = + flashrt::models::pi05::load_native_source_tensor(file, key, out); + if (!st.ok_status()) std::cerr << st.message << '\n'; + return st.ok_status(); +} + bool load(const SafetensorsFile& file, const std::string& key, NativeFloatTensor* out) { const Status st = @@ -30,118 +38,87 @@ bool load(const SafetensorsFile& file, const std::string& key, 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); + NativeSourceTensorView source; + return source_view(file, std::string(kVision) + + ".embeddings.patch_embedding.weight", + &source) && + flashrt::models::pi05::native_source_patch_oihw_to_hwio_bf16( + source, out).ok_status(); } 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()) { + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + if (!source_view(file, prefix + ".self_attn.q_proj.weight", &q) || + !source_view(file, prefix + ".self_attn.k_proj.weight", &k) || + !source_view(file, prefix + ".self_attn.v_proj.weight", &v)) { return false; } + NativeFloatTensor norm; + const NativeFloatTensor* norm_ptr = nullptr; 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); + if (!load(file, prefix + ".input_layernorm.weight", &norm)) return false; + norm_ptr = &norm; } - NativeFloatTensor joined; - return flashrt::models::pi05::native_concat_rows_transpose( - {&qi, &ki, &v_input}, &joined).ok_status() && - finish(joined, out); + return flashrt::models::pi05::native_source_qkv_to_bf16( + q, k, v, 8, 1, norm_ptr, out).ok_status(); } 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); + NativeSourceTensorView gate; + NativeSourceTensorView up; + return source_view(file, std::string(kDecoder) + + ".mlp.gate_proj.weight", &gate) && + source_view(file, std::string(kDecoder) + + ".mlp.up_proj.weight", &up) && + flashrt::models::pi05::native_source_pair_transpose_concat_bf16( + gate, up, out).ok_status(); } 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); + NativeSourceTensorView source; + return source_view(file, "action_out_proj.weight", &source) && + flashrt::models::pi05::native_source_round_scale_to_bf16( + source, -1.0f / static_cast(num_steps), + true, out).ok_status(); +} + +bool rounded_transpose(const SafetensorsFile& file, + const std::string& key, + NativeBf16Tensor* out) { + NativeSourceTensorView source; + return source_view(file, key, &source) && + flashrt::models::pi05::native_source_to_bf16(source, true, out) + .ok_status(); +} + +bool rounded_copy(const SafetensorsFile& file, + const std::string& key, + NativeBf16Tensor* out) { + NativeSourceTensorView source; + return source_view(file, key, &source) && + flashrt::models::pi05::native_source_to_bf16(source, false, out) + .ok_status(); +} + +bool folded_transpose(const SafetensorsFile& file, + const std::string& key, + const std::string& norm_key, + NativeBf16Tensor* out) { + NativeSourceTensorView source; + NativeFloatTensor norm; + return source_view(file, key, &source) && load(file, norm_key, &norm) && + flashrt::models::pi05::native_source_fold_rms_columns_transpose( + source, norm, out) + .ok_status(); } bool time_embeds(int num_steps, NativeBf16Tensor* out) { @@ -185,6 +162,20 @@ 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 == "encoder_o0_fast") { + ok = rounded_transpose( + file, std::string(kEncoder) + ".self_attn.o_proj.weight", + &output); + } else if (op == "encoder_gate0_fast") { + ok = folded_transpose( + file, std::string(kEncoder) + ".mlp.gate_proj.weight", + std::string(kEncoder) + ".post_attention_layernorm.weight", + &output); + } else if (op == "decoder_mod_bias0_fast") { + ok = rounded_copy( + file, std::string(kDecoder) + + ".input_layernorm.dense.bias", + &output); } else if (op == "action_out10") { ok = action_out(file, 10, &output); } else if (op == "action_out5") { diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp index b3f22712..9703229b 100644 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ b/cpp/tests/test_pi05_native_weight_materializer.cpp @@ -181,7 +181,12 @@ int main() { flashrt::models::pi05::NativeDeviceWeightStore destination(ctx); flashrt::models::pi05::NativeWeightMaterializer materializer( source, &destination); - assert(materializer.materialize_encoder_layer(0).ok_status()); + const auto encoder_status = materializer.materialize_encoder_layer(0); + if (!encoder_status.ok_status()) { + std::fprintf(stderr, "encoder materialization failed: %s\n", + encoder_status.message.c_str()); + } + assert(encoder_status.ok_status()); assert(destination.size() == 5); const auto* qkv = destination.find("encoder_attn_qkv_w_0"); assert(qkv && qkv->shape == std::vector({4, 24})); diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp index 545040ac..21b78f3c 100644 --- a/cpp/tests/test_pi05_native_weight_ops.cpp +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -12,6 +12,8 @@ namespace { using flashrt::models::pi05::NativeBf16Tensor; using flashrt::models::pi05::NativeFloatTensor; +using flashrt::models::pi05::NativeSourceDType; +using flashrt::models::pi05::NativeSourceTensorView; void expect(const NativeFloatTensor& tensor, std::initializer_list shape, @@ -24,6 +26,18 @@ void expect(const NativeFloatTensor& tensor, } } +void expect_bf16(const NativeBf16Tensor& 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(tensor.values[i++] == + flashrt::modalities::float_to_bfloat16(value)); + } +} + std::string temp_path() { char path[] = "/tmp/frt_pi05_weight_ops_XXXXXX"; const int fd = ::mkstemp(path); @@ -76,6 +90,21 @@ int main() { expect(loaded, {2}, {3.0f, -4.0f}); assert(load_native_float_tensor(file, "f16", &loaded).ok_status()); expect(loaded, {2}, {5.0f, -6.0f}); + + NativeSourceTensorView mapped; + NativeBf16Tensor mapped_bf16; + assert(load_native_source_tensor(file, "f32", &mapped).ok_status()); + assert(mapped.dtype == NativeSourceDType::kF32); + assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); + expect_bf16(mapped_bf16, {2}, {1.25f, -2.5f}); + assert(load_native_source_tensor(file, "bf16", &mapped).ok_status()); + assert(mapped.dtype == NativeSourceDType::kBf16); + assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); + expect_bf16(mapped_bf16, {2}, {3.0f, -4.0f}); + assert(load_native_source_tensor(file, "f16", &mapped).ok_status()); + assert(mapped.dtype == NativeSourceDType::kF16); + assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); + expect_bf16(mapped_bf16, {2}, {5.0f, -6.0f}); assert(::unlink(path.c_str()) == 0); NativeFloatTensor matrix{{2, 3}, {1, 2, 3, 4, 5, 6}}; @@ -139,37 +168,37 @@ int main() { } NativeBf16Tensor direct; - assert(native_f32_to_bf16(matrix.values.data(), matrix.shape, false, - &direct).ok_status()); - assert(direct.shape == converted.shape && - direct.values == converted.values); - assert(native_f32_to_bf16(matrix.values.data(), matrix.shape, true, - &direct).ok_status()); - NativeFloatTensor transposed; - NativeBf16Tensor transposed_bf16; - assert(native_transpose_2d(matrix, &transposed).ok_status()); - assert(native_to_bf16(transposed, &transposed_bf16).ok_status()); - assert(direct.shape == transposed_bf16.shape && - direct.values == transposed_bf16.values); - - assert(native_f32_fold_rms_columns_transpose( - matrix.values.data(), 2, 3, norm, &direct).ok_status()); - NativeFloatTensor folded; - assert(native_fold_rms_columns(matrix, norm, &folded).ok_status()); - assert(native_transpose_2d(folded, &transposed).ok_status()); - assert(native_to_bf16(transposed, &transposed_bf16).ok_status()); - assert(direct.shape == transposed_bf16.shape && - direct.values == transposed_bf16.values); - - constexpr float kScale = -0.1f; - assert(native_f32_round_scale_to_bf16( - unrounded.values.data(), unrounded.shape, kScale, false, - &direct).ok_status()); - NativeFloatTensor rounded_scaled; - assert(native_scale(result, kScale, &rounded_scaled).ok_status()); - assert(native_to_bf16(rounded_scaled, &converted).ok_status()); - assert(direct.shape == converted.shape && - direct.values == converted.values); + const float source_f32[] = {1, 2, 3, 4, 5, 6, 7, 8}; + std::uint16_t source_bf16[8]; + std::uint16_t source_f16[8]; + for (std::size_t i = 0; i < 8; ++i) { + source_bf16[i] = flashrt::modalities::float_to_bfloat16(source_f32[i]); + source_f16[i] = flashrt::modalities::float_to_float16(source_f32[i]); + } + const NativeSourceTensorView source_views[] = { + {source_f32, {2, 4}, NativeSourceDType::kF32}, + {source_bf16, {2, 4}, NativeSourceDType::kBf16}, + {source_f16, {2, 4}, NativeSourceDType::kF16}, + }; + NativeFloatTensor source_norm{{4}, {0, 0, 0, 0}}; + for (const NativeSourceTensorView& source_view : source_views) { + assert(native_source_to_bf16(source_view, true, &direct).ok_status()); + expect_bf16(direct, {4, 2}, {1, 5, 2, 6, 3, 7, 4, 8}); + assert(native_source_fold_rms_columns_transpose( + source_view, source_norm, &direct).ok_status()); + expect_bf16(direct, {4, 2}, {1, 5, 2, 6, 3, 7, 4, 8}); + assert(native_source_round_scale_to_bf16( + source_view, 2.0f, true, &direct).ok_status()); + expect_bf16(direct, {4, 2}, {2, 10, 4, 12, 6, 14, 8, 16}); + assert(native_source_qkv_to_bf16( + source_view, source_view, source_view, 1, 1, nullptr, + &direct).ok_status()); + expect_bf16(direct, {4, 6}, + {1, 5, 1, 5, 1, 5, + 2, 6, 2, 6, 2, 6, + 3, 7, 3, 7, 3, 7, + 4, 8, 4, 8, 4, 8}); + } assert(!native_interleave_qk_rows(matrix, 2, &result).ok_status()); assert(!native_concat_columns(matrix, k, &result).ok_status()); diff --git a/docs/cpp_runtime_modalities.md b/docs/cpp_runtime_modalities.md index e53e5856..f3554e41 100644 --- a/docs/cpp_runtime_modalities.md +++ b/docs/cpp_runtime_modalities.md @@ -82,8 +82,10 @@ Current Pi0.5 status: resize/normalize/cast directly into export device buffers; - conservative action staging path: device action buffer -> D2H -> CPU reference postprocess; -- native checkpoint loader/tokenizer/capture is not implemented yet. It will - become a producer for the same `frt_runtime_export_v1`, not a Nexus feature. +- native SM120 checkpoint loading, tokenizer/prompt staging, weight + materialization, and graph capture are implemented by the optional + `frt_model_runtime_open_v1` producer. They present the same model-runtime ABI + and remain FlashRT responsibilities, not Nexus features. ## CPU Reference First diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 7788a62b..38ed51c7 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -95,6 +95,14 @@ portable implementation with identical identity bytes. The model hash and weight materialization execute concurrently, but the factory publishes no runtime until both have succeeded. +The native loader maps the checkpoint read-only and directly emits final BF16 +device layouts from F32, BF16, or F16 source tensors. QKV interleave/concat, +RMS folding, patch permutation, transpose, and output scaling are fused into +that pass; there is no checkpoint-sized Python dictionary or chain of float +intermediates. `FLASHRT_PROFILE_NATIVE_SETUP=1` reports header, materialization, +workspace/style, input initialization, capture, stream, and total setup time. +This diagnostic is setup-only and does not change the runtime contract. + ## 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 c865a134..e48cdfc2 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -247,12 +247,15 @@ gates compare the resulting BF16 bytes against PyTorch for both bare OpenPI keys and LeRobot `model.`-prefixed keys. Build native producers with `CMAKE_BUILD_TYPE=Release` for deployment. The -Pi0.5 checkpoint stores F32 tensors, so setup converts and lays out independent -weight ranges in parallel and writes the final BF16 payload directly instead -of creating intermediate rounded and transposed F32 copies. These direct -transforms are byte-compared with the reference multi-step transforms in the -weight-op tests. Debug builds intentionally retain unoptimized setup code and -are not a startup-latency reference. +native source view accepts F32, BF16, and F16 safetensors without materializing +a checkpoint-sized float dictionary. Type dispatch occurs once per tensor; +aligned payloads use typed loops and valid unaligned safetensors ranges use +safe scalar reads. Plain copy/transpose preserves source BF16 bits. Conversion, +RMS fold, Q/K interleave, QKV concat/transpose, patch OIHW-to-HWIO layout, and +action scaling write the final BF16 payload directly. They do not create +rounded, permuted, or concatenated F32 intermediates. Tests cover every source +dtype, and real-checkpoint gates byte-compare all fused transform families with +the PyTorch multi-step reference. Checkpoint identity remains a full-file SHA-256, not a path, timestamp, or partial-content surrogate. When OpenSSL is available at configure time the @@ -261,13 +264,16 @@ remains the build fallback. Model hashing runs concurrently with independent weight materialization, and both must complete successfully before the runtime descriptor is published. -For a 14.47 GB F32 Pi0.5 safetensors checkpoint on RTX 5090 SM120, the Release -`pi05_native_open_probe` full lifecycle measured 9.47 seconds with a warm file -cache and 10.91 seconds after evicting that file from the page cache. This -scope includes identity hashing, weight conversion/upload, workspace setup, -CUDA graph capture, one inference, and teardown. Compare producer startup -using this complete scope; a Python `load_model` timer that excludes graph -capture and identity construction is not the same metric. +For a 14.47 GB F32 Pi0.5 safetensors checkpoint on RTX 5090 SM120, a Release +validation run measured 3.52 seconds inside native setup and 4.54 seconds for +the `pi05_native_open_probe` full lifecycle. The setup breakdown was 3.38 +seconds for weight materialization/upload, 91 ms for workspace/style setup, +and 52 ms for CUDA graph capture. The full lifecycle additionally includes +identity completion, one inference, output handling, and teardown. These are +reference measurements, not a latency ABI. Use +`FLASHRT_PROFILE_NATIVE_SETUP=1` to print the same setup phase breakdown on a +deployment system. Compare producer startup using the complete scope; a Python +timer that stops after safetensors conversion is not the same metric. 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 diff --git a/docs/pr_review_checklist.md b/docs/pr_review_checklist.md index 6efc0dd7..3d9af35e 100644 --- a/docs/pr_review_checklist.md +++ b/docs/pr_review_checklist.md @@ -140,6 +140,20 @@ Blockers: - A hot-path verb (`set_input`/`get_output`, SWAP writes, tick) allocates, recaptures, or rebinds graph pointers. +Native checkpoint-loading changes additionally require: + +- A phase profile that separates file/header work, materialization and upload, + workspace setup, graph capture, and complete producer startup. +- Byte-level transform parity against the established producer, including the + exact order of source rounding, scaling, folding, and final dtype conversion. +- Coverage for every declared source dtype and for valid unaligned safetensors + payload offsets; a fast path must not silently become the only valid path. +- A real-input end-to-end numerical gate after transform fusion or upload-order + changes. A lower startup time is not evidence of numerical equivalence. +- No checkpoint-sized duplicate tensor dictionary or transformed layout cache + unless measurements show that its lifecycle, invalidation key, disk cost, + and restart benefit justify the added mechanism. + ## 6. Public API And Import Boundaries Required: From 7dc7e1f91a5f5c5a4b627685800d23eea021d982 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 11:38:44 -0400 Subject: [PATCH 64/83] feat(cpp): extend native weight loading primitives --- .../include/flashrt/cpp/loader/safetensors.h | 4 + cpp/loader/src/safetensors.cpp | 51 ++- cpp/modalities/src/types.cpp | 52 ++- cpp/modalities/src/vision_cpu.cpp | 7 + .../cpp/models/pi05/native_device_weights.h | 5 + .../cpp/models/pi05/native_quantization.h | 5 + .../cpp/models/pi05/native_weight_ops.h | 43 ++ cpp/models/pi05/src/native_device_weights.cpp | 10 + cpp/models/pi05/src/native_quantization.cu | 58 ++- .../src/native_quantization_unavailable.cpp | 9 + cpp/models/pi05/src/native_weight_ops.cpp | 372 ++++++++++++++++-- cpp/tests/pi05_native_quant_probe.cpp | 53 ++- cpp/tests/test_device_staging.cpp | 10 +- cpp/tests/test_modalities.cpp | 27 ++ cpp/tests/test_pi05_native_quantization.cpp | 21 + cpp/tests/test_pi05_native_weight_ops.cpp | 109 +++++ cpp/tests/test_safetensors_loader.cpp | 20 + 17 files changed, 801 insertions(+), 55 deletions(-) diff --git a/cpp/loader/include/flashrt/cpp/loader/safetensors.h b/cpp/loader/include/flashrt/cpp/loader/safetensors.h index e0fa3c5e..e63d7803 100644 --- a/cpp/loader/include/flashrt/cpp/loader/safetensors.h +++ b/cpp/loader/include/flashrt/cpp/loader/safetensors.h @@ -37,6 +37,9 @@ class SafetensorsFile { const std::map& tensors() const { return tensors_; } + const std::map& metadata() const { + return metadata_; + } const SafetensorInfo* find(const std::string& name) const; const void* data(const SafetensorInfo& tensor) const; @@ -50,6 +53,7 @@ class SafetensorsFile { std::string path_; std::string error_; std::map tensors_; + std::map metadata_; }; } // namespace loader diff --git a/cpp/loader/src/safetensors.cpp b/cpp/loader/src/safetensors.cpp index 23eef574..c8c383aa 100644 --- a/cpp/loader/src/safetensors.cpp +++ b/cpp/loader/src/safetensors.cpp @@ -49,11 +49,12 @@ class HeaderParser { HeaderParser(const char* begin, const char* end) : begin_(begin), cur_(begin), end_(end) {} - bool parse(std::map* tensors) { + bool parse(std::map* tensors, + std::map* metadata) { skip_ws(); if (!consume('{')) return fail("safetensors header must be an object"); skip_ws(); - if (consume('}')) return finish(tensors); + if (consume('}')) return finish(tensors, metadata); while (cur_ < end_) { std::string key; if (!parse_string(&key)) return false; @@ -61,7 +62,11 @@ class HeaderParser { if (!consume(':')) return fail("expected ':' after tensor name"); skip_ws(); if (key == "__metadata__") { - if (!skip_value()) return false; + if (metadata_seen_) { + return fail("duplicate safetensors metadata object"); + } + metadata_seen_ = true; + if (!parse_metadata()) return false; } else { SafetensorInfo tensor; if (!parse_tensor(&tensor)) return false; @@ -70,7 +75,7 @@ class HeaderParser { } } skip_ws(); - if (consume('}')) return finish(tensors); + if (consume('}')) return finish(tensors, metadata); if (!consume(',')) return fail("expected ',' or '}' in header"); skip_ws(); } @@ -80,14 +85,44 @@ class HeaderParser { const std::string& error() const { return error_; } private: - bool finish(std::map* tensors) { + bool finish(std::map* tensors, + std::map* metadata) { 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_); + if (metadata) *metadata = std::move(metadata_); return true; } + bool parse_metadata() { + if (!consume('{')) return fail("safetensors metadata must be an object"); + skip_ws(); + if (consume('}')) return true; + while (cur_ < end_) { + std::string key; + if (!parse_string(&key)) return false; + skip_ws(); + if (!consume(':')) return fail("expected ':' in safetensors metadata"); + skip_ws(); + std::string value; + if (cur_ >= end_ || *cur_ != '"') { + return fail("safetensors metadata values must be strings"); + } + if (!parse_string(&value)) return false; + if (!metadata_.emplace(std::move(key), std::move(value)).second) { + return fail("duplicate safetensors metadata key"); + } + skip_ws(); + if (consume('}')) return true; + if (!consume(',')) { + return fail("expected ',' in safetensors metadata"); + } + skip_ws(); + } + return fail("unterminated safetensors metadata"); + } + bool parse_tensor(SafetensorInfo* tensor) { if (!consume('{')) return fail("tensor metadata must be an object"); bool have_dtype = false; @@ -320,6 +355,8 @@ class HeaderParser { const char* end_; std::string error_; std::map parsed_; + std::map metadata_; + bool metadata_seen_ = false; }; } // namespace @@ -346,6 +383,7 @@ void SafetensorsFile::move_from(SafetensorsFile&& other) noexcept { path_ = std::move(other.path_); error_ = std::move(other.error_); tensors_ = std::move(other.tensors_); + metadata_ = std::move(other.metadata_); other.fd_ = -1; other.mapping_ = nullptr; other.mapping_bytes_ = 0; @@ -391,7 +429,7 @@ bool SafetensorsFile::open(const std::string& path) { data_offset_ = 8 + header_bytes; const char* header = static_cast(mapping_) + 8; HeaderParser parser(header, header + header_bytes); - if (!parser.parse(&tensors_)) { + if (!parser.parse(&tensors_, &metadata_)) { error_ = parser.error(); close(); return false; @@ -436,6 +474,7 @@ void SafetensorsFile::close() { data_offset_ = 0; path_.clear(); tensors_.clear(); + metadata_.clear(); } const SafetensorInfo* SafetensorsFile::find(const std::string& name) const { diff --git a/cpp/modalities/src/types.cpp b/cpp/modalities/src/types.cpp index 95f64739..85f21ac0 100644 --- a/cpp/modalities/src/types.cpp +++ b/cpp/modalities/src/types.cpp @@ -74,25 +74,47 @@ std::uint16_t float_to_float16(float value) { std::uint32_t x = 0; std::memcpy(&x, &value, sizeof(x)); const std::uint32_t sign = (x >> 16) & 0x8000u; - std::int32_t exp = static_cast((x >> 23) & 0xffu) - 127 + 15; + const std::uint32_t exponent_bits = (x >> 23) & 0xffu; std::uint32_t mant = x & 0x7fffffu; - if (exp <= 0) { - if (exp < -10) return static_cast(sign); - mant |= 0x800000u; - const std::uint32_t shift = static_cast(14 - exp); - std::uint32_t half = mant >> shift; - if ((mant >> (shift - 1)) & 1u) half += 1; - return static_cast(sign | half); + if (exponent_bits == 0xffu) { + if (!mant) return static_cast(sign | 0x7c00u); + return static_cast(sign | 0x7e00u); } - if (exp >= 31) { - if (mant == 0) return static_cast(sign | 0x7c00u); - return static_cast(sign | 0x7c00u | (mant >> 13) | 1u); + const std::int32_t exponent = + static_cast(exponent_bits) - 127; + if (exponent > 15) { + return static_cast(sign | 0x7c00u); } - - std::uint32_t half = (static_cast(exp) << 10) | (mant >> 13); - if (mant & 0x1000u) half += 1; - return static_cast(sign | half); + if (exponent >= -14) { + std::uint32_t half_exponent = + static_cast(exponent + 15); + std::uint32_t half_mantissa = mant >> 13; + const std::uint32_t remainder = mant & 0x1fffu; + if (remainder > 0x1000u || + (remainder == 0x1000u && (half_mantissa & 1u))) { + if (++half_mantissa == 0x400u) { + half_mantissa = 0; + if (++half_exponent == 31u) { + return static_cast(sign | 0x7c00u); + } + } + } + return static_cast( + sign | (half_exponent << 10) | half_mantissa); + } + if (exponent < -25) return static_cast(sign); + + mant |= 0x800000u; + const std::uint32_t shift = static_cast(-exponent - 1); + std::uint32_t half_mantissa = mant >> shift; + const std::uint32_t remainder = mant & ((1u << shift) - 1u); + const std::uint32_t halfway = 1u << (shift - 1u); + if (remainder > halfway || + (remainder == halfway && (half_mantissa & 1u))) { + ++half_mantissa; + } + return static_cast(sign | half_mantissa); } float float16_to_float(std::uint16_t value) { diff --git a/cpp/modalities/src/vision_cpu.cpp b/cpp/modalities/src/vision_cpu.cpp index cc217d4e..3a93098d 100644 --- a/cpp/modalities/src/vision_cpu.cpp +++ b/cpp/modalities/src/vision_cpu.cpp @@ -28,6 +28,13 @@ Status vision_staging_create(VisionStaging* out, std::uint32_t n_views, return Status::error(StatusCode::kInvalidArgument, "invalid vision staging capacity"); } + if (max_frame_bytes > + std::numeric_limits::max() / n_views || + max_frame_bytes > + std::numeric_limits::max() / n_views) { + return Status::error(StatusCode::kInvalidArgument, + "vision staging capacity overflows size_t"); + } *out = VisionStaging{}; const std::uint64_t total = max_frame_bytes * n_views; cudaError_t rc = cudaMalloc(&out->device, total); 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 2cee6402..741cc1e2 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 @@ -6,6 +6,7 @@ #include #include +#include #include #include @@ -15,6 +16,7 @@ namespace pi05 { enum class NativeWeightDType { kBf16, + kFloat16, kFp8E4M3, kInt8, kFloat32, @@ -35,6 +37,8 @@ class NativeDeviceWeightStore { modalities::Status upload(const std::string& name, const NativeBf16Tensor& tensor); + modalities::Status upload(const std::string& name, + const NativeF16Tensor& tensor); modalities::Status upload_bytes( const std::string& name, const std::vector& shape, @@ -50,6 +54,7 @@ class NativeDeviceWeightStore { private: frt_ctx ctx_ = nullptr; // borrowed; the context owns every buffer std::map weights_; + std::mutex upload_mutex_; }; } // namespace pi05 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 index 96b31868..bc3c1e9d 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h @@ -27,6 +27,11 @@ modalities::Status native_quantize_fp8_e4m3( bool transpose, NativeFp8Tensor* out); +modalities::Status native_quantize_fp8_e4m3( + const NativeF16Tensor& fp16_weight, + bool transpose, + NativeFp8Tensor* out); + modalities::Status native_quantize_int8_per_output( const NativeFloatTensor& bf16_weight, NativeInt8Tensor* 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 0f7dc4b5..128138bb 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 @@ -22,6 +22,11 @@ struct NativeBf16Tensor { std::vector values; }; +struct NativeF16Tensor { + std::vector shape; + std::vector values; +}; + enum class NativeSourceDType { kF32, kBf16, @@ -44,6 +49,36 @@ modalities::Status native_source_to_bf16( bool transpose, NativeBf16Tensor* out); +modalities::Status native_source_to_f16( + const NativeSourceTensorView& input, + bool transpose, + NativeF16Tensor* out); + +modalities::Status native_source_qkv_to_f16( + const NativeSourceTensorView& q, + const NativeSourceTensorView& k, + const NativeSourceTensorView& v, + std::uint64_t q_heads, + std::uint64_t k_heads, + const NativeFloatTensor* norm, + bool transpose, + NativeF16Tensor* out); + +modalities::Status native_source_pair_to_f16( + const NativeSourceTensorView& left, + const NativeSourceTensorView& right, + const NativeFloatTensor* norm, + bool transpose, + NativeF16Tensor* out); + +modalities::Status native_source_concat_vectors_to_f16( + const std::vector& inputs, + NativeF16Tensor* out); + +modalities::Status native_source_patch_oihw_to_hwio_f16( + const NativeSourceTensorView& input, + NativeF16Tensor* out); + modalities::Status native_source_fold_rms_columns_transpose( const NativeSourceTensorView& weight, const NativeFloatTensor& norm, @@ -85,6 +120,9 @@ modalities::Status load_native_float_tensor( modalities::Status native_to_bf16(const NativeFloatTensor& input, NativeBf16Tensor* out); +modalities::Status native_to_f16(const NativeFloatTensor& input, + NativeF16Tensor* out); + modalities::Status native_round_to_bf16_float( const NativeFloatTensor& input, NativeFloatTensor* out); @@ -128,6 +166,11 @@ modalities::Status native_pi05_time_embeddings( std::uint64_t embedding_dim, NativeFloatTensor* out); +modalities::Status native_pi05_time_embeddings_f16( + int num_steps, + std::uint64_t embedding_dim, + NativeF16Tensor* out); + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/models/pi05/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp index a6289b55..8e70dcc9 100644 --- a/cpp/models/pi05/src/native_device_weights.cpp +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -36,6 +36,7 @@ bool element_count(const std::vector& shape, std::size_t element_bytes(NativeWeightDType dtype) { switch (dtype) { case NativeWeightDType::kBf16: return sizeof(std::uint16_t); + case NativeWeightDType::kFloat16: 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); @@ -53,12 +54,21 @@ modalities::Status NativeDeviceWeightStore::upload( tensor.values.size() * sizeof(std::uint16_t)); } +modalities::Status NativeDeviceWeightStore::upload( + const std::string& name, + const NativeF16Tensor& tensor) { + return upload_bytes(name, tensor.shape, NativeWeightDType::kFloat16, + 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) { + std::lock_guard lock(upload_mutex_); if (!ctx_ || name.empty()) return invalid("invalid device weight store"); if (weights_.find(name) != weights_.end()) { return invalid("duplicate device weight name"); diff --git a/cpp/models/pi05/src/native_quantization.cu b/cpp/models/pi05/src/native_quantization.cu index e53ec458..d2a6366c 100644 --- a/cpp/models/pi05/src/native_quantization.cu +++ b/cpp/models/pi05/src/native_quantization.cu @@ -34,6 +34,16 @@ bool finite_values(const NativeFloatTensor& tensor) { return true; } +bool valid_matrix(const NativeF16Tensor& 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(); +} + } // namespace modalities::Status native_quantize_fp8_e4m3( @@ -58,17 +68,63 @@ modalities::Status native_quantize_fp8_e4m3( NativeFp8Tensor result; result.shape = arranged.shape; result.scale = std::max(amax / 448.0f, 1.0e-12f); + // Producer-side CUDA tensor division lowers to one FP32 reciprocal + // followed by multiplication. Preserve that rounding at FP8 boundaries. + const float inverse_scale = 1.0f / result.scale; 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)); + std::min(448.0f, arranged.values[i] * inverse_scale)); result.values[i] = __nv_fp8_e4m3(value).__x; } *out = std::move(result); return modalities::Status::ok(); } +modalities::Status native_quantize_fp8_e4m3( + const NativeF16Tensor& fp16_weight, + bool transpose, + NativeFp8Tensor* out) { + if (!out || !valid_matrix(fp16_weight)) { + return invalid("FP8 weight must be a valid FP16 matrix"); + } + const std::size_t rows = static_cast(fp16_weight.shape[0]); + const std::size_t columns = + static_cast(fp16_weight.shape[1]); + float amax = 0.0f; + for (std::uint16_t bits : fp16_weight.values) { + const float value = modalities::float16_to_float(bits); + if (!std::isfinite(value)) { + return invalid("FP8 weight must contain finite FP16 values"); + } + amax = std::max(amax, std::fabs(value)); + } + NativeFp8Tensor result; + result.shape = transpose + ? std::vector{fp16_weight.shape[1], + fp16_weight.shape[0]} + : fp16_weight.shape; + result.scale = std::max(amax / 448.0f, 1.0e-12f); + const float inverse_scale = 1.0f / result.scale; + result.values.resize(fp16_weight.values.size()); + for (std::size_t row = 0; row < rows; ++row) { + for (std::size_t column = 0; column < columns; ++column) { + const std::size_t source = row * columns + column; + const std::size_t destination = transpose + ? column * rows + row + : source; + const float value = modalities::float16_to_float( + fp16_weight.values[source]); + const float scaled = std::max( + -448.0f, std::min(448.0f, value * inverse_scale)); + result.values[destination] = __nv_fp8_e4m3(scaled).__x; + } + } + *out = std::move(result); + return modalities::Status::ok(); +} + modalities::Status native_quantize_int8_per_output( const NativeFloatTensor& bf16_weight, NativeInt8Tensor* out) { diff --git a/cpp/models/pi05/src/native_quantization_unavailable.cpp b/cpp/models/pi05/src/native_quantization_unavailable.cpp index 94d6957c..d043b01d 100644 --- a/cpp/models/pi05/src/native_quantization_unavailable.cpp +++ b/cpp/models/pi05/src/native_quantization_unavailable.cpp @@ -20,6 +20,15 @@ modalities::Status native_quantize_fp8_e4m3( return unavailable(); } +modalities::Status native_quantize_fp8_e4m3( + const NativeF16Tensor&, + bool, + NativeFp8Tensor*) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "FP8 quantization requires the CUDA build"); +} + modalities::Status native_quantize_int8_per_output( const NativeFloatTensor&, NativeInt8Tensor*) { diff --git a/cpp/models/pi05/src/native_weight_ops.cpp b/cpp/models/pi05/src/native_weight_ops.cpp index 11f672b6..85388399 100644 --- a/cpp/models/pi05/src/native_weight_ops.cpp +++ b/cpp/models/pi05/src/native_weight_ops.cpp @@ -116,6 +116,9 @@ struct F32Reader { std::uint16_t bf16(std::size_t index) const { return modalities::float_to_bfloat16(data[index]); } + std::uint16_t f16(std::size_t index) const { + return modalities::float_to_float16(data[index]); + } }; struct Bf16Reader { @@ -124,6 +127,9 @@ struct Bf16Reader { return modalities::bfloat16_to_float(data[index]); } std::uint16_t bf16(std::size_t index) const { return data[index]; } + std::uint16_t f16(std::size_t index) const { + return modalities::float_to_float16((*this)[index]); + } }; struct F16Reader { @@ -135,6 +141,7 @@ struct F16Reader { return modalities::float_to_bfloat16( modalities::float16_to_float(data[index])); } + std::uint16_t f16(std::size_t index) const { return data[index]; } }; struct UnalignedF32Reader { @@ -147,6 +154,9 @@ struct UnalignedF32Reader { std::uint16_t bf16(std::size_t index) const { return modalities::float_to_bfloat16((*this)[index]); } + std::uint16_t f16(std::size_t index) const { + return modalities::float_to_float16((*this)[index]); + } }; template @@ -165,6 +175,10 @@ struct UnalignedU16Reader { return IsBf16 ? bits(index) : modalities::float_to_bfloat16((*this)[index]); } + std::uint16_t f16(std::size_t index) const { + return IsBf16 ? modalities::float_to_float16((*this)[index]) + : bits(index); + } }; template @@ -283,6 +297,267 @@ modalities::Status native_source_to_bf16( return modalities::Status::ok(); } +modalities::Status native_source_to_f16( + const NativeSourceTensorView& input, + bool transpose, + NativeF16Tensor* out) { + std::size_t count = 0; + if (!out || !element_count(input.shape, &count) || + (transpose && input.shape.size() != 2)) { + return invalid("invalid direct source to FP16 input"); + } + NativeF16Tensor converted; + converted.shape = transpose + ? std::vector{input.shape[1], input.shape[0]} + : input.shape; + converted.values.resize(count); + modalities::Status st = dispatch_source(input, [&](const auto& reader) { + if (!transpose) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + converted.values[i] = reader.f16(i); + } + }); + } else { + const std::size_t rows = static_cast(input.shape[0]); + const std::size_t cols = static_cast(input.shape[1]); + tiled_transform_transpose( + rows, cols, rows, 0, converted.values.data(), + [&](std::size_t row, std::size_t col) { + return reader.f16(row * cols + col); + }); + } + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *out = std::move(converted); + return modalities::Status::ok(); +} + +modalities::Status native_source_qkv_to_f16( + const NativeSourceTensorView& q, + const NativeSourceTensorView& k, + const NativeSourceTensorView& v, + std::uint64_t q_heads, + std::uint64_t k_heads, + const NativeFloatTensor* norm, + bool transpose, + NativeF16Tensor* out) { + std::size_t q_count = 0; + std::size_t k_count = 0; + std::size_t v_count = 0; + if (!out || q.shape.size() != 2 || k.shape.size() != 2 || + v.shape.size() != 2 || q.shape[1] != k.shape[1] || + q.shape[1] != v.shape[1] || !element_count(q.shape, &q_count) || + !element_count(k.shape, &k_count) || + !element_count(v.shape, &v_count) || + q.shape[0] > std::numeric_limits::max() - k.shape[0] || + q.shape[0] + k.shape[0] > + std::numeric_limits::max() - v.shape[0] || + (q_heads && (q.shape[0] % q_heads || + (q.shape[0] / q_heads) % 2)) || + (k_heads && (k.shape[0] % k_heads || + (k.shape[0] / k_heads) % 2)) || + q.shape[1] > std::numeric_limits::max() || + q.shape[0] + k.shape[0] + v.shape[0] > + std::numeric_limits::max() || + (norm && (!valid_tensor(*norm) || norm->shape.size() != 1 || + norm->shape[0] != q.shape[1]))) { + return invalid("FP16 QKV source tensors have incompatible shapes"); + } + const std::size_t cols = static_cast(q.shape[1]); + const std::size_t total_rows = static_cast( + q.shape[0] + k.shape[0] + v.shape[0]); + NativeF16Tensor joined; + joined.shape = transpose + ? std::vector{q.shape[1], + static_cast(total_rows)} + : std::vector{static_cast(total_rows), + q.shape[1]}; + if (total_rows && cols > std::numeric_limits::max() / + total_rows) { + return invalid("FP16 QKV output shape overflows size_t"); + } + joined.values.resize(total_rows * cols); + const auto write = [&](const NativeSourceTensorView& source, + std::size_t heads, + bool interleave, + std::size_t row_offset) { + return dispatch_source(source, [&](const auto& reader) { + const std::size_t rows = static_cast(source.shape[0]); + const auto value = [&](std::size_t output_row, std::size_t col) { + const std::size_t source_row = interleave + ? interleaved_source_row(output_row, rows, heads) + : output_row; + const float weight = reader[source_row * cols + col]; + const float folded = norm + ? weight * (1.0f + norm->values[col]) + : weight; + return modalities::float_to_float16(folded); + }; + if (transpose) { + tiled_transform_transpose( + rows, cols, total_rows, row_offset, + joined.values.data(), value); + } else { + parallel_ranges(rows, 16, + [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++row) { + std::uint16_t* destination = joined.values.data() + + (row_offset + row) * cols; + for (std::size_t col = 0; col < cols; ++col) { + destination[col] = value(row, col); + } + } + }); + } + return modalities::Status::ok(); + }); + }; + modalities::Status st = write(q, q_heads, q_heads != 0, 0); + if (!st.ok_status()) return st; + st = write(k, k_heads, k_heads != 0, + static_cast(q.shape[0])); + if (!st.ok_status()) return st; + st = write(v, 1, false, + static_cast(q.shape[0] + k.shape[0])); + if (!st.ok_status()) return st; + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_source_pair_to_f16( + const NativeSourceTensorView& left, + const NativeSourceTensorView& right, + const NativeFloatTensor* norm, + bool transpose, + NativeF16Tensor* out) { + std::size_t source_count = 0; + if (!out || left.shape.size() != 2 || right.shape != left.shape || + !element_count(left.shape, &source_count) || + left.shape[0] > std::numeric_limits::max() / 2 || + source_count > std::numeric_limits::max() / 2 || + (norm && (!valid_tensor(*norm) || norm->shape.size() != 1 || + norm->shape[0] != left.shape[1]))) { + return invalid("FP16 paired source tensors have incompatible shapes"); + } + const std::size_t rows = static_cast(left.shape[0]); + const std::size_t cols = static_cast(left.shape[1]); + const std::size_t total_rows = rows * 2; + NativeF16Tensor joined; + joined.shape = transpose + ? std::vector{left.shape[1], left.shape[0] * 2} + : std::vector{left.shape[0] * 2, left.shape[1]}; + joined.values.resize(source_count * 2); + const auto write = [&](const NativeSourceTensorView& source, + std::size_t row_offset) { + return dispatch_source(source, [&](const auto& reader) { + const auto value = [&](std::size_t row, std::size_t col) { + const float weight = reader[row * cols + col]; + return modalities::float_to_float16( + norm ? weight * (1.0f + norm->values[col]) : weight); + }; + if (transpose) { + tiled_transform_transpose( + rows, cols, total_rows, row_offset, + joined.values.data(), value); + } else { + parallel_ranges(rows, 16, + [&](std::size_t begin, std::size_t end) { + for (std::size_t row = begin; row < end; ++row) { + std::uint16_t* destination = joined.values.data() + + (row_offset + row) * cols; + for (std::size_t col = 0; col < cols; ++col) { + destination[col] = value(row, col); + } + } + }); + } + return modalities::Status::ok(); + }); + }; + modalities::Status st = write(left, 0); + if (!st.ok_status()) return st; + st = write(right, rows); + if (!st.ok_status()) return st; + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_source_concat_vectors_to_f16( + const std::vector& inputs, + NativeF16Tensor* out) { + if (!out || inputs.empty()) return invalid("FP16 vector concat has no inputs"); + std::size_t total = 0; + for (const NativeSourceTensorView* input : inputs) { + if (!input || input->shape.size() != 1 || + input->shape[0] > std::numeric_limits::max() || + input->shape[0] > std::numeric_limits::max() - total) { + return invalid("FP16 vector concat tensors have incompatible shapes"); + } + total += static_cast(input->shape[0]); + } + NativeF16Tensor joined; + joined.shape = {static_cast(total)}; + joined.values.resize(total); + std::size_t offset = 0; + for (const NativeSourceTensorView* input : inputs) { + const std::size_t count = static_cast(input->shape[0]); + modalities::Status st = dispatch_source(*input, [&](const auto& reader) { + parallel_ranges(count, 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + joined.values[offset + i] = reader.f16(i); + } + }); + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + offset += count; + } + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_source_patch_oihw_to_hwio_f16( + const NativeSourceTensorView& input, + NativeF16Tensor* out) { + std::size_t count = 0; + if (!out || input.shape.size() != 4 || + !element_count(input.shape, &count)) { + return invalid("FP16 patch permutation requires a 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]); + NativeF16Tensor result; + result.shape = {input.shape[2], input.shape[3], input.shape[1], + input.shape[0]}; + result.values.resize(count); + modalities::Status st = dispatch_source(input, [&](const auto& reader) { + parallel_ranges(height * width * channels, 32, + [&](std::size_t begin, std::size_t end) { + for (std::size_t hwc = begin; hwc < end; ++hwc) { + const std::size_t c = hwc % channels; + const std::size_t hw = hwc / channels; + const std::size_t h = hw / width; + const std::size_t w = hw % width; + for (std::size_t o = 0; o < outputs; ++o) { + const std::size_t src = + ((o * channels + c) * height + h) * width + w; + result.values[hwc * outputs + o] = reader.f16(src); + } + } + }); + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *out = std::move(result); + return modalities::Status::ok(); +} + modalities::Status native_source_fold_rms_columns_transpose( const NativeSourceTensorView& weight, const NativeFloatTensor& norm, @@ -545,46 +820,28 @@ 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); - } + if (!out) return invalid("invalid native tensor load"); + NativeSourceTensorView source; + modalities::Status st = load_native_source_tensor(file, key, &source); + if (!st.ok_status()) return st; std::size_t count = 0; - if (!element_count(tensor->shape, &count)) { + if (!element_count(source.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.shape = source.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); - parallel_ranges(count, 1 << 18, [&](std::size_t begin, - std::size_t end) { - for (std::size_t i = begin; i < end; ++i) { - loaded.values[i] = modalities::bfloat16_to_float(src[i]); - } - }); - } else if (tensor->dtype == "F16") { - const auto* src = static_cast(data); + st = dispatch_source(source, [&](const auto& reader) { parallel_ranges(count, 1 << 18, [&](std::size_t begin, std::size_t end) { for (std::size_t i = begin; i < end; ++i) { - loaded.values[i] = modalities::float16_to_float(src[i]); + loaded.values[i] = reader[i]; } }); - } else { - return modalities::Status::error( - modalities::StatusCode::kUnsupported, - "native tensor dtype is not a floating-point weight: " + - tensor->dtype); - } + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; *out = std::move(loaded); return modalities::Status::ok(); } @@ -606,6 +863,23 @@ modalities::Status native_to_bf16(const NativeFloatTensor& input, return modalities::Status::ok(); } +modalities::Status native_to_f16(const NativeFloatTensor& input, + NativeF16Tensor* out) { + if (!out || !valid_tensor(input)) return invalid("invalid FP16 input"); + NativeF16Tensor converted; + converted.shape = input.shape; + converted.values.resize(input.values.size()); + parallel_ranges(input.values.size(), 1 << 18, + [&](std::size_t begin, std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + converted.values[i] = + modalities::float_to_float16(input.values[i]); + } + }); + *out = std::move(converted); + return modalities::Status::ok(); +} + modalities::Status native_round_to_bf16_float( const NativeFloatTensor& input, NativeFloatTensor* out) { @@ -885,6 +1159,46 @@ modalities::Status native_pi05_time_embeddings( return modalities::Status::ok(); } +modalities::Status native_pi05_time_embeddings_f16( + int num_steps, + std::uint64_t embedding_dim, + NativeF16Tensor* out) { + if (!out || num_steps <= 0 || embedding_dim < 2 || + embedding_dim % 2 != 0 || + embedding_dim > std::numeric_limits::max() / + static_cast(num_steps)) { + return invalid("Pi0.5 FP16 time embedding shape is invalid"); + } + const std::uint64_t half = embedding_dim / 2; + NativeF16Tensor result; + result.shape = {static_cast(num_steps), embedding_dim}; + result.values.resize(static_cast(num_steps) * embedding_dim); + constexpr double kMinPeriod = 4.0e-3; + constexpr double kPeriodRatio = 1000.0; + constexpr double kTwoPi = 6.283185307179586476925286766559; + for (int step = 0; step < num_steps; ++step) { + const float t_f32 = static_cast( + 1.0 - static_cast(step) / + static_cast(num_steps)); + const double t = static_cast(t_f32); + const std::size_t row = static_cast(step) * embedding_dim; + for (std::uint64_t i = 0; i < half; ++i) { + const double fraction = half == 1 + ? 0.0 + : static_cast(i) / static_cast(half - 1); + const double period = + kMinPeriod * std::pow(kPeriodRatio, fraction); + const double angle = t * kTwoPi / period; + result.values[row + i] = modalities::float_to_float16( + static_cast(std::sin(angle))); + result.values[row + half + i] = modalities::float_to_float16( + static_cast(std::cos(angle))); + } + } + *out = std::move(result); + return modalities::Status::ok(); +} + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/tests/pi05_native_quant_probe.cpp b/cpp/tests/pi05_native_quant_probe.cpp index 8a5e12f7..b7a59677 100644 --- a/cpp/tests/pi05_native_quant_probe.cpp +++ b/cpp/tests/pi05_native_quant_probe.cpp @@ -1,6 +1,7 @@ #include "flashrt/cpp/models/pi05/native_quantization.h" #include +#include #include #include #include @@ -9,13 +10,17 @@ namespace { using flashrt::loader::SafetensorsFile; +using flashrt::models::pi05::NativeF16Tensor; using flashrt::models::pi05::NativeFloatTensor; using flashrt::models::pi05::NativeFp8Tensor; using flashrt::models::pi05::NativeInt8Tensor; +using flashrt::models::pi05::NativeSourceTensorView; using flashrt::modalities::Status; constexpr const char* kDecoder = "paligemma_with_expert.gemma_expert.model.layers.0"; +constexpr const char* kEncoder = + "paligemma_with_expert.paligemma.model.language_model.layers.0"; bool load(const SafetensorsFile& file, const std::string& key, NativeFloatTensor* out) { @@ -55,6 +60,22 @@ bool decoder_qkv(const SafetensorsFile& file, NativeFloatTensor* out) { .ok_status(); } +bool encoder_gate_up(const SafetensorsFile& file, NativeF16Tensor* out) { + NativeSourceTensorView gate; + NativeSourceTensorView up; + NativeFloatTensor norm; + return flashrt::models::pi05::load_native_source_tensor( + file, std::string(kEncoder) + ".mlp.gate_proj.weight", &gate) + .ok_status() && + flashrt::models::pi05::load_native_source_tensor( + file, std::string(kEncoder) + ".mlp.up_proj.weight", &up) + .ok_status() && + load(file, std::string(kEncoder) + + ".post_attention_layernorm.weight", &norm) && + flashrt::models::pi05::native_source_pair_to_f16( + gate, up, &norm, false, 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); @@ -66,6 +87,7 @@ std::uint64_t fnv1a(const void* data, std::size_t bytes) { } void print_shape(const std::vector& shape) { + std::cout << std::dec; for (std::size_t i = 0; i < shape.size(); ++i) { if (i) std::cout << ','; std::cout << shape[i]; @@ -87,8 +109,8 @@ void print_result(const std::vector& shape, } // namespace int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_quant_probe CHECKPOINT OP\n"; + if (argc != 3 && argc != 4) { + std::cerr << "usage: pi05_native_quant_probe CHECKPOINT OP [OUTPUT]\n"; return 2; } SafetensorsFile file; @@ -96,9 +118,34 @@ int main(int argc, char** argv) { std::cerr << file.error() << '\n'; return 2; } + const std::string op = argv[2]; + if (op == "encoder_gate_up0_fp8") { + NativeF16Tensor weight; + NativeFp8Tensor output; + if (!encoder_gate_up(file, &weight)) return 1; + std::cout << "input_fnv=" << std::hex << std::setw(16) + << std::setfill('0') + << fnv1a(weight.values.data(), + weight.values.size() * sizeof(std::uint16_t)) + << ' '; + const Status st = flashrt::models::pi05::native_quantize_fp8_e4m3( + weight, false, &output); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + return 1; + } + if (argc == 4) { + std::ofstream file(argv[3], std::ios::binary | std::ios::trunc); + file.write(reinterpret_cast(output.values.data()), + static_cast(output.values.size())); + if (!file) return 1; + } + print_result(output.shape, output.values.data(), output.values.size(), + {output.scale}); + return 0; + } 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'; diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index f8def582..57855ef4 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include using flashrt::modalities::DType; @@ -57,6 +58,13 @@ std::uint32_t bf16_ulp_distance(std::uint16_t a, std::uint16_t b) { } void test_vision_h2d_staging() { + flashrt::modalities::VisionStaging overflow; + auto st = flashrt::modalities::vision_staging_create( + &overflow, 2, + std::numeric_limits::max() / 2 + 1); + assert(!st.ok_status()); + assert(!overflow.device && !overflow.host_pinned); + const auto spec = flashrt::models::pi05::vision_preprocess_spec(1); const std::uint64_t bytes = required_vision_output_bytes(spec); @@ -78,7 +86,7 @@ void test_vision_h2d_staging() { TensorView dst{device, bytes, DType::kBFloat16, MemoryPlace::kDevice, Layout::kNHWC, Shape{1, 224, 224, 3}}; - auto st = preprocess_vision(spec, {frame}, dst); + st = preprocess_vision(spec, {frame}, dst); assert(st.ok_status()); std::vector got(bytes / 2); diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index 0e264a06..f4d012e1 100644 --- a/cpp/tests/test_modalities.cpp +++ b/cpp/tests/test_modalities.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include using flashrt::modalities::DType; @@ -19,12 +20,37 @@ using flashrt::modalities::TensorView; using flashrt::modalities::VisionFrame; using flashrt::modalities::bfloat16_to_float; using flashrt::modalities::float_to_bfloat16; +using flashrt::modalities::float_to_float16; using flashrt::modalities::postprocess_action_cpu; using flashrt::modalities::preprocess_vision_cpu; using flashrt::modalities::required_vision_output_bytes; namespace { +void test_float16_round_to_nearest_even() { + assert(float_to_float16(1.00048828125f) == 0x3c00u); + assert(float_to_float16(1.00146484375f) == 0x3c02u); + assert(float_to_float16(-1.00048828125f) == 0xbc00u); + assert(float_to_float16(-1.00146484375f) == 0xbc02u); + + const float half_min_subnormal = std::ldexp(1.0f, -25); + assert(float_to_float16(half_min_subnormal) == 0x0000u); + assert(float_to_float16(std::nextafter( + half_min_subnormal, std::numeric_limits::infinity())) == + 0x0001u); + assert(float_to_float16(3.0f * std::ldexp(1.0f, -25)) == 0x0002u); + assert(float_to_float16(std::ldexp(1.0f, -14) - + std::ldexp(1.0f, -25)) == 0x0400u); + + assert(float_to_float16(std::numeric_limits::infinity()) == + 0x7c00u); + assert(float_to_float16(-std::numeric_limits::infinity()) == + 0xfc00u); + const std::uint16_t nan = + float_to_float16(std::numeric_limits::quiet_NaN()); + assert((nan & 0x7c00u) == 0x7c00u && (nan & 0x03ffu) != 0); +} + void test_pi05_vision_spec_and_preprocess() { const auto spec = flashrt::models::pi05::vision_preprocess_spec(2); assert(spec.view_order.size() == 2); @@ -201,6 +227,7 @@ void test_pi05_runtime_io_adapter() { } // namespace int main() { + test_float16_round_to_nearest_even(); test_pi05_vision_spec_and_preprocess(); test_view_order_guard(); test_action_postprocess(); diff --git a/cpp/tests/test_pi05_native_quantization.cpp b/cpp/tests/test_pi05_native_quantization.cpp index 6ae4dfaa..9340c230 100644 --- a/cpp/tests/test_pi05_native_quantization.cpp +++ b/cpp/tests/test_pi05_native_quantization.cpp @@ -16,6 +16,27 @@ int main() { assert(fp8.values == std::vector( {0xfe, 0xb8, 0x00, 0x38, 0x7e})); + NativeF16Tensor fp16_input{{2, 3}, {}}; + for (float value : {-448.0f, -1.0f, 0.0f, 1.0f, 224.0f, 448.0f}) { + fp16_input.values.push_back( + flashrt::modalities::float_to_float16(value)); + } + assert(native_quantize_fp8_e4m3(fp16_input, true, &fp8).ok_status()); + assert(fp8.shape == std::vector({3, 2})); + assert(fp8.scale == 1.0f); + assert(fp8.values == std::vector( + {0xfe, 0x38, 0xb8, 0x76, 0x00, 0x7e})); + + NativeF16Tensor reciprocal_input{{1, 2}, {}}; + for (float value : {-0.10406494140625f, 3.0078125f}) { + reciprocal_input.values.push_back( + flashrt::modalities::float_to_float16(value)); + } + assert(native_quantize_fp8_e4m3( + reciprocal_input, false, &fp8).ok_status()); + assert(fp8.scale == 0.0067138671875f); + assert(fp8.values == std::vector({0xd7, 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()); diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp index 21b78f3c..0072916b 100644 --- a/cpp/tests/test_pi05_native_weight_ops.cpp +++ b/cpp/tests/test_pi05_native_weight_ops.cpp @@ -11,6 +11,7 @@ namespace { using flashrt::models::pi05::NativeBf16Tensor; +using flashrt::models::pi05::NativeF16Tensor; using flashrt::models::pi05::NativeFloatTensor; using flashrt::models::pi05::NativeSourceDType; using flashrt::models::pi05::NativeSourceTensorView; @@ -38,6 +39,17 @@ void expect_bf16(const NativeBf16Tensor& tensor, } } +void expect_f16(const NativeF16Tensor& 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(tensor.values[i++] == flashrt::modalities::float_to_float16(value)); + } +} + std::string temp_path() { char path[] = "/tmp/frt_pi05_weight_ops_XXXXXX"; const int fd = ::mkstemp(path); @@ -93,18 +105,25 @@ int main() { NativeSourceTensorView mapped; NativeBf16Tensor mapped_bf16; + NativeF16Tensor mapped_f16; assert(load_native_source_tensor(file, "f32", &mapped).ok_status()); assert(mapped.dtype == NativeSourceDType::kF32); assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); expect_bf16(mapped_bf16, {2}, {1.25f, -2.5f}); + assert(native_source_to_f16(mapped, false, &mapped_f16).ok_status()); + expect_f16(mapped_f16, {2}, {1.25f, -2.5f}); assert(load_native_source_tensor(file, "bf16", &mapped).ok_status()); assert(mapped.dtype == NativeSourceDType::kBf16); assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); expect_bf16(mapped_bf16, {2}, {3.0f, -4.0f}); + assert(native_source_to_f16(mapped, false, &mapped_f16).ok_status()); + expect_f16(mapped_f16, {2}, {3.0f, -4.0f}); assert(load_native_source_tensor(file, "f16", &mapped).ok_status()); assert(mapped.dtype == NativeSourceDType::kF16); assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); expect_bf16(mapped_bf16, {2}, {5.0f, -6.0f}); + assert(native_source_to_f16(mapped, false, &mapped_f16).ok_status()); + expect_f16(mapped_f16, {2}, {5.0f, -6.0f}); assert(::unlink(path.c_str()) == 0); NativeFloatTensor matrix{{2, 3}, {1, 2, 3, 4, 5, 6}}; @@ -150,6 +169,25 @@ int main() { assert(!native_pi05_time_embeddings(0, 4, &result).ok_status()); assert(!native_pi05_time_embeddings(2, 3, &result).ok_status()); + NativeF16Tensor time_f16; + assert(native_pi05_time_embeddings_f16(2, 4, &time_f16).ok_status()); + assert(time_f16.shape == std::vector({2, 4})); + assert(time_f16.values.size() == 8); + const double two_pi = 6.2831853071795864769; + const double first_angles[] = {two_pi / 4.0e-3, two_pi / 4.0}; + for (std::size_t i = 0; i < 2; ++i) { + assert(time_f16.values[i] == flashrt::modalities::float_to_float16( + static_cast(std::sin(first_angles[i])))); + assert(time_f16.values[2 + i] == + flashrt::modalities::float_to_float16( + static_cast(std::cos(first_angles[i])))); + } + assert(native_pi05_time_embeddings_f16(10, 1024, &time_f16).ok_status()); + assert(time_f16.shape == std::vector({10, 1024})); + assert(time_f16.values[6 * 1024 + 7] == 0xb1c0u); + assert(time_f16.values[8 * 1024 + 7] == 0x2dc6u); + assert(time_f16.values[9 * 1024 + 3] == 0x292au); + 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( @@ -167,6 +205,10 @@ int main() { flashrt::modalities::float_to_bfloat16(matrix.values[i])); } + NativeF16Tensor converted_f16; + assert(native_to_f16(matrix, &converted_f16).ok_status()); + expect_f16(converted_f16, {2, 3}, {1, 2, 3, 4, 5, 6}); + NativeBf16Tensor direct; const float source_f32[] = {1, 2, 3, 4, 5, 6, 7, 8}; std::uint16_t source_bf16[8]; @@ -198,8 +240,75 @@ int main() { 2, 6, 2, 6, 2, 6, 3, 7, 3, 7, 3, 7, 4, 8, 4, 8, 4, 8}); + + NativeF16Tensor direct_f16; + assert(native_source_to_f16(source_view, true, &direct_f16).ok_status()); + expect_f16(direct_f16, {4, 2}, {1, 5, 2, 6, 3, 7, 4, 8}); + assert(native_source_qkv_to_f16( + source_view, source_view, source_view, 1, 1, nullptr, + true, &direct_f16).ok_status()); + expect_f16(direct_f16, {4, 6}, + {1, 5, 1, 5, 1, 5, + 2, 6, 2, 6, 2, 6, + 3, 7, 3, 7, 3, 7, + 4, 8, 4, 8, 4, 8}); + assert(native_source_qkv_to_f16( + source_view, source_view, source_view, 1, 1, nullptr, + false, &direct_f16).ok_status()); + expect_f16(direct_f16, {6, 4}, + {1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8}); } + const float fold_source[] = {1.003f, -2.007f, 3.011f, -4.015f}; + const NativeSourceTensorView fold_view{ + fold_source, {2, 2}, NativeSourceDType::kF32}; + NativeFloatTensor fold_norm{{2}, {0.125f, -0.375f}}; + NativeF16Tensor direct_f16; + assert(native_source_qkv_to_f16( + fold_view, fold_view, fold_view, 1, 1, &fold_norm, + false, &direct_f16).ok_status()); + expect_f16(direct_f16, {6, 2}, + {1.003f * 1.125f, -2.007f * 0.625f, + 3.011f * 1.125f, -4.015f * 0.625f, + 1.003f * 1.125f, -2.007f * 0.625f, + 3.011f * 1.125f, -4.015f * 0.625f, + 1.003f * 1.125f, -2.007f * 0.625f, + 3.011f * 1.125f, -4.015f * 0.625f}); + + assert(native_source_pair_to_f16( + source_views[0], source_views[0], nullptr, false, + &direct_f16).ok_status()); + expect_f16(direct_f16, {4, 4}, + {1, 2, 3, 4, 5, 6, 7, 8, + 1, 2, 3, 4, 5, 6, 7, 8}); + assert(native_source_pair_to_f16( + source_views[0], source_views[0], nullptr, true, + &direct_f16).ok_status()); + expect_f16(direct_f16, {4, 4}, + {1, 5, 1, 5, 2, 6, 2, 6, + 3, 7, 3, 7, 4, 8, 4, 8}); + + const float vector_a[] = {1, 2}; + const std::uint16_t vector_b[] = { + flashrt::modalities::float_to_bfloat16(3), + flashrt::modalities::float_to_bfloat16(4)}; + const NativeSourceTensorView vector_a_view{ + vector_a, {2}, NativeSourceDType::kF32}; + const NativeSourceTensorView vector_b_view{ + vector_b, {2}, NativeSourceDType::kBf16}; + assert(native_source_concat_vectors_to_f16( + {&vector_a_view, &vector_b_view}, &direct_f16).ok_status()); + expect_f16(direct_f16, {4}, {1, 2, 3, 4}); + + const float patch_source[] = {0, 1, 2, 3, 4, 5, 6, 7}; + const NativeSourceTensorView patch_view{ + patch_source, {2, 2, 2, 1}, NativeSourceDType::kF32}; + assert(native_source_patch_oihw_to_hwio_f16( + patch_view, &direct_f16).ok_status()); + expect_f16(direct_f16, {2, 1, 2, 2}, {0, 4, 2, 6, 1, 5, 3, 7}); + 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"); diff --git a/cpp/tests/test_safetensors_loader.cpp b/cpp/tests/test_safetensors_loader.cpp index a24c83ba..7d5c23a5 100644 --- a/cpp/tests/test_safetensors_loader.cpp +++ b/cpp/tests/test_safetensors_loader.cpp @@ -53,6 +53,8 @@ int main() { assert(file.open(path)); assert(file.is_open()); assert(file.tensors().size() == 2); + assert(file.metadata().size() == 1); + assert(file.metadata().at("format") == "pt"); const auto* values = file.find("values"); assert(values); assert(values->dtype == "F32"); @@ -63,6 +65,7 @@ int main() { SafetensorsFile moved(std::move(file)); assert(!file.is_open()); assert(moved.find("u8")); + assert(moved.metadata().at("format") == "pt"); assert(::unlink(path.c_str()) == 0); assert(std::memcmp(moved.data(*moved.find("values")), expected, sizeof(expected)) == 0); @@ -102,6 +105,23 @@ int main() { assert(!file.open(invalid)); assert(file.error().find("duplicate") != std::string::npos); + write_file(invalid, + "{\"__metadata__\":{\"format\":7}," + "\"x\":{\"dtype\":\"U8\",\"shape\":[1]," + "\"data_offsets\":[0,1]}}", + std::string(1, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("must be strings") != std::string::npos); + + write_file(invalid, + "{\"__metadata__\":{\"format\":\"pt\"}," + "\"__metadata__\":{\"source\":\"test\"}," + "\"x\":{\"dtype\":\"U8\",\"shape\":[1]," + "\"data_offsets\":[0,1]}}", + std::string(1, '\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]}}", From ed7ba15440cc1d931b15a685a41be28c238d30ff Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 11:39:18 -0400 Subject: [PATCH 65/83] feat(pi05): add native Thor FP8 calibration pipeline --- cpp/CMakeLists.txt | 126 +- .../include/flashrt/cpp/models/pi05/c_api.h | 42 + .../cpp/models/pi05/native_calibration.h | 50 + .../cpp/models/pi05/native_graph_owner.h | 31 +- .../cpp/models/pi05/native_graph_runtime.h | 38 + .../flashrt/cpp/models/pi05/native_rope.h | 19 + .../pi05/native_thor_calibration_session.h | 65 + .../cpp/models/pi05/native_thor_fp8_forward.h | 53 + .../cpp/models/pi05/native_thor_graph_owner.h | 68 + .../models/pi05/native_thor_kernel_driver.h | 163 +++ .../pi05/native_thor_style_precompute.h | 31 + .../pi05/native_thor_weight_materializer.h | 64 + .../cpp/models/pi05/native_workspace.h | 17 + cpp/models/pi05/src/native_calibration.cpp | 394 +++++ cpp/models/pi05/src/native_graph_owner.cpp | 6 + cpp/models/pi05/src/native_model_runtime.cpp | 159 +- cpp/models/pi05/src/native_open.cpp | 70 +- cpp/models/pi05/src/native_open_internal.h | 8 + cpp/models/pi05/src/native_rope.cu | 62 + .../pi05/src/native_rope_unavailable.cpp | 16 + .../src/native_thor_calibration_c_api.cpp | 254 ++++ .../src/native_thor_calibration_session.cpp | 499 +++++++ .../pi05/src/native_thor_fp8_forward.cpp | 1297 +++++++++++++++++ .../pi05/src/native_thor_graph_owner.cpp | 318 ++++ .../pi05/src/native_thor_kernel_driver.cu | 622 ++++++++ cpp/models/pi05/src/native_thor_setup_ops.cu | 28 + .../pi05/src/native_thor_style_precompute.cpp | 210 +++ .../src/native_thor_weight_materializer.cpp | 545 +++++++ cpp/models/pi05/src/native_workspace.cpp | 240 ++- cpp/tests/gate_pi05_native_thor_fp8.py | 385 +++++ cpp/tests/pi05_native_open_probe.cpp | 25 +- cpp/tests/pi05_native_rope_probe.cpp | 27 +- cpp/tests/pi05_native_thor_fp8_probe.cpp | 357 +++++ cpp/tests/test_pi05_c_api.cpp | 5 + cpp/tests/test_pi05_native_calibration.cpp | 91 ++ cpp/tests/test_pi05_native_open.cpp | 26 + cpp/tests/test_pi05_native_workspace.cpp | 107 +- 37 files changed, 6437 insertions(+), 81 deletions(-) create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_calibration.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_runtime.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rope.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h create mode 100644 cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h create mode 100644 cpp/models/pi05/src/native_calibration.cpp create mode 100644 cpp/models/pi05/src/native_rope.cu create mode 100644 cpp/models/pi05/src/native_rope_unavailable.cpp create mode 100644 cpp/models/pi05/src/native_thor_calibration_c_api.cpp create mode 100644 cpp/models/pi05/src/native_thor_calibration_session.cpp create mode 100644 cpp/models/pi05/src/native_thor_fp8_forward.cpp create mode 100644 cpp/models/pi05/src/native_thor_graph_owner.cpp create mode 100644 cpp/models/pi05/src/native_thor_kernel_driver.cu create mode 100644 cpp/models/pi05/src/native_thor_setup_ops.cu create mode 100644 cpp/models/pi05/src/native_thor_style_precompute.cpp create mode 100644 cpp/models/pi05/src/native_thor_weight_materializer.cpp create mode 100644 cpp/tests/gate_pi05_native_thor_fp8.py create mode 100644 cpp/tests/pi05_native_thor_fp8_probe.cpp create mode 100644 cpp/tests/test_pi05_native_calibration.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 964db7b0..cc8f2da4 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -24,8 +24,13 @@ option(FLASHRT_CPP_WITH_CUDA_STAGING "Enable conservative CUDA H2D/D2H modality 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) +option(FLASHRT_CPP_WITH_THOR_FP8 + "Enable the native Thor SM110 FP8 kernel backend" OFF) set(FLASHRT_CPP_FA2_LIBRARY "" CACHE FILEPATH "Path to the Python-free libflashrt_fa2_raw shared library") +set(FLASHRT_CPP_CUTLASS_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/../third_party/cutlass" CACHE PATH + "Path to CUTLASS for the native Thor SM110 FP8 backend") if(FLASHRT_CPP_WITH_CUDA_STAGING) find_package(CUDAToolkit REQUIRED) endif() @@ -53,6 +58,26 @@ if(FLASHRT_CPP_WITH_FA2) set_target_properties(flashrt_cpp_fa2_external PROPERTIES IMPORTED_LOCATION ${FLASHRT_CPP_FA2_LIBRARY}) endif() +if(FLASHRT_CPP_WITH_THOR_FP8) + if(NOT FLASHRT_CPP_WITH_CUDA_KERNELS) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_THOR_FP8 requires FLASHRT_CPP_WITH_CUDA_KERNELS") + endif() + if(NOT FLASHRT_CPP_WITH_CUDA_STAGING) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_THOR_FP8 requires FLASHRT_CPP_WITH_CUDA_STAGING") + endif() + if(NOT FLASHRT_CPP_WITH_EXEC) + message(FATAL_ERROR + "FLASHRT_CPP_WITH_THOR_FP8 requires FLASHRT_CPP_WITH_EXEC") + endif() + if(NOT EXISTS + "${FLASHRT_CPP_CUTLASS_DIR}/include/cutlass/cutlass.h") + message(FATAL_ERROR + "FLASHRT_CPP_WITH_THOR_FP8 requires CUTLASS; set " + "FLASHRT_CPP_CUTLASS_DIR to a compatible checkout") + endif() +endif() if(FLASHRT_CPP_WITH_SENTENCEPIECE) find_path(FLASHRT_SENTENCEPIECE_INCLUDE_DIR sentencepiece_processor.h) find_library(FLASHRT_SENTENCEPIECE_LIBRARY sentencepiece) @@ -114,6 +139,16 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) endif() add_library(flashrt_cpp_modalities STATIC ${FLASHRT_CPP_MODALITY_SRCS}) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + # Match the CUDA preprocess kernels' explicit round-to-nearest operations in + # optimized builds; otherwise the host compiler may contract lerps to FMAs. + set_property(SOURCE modalities/src/vision_cpu.cpp APPEND PROPERTY + COMPILE_OPTIONS -ffp-contract=off) + # Weight fusion mirrors producer-side add/cast/multiply tensor operations. + # Keep each float operation rounded before the final FP16 conversion. + set_property(SOURCE models/pi05/src/native_weight_ops.cpp APPEND PROPERTY + COMPILE_OPTIONS -ffp-contract=off) +endif() target_include_directories(flashrt_cpp_modalities PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/modalities/include) if(FLASHRT_CPP_WITH_CUDA_STAGING) @@ -164,6 +199,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_calibration.cpp models/pi05/src/native_workspace.cpp models/pi05/src/native_rtx_attention.cpp models/pi05/src/prompt_format.cpp @@ -172,10 +208,12 @@ set(FLASHRT_CPP_PI05_SRCS models/pi05/src/runtime.cpp) if(FLASHRT_CPP_WITH_CUDA_KERNELS) list(APPEND FLASHRT_CPP_PI05_SRCS - models/pi05/src/native_quantization.cu) + models/pi05/src/native_quantization.cu + models/pi05/src/native_rope.cu) else() list(APPEND FLASHRT_CPP_PI05_SRCS - models/pi05/src/native_quantization_unavailable.cpp) + models/pi05/src/native_quantization_unavailable.cpp + models/pi05/src/native_rope_unavailable.cpp) endif() add_library(flashrt_cpp_pi05 STATIC ${FLASHRT_CPP_PI05_SRCS}) @@ -209,6 +247,52 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) $<$: --expt-relaxed-constexpr -O3 --ftz=true --prec-div=false --prec-sqrt=false>) + if(FLASHRT_CPP_WITH_THOR_FP8) + add_library(flashrt_cpp_pi05_thor_precise OBJECT + models/pi05/src/native_thor_setup_ops.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/attention/fmha_fp16_strided.cu) + target_include_directories(flashrt_cpp_pi05_thor_precise PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/attention + ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha + ${FLASHRT_CPP_CUTLASS_DIR}/include + ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) + target_compile_options(flashrt_cpp_pi05_thor_precise PRIVATE + $<$: + --expt-relaxed-constexpr --expt-extended-lambda -O3>) + set_target_properties(flashrt_cpp_pi05_thor_precise PROPERTIES + CUDA_ARCHITECTURES 110a + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + POSITION_INDEPENDENT_CODE ON) + target_sources(flashrt_cpp_pi05_kernels PRIVATE + $) + target_sources(flashrt_cpp_pi05_kernels PRIVATE + models/pi05/src/native_thor_kernel_driver.cu + models/pi05/src/native_thor_weight_materializer.cpp + models/pi05/src/native_thor_fp8_forward.cpp + models/pi05/src/native_thor_graph_owner.cpp + models/pi05/src/native_thor_calibration_session.cpp + models/pi05/src/native_thor_style_precompute.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/cutlass_sm100.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/attention_cublas.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/decoder_fused.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/quantize.cu + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/softmax.cu) + target_include_directories(flashrt_cpp_pi05_kernels PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/attention + ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha + ${FLASHRT_CPP_CUTLASS_DIR}/include + ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) + target_compile_definitions(flashrt_cpp_pi05_kernels + PUBLIC FLASHRT_CPP_WITH_THOR_FP8=1) + target_compile_options(flashrt_cpp_pi05_kernels PRIVATE + $<$: + --expt-extended-lambda --use_fast_math + -gencode=arch=compute_110a,code=sm_110a>) + set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES + CUDA_ARCHITECTURES 110a + CUDA_RESOLVE_DEVICE_SYMBOLS ON) + endif() if(FLASHRT_CPP_WITH_FA2) target_sources(flashrt_cpp_pi05_kernels PRIVATE models/pi05/src/native_graph_owner.cpp @@ -225,6 +309,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_thor_calibration_c_api.cpp models/pi05/src/native_model_runtime.cpp models/pi05/src/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c @@ -237,6 +322,15 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) endif() if(BUILD_TESTING) + # The test executables use assert() as their fail-fast check, including + # around setup calls whose side effects are required by later checks. + # Keep those checks active for Release test builds as well. + if(MSVC) + add_compile_options($<$:/UNDEBUG>) + else() + add_compile_options($<$:-UNDEBUG>) + endif() + 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) @@ -268,6 +362,11 @@ if(BUILD_TESTING) 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(test_pi05_native_calibration + tests/test_pi05_native_calibration.cpp) + target_link_libraries(test_pi05_native_calibration PRIVATE flashrt_cpp_pi05) + add_test(NAME pi05_native_calibration COMMAND test_pi05_native_calibration) + add_executable(pi05_native_weight_probe tests/pi05_native_weight_probe.cpp) target_link_libraries(pi05_native_weight_probe PRIVATE flashrt_cpp_pi05) @@ -369,11 +468,6 @@ if(BUILD_TESTING) PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) if(FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_open_probe - tests/pi05_native_open_probe.cpp) - target_link_libraries(pi05_native_open_probe - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - add_executable(pi05_native_e2e_probe tests/pi05_native_e2e_probe.cpp) target_link_libraries(pi05_native_e2e_probe @@ -396,6 +490,14 @@ if(BUILD_TESTING) COMMAND test_pi05_native_rtx_attention_driver) endif() + if(FLASHRT_CPP_WITH_SENTENCEPIECE AND + (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) + 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() + if(FLASHRT_CPP_WITH_SENTENCEPIECE) add_executable(pi05_tokenizer_corpus_probe tests/pi05_tokenizer_corpus_probe.cpp) @@ -403,6 +505,13 @@ if(BUILD_TESTING) PRIVATE flashrt_cpp_pi05) endif() + if(FLASHRT_CPP_WITH_THOR_FP8 AND FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_native_thor_fp8_probe + tests/pi05_native_thor_fp8_probe.cpp) + target_link_libraries(pi05_native_thor_fp8_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + endif() + add_executable(pi05_native_rope_probe tests/pi05_native_rope_probe.cpp) target_link_libraries(pi05_native_rope_probe @@ -450,7 +559,8 @@ 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) + if((FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8) AND + FLASHRT_CPP_WITH_SENTENCEPIECE) target_compile_definitions(test_pi05_native_open PRIVATE FLASHRT_CPP_PI05_NATIVE_OPEN_ENABLED=1) endif() 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 3eb59ea9..d4011c9a 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 @@ -10,6 +10,8 @@ extern "C" { #endif typedef struct frt_pi05_runtime_s frt_pi05_runtime; +typedef struct frt_pi05_calibration_session_s + frt_pi05_calibration_session; enum frt_pi05_pixel_format { FRT_PI05_PIXEL_RGB8 = 0, @@ -100,6 +102,25 @@ typedef struct frt_pi05_vision_frame { uint64_t timestamp_ns; } frt_pi05_vision_frame; +/* One calibration observation. `frames` is the complete named camera set for + * one observation. Dataset and multi-timestamp calibration call observe + * repeatedly; synchronization, iteration, and decoding remain host policy. + * + * `noise` is optional f32 [chunk, 32]. Supplying it makes oracle comparison + * exact. When omitted, FlashRT generates deterministic normal noise from + * noise_seed and the current sample index. */ +typedef struct frt_pi05_calibration_sample_v1 { + uint32_t struct_size; + const char* prompt; + const float* state; + uint64_t n_state; + const frt_pi05_vision_frame* frames; + uint64_t n_frames; + const float* noise; + uint64_t n_noise; + uint64_t noise_seed; +} frt_pi05_calibration_sample_v1; + int frt_pi05_runtime_create(const frt_runtime_export_v1* exp, const frt_pi05_runtime_config* config, frt_pi05_runtime** out); @@ -120,6 +141,27 @@ int frt_pi05_runtime_read_actions(frt_pi05_runtime*, const frt_runtime_export_v1* frt_pi05_runtime_export(frt_pi05_runtime*); const char* frt_pi05_runtime_last_error(frt_pi05_runtime*); +/* Native SM110 FP8 calibration. `config_json` uses the same checkpoint, + * tokenizer, shape, and normalization fields as frt_model_runtime_open_v1. + * The precision must resolve to fp8_e4m3fn. A session accepts one or many + * observations and publishes one identity-bound safetensors artifact. */ +int frt_pi05_calibration_create_v1( + const char* config_json, + double percentile, + frt_pi05_calibration_session** out); +int frt_pi05_calibration_observe_v1( + frt_pi05_calibration_session*, + const frt_pi05_calibration_sample_v1* sample); +int frt_pi05_calibration_finalize_v1( + frt_pi05_calibration_session*, + const char* artifact_path); +uint64_t frt_pi05_calibration_sample_count_v1( + const frt_pi05_calibration_session*); +const char* frt_pi05_calibration_last_error_v1( + const frt_pi05_calibration_session*); +const char* frt_pi05_calibration_create_last_error_v1(void); +void frt_pi05_calibration_destroy_v1(frt_pi05_calibration_session*); + #ifdef __cplusplus } #endif diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_calibration.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_calibration.h new file mode 100644 index 00000000..a1d2e600 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_calibration.h @@ -0,0 +1,50 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_CALIBRATION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_CALIBRATION_H + +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeCalibrationArtifact { + std::string hardware; + std::string weights_sha256; + std::string tokenizer_sha256; + int num_views = 0; + int max_prompt_tokens = 0; + int state_dim = 0; + int chunk_size = 0; + int num_steps = 0; + int vision_pool_factor = 0; + std::uint64_t sample_count = 0; + double percentile = 100.0; + std::vector encoder_scales; + std::vector decoder_scales; +}; + +modalities::Status validate_native_calibration_artifact( + const NativeCalibrationArtifact& artifact); + +modalities::Status save_native_calibration_artifact( + const std::string& path, + const NativeCalibrationArtifact& artifact); + +modalities::Status load_native_calibration_artifact( + const std::string& path, + NativeCalibrationArtifact* artifact); + +modalities::Status reduce_native_calibration_samples( + const std::vector>& samples, + double percentile, + std::vector* reduced); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_CALIBRATION_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h index 0bfe18ba..2321254e 100644 --- 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 @@ -2,6 +2,7 @@ #define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H #include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_graph_runtime.h" #include #include @@ -10,38 +11,30 @@ 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 { +class NativeGraphOwner final : public NativeGraphRuntime { public: static std::unique_ptr create( const std::string& checkpoint_path, const NativeGraphConfig& config, modalities::Status* status); - ~NativeGraphOwner(); + ~NativeGraphOwner() override; 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_; } + frt_ctx context() const override { return ctx_; } + frt_graph infer_graph() const override { return infer_graph_; } + int stream_id() const override { return stream_id_; } + void* native_stream() const override { 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_; } + NativeDeviceWeightStore& weights() override { return weights_; } + const NativeDeviceWeightStore& weights() const override { return weights_; } + NativeWorkspace& workspace() override { return workspace_; } + const NativeWorkspace& workspace() const override { return workspace_; } NativeRtxAttentionWorkspace& attention() { return attention_; } const NativeRtxAttentionWorkspace& attention() const { return attention_; } - modalities::Status set_prompt_length(int prompt_tokens); + modalities::Status set_prompt_length(int prompt_tokens) override; int replay() const; modalities::Status synchronize() const; diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_runtime.h new file mode 100644 index 00000000..a2776179 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_runtime.h @@ -0,0 +1,38 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H + +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_workspace.h" + +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 NativeGraphRuntime { +public: + virtual ~NativeGraphRuntime() = default; + + virtual frt_ctx context() const = 0; + virtual frt_graph infer_graph() const = 0; + virtual int stream_id() const = 0; + virtual void* native_stream() const = 0; + virtual NativeDeviceWeightStore& weights() = 0; + virtual const NativeDeviceWeightStore& weights() const = 0; + virtual NativeWorkspace& workspace() = 0; + virtual const NativeWorkspace& workspace() const = 0; + virtual modalities::Status set_prompt_length(int prompt_tokens) = 0; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rope.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rope.h new file mode 100644 index 00000000..a7bd0a95 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rope.h @@ -0,0 +1,19 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_ROPE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_ROPE_H + +#include "flashrt/cpp/modalities/types.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +modalities::Status generate_native_thor_rope_f16( + void* output, int start_position, int positions, std::uintptr_t stream); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_ROPE_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h new file mode 100644 index 00000000..352fb8ff --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h @@ -0,0 +1,65 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H + +#include "flashrt/cpp/modalities/vision.h" +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeThorCalibrationConfig { + std::string checkpoint_path; + std::string tokenizer_model_path; + int max_prompt_tokens = 200; + int state_dim = 0; + int num_views = 2; + int chunk_size = 10; + int num_steps = 10; + int vision_pool_factor = 1; + int max_frame_width = 1280; + int max_frame_height = 720; + std::vector state_q01; + std::vector state_q99; +}; + +class NativeThorCalibrationSession { +public: + static std::unique_ptr create( + const NativeThorCalibrationConfig& config, + double percentile, + modalities::Status* status); + + ~NativeThorCalibrationSession(); + + NativeThorCalibrationSession(const NativeThorCalibrationSession&) = delete; + NativeThorCalibrationSession& operator=( + const NativeThorCalibrationSession&) = delete; + + modalities::Status observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed); + modalities::Status finalize(const std::string& artifact_path) const; + std::uint64_t sample_count() const; + +private: + struct Impl; + explicit NativeThorCalibrationSession(std::unique_ptr impl); + std::unique_ptr impl_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h new file mode 100644 index 00000000..7efd7126 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h @@ -0,0 +1,53 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_FP8_FORWARD_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_FP8_FORWARD_H + +#include "flashrt/cpp/models/pi05/native_thor_kernel_driver.h" +#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeThorFp8Forward { +public: + explicit NativeThorFp8Forward(const NativeThorKernelDriver* driver) + : driver_(driver) {} + + modalities::Status vision(const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const NativeThorWeightScales& weight_scales, + std::uintptr_t stream) const; + modalities::Status encoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const std::vector& activation_weight_alphas, + std::uintptr_t stream) const; + modalities::Status diffusion(const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const; + + modalities::Status calibrate_encoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const NativeThorWeightScales& weight_scales, + std::vector* sample_scales, + std::uintptr_t stream) const; + modalities::Status calibrate_decoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::vector* sample_scales, + std::uintptr_t stream) const; + +private: + const NativeThorKernelDriver* driver_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_FP8_FORWARD_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h new file mode 100644 index 00000000..b3393739 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h @@ -0,0 +1,68 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H + +#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/native_thor_fp8_forward.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeThorGraphOwner final : public NativeGraphRuntime { +public: + static std::unique_ptr create( + const std::string& checkpoint_path, + const NativeGraphConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status); + + ~NativeThorGraphOwner() override; + + NativeThorGraphOwner(const NativeThorGraphOwner&) = delete; + NativeThorGraphOwner& operator=(const NativeThorGraphOwner&) = delete; + + frt_ctx context() const override { return ctx_; } + frt_graph infer_graph() const override { return infer_graph_; } + int stream_id() const override { return stream_id_; } + void* native_stream() const override { return replay_stream_; } + NativeDeviceWeightStore& weights() override { return weights_; } + const NativeDeviceWeightStore& weights() const override { return weights_; } + NativeWorkspace& workspace() override { return workspace_; } + const NativeWorkspace& workspace() const override { return workspace_; } + + modalities::Status set_prompt_length(int prompt_tokens) override; + int replay() const; + modalities::Status synchronize() const; + +private: + NativeThorGraphOwner(frt_ctx ctx, const NativeGraphConfig& config); + modalities::Status initialize( + const std::string& checkpoint_path, + const NativeCalibrationArtifact& calibration); + modalities::Status record(void* stream); + static void record_graph(void* user, void* stream); + + frt_ctx ctx_ = nullptr; + NativeGraphConfig config_; + NativeDeviceWeightStore weights_; + NativeWorkspace workspace_; + NativeThorKernelDriver driver_; + NativeThorFp8Forward forward_; + NativeThorWeightScales weight_scales_; + std::vector encoder_alphas_; + 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_THOR_GRAPH_OWNER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h new file mode 100644 index 00000000..46c1626c --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h @@ -0,0 +1,163 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_KERNEL_DRIVER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_KERNEL_DRIVER_H + +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeThorFp8Tactic { + kSquare, + kT1, + kWide, + kPlain, +}; + +class NativeThorKernelDriver { +public: + NativeThorKernelDriver() noexcept; + ~NativeThorKernelDriver(); + + NativeThorKernelDriver(const NativeThorKernelDriver&) = delete; + NativeThorKernelDriver& operator=(const NativeThorKernelDriver&) = delete; + + modalities::Status status() const; + modalities::Status fp16_nn(void* a, void* b, void* output, + int m, int n, int k, + std::uintptr_t stream) const; + modalities::Status fp8_nn_bias( + void* a, void* b, void* output, void* bias, + int m, int n, int k, float alpha, std::uintptr_t stream) const; + modalities::Status fp8_nn_bias_residual( + void* a, void* b, void* output, void* bias, + int m, int n, int k, float alpha, std::uintptr_t stream) const; + modalities::Status fp8_nn_gelu_bias( + void* a, void* b, void* output, void* bias, + int m, int n, int k, float alpha, std::uintptr_t stream) const; + modalities::Status fp8_cutlass( + void* a, void* b, void* output, int m, int n, int k, + float alpha, float beta, NativeThorFp8Tactic tactic, + std::uintptr_t stream) const; + modalities::Status fp8_descale( + void* a, void* b, void* output, int m, int n, int k, + const float* activation_scale, const float* weight_scale, + std::uintptr_t stream) const; + + modalities::Status add_bias_fp16(void* values, const void* bias, + int rows, int columns, + std::uintptr_t stream) const; + modalities::Status layer_norm_fp16( + const void* values, const void* weight, const void* bias, void* output, + int rows, int columns, float epsilon, std::uintptr_t stream) const; + modalities::Status layer_norm_fp8( + const void* values, void* output, const void* weight, const void* bias, + int rows, int columns, float epsilon, std::uintptr_t stream) const; + modalities::Status rms_norm_fp16( + const void* values, const void* weight, void* output, + int rows, int columns, float epsilon, std::uintptr_t stream) const; + modalities::Status rms_norm_fp8_noweight( + const void* values, void* output, int rows, int columns, + const float* scale, std::uintptr_t stream) const; + modalities::Status residual_rms_norm_fp8_noweight( + void* residual, const void* values, void* output, + int rows, int columns, const float* scale, + std::uintptr_t stream) const; + modalities::Status quantize_fp8_static( + const void* values, void* output, const float* scale, + std::size_t elements, std::uintptr_t stream) const; + modalities::Status quantize_fp8_dynamic( + const void* values, void* output, float* scale, + std::size_t elements, std::uintptr_t stream) const; + + modalities::Status gelu_fp16(void* values, std::size_t elements, + std::uintptr_t stream) const; + modalities::Status silu_fp16(void* values, std::size_t elements, + std::uintptr_t stream) const; + modalities::Status precise_silu_fp16( + void* values, std::size_t elements, std::uintptr_t stream) const; + modalities::Status gate_gelu_fp16( + const void* merged, void* output, int rows, int hidden, + std::uintptr_t stream) const; + modalities::Status gate_gelu_fp8( + const void* merged, void* output, int rows, int hidden, + const float* scale, std::uintptr_t stream) const; + modalities::Status residual_add_fp16( + void* residual, const void* values, std::size_t elements, + std::uintptr_t stream) const; + + modalities::Status fused_adarms_fp8( + const void* values, const void* style, void* output, void* gate, + int rows, int columns, const float* scale, + std::uintptr_t stream) const; + modalities::Status gate_res_adarms_fp8( + const void* gemm_output, const void* previous_gate, void* residual, + const void* style, void* output, void* gate, + int rows, int columns, const float* scale, + std::uintptr_t stream) const; + modalities::Status gate_res_fp16( + const void* gemm_output, const void* gate, void* residual, + std::size_t elements, std::uintptr_t stream) const; + modalities::Status adarms_fp16( + const void* values, const void* style, void* output, void* gate, + int rows, int columns, std::uintptr_t stream) const; + + modalities::Status qkv_rope_cache_fp16( + const void* qkv, const void* rope, void* query, void* key_cache, + void* value_cache, int rows, int query_columns, int key_columns, + int head_dimension, int qkv_stride, int cache_offset, + int cache_stride, std::uintptr_t stream) const; + modalities::Status qkv_rope_cache_devpos_fp16( + const void* qkv, const void* rope, void* query, void* key_cache, + void* value_cache, const int* device_position, int rows, + int query_columns, int key_columns, int head_dimension, + int qkv_stride, int cache_offset, int cache_stride, + std::uintptr_t stream) const; + modalities::Status attention_seqused_fp16( + const void* query, const void* key, const void* value, void* logits, + void* output, int query_rows, int max_key_rows, int heads, + int head_dimension, const int* valid_key_rows, float scale, + std::uintptr_t stream) const; + modalities::Status vision_fmha_fp16( + const void* query, const void* key, const void* value, void* output, + int batch, int query_rows, int key_rows, int query_heads, + int key_heads, int head_dimension, int query_stride, int key_stride, + std::uintptr_t stream) const; + + modalities::Status patch_im2col_fp16( + const void* images, void* patches, int num_views, + std::uintptr_t stream) const; + modalities::Status patch_bias_position_fp16( + void* output, const void* bias, const void* position, + int sequence, int columns, int positions_per_view, + std::uintptr_t stream) const; + modalities::Status bias_residual_fp16( + void* residual, const void* values, const void* bias, + int rows, int columns, std::uintptr_t stream) const; + + modalities::Status gmm_fp16( + const void* a, const void* b, void* output, + int m, int n, int k, float beta, std::uintptr_t stream) const; + modalities::Status gmm_fp16_out_fp32( + const void* a, const void* b, float* output, + int m, int n, int k, std::uintptr_t stream) const; + modalities::Status action_update_fp16( + const float* delta, const void* bias, void* noise, + int rows, int columns, float dt, 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_THOR_KERNEL_DRIVER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h new file mode 100644 index 00000000..b152321d --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h @@ -0,0 +1,31 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_STYLE_PRECOMPUTE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_STYLE_PRECOMPUTE_H + +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_thor_kernel_driver.h" +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeThorStylePrecomputer { +public: + explicit NativeThorStylePrecomputer(const NativeThorKernelDriver* driver) + : driver_(driver) {} + + modalities::Status run(const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const; + +private: + const NativeThorKernelDriver* driver_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_STYLE_PRECOMPUTE_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h new file mode 100644 index 00000000..8bedf5e8 --- /dev/null +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h @@ -0,0 +1,64 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_WEIGHT_MATERIALIZER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_WEIGHT_MATERIALIZER_H + +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_quantization.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeThorMaterializationOptions { + int num_steps = 10; + bool include_embedding = true; +}; + +struct NativeThorWeightScales { + std::vector vision; + std::vector encoder; + std::vector decoder; +}; + +class NativeThorWeightMaterializer { +public: + NativeThorWeightMaterializer(const loader::SafetensorsFile& source, + NativeDeviceWeightStore* destination) + : source_(source), destination_(destination) {} + + modalities::Status materialize_all( + const NativeThorMaterializationOptions& options, + NativeThorWeightScales* scales); + +private: + modalities::Status upload_f16(const std::string& source_key, + const std::string& destination_name, + bool transpose); + modalities::Status upload_f16(const std::string& destination_name, + const NativeF16Tensor& tensor); + modalities::Status upload_fp8(const std::string& destination_name, + const NativeF16Tensor& tensor, + std::vector* scales); + modalities::Status materialize_vision_globals(); + modalities::Status materialize_vision_layer(int layer, + std::vector* scales); + modalities::Status materialize_encoder_layer(int layer, + std::vector* scales); + modalities::Status materialize_decoder_layer(int layer, + std::vector* scales); + modalities::Status materialize_decoder_globals(int num_steps); + modalities::Status materialize_embedding(); + modalities::Status upload_scale_vector(const std::string& name, + const std::vector& values); + + const loader::SafetensorsFile& source_; + NativeDeviceWeightStore* destination_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_WEIGHT_MATERIALIZER_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h index 1098bb24..2d2aaebb 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 @@ -15,12 +15,19 @@ namespace flashrt { namespace models { namespace pi05 { +enum class NativeWorkspaceFlavor { + kBf16, + kThorFp8, +}; + struct NativeWorkspaceConfig { int num_views = 2; int max_prompt_tokens = 200; int chunk_size = 10; int num_steps = 10; int vision_pool_factor = 1; + NativeWorkspaceFlavor flavor = NativeWorkspaceFlavor::kBf16; + bool enable_calibration = false; }; struct NativeWorkspaceBuffer { @@ -39,6 +46,7 @@ class NativeWorkspace { modalities::Status allocate(const NativeWorkspaceConfig& config); modalities::Status update_decoder_rope(int prompt_tokens); + modalities::Status set_fixed_prompt_length(int prompt_tokens); modalities::Status expand_vision_position_embedding( const NativeDeviceWeightStore& weights); const NativeWorkspaceBuffer* find(const std::string& name) const; @@ -49,6 +57,11 @@ class NativeWorkspace { int vision_sequence() const { return vision_sequence_; } int encoder_vision_sequence() const { return encoder_vision_sequence_; } int encoder_sequence() const { return encoder_sequence_; } + int total_keys() const { return encoder_sequence_ + chunk_size_; } + int num_views() const { return num_views_; } + int chunk_size() const { return chunk_size_; } + int num_steps() const { return num_steps_; } + NativeWorkspaceFlavor flavor() const { return flavor_; } private: modalities::Status add(const std::string& name, @@ -70,7 +83,11 @@ class NativeWorkspace { int num_views_ = 0; int max_prompt_tokens_ = 0; int chunk_size_ = 0; + int num_steps_ = 0; + NativeWorkspaceFlavor flavor_ = NativeWorkspaceFlavor::kBf16; frt_buffer decoder_rope_buffer_ = nullptr; + frt_buffer prompt_embedding_buffer_ = nullptr; + frt_buffer prompt_length_buffers_[3] = {}; std::vector rope_table_; }; diff --git a/cpp/models/pi05/src/native_calibration.cpp b/cpp/models/pi05/src/native_calibration.cpp new file mode 100644 index 00000000..6da57b76 --- /dev/null +++ b/cpp/models/pi05/src/native_calibration.cpp @@ -0,0 +1,394 @@ +#include "flashrt/cpp/models/pi05/native_calibration.h" + +#include "flashrt/cpp/loader/safetensors.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +constexpr char kSchema[] = "flashrt.pi05.fp8_calibration.v1"; +constexpr char kModel[] = "pi05"; +constexpr char kPrecision[] = "fp8_e4m3fn"; +constexpr char kTensorDtype[] = "float16"; +constexpr char kReducer[] = "linear_percentile"; +constexpr std::size_t kEncoderScaleCount = 18 * 4; + +modalities::Status invalid(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status not_found(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kNotFound, + message); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +bool valid_digest(const std::string& value) { + if (value.size() != 64) return false; + return std::all_of(value.begin(), value.end(), [](unsigned char c) { + return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'); + }); +} + +bool valid_scales(const std::vector& scales) { + return std::all_of(scales.begin(), scales.end(), [](float value) { + return std::isfinite(value) && value > 0.0f; + }); +} + +std::string json_string(const std::string& value) { + std::string escaped; + escaped.reserve(value.size() + 2); + escaped.push_back('"'); + constexpr char hex[] = "0123456789abcdef"; + for (unsigned char c : value) { + switch (c) { + case '"': escaped += "\\\""; break; + case '\\': escaped += "\\\\"; break; + case '\b': escaped += "\\b"; break; + case '\f': escaped += "\\f"; break; + case '\n': escaped += "\\n"; break; + case '\r': escaped += "\\r"; break; + case '\t': escaped += "\\t"; break; + default: + if (c < 0x20) { + escaped += "\\u00"; + escaped.push_back(hex[c >> 4]); + escaped.push_back(hex[c & 0xf]); + } else { + escaped.push_back(static_cast(c)); + } + } + } + escaped.push_back('"'); + return escaped; +} + +std::string decimal(double value) { + std::ostringstream out; + out << std::setprecision(17) << value; + return out.str(); +} + +std::map artifact_metadata( + const NativeCalibrationArtifact& artifact) { + return { + {"schema", kSchema}, + {"model", kModel}, + {"precision", kPrecision}, + {"tensor_dtype", kTensorDtype}, + {"reducer", kReducer}, + {"hardware", artifact.hardware}, + {"weights_sha256", artifact.weights_sha256}, + {"tokenizer_sha256", artifact.tokenizer_sha256}, + {"num_views", std::to_string(artifact.num_views)}, + {"max_prompt_tokens", std::to_string(artifact.max_prompt_tokens)}, + {"state_dim", std::to_string(artifact.state_dim)}, + {"chunk_size", std::to_string(artifact.chunk_size)}, + {"num_steps", std::to_string(artifact.num_steps)}, + {"vision_pool_factor", std::to_string(artifact.vision_pool_factor)}, + {"sample_count", std::to_string(artifact.sample_count)}, + {"percentile", decimal(artifact.percentile)}, + }; +} + +std::string make_header(const NativeCalibrationArtifact& artifact) { + const auto metadata = artifact_metadata(artifact); + const std::uint64_t encoder_bytes = + artifact.encoder_scales.size() * sizeof(float); + const std::uint64_t decoder_bytes = + artifact.decoder_scales.size() * sizeof(float); + std::ostringstream out; + out << "{\"__metadata__\":{"; + bool first = true; + for (const auto& entry : metadata) { + if (!first) out << ','; + first = false; + out << json_string(entry.first) << ':' << json_string(entry.second); + } + out << "},\"encoder_scales\":{\"dtype\":\"F32\",\"shape\":[" + << artifact.encoder_scales.size() + << "],\"data_offsets\":[0," << encoder_bytes + << "]},\"decoder_scales\":{\"dtype\":\"F32\",\"shape\":[" + << artifact.decoder_scales.size() << "],\"data_offsets\":[" + << encoder_bytes << ',' << encoder_bytes + decoder_bytes << "]}}"; + std::string header = out.str(); + while (header.size() % 8) header.push_back(' '); + return header; +} + +bool write_all(int fd, const void* data, std::size_t bytes) { + const auto* cursor = static_cast(data); + while (bytes) { + const ssize_t written = ::write(fd, cursor, bytes); + if (written < 0) { + if (errno == EINTR) continue; + return false; + } + if (!written) return false; + cursor += written; + bytes -= static_cast(written); + } + return true; +} + +std::string parent_directory(const std::string& path) { + const std::size_t slash = path.find_last_of('/'); + if (slash == std::string::npos) return "."; + if (slash == 0) return "/"; + return path.substr(0, slash); +} + +bool parse_int(const std::map& metadata, + const char* key, + int* value) { + const auto it = metadata.find(key); + if (it == metadata.end() || it->second.empty()) return false; + errno = 0; + char* end = nullptr; + const long parsed = std::strtol(it->second.c_str(), &end, 10); + if (errno || end != it->second.c_str() + it->second.size() || + parsed < std::numeric_limits::min() || + parsed > std::numeric_limits::max()) { + return false; + } + if (value) *value = static_cast(parsed); + return true; +} + +bool parse_u64(const std::map& metadata, + const char* key, + std::uint64_t* value) { + const auto it = metadata.find(key); + if (it == metadata.end() || it->second.empty() || it->second[0] == '-') { + return false; + } + errno = 0; + char* end = nullptr; + const unsigned long long parsed = + std::strtoull(it->second.c_str(), &end, 10); + if (errno || end != it->second.c_str() + it->second.size()) return false; + if (value) *value = static_cast(parsed); + return true; +} + +bool parse_double(const std::map& metadata, + const char* key, + double* value) { + const auto it = metadata.find(key); + if (it == metadata.end() || it->second.empty()) return false; + errno = 0; + char* end = nullptr; + const double parsed = std::strtod(it->second.c_str(), &end); + if (errno || end != it->second.c_str() + it->second.size() || + !std::isfinite(parsed)) { + return false; + } + if (value) *value = parsed; + return true; +} + +bool metadata_is(const std::map& metadata, + const char* key, + const char* expected) { + const auto it = metadata.find(key); + return it != metadata.end() && it->second == expected; +} + +bool copy_f32_tensor(const loader::SafetensorsFile& file, + const char* name, + std::vector* values) { + if (!values) return false; + const loader::SafetensorInfo* tensor = file.find(name); + if (!tensor || tensor->dtype != "F32" || tensor->shape.size() != 1 || + tensor->shape[0] > + std::numeric_limits::max() / sizeof(float) || + tensor->bytes != tensor->shape[0] * sizeof(float)) { + return false; + } + std::vector copied(static_cast(tensor->shape[0])); + if (!copied.empty()) { + const void* data = file.data(*tensor); + if (!data) return false; + std::memcpy(copied.data(), data, tensor->bytes); + } + *values = std::move(copied); + return true; +} + +} // namespace + +modalities::Status validate_native_calibration_artifact( + const NativeCalibrationArtifact& artifact) { + if (artifact.hardware.empty() || !valid_digest(artifact.weights_sha256) || + !valid_digest(artifact.tokenizer_sha256) || artifact.num_views < 1 || + artifact.num_views > 3 || artifact.max_prompt_tokens < 1 || + artifact.state_dim < 1 || artifact.chunk_size < 1 || + artifact.num_steps < 1 || + (artifact.vision_pool_factor != 1 && + artifact.vision_pool_factor != 2 && + artifact.vision_pool_factor != 4) || + !artifact.sample_count || !std::isfinite(artifact.percentile) || + artifact.percentile < 0.0 || artifact.percentile > 100.0) { + return invalid("native calibration metadata is invalid"); + } + if (artifact.encoder_scales.size() != kEncoderScaleCount || + artifact.decoder_scales.size() != + static_cast(artifact.num_steps) * 18 * 4 || + !valid_scales(artifact.encoder_scales) || + !valid_scales(artifact.decoder_scales)) { + return invalid("native calibration scale payload is invalid"); + } + return modalities::Status::ok(); +} + +modalities::Status save_native_calibration_artifact( + const std::string& path, + const NativeCalibrationArtifact& artifact) { + if (path.empty()) return invalid("native calibration path is empty"); + modalities::Status st = validate_native_calibration_artifact(artifact); + if (!st.ok_status()) return st; + + const std::uint16_t endian_probe = 1; + if (*reinterpret_cast(&endian_probe) != 1) { + return backend("safetensors writing requires a little-endian host"); + } + const std::string header = make_header(artifact); + std::string temporary = path + ".tmp.XXXXXX"; + std::vector writable(temporary.begin(), temporary.end()); + writable.push_back('\0'); + const int fd = ::mkstemp(writable.data()); + if (fd < 0) { + return backend(std::string("unable to create calibration artifact: ") + + std::strerror(errno)); + } + temporary = writable.data(); + bool ok = true; + unsigned char header_size[8]; + std::uint64_t size = header.size(); + for (int i = 0; i < 8; ++i) { + header_size[i] = static_cast((size >> (8 * i)) & 0xffu); + } + ok = write_all(fd, header_size, sizeof(header_size)) && + write_all(fd, header.data(), header.size()) && + write_all(fd, artifact.encoder_scales.data(), + artifact.encoder_scales.size() * sizeof(float)) && + write_all(fd, artifact.decoder_scales.data(), + artifact.decoder_scales.size() * sizeof(float)) && + ::fsync(fd) == 0; + const int close_rc = ::close(fd); + if (!ok || close_rc != 0 || ::rename(temporary.c_str(), path.c_str()) != 0) { + const int saved_errno = errno; + ::unlink(temporary.c_str()); + return backend(std::string("unable to publish calibration artifact: ") + + std::strerror(saved_errno)); + } + const std::string parent = parent_directory(path); + const int dir_fd = ::open(parent.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC); + if (dir_fd >= 0) { + ::fsync(dir_fd); + ::close(dir_fd); + } + return modalities::Status::ok(); +} + +modalities::Status load_native_calibration_artifact( + const std::string& path, + NativeCalibrationArtifact* artifact) { + if (path.empty() || !artifact) { + return invalid("native calibration load arguments are invalid"); + } + loader::SafetensorsFile file; + if (!file.open(path)) return not_found(file.error()); + const auto& metadata = file.metadata(); + if (!metadata_is(metadata, "schema", kSchema) || + !metadata_is(metadata, "model", kModel) || + !metadata_is(metadata, "precision", kPrecision) || + !metadata_is(metadata, "tensor_dtype", kTensorDtype) || + !metadata_is(metadata, "reducer", kReducer)) { + return invalid("native calibration artifact schema is incompatible"); + } + NativeCalibrationArtifact loaded; + const auto hardware = metadata.find("hardware"); + const auto weights = metadata.find("weights_sha256"); + const auto tokenizer = metadata.find("tokenizer_sha256"); + if (hardware == metadata.end() || weights == metadata.end() || + tokenizer == metadata.end() || + !parse_int(metadata, "num_views", &loaded.num_views) || + !parse_int(metadata, "max_prompt_tokens", &loaded.max_prompt_tokens) || + !parse_int(metadata, "state_dim", &loaded.state_dim) || + !parse_int(metadata, "chunk_size", &loaded.chunk_size) || + !parse_int(metadata, "num_steps", &loaded.num_steps) || + !parse_int(metadata, "vision_pool_factor", &loaded.vision_pool_factor) || + !parse_u64(metadata, "sample_count", &loaded.sample_count) || + !parse_double(metadata, "percentile", &loaded.percentile) || + !copy_f32_tensor(file, "encoder_scales", &loaded.encoder_scales) || + !copy_f32_tensor(file, "decoder_scales", &loaded.decoder_scales)) { + return invalid("native calibration artifact metadata is incomplete"); + } + loaded.hardware = hardware->second; + loaded.weights_sha256 = weights->second; + loaded.tokenizer_sha256 = tokenizer->second; + modalities::Status st = validate_native_calibration_artifact(loaded); + if (!st.ok_status()) return st; + *artifact = std::move(loaded); + return modalities::Status::ok(); +} + +modalities::Status reduce_native_calibration_samples( + const std::vector>& samples, + double percentile, + std::vector* reduced) { + if (!reduced || samples.empty() || !std::isfinite(percentile) || + percentile < 0.0 || percentile > 100.0 || samples.front().empty()) { + return invalid("native calibration reduction input is invalid"); + } + const std::size_t points = samples.front().size(); + for (const auto& sample : samples) { + if (sample.size() != points || !valid_scales(sample)) { + return invalid("native calibration samples are incompatible"); + } + } + std::vector result(points); + std::vector column(samples.size()); + const double rank = + percentile * static_cast(samples.size() - 1) / 100.0; + const std::size_t lower = static_cast(std::floor(rank)); + const std::size_t upper = static_cast(std::ceil(rank)); + const double fraction = rank - static_cast(lower); + for (std::size_t point = 0; point < points; ++point) { + for (std::size_t sample = 0; sample < samples.size(); ++sample) { + column[sample] = static_cast(samples[sample][point]); + } + std::sort(column.begin(), column.end()); + const double value = column[lower] + + (column[upper] - column[lower]) * fraction; + result[point] = static_cast(value); + } + *reduced = std::move(result); + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index da678126..f9af08b2 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -6,8 +6,10 @@ #include #include +#include #include #include +#include #include namespace flashrt { @@ -57,6 +59,10 @@ std::unique_ptr NativeGraphOwner::create( if (config.num_views < 1 || config.num_views > 3 || config.max_prompt_tokens < 1 || config.chunk_size < 1 || config.num_steps < 1 || + static_cast(config.max_prompt_tokens) + + static_cast(config.chunk_size) + + static_cast(config.num_views) * 256 > + static_cast(std::numeric_limits::max()) || (config.vision_pool_factor != 1 && config.vision_pool_factor != 2 && config.vision_pool_factor != 4)) { if (status) *status = invalid("native graph configuration is invalid"); diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 4da75e70..706d1bea 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -1,11 +1,19 @@ #include "native_open_internal.h" -#if defined(FLASHRT_CPP_WITH_FA2) && defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_SENTENCEPIECE) && \ + (defined(FLASHRT_CPP_WITH_FA2) || defined(FLASHRT_CPP_WITH_THOR_FP8)) #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_runtime.h" +#if defined(FLASHRT_CPP_WITH_FA2) #include "flashrt/cpp/models/pi05/native_graph_owner.h" +#endif +#if defined(FLASHRT_CPP_WITH_THOR_FP8) +#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/native_thor_graph_owner.h" +#endif #include @@ -21,17 +29,17 @@ namespace pi05 { namespace { const NativeWorkspaceBuffer* workspace_buffer( - const NativeGraphOwner& owner, + const NativeGraphRuntime& owner, const char* name) { return owner.workspace().find(name); } void release_graph_owner(void* owner) { - delete static_cast(owner); + delete static_cast(owner); } int update_prompt_length(void* owner, std::uint64_t prompt_len) { - auto* graph = static_cast(owner); + 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))); @@ -84,12 +92,51 @@ int build_native_model_runtime(const NativeOpenConfig& config, if (error) *error = cudaGetErrorString(cuda_rc); return -6; } - if (properties.major != 12 || properties.minor != 0) { - if (error) *error = "Pi0.5 native_v2 requires RTX SM120"; - return -3; - } const std::string hardware_id = "sm" + std::to_string(properties.major * 10 + properties.minor); + enum class Precision { kBf16, kFp8E4M3Fn }; + Precision precision; + if (config.precision == "auto") { + if (properties.major == 12 && properties.minor == 0) { + precision = Precision::kBf16; + } else if (properties.major == 11 && properties.minor == 0) { + precision = Precision::kFp8E4M3Fn; + } else { + if (error) { + *error = "Pi0.5 native_v2 has no backend for " + hardware_id; + } + return -3; + } + } else if (config.precision == "bf16") { + precision = Precision::kBf16; + } else if (config.precision == "fp8_e4m3fn") { + precision = Precision::kFp8E4M3Fn; + } else { + if (error) *error = "Pi0.5 native precision is invalid"; + return -1; + } + if (precision == Precision::kBf16 && + (properties.major != 12 || properties.minor != 0)) { + if (error) *error = "Pi0.5 native BF16 requires SM120"; + return -3; + } + if (precision == Precision::kFp8E4M3Fn && + (properties.major != 11 || properties.minor != 0)) { + if (error) *error = "Pi0.5 native FP8 requires SM110"; + return -3; + } +#if !defined(FLASHRT_CPP_WITH_FA2) + if (precision == Precision::kBf16) { + if (error) *error = "Pi0.5 native BF16 backend is not built"; + return -3; + } +#endif +#if !defined(FLASHRT_CPP_WITH_THOR_FP8) + if (precision == Precision::kFp8E4M3Fn) { + if (error) *error = "Pi0.5 native Thor FP8 backend is not built"; + return -3; + } +#endif struct HashResult { bool ok = false; @@ -113,6 +160,40 @@ int build_native_model_runtime(const NativeOpenConfig& config, return -2; } + NativeCalibrationArtifact calibration; + std::string calibration_sha256; + if (precision == Precision::kFp8E4M3Fn) { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) + if (config.calibration_path.empty()) { + if (error) *error = "Pi0.5 native FP8 requires calibration_path"; + return -1; + } + modalities::Status calibration_status = + load_native_calibration_artifact(config.calibration_path, + &calibration); + if (!calibration_status.ok_status()) { + if (error) *error = calibration_status.message; + return cface::status_code(calibration_status); + } + if (calibration.hardware != hardware_id || + calibration.tokenizer_sha256 != tokenizer_sha256 || + calibration.num_views != config.num_views || + calibration.max_prompt_tokens != config.max_prompt_tokens || + calibration.state_dim != config.state_dim || + calibration.chunk_size != config.chunk || + calibration.num_steps != config.num_steps || + calibration.vision_pool_factor != config.vision_pool_factor) { + if (error) *error = "Pi0.5 calibration identity does not match config"; + return -2; + } + if (!loader::sha256_file(config.calibration_path, + &calibration_sha256, &hash_error)) { + if (error) *error = hash_error; + return -2; + } +#endif + } + NativeGraphConfig graph_config; graph_config.num_views = config.num_views; graph_config.max_prompt_tokens = config.max_prompt_tokens; @@ -120,8 +201,18 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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); + std::unique_ptr graph; + if (precision == Precision::kBf16) { +#if defined(FLASHRT_CPP_WITH_FA2) + graph = NativeGraphOwner::create( + config.checkpoint_path, graph_config, &st); +#endif + } else { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) + graph = NativeThorGraphOwner::create( + config.checkpoint_path, graph_config, calibration, &st); +#endif + } if (!graph) { if (error) *error = st.message; return cface::status_code(st); @@ -131,6 +222,11 @@ int build_native_model_runtime(const NativeOpenConfig& config, if (error) *error = weights_sha256.error; return -2; } + if (precision == Precision::kFp8E4M3Fn && + calibration.weights_sha256 != weights_sha256.digest) { + if (error) *error = "Pi0.5 calibration checkpoint digest mismatch"; + return -2; + } const NativeWorkspaceBuffer* images = workspace_buffer(*graph, "observation_images_normalized"); @@ -150,7 +246,9 @@ int build_native_model_runtime(const NativeOpenConfig& config, "embedding_weight"); if (!images || !noise || !encoder || !previous || !prefix_weights || !guidance || !prompt || !embedding || - embedding->dtype != NativeWeightDType::kBf16 || + embedding->dtype != (precision == Precision::kBf16 + ? NativeWeightDType::kBf16 + : NativeWeightDType::kFloat16) || embedding->shape.size() != 2 || embedding->shape[1] != 2048) { if (error) *error = "native graph export buffers are incomplete"; return -6; @@ -201,11 +299,16 @@ int build_native_model_runtime(const NativeOpenConfig& config, FRT_RT_REGION_SNAPSHOT | FRT_RT_REGION_RESTORE) == 0; if (!ok) return fail_builder(builder, error, "native region build failed"); + const bool thor_fp8 = precision == Precision::kFp8E4M3Fn; + const std::string precision_id = thor_fp8 ? "fp8_e4m3fn" : "bf16"; + const std::string pipeline_id = thor_fp8 ? "NativeThorFp8" : "NativeBf16"; + const std::string tensor_dtype = thor_fp8 ? "float16" : "bf16"; ok = add_identity(builder, "model", "pi05") && add_identity(builder, "producer", "native") && - add_identity(builder, "pipeline", "NativeBf16") && + add_identity(builder, "pipeline", pipeline_id) && add_identity(builder, "hardware", hardware_id) && - add_identity(builder, "tensor_dtype", "bf16") && + add_identity(builder, "precision", precision_id) && + add_identity(builder, "tensor_dtype", tensor_dtype) && add_identity(builder, "weights_sha256", weights_sha256.digest) && add_identity(builder, "tokenizer_sha256", tokenizer_sha256) && add_identity(builder, "io", "native_v2") && @@ -221,11 +324,15 @@ int build_native_model_runtime(const NativeOpenConfig& config, add_identity(builder, "model_action_dim", "32") && add_identity(builder, "robot_action_dim", std::to_string(config.action_q01.size())); + if (ok && thor_fp8) { + ok = add_identity(builder, "calibration_sha256", calibration_sha256); + } if (!ok) return fail_builder(builder, error, "native identity build failed"); std::ostringstream manifest; manifest << "{\"model\":\"pi05\",\"producer\":\"native\"," << "\"hardware\":\"" << hardware_id + << "\",\"precision\":\"" << precision_id << "\",\"io\":\"native_v2\"," << "\"stage_plan\":{\"name\":\"full\"," << "\"stages\":[{\"name\":\"infer\"," @@ -243,6 +350,8 @@ int build_native_model_runtime(const NativeOpenConfig& config, const std::uint64_t action_bytes = static_cast(config.chunk) * config.action_q01.size() * sizeof(float); + const uint32_t io_dtype = + thor_fp8 ? FRT_RT_DTYPE_F16 : FRT_RT_DTYPE_BF16; 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, @@ -252,12 +361,12 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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, + builder, "images", FRT_RT_MOD_IMAGE, io_dtype, 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, + builder, "noise", FRT_RT_MOD_TENSOR, io_dtype, 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 && @@ -266,14 +375,14 @@ int build_native_model_runtime(const NativeOpenConfig& config, FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, action_shape, 2, 0, nullptr, 0, action_bytes) == 0 && frt_runtime_builder_add_port( - builder, "actions_raw", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_BF16, + builder, "actions_raw", FRT_RT_MOD_TENSOR, io_dtype, 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(); + NativeGraphRuntime* 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(); @@ -306,20 +415,24 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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; + const int runtime_dtype = thor_fp8 ? FRT_PI05_DTYPE_FLOAT16 + : FRT_PI05_DTYPE_BFLOAT16; + runtime_config.image_dtype = runtime_dtype; + runtime_config.action_dtype = runtime_dtype; + runtime_config.max_frame_width = config.max_frame_width; + runtime_config.max_frame_height = config.max_frame_height; 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_table_dtype = runtime_dtype; 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.prompt_embedding_dtype = runtime_dtype; 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(); @@ -357,7 +470,9 @@ 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"; + if (error) { + *error = "native graph backend and SentencePiece are unavailable"; + } return -3; } diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index fb078cf2..6a803f1e 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -602,11 +603,15 @@ int validate_config( std::string checkpoint_path; std::string tokenizer_model_path; std::string state_prompt_mode; + std::string precision; + std::string calibration_path; 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)) { + !string_field(obj, "state_prompt_mode", &state_prompt_mode, true) || + !string_field(obj, "precision", &precision, false) || + !string_field(obj, "calibration_path", &calibration_path, false)) { return -1; } if (io != "native_v2") { @@ -618,6 +623,13 @@ int validate_config( "Pi0.5 native_v2 requires state_prompt_mode='fixed'"; return -1; } + if (precision.empty()) precision = "auto"; + if (precision != "auto" && precision != "bf16" && + precision != "fp8_e4m3fn") { + g_last_error = + "precision must be 'auto', 'bf16', or 'fp8_e4m3fn'"; + return -1; + } if (!path_exists(checkpoint_path)) { g_last_error = "checkpoint_path does not exist"; return -2; @@ -626,6 +638,11 @@ int validate_config( g_last_error = "tokenizer_model_path does not name a file"; return -2; } + if (!calibration_path.empty() && + !regular_file_exists(calibration_path)) { + g_last_error = "calibration_path does not name a file"; + return -2; + } if (!validate_pi05_safetensors(checkpoint_path)) { return -2; } @@ -636,12 +653,16 @@ int validate_config( int64_t chunk = 0; int64_t num_steps = 10; int64_t vision_pool_factor = 1; + int64_t max_frame_width = 1280; + int64_t max_frame_height = 720; if (!integer_field(obj, "max_prompt_tokens", &max_prompt_tokens) || !integer_field(obj, "state_dim", &state_dim) || !integer_field(obj, "num_views", &num_views) || !integer_field(obj, "chunk", &chunk) || !integer_field(obj, "num_steps", &num_steps) || - !integer_field(obj, "vision_pool_factor", &vision_pool_factor)) { + !integer_field(obj, "vision_pool_factor", &vision_pool_factor) || + !integer_field(obj, "max_frame_width", &max_frame_width) || + !integer_field(obj, "max_frame_height", &max_frame_height)) { return -1; } if (max_prompt_tokens < 200 || max_prompt_tokens > INT_MAX) { @@ -669,15 +690,24 @@ int validate_config( g_last_error = "vision_pool_factor must be one of 1, 2, or 4"; return -1; } + if (max_frame_width <= 0 || max_frame_width > INT_MAX || + max_frame_height <= 0 || max_frame_height > INT_MAX) { + g_last_error = "max_frame_width/height must be in [1, INT_MAX]"; + return -1; + } flashrt::models::pi05::NativeOpenConfig config; config.checkpoint_path = checkpoint_path; config.tokenizer_model_path = tokenizer_model_path; + config.precision = precision; + config.calibration_path = calibration_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); + config.max_frame_width = static_cast(max_frame_width); + config.max_frame_height = static_cast(max_frame_height); if (!validate_norm_stats(checkpoint_path, state_dim, &config)) { return -2; } @@ -689,25 +719,51 @@ int validate_config( } // namespace +namespace flashrt { +namespace models { +namespace pi05 { + +int parse_native_open_config(const char* config_json, + NativeOpenConfig* out, + std::string* error) { + const int rc = validate_config(config_json, out); + if (error) *error = g_last_error; + return rc; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt + extern "C" int frt_model_runtime_open_v1(const char* config_json, - frt_model_runtime_v1** out) { + frt_model_runtime_v1** out) try { if (!out) { g_last_error = "out is null"; return -1; } *out = nullptr; flashrt::models::pi05::NativeOpenConfig config; - const int rc = validate_config(config_json, &config); + const int rc = flashrt::models::pi05::parse_native_open_config( + config_json, &config, &g_last_error); if (rc != 0) return rc; -#if defined(FLASHRT_CPP_WITH_FA2) && defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_SENTENCEPIECE) && \ + (defined(FLASHRT_CPP_WITH_FA2) || defined(FLASHRT_CPP_WITH_THOR_FP8)) 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"; + "Pi0.5 native open validated config; this build requires " + "SentencePiece and a native graph backend"; return -3; #endif +} catch (const std::exception& error) { + if (out) *out = nullptr; + g_last_error = error.what(); + return -6; +} catch (...) { + if (out) *out = nullptr; + g_last_error = "Pi0.5 native open failed"; + return -6; } 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 index 21683dfe..a32e6747 100644 --- a/cpp/models/pi05/src/native_open_internal.h +++ b/cpp/models/pi05/src/native_open_internal.h @@ -13,18 +13,26 @@ namespace pi05 { struct NativeOpenConfig { std::string checkpoint_path; std::string tokenizer_model_path; + std::string precision = "auto"; + std::string calibration_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; + int max_frame_width = 1280; + int max_frame_height = 720; std::vector state_q01; std::vector state_q99; std::vector action_q01; std::vector action_q99; }; +int parse_native_open_config(const char* config_json, + NativeOpenConfig* out, + std::string* error); + int build_native_model_runtime(const NativeOpenConfig& config, frt_model_runtime_v1** out, std::string* error); diff --git a/cpp/models/pi05/src/native_rope.cu b/cpp/models/pi05/src/native_rope.cu new file mode 100644 index 00000000..ca57bfe0 --- /dev/null +++ b/cpp/models/pi05/src/native_rope.cu @@ -0,0 +1,62 @@ +#include "flashrt/cpp/models/pi05/native_rope.h" + +#include +#include +#include + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +constexpr int kFrequencies = 128; + +__global__ void generate_rope_kernel(__half* output, int start_position, + int positions) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= positions * kFrequencies) return; + const int row = index / kFrequencies; + const int frequency = index - row * kFrequencies; + const float exponent = static_cast(frequency) / kFrequencies; + const float denominator = powf(10000.0f, exponent); + float inverse_frequency = __fdiv_rn(1.0f, denominator); + inverse_frequency = __bfloat162float( + __float2bfloat16_rn(inverse_frequency)); + const float phase = __fmul_rn( + static_cast(start_position + row), inverse_frequency); + output[2 * index] = __float2half_rn(cosf(phase)); + output[2 * index + 1] = __float2half_rn(sinf(phase)); +} + +} // namespace + +modalities::Status generate_native_thor_rope_f16( + void* output, int start_position, int positions, std::uintptr_t stream) { + constexpr int kMax = std::numeric_limits::max(); + if (!output || start_position < 0 || positions <= 0 || + positions > kMax / kFrequencies || + start_position > kMax - (positions - 1)) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "Thor RoPE generation arguments are invalid"); + } + cudaStream_t cuda_stream = reinterpret_cast(stream); + const int elements = positions * kFrequencies; + generate_rope_kernel<<<(elements + 255) / 256, 256, 0, cuda_stream>>>( + static_cast<__half*>(output), start_position, positions); + cudaError_t rc = cudaGetLastError(); + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + return rc == cudaSuccess + ? modalities::Status::ok() + : modalities::Status::error( + modalities::StatusCode::kBackend, + std::string("Thor RoPE generation failed: ") + + cudaGetErrorString(rc)); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_rope_unavailable.cpp b/cpp/models/pi05/src/native_rope_unavailable.cpp new file mode 100644 index 00000000..84655921 --- /dev/null +++ b/cpp/models/pi05/src/native_rope_unavailable.cpp @@ -0,0 +1,16 @@ +#include "flashrt/cpp/models/pi05/native_rope.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +modalities::Status generate_native_thor_rope_f16( + void*, int, int, std::uintptr_t) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Thor RoPE generation requires the CUDA kernels build"); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_calibration_c_api.cpp b/cpp/models/pi05/src/native_thor_calibration_c_api.cpp new file mode 100644 index 00000000..e55d8440 --- /dev/null +++ b/cpp/models/pi05/src/native_thor_calibration_c_api.cpp @@ -0,0 +1,254 @@ +#include "flashrt/cpp/models/pi05/c_api.h" + +#include "config_map.h" +#include "flashrt/cpp/models/pi05/spec.h" +#include "native_open_internal.h" + +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#include "flashrt/cpp/models/pi05/native_thor_calibration_session.h" +#endif + +#include +#include +#include +#include +#include +#include + +using flashrt::modalities::DType; +using flashrt::modalities::Layout; +using flashrt::modalities::MemoryPlace; +using flashrt::modalities::PixelFormat; +using flashrt::modalities::Shape; +using flashrt::modalities::VisionFrame; + +thread_local std::string g_calibration_create_error; + +struct frt_pi05_calibration_session_s { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + std::unique_ptr impl; +#endif + std::string last_error; + std::vector view_names; + std::vector frames; +}; + +namespace { + +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +int set_error(frt_pi05_calibration_session* session, + const flashrt::modalities::Status& status) { + if (session) session->last_error = status.message; + return flashrt::models::pi05::cface::status_code(status); +} + +bool convert_frames(frt_pi05_calibration_session* session, + const frt_pi05_calibration_sample_v1& sample) { + if (!session || session->view_names.empty() || !sample.frames || + sample.n_frames != session->view_names.size()) { + return false; + } + const std::size_t count = static_cast(sample.n_frames); + session->frames.clear(); + session->frames.resize(count); + bool seen[3]{}; + for (std::uint64_t i = 0; i < sample.n_frames; ++i) { + const frt_pi05_vision_frame& source = sample.frames[i]; + if (source.struct_size < sizeof(frt_pi05_vision_frame) || + !source.name || !source.data || source.width <= 0 || + source.height <= 0 || source.stride_bytes <= 0 || + source.pixel_format != FRT_PI05_PIXEL_RGB8) { + return false; + } + std::size_t slot = count; + for (std::size_t candidate = 0; candidate < count; ++candidate) { + if (source.name == session->view_names[candidate]) { + slot = candidate; + break; + } + } + if (slot == count || seen[slot]) return false; + seen[slot] = true; + VisionFrame& frame = session->frames[slot]; + frame.name = source.name; + frame.image.data = const_cast(source.data); + frame.image.bytes = source.bytes; + frame.image.dtype = DType::kUInt8; + frame.image.place = MemoryPlace::kHost; + frame.image.layout = Layout::kHWC; + frame.image.shape = + Shape{static_cast(source.height), + static_cast(source.width), 3}; + frame.format = PixelFormat::kRGB8; + frame.width = source.width; + frame.height = source.height; + frame.stride_bytes = source.stride_bytes; + frame.timestamp_ns = source.timestamp_ns; + } + for (std::size_t slot = 0; slot < count; ++slot) { + if (!seen[slot]) return false; + } + return true; +} +#endif + +} // namespace + +extern "C" int frt_pi05_calibration_create_v1( + const char* config_json, + double percentile, + frt_pi05_calibration_session** out) try { + g_calibration_create_error.clear(); + if (!out) { + g_calibration_create_error = "calibration out is null"; + return -1; + } + *out = nullptr; +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + flashrt::models::pi05::NativeOpenConfig open_config; + int rc = flashrt::models::pi05::parse_native_open_config( + config_json, &open_config, &g_calibration_create_error); + if (rc != 0) return rc; + if (open_config.precision == "bf16") { + g_calibration_create_error = + "Pi0.5 calibration precision must resolve to fp8_e4m3fn"; + return -1; + } + flashrt::models::pi05::NativeThorCalibrationConfig config; + config.checkpoint_path = open_config.checkpoint_path; + config.tokenizer_model_path = open_config.tokenizer_model_path; + config.max_prompt_tokens = open_config.max_prompt_tokens; + config.state_dim = open_config.state_dim; + config.num_views = open_config.num_views; + config.chunk_size = open_config.chunk; + config.num_steps = open_config.num_steps; + config.vision_pool_factor = open_config.vision_pool_factor; + config.max_frame_width = open_config.max_frame_width; + config.max_frame_height = open_config.max_frame_height; + config.state_q01 = std::move(open_config.state_q01); + config.state_q99 = std::move(open_config.state_q99); + flashrt::modalities::Status status; + auto impl = + flashrt::models::pi05::NativeThorCalibrationSession::create( + config, percentile, &status); + if (!impl) { + g_calibration_create_error = status.message; + return flashrt::models::pi05::cface::status_code(status); + } + std::unique_ptr session( + new (std::nothrow) frt_pi05_calibration_session); + if (!session) { + g_calibration_create_error = "calibration handle allocation failed"; + return -6; + } + session->impl = std::move(impl); + session->view_names = + flashrt::models::pi05::vision_preprocess_spec(config.num_views) + .view_order; + *out = session.release(); + return 0; +#else + (void)config_json; + (void)percentile; + g_calibration_create_error = + "Pi0.5 calibration requires Thor FP8 and SentencePiece"; + return -3; +#endif +} catch (const std::exception& error) { + if (out) *out = nullptr; + g_calibration_create_error = error.what(); + return -6; +} catch (...) { + if (out) *out = nullptr; + g_calibration_create_error = "calibration creation failed"; + return -6; +} + +extern "C" int frt_pi05_calibration_observe_v1( + frt_pi05_calibration_session* session, + const frt_pi05_calibration_sample_v1* sample) try { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + if (!session || !session->impl || !sample || + sample->struct_size < sizeof(frt_pi05_calibration_sample_v1) || + !sample->prompt || (!sample->state && sample->n_state) || + (!sample->noise && sample->n_noise) || + !convert_frames(session, *sample)) { + if (session) session->last_error = "calibration sample is invalid"; + return -1; + } + const flashrt::modalities::Status status = session->impl->observe( + sample->prompt, sample->state, sample->n_state, session->frames, + sample->noise, sample->n_noise, sample->noise_seed); + if (!status.ok_status()) return set_error(session, status); + session->last_error.clear(); + return 0; +#else + (void)session; + (void)sample; + return -3; +#endif +} catch (const std::exception& error) { + if (session) session->last_error = error.what(); + return -6; +} catch (...) { + if (session) session->last_error = "calibration observation failed"; + return -6; +} + +extern "C" int frt_pi05_calibration_finalize_v1( + frt_pi05_calibration_session* session, + const char* artifact_path) try { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + if (!session || !session->impl || !artifact_path || !artifact_path[0]) { + if (session) session->last_error = "calibration output path is invalid"; + return -1; + } + const flashrt::modalities::Status status = + session->impl->finalize(artifact_path); + if (!status.ok_status()) return set_error(session, status); + session->last_error.clear(); + return 0; +#else + (void)session; + (void)artifact_path; + return -3; +#endif +} catch (const std::exception& error) { + if (session) session->last_error = error.what(); + return -6; +} catch (...) { + if (session) session->last_error = "calibration finalization failed"; + return -6; +} + +extern "C" uint64_t frt_pi05_calibration_sample_count_v1( + const frt_pi05_calibration_session* session) { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ + defined(FLASHRT_CPP_HAS_SENTENCEPIECE) + return session && session->impl ? session->impl->sample_count() : 0; +#else + (void)session; + return 0; +#endif +} + +extern "C" const char* frt_pi05_calibration_last_error_v1( + const frt_pi05_calibration_session* session) { + return session ? session->last_error.c_str() + : g_calibration_create_error.c_str(); +} + +extern "C" const char* frt_pi05_calibration_create_last_error_v1() { + return g_calibration_create_error.c_str(); +} + +extern "C" void frt_pi05_calibration_destroy_v1( + frt_pi05_calibration_session* session) { + delete session; +} diff --git a/cpp/models/pi05/src/native_thor_calibration_session.cpp b/cpp/models/pi05/src/native_thor_calibration_session.cpp new file mode 100644 index 00000000..93e56241 --- /dev/null +++ b/cpp/models/pi05/src/native_thor_calibration_session.cpp @@ -0,0 +1,499 @@ +#include "flashrt/cpp/models/pi05/native_thor_calibration_session.h" + +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/loader/sha256.h" +#include "flashrt/cpp/models/pi05/io.h" +#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/native_thor_fp8_forward.h" +#include "flashrt/cpp/models/pi05/native_thor_style_precompute.h" +#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" +#include "flashrt/cpp/models/pi05/prompt_embed.h" + +#include + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +struct HashResult { + bool ok = false; + std::string digest; + std::string error; +}; + +std::uint64_t splitmix64(std::uint64_t* state) { + std::uint64_t value = (*state += 0x9e3779b97f4a7c15ull); + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ull; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebull; + return value ^ (value >> 31); +} + +double uniform_open(std::uint64_t* state) { + constexpr double kDenominator = 9007199254740993.0; + return (static_cast(splitmix64(state) >> 11) + 1.0) / + kDenominator; +} + +void deterministic_normal_f16(std::uint64_t seed, + std::vector* output) { + constexpr double kTwoPi = 6.283185307179586476925286766559; + std::uint64_t state = seed ^ 0x243f6a8885a308d3ull; + for (std::size_t i = 0; i < output->size(); i += 2) { + const double radius = std::sqrt(-2.0 * std::log(uniform_open(&state))); + const double angle = kTwoPi * uniform_open(&state); + (*output)[i] = modalities::float_to_float16( + static_cast(radius * std::cos(angle))); + if (i + 1 < output->size()) { + (*output)[i + 1] = modalities::float_to_float16( + static_cast(radius * std::sin(angle))); + } + } +} + +bool valid_config(const NativeThorCalibrationConfig& config) { + const std::uint64_t width = + static_cast(config.max_frame_width); + const std::uint64_t height = + static_cast(config.max_frame_height); + bool valid_quantiles = + config.state_q01.size() == + static_cast(config.state_dim) && + config.state_q99.size() == config.state_q01.size(); + for (std::size_t i = 0; + valid_quantiles && i < config.state_q01.size(); ++i) { + valid_quantiles = std::isfinite(config.state_q01[i]) && + std::isfinite(config.state_q99[i]) && + config.state_q99[i] > config.state_q01[i]; + } + return !config.checkpoint_path.empty() && + !config.tokenizer_model_path.empty() && config.state_dim > 0 && + config.num_views >= 1 && config.num_views <= 3 && + config.max_prompt_tokens >= 1 && + !(config.max_prompt_tokens & 1) && config.chunk_size > 0 && + config.num_steps > 0 && config.vision_pool_factor == 1 && + static_cast(config.max_prompt_tokens) + + static_cast(config.chunk_size) + + static_cast(config.num_views) * 256 <= + static_cast( + std::numeric_limits::max()) && + config.max_frame_width > 0 && config.max_frame_height > 0 && + width <= std::numeric_limits::max() / height / 4 && + valid_quantiles; +} + +modalities::TensorView device_view(const NativeWorkspaceBuffer* buffer, + modalities::DType dtype, + modalities::Layout layout, + modalities::Shape shape) { + modalities::TensorView view; + if (!buffer) return view; + view.data = frt_buffer_dptr(buffer->buffer); + view.bytes = frt_buffer_bytes(buffer->buffer); + view.dtype = dtype; + view.place = modalities::MemoryPlace::kDevice; + view.layout = layout; + view.shape = shape; + return view; +} + +} // namespace + +struct NativeThorCalibrationSession::Impl { + explicit Impl(frt_ctx context, NativeThorCalibrationConfig value, + double requested_percentile) + : config(std::move(value)), + percentile(requested_percentile), + ctx(context), + weights(context), + workspace(context), + forward(&driver) {} + + ~Impl() { + if (stream) { + cudaStreamSynchronize(stream); + cudaStreamDestroy(stream); + stream = nullptr; + } + modalities::text_embedding_staging_destroy(&text_staging); + modalities::vision_staging_destroy(&vision_staging); + if (ctx) { + frt_ctx_destroy(ctx); + ctx = nullptr; + } + } + + modalities::Status initialize() { + int device = 0; + cudaDeviceProp properties{}; + cudaError_t rc = cudaGetDevice(&device); + if (rc == cudaSuccess) { + rc = cudaGetDeviceProperties(&properties, device); + } + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + if (properties.major != 11 || properties.minor != 0) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Pi0.5 FP8 calibration requires SM110"); + } + hardware = + "sm" + std::to_string(properties.major * 10 + properties.minor); + + const std::string weights_path = + config.checkpoint_path + "/model.safetensors"; + std::future weights_hash = std::async( + std::launch::async, [weights_path] { + HashResult result; + result.ok = loader::sha256_file( + weights_path, &result.digest, &result.error); + return result; + }); + std::string hash_error; + if (!loader::sha256_file(config.tokenizer_model_path, + &tokenizer_sha256, &hash_error)) { + return modalities::Status::error( + modalities::StatusCode::kNotFound, hash_error); + } + loader::SafetensorsFile source; + if (!source.open(weights_path)) { + return modalities::Status::error( + modalities::StatusCode::kNotFound, source.error()); + } + NativeThorWeightMaterializer materializer(source, &weights); + NativeThorMaterializationOptions options; + options.num_steps = config.num_steps; + options.include_embedding = true; + modalities::Status st = + materializer.materialize_all(options, &weight_scales); + if (!st.ok_status()) return st; + HashResult digest = weights_hash.get(); + if (!digest.ok) { + return modalities::Status::error( + modalities::StatusCode::kNotFound, digest.error); + } + weights_sha256 = std::move(digest.digest); + + 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; + workspace_config.flavor = NativeWorkspaceFlavor::kThorFp8; + workspace_config.enable_calibration = true; + st = workspace.allocate(workspace_config); + if (!st.ok_status()) return st; + st = workspace.expand_vision_position_embedding(weights); + if (!st.ok_status()) return st; + st = workspace.set_fixed_prompt_length(0); + if (!st.ok_status()) return st; + + rc = cudaStreamCreate(&stream); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + NativeThorStylePrecomputer precomputer(&driver); + st = precomputer.run( + weights, &workspace, reinterpret_cast(stream)); + if (!st.ok_status()) return st; + + const std::uint64_t frame_capacity = + static_cast(config.max_frame_width) * + static_cast(config.max_frame_height) * 4; + st = modalities::vision_staging_create( + &vision_staging, static_cast(config.num_views), + frame_capacity); + if (!st.ok_status()) return st; + st = modalities::text_embedding_staging_create( + &text_staging, config.max_prompt_tokens); + if (!st.ok_status()) return st; + st = tokenizer.load_model(config.tokenizer_model_path); + if (!st.ok_status()) return st; + tokenizer.reserve(config.max_prompt_tokens); + + const NativeWorkspaceBuffer* image_buffer = + workspace.find("observation_images_normalized"); + const NativeWorkspaceBuffer* noise_buffer = + workspace.find("diffusion_noise"); + image_output = device_view( + image_buffer, modalities::DType::kFloat16, + modalities::Layout::kNHWC, + {static_cast(config.num_views), 224, 224, 3}); + action_output = device_view( + noise_buffer, modalities::DType::kFloat16, + modalities::Layout::kFlat, + {static_cast(config.chunk_size), 32}); + if (!image_output.data || !action_output.data) { + return invalid("Thor calibration IO buffers are incomplete"); + } + io.reset(new (std::nothrow) RuntimeIo( + config.num_views, image_output, action_output, {}, {}, stream, + config.chunk_size, 32, 32, modalities::DType::kFloat16, + &vision_staging, nullptr, true)); + if (!io) return backend("Thor calibration IO allocation failed"); + + const NativeDeviceWeight* embedding = + weights.find("embedding_weight"); + const NativeWorkspaceBuffer* prompt = + workspace.find("prompt_embedding"); + if (!embedding || !prompt || + embedding->dtype != NativeWeightDType::kFloat16 || + embedding->shape.size() != 2 || embedding->shape[1] != 2048) { + return invalid("Thor calibration embedding buffers are invalid"); + } + embedding_table.data = frt_buffer_dptr(embedding->buffer); + embedding_table.bytes = frt_buffer_bytes(embedding->buffer); + embedding_table.dtype = modalities::DType::kFloat16; + embedding_table.place = modalities::MemoryPlace::kDevice; + embedding_table.layout = modalities::Layout::kFlat; + embedding_table.shape = + {embedding->shape[0], embedding->shape[1]}; + prompt_output = device_view( + prompt, modalities::DType::kFloat16, + modalities::Layout::kFlat, + {static_cast(config.max_prompt_tokens), 2048}); + prompt_spec.vocab_size = embedding->shape[0]; + prompt_spec.hidden_dim = 2048; + prompt_spec.max_tokens = config.max_prompt_tokens; + prompt_spec.scale = std::sqrt(2048.0f); + normalized_state.resize(config.state_dim); + token_ids.reserve(static_cast(config.max_prompt_tokens) + 1); + const std::size_t max_prompt_bytes = + static_cast(config.max_prompt_tokens) * 8; + formatted_prompt.reserve( + max_prompt_bytes + static_cast(config.state_dim) * 5 + + 32); + noise_f16.resize(static_cast(config.chunk_size) * 32); + return modalities::Status::ok(); + } + + modalities::Status observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed) { + if (!state || n_state != static_cast(config.state_dim)) { + return invalid("Thor calibration state shape is invalid"); + } + if ((noise && n_noise != noise_f16.size()) || + (!noise && n_noise != 0)) { + return invalid("Thor calibration noise shape is invalid"); + } + for (std::size_t i = 0; i < normalized_state.size(); ++i) { + if (!std::isfinite(state[i])) { + return invalid("Thor calibration state contains non-finite data"); + } + 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; + } + std::uint64_t prompt_len = 0; + modalities::Status st = embed_prompt( + tokenizer, prompt_spec, prompt, normalized_state.data(), + normalized_state.size(), embedding_table, prompt_output, + &token_ids, &prompt_len, stream, &text_staging, + &formatted_prompt); + if (!st.ok_status()) return st; + rc_check = cudaStreamSynchronize(stream); + if (rc_check != cudaSuccess) return backend(cudaGetErrorString(rc_check)); + st = workspace.set_fixed_prompt_length(static_cast(prompt_len)); + if (!st.ok_status()) return st; + st = io->prepare_vision(frames); + if (!st.ok_status()) return st; + + const NativeWorkspaceBuffer* encoder = workspace.find("encoder_x"); + const NativeWorkspaceBuffer* prompt_buffer = + workspace.find("prompt_embedding"); + if (!encoder || !prompt_buffer) { + return invalid("Thor calibration prompt window is missing"); + } + const std::size_t prompt_offset = + static_cast(workspace.encoder_vision_sequence()) * + 2048 * sizeof(std::uint16_t); + rc_check = cudaMemcpyAsync( + static_cast(frt_buffer_dptr(encoder->buffer)) + + prompt_offset, + frt_buffer_dptr(prompt_buffer->buffer), + frt_buffer_bytes(prompt_buffer->buffer), cudaMemcpyDeviceToDevice, + stream); + if (rc_check != cudaSuccess) return backend(cudaGetErrorString(rc_check)); + + if (noise) { + for (std::size_t i = 0; i < noise_f16.size(); ++i) { + if (!std::isfinite(noise[i])) { + return invalid( + "Thor calibration noise contains non-finite data"); + } + noise_f16[i] = modalities::float_to_float16(noise[i]); + } + } else { + deterministic_normal_f16( + noise_seed + static_cast(encoder_samples.size()), + &noise_f16); + } + rc_check = cudaMemcpyAsync( + action_output.data, noise_f16.data(), + noise_f16.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice, + stream); + if (rc_check != cudaSuccess) return backend(cudaGetErrorString(rc_check)); + + const std::uintptr_t native_stream = + reinterpret_cast(stream); + st = forward.vision(weights, &workspace, weight_scales, native_stream); + if (!st.ok_status()) return st; + std::vector encoder_scale; + st = forward.calibrate_encoder( + weights, &workspace, weight_scales, &encoder_scale, native_stream); + if (!st.ok_status()) return st; + std::vector decoder_scale; + st = forward.calibrate_decoder( + weights, &workspace, &decoder_scale, native_stream); + if (!st.ok_status()) return st; + encoder_samples.push_back(std::move(encoder_scale)); + decoder_samples.push_back(std::move(decoder_scale)); + return modalities::Status::ok(); + } + + modalities::Status finalize(const std::string& artifact_path) const { + if (encoder_samples.empty() || + encoder_samples.size() != decoder_samples.size()) { + return invalid("Thor calibration has no complete samples"); + } + NativeCalibrationArtifact artifact; + artifact.hardware = hardware; + artifact.weights_sha256 = weights_sha256; + artifact.tokenizer_sha256 = tokenizer_sha256; + artifact.num_views = config.num_views; + artifact.max_prompt_tokens = config.max_prompt_tokens; + artifact.state_dim = config.state_dim; + artifact.chunk_size = config.chunk_size; + artifact.num_steps = config.num_steps; + artifact.vision_pool_factor = config.vision_pool_factor; + artifact.sample_count = encoder_samples.size(); + artifact.percentile = percentile; + modalities::Status st = reduce_native_calibration_samples( + encoder_samples, percentile, &artifact.encoder_scales); + if (!st.ok_status()) return st; + st = reduce_native_calibration_samples( + decoder_samples, percentile, &artifact.decoder_scales); + if (!st.ok_status()) return st; + return save_native_calibration_artifact(artifact_path, artifact); + } + + NativeThorCalibrationConfig config; + double percentile = 99.9; + std::string hardware; + std::string weights_sha256; + std::string tokenizer_sha256; + frt_ctx ctx = nullptr; + NativeDeviceWeightStore weights; + NativeWorkspace workspace; + NativeThorKernelDriver driver; + NativeThorFp8Forward forward; + NativeThorWeightScales weight_scales; + cudaStream_t stream = nullptr; + cudaError_t rc_check = cudaSuccess; + modalities::VisionStaging vision_staging; + modalities::TextEmbeddingStaging text_staging; + modalities::SentencePieceTokenizer tokenizer; + std::unique_ptr io; + modalities::TensorView image_output; + modalities::TensorView action_output; + modalities::TensorView embedding_table; + modalities::TensorView prompt_output; + PromptEmbeddingSpec prompt_spec; + std::vector token_ids; + std::vector normalized_state; + std::string formatted_prompt; + std::vector noise_f16; + std::vector> encoder_samples; + std::vector> decoder_samples; +}; + +NativeThorCalibrationSession::NativeThorCalibrationSession( + std::unique_ptr impl) + : impl_(std::move(impl)) {} + +NativeThorCalibrationSession::~NativeThorCalibrationSession() = default; + +std::unique_ptr +NativeThorCalibrationSession::create( + const NativeThorCalibrationConfig& config, + double percentile, + modalities::Status* status) { + if (!valid_config(config) || !std::isfinite(percentile) || + percentile < 0.0 || percentile > 100.0) { + if (status) *status = invalid("Thor calibration config is invalid"); + return nullptr; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) { + if (status) *status = backend("Thor calibration context creation failed"); + return nullptr; + } + std::unique_ptr impl( + new (std::nothrow) Impl(ctx, config, percentile)); + if (!impl) { + frt_ctx_destroy(ctx); + if (status) *status = backend("Thor calibration allocation failed"); + return nullptr; + } + modalities::Status st = impl->initialize(); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + std::unique_ptr session( + new (std::nothrow) NativeThorCalibrationSession(std::move(impl))); + if (!session) { + if (status) *status = backend("Thor calibration session allocation failed"); + return nullptr; + } + if (status) *status = modalities::Status::ok(); + return session; +} + +modalities::Status NativeThorCalibrationSession::observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed) { + return impl_ ? impl_->observe(prompt, state, n_state, frames, noise, + n_noise, noise_seed) + : invalid("Thor calibration session is invalid"); +} + +modalities::Status NativeThorCalibrationSession::finalize( + const std::string& artifact_path) const { + return impl_ ? impl_->finalize(artifact_path) + : invalid("Thor calibration session is invalid"); +} + +std::uint64_t NativeThorCalibrationSession::sample_count() const { + return impl_ ? impl_->encoder_samples.size() : 0; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_fp8_forward.cpp b/cpp/models/pi05/src/native_thor_fp8_forward.cpp new file mode 100644 index 00000000..1db58508 --- /dev/null +++ b/cpp/models/pi05/src/native_thor_fp8_forward.cpp @@ -0,0 +1,1297 @@ +#include "flashrt/cpp/models/pi05/native_thor_fp8_forward.h" + +#include + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +constexpr int kVisionWidth = 1152; +constexpr int kVisionHidden = 4304; +constexpr int kVisionHeads = 16; +constexpr int kVisionHeadDimension = 72; +constexpr int kEncoderWidth = 2048; +constexpr int kEncoderHidden = 16384; +constexpr int kDecoderWidth = 1024; +constexpr int kDecoderHidden = 4096; +constexpr int kHeads = 8; +constexpr int kHeadDimension = 256; +constexpr int kLayers = 18; + +static_assert(kVisionHeads * kVisionHeadDimension == kVisionWidth, + "Pi0.5 vision attention shape must cover the hidden width"); + +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); +} + +const NativeWorkspaceBuffer* buffer( + const NativeWorkspace& workspace, const char* name, + modalities::DType dtype, + std::initializer_list shape) { + const NativeWorkspaceBuffer* value = workspace.find(name); + return value && value->dtype == dtype && + value->shape == std::vector(shape) + ? value + : nullptr; +} + +const NativeDeviceWeight* weight( + const NativeDeviceWeightStore& weights, const std::string& name, + NativeWeightDType dtype, + std::initializer_list shape) { + const NativeDeviceWeight* value = weights.find(name); + return value && value->dtype == dtype && + value->shape == std::vector(shape) + ? value + : nullptr; +} + +void* dptr(const NativeWorkspaceBuffer* value) { + return value ? frt_buffer_dptr(value->buffer) : nullptr; +} + +void* dptr(const NativeDeviceWeight* value) { + return value ? frt_buffer_dptr(value->buffer) : nullptr; +} + +void* offset_bytes(void* base, std::size_t bytes) { + return static_cast(base) + bytes; +} + +const void* offset_bytes(const void* base, std::size_t bytes) { + return static_cast(base) + bytes; +} + +float* scale_ptr(const NativeWorkspaceBuffer* scales, std::size_t index) { + return static_cast(dptr(scales)) + index; +} + +float* scale_ptr(const NativeDeviceWeight* scales, std::size_t index) { + return static_cast(dptr(scales)) + index; +} + +modalities::Status copy_device(void* destination, const void* source, + std::size_t bytes, std::uintptr_t stream) { + const cudaError_t rc = cudaMemcpyAsync( + destination, source, bytes, cudaMemcpyDeviceToDevice, + reinterpret_cast(stream)); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +modalities::Status synchronize(std::uintptr_t stream) { + const cudaError_t rc = cudaStreamSynchronize( + reinterpret_cast(stream)); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +modalities::Status measure_scale( + const NativeThorKernelDriver& driver, const void* values, void* fp8_scratch, + float* dynamic_scale, float* destination_scale, std::size_t elements, + std::uintptr_t stream, float* host_value = nullptr) { + modalities::Status st = driver.quantize_fp8_dynamic( + values, fp8_scratch, dynamic_scale, elements, stream); + if (!st.ok_status()) return st; + st = synchronize(stream); + if (!st.ok_status()) return st; + if (host_value) { + const cudaError_t rc = cudaMemcpy( + host_value, dynamic_scale, sizeof(*host_value), + cudaMemcpyDeviceToHost); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + } + return copy_device(destination_scale, dynamic_scale, sizeof(float), stream); +} + +} // namespace + +modalities::Status NativeThorFp8Forward::vision( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const NativeThorWeightScales& weight_scales, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace || + workspace->flavor() != NativeWorkspaceFlavor::kThorFp8 || + weight_scales.vision.size() != 27 * 4) { + return invalid("Thor vision forward owner is invalid"); + } + const std::uint64_t sequence = workspace->vision_sequence(); + const std::uint64_t views = workspace->num_views(); + const NativeWorkspaceBuffer* images = buffer( + *workspace, "observation_images_normalized", + modalities::DType::kFloat16, {views, 224, 224, 3}); + const NativeWorkspaceBuffer* patches = buffer( + *workspace, "vision_patches", modalities::DType::kFloat16, + {sequence, 588}); + const NativeWorkspaceBuffer* position = buffer( + *workspace, "vision_pos_embed_expanded", + modalities::DType::kFloat16, {sequence, kVisionWidth}); + const NativeWorkspaceBuffer* x = buffer( + *workspace, "vision_x", modalities::DType::kFloat16, + {sequence, kVisionWidth}); + const NativeWorkspaceBuffer* x_fp8 = buffer( + *workspace, "vision_x_fp8", modalities::DType::kUInt8, + {sequence, kVisionWidth}); + const NativeWorkspaceBuffer* qkv = buffer( + *workspace, "vision_QKV", modalities::DType::kFloat16, + {sequence, 3 * kVisionWidth}); + const NativeWorkspaceBuffer* attention = buffer( + *workspace, "vision_attn", modalities::DType::kFloat16, + {sequence, kVisionWidth}); + const NativeWorkspaceBuffer* hidden = buffer( + *workspace, "vision_hidden", modalities::DType::kFloat16, + {sequence, kVisionHidden}); + const NativeWorkspaceBuffer* hidden_fp8 = buffer( + *workspace, "vision_hidden_fp8", modalities::DType::kUInt8, + {sequence, kVisionHidden}); + const NativeWorkspaceBuffer* unit_scale = buffer( + *workspace, "vision_unit_scale", modalities::DType::kFloat32, {1}); + const NativeWorkspaceBuffer* encoder_x = buffer( + *workspace, "encoder_x", modalities::DType::kFloat16, + {static_cast(workspace->encoder_sequence()), + kEncoderWidth}); + const NativeDeviceWeight* patch_w = weight( + weights, "vision_patch_embedding_w", NativeWeightDType::kFloat16, + {14, 14, 3, kVisionWidth}); + const NativeDeviceWeight* patch_b = weight( + weights, "vision_patch_embedding_b", NativeWeightDType::kFloat16, + {kVisionWidth}); + if (!images || !patches || !position || !x || !x_fp8 || !qkv || + !attention || !hidden || !hidden_fp8 || !unit_scale || !encoder_x || + !patch_w || !patch_b) { + return invalid("Thor vision buffers or weights are incomplete"); + } + + modalities::Status st = driver_->patch_im2col_fp16( + dptr(images), dptr(patches), static_cast(views), stream); + if (!st.ok_status()) return st; + st = driver_->fp16_nn(dptr(patches), dptr(patch_w), dptr(x), + static_cast(sequence), kVisionWidth, 588, + stream); + if (!st.ok_status()) return st; + st = driver_->patch_bias_position_fp16( + dptr(x), dptr(patch_b), dptr(position), static_cast(sequence), + kVisionWidth, 256, stream); + if (!st.ok_status()) return st; + + for (int layer = 0; layer < 27; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* ln_attn_w = weight( + weights, "vision_pre_attn_norm_w_" + suffix, + NativeWeightDType::kFloat16, {kVisionWidth}); + const NativeDeviceWeight* ln_attn_b = weight( + weights, "vision_pre_attn_norm_b_" + suffix, + NativeWeightDType::kFloat16, {kVisionWidth}); + const NativeDeviceWeight* qkv_w = weight( + weights, "vision_attn_qkv_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kVisionWidth, 3 * kVisionWidth}); + const NativeDeviceWeight* qkv_b = weight( + weights, "vision_attn_qkv_b_" + suffix, + NativeWeightDType::kFloat16, {3 * kVisionWidth}); + const NativeDeviceWeight* o_w = weight( + weights, "vision_attn_o_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kVisionWidth, kVisionWidth}); + const NativeDeviceWeight* o_b = weight( + weights, "vision_attn_o_b_" + suffix, + NativeWeightDType::kFloat16, {kVisionWidth}); + const NativeDeviceWeight* ln_ffn_w = weight( + weights, "vision_pre_ffn_norm_w_" + suffix, + NativeWeightDType::kFloat16, {kVisionWidth}); + const NativeDeviceWeight* ln_ffn_b = weight( + weights, "vision_pre_ffn_norm_b_" + suffix, + NativeWeightDType::kFloat16, {kVisionWidth}); + const NativeDeviceWeight* up_w = weight( + weights, "vision_ffn_up_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kVisionWidth, kVisionHidden}); + const NativeDeviceWeight* up_b = weight( + weights, "vision_ffn_up_b_" + suffix, + NativeWeightDType::kFloat16, {kVisionHidden}); + const NativeDeviceWeight* down_w = weight( + weights, "vision_ffn_down_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kVisionHidden, kVisionWidth}); + const NativeDeviceWeight* down_b = weight( + weights, "vision_ffn_down_b_" + suffix, + NativeWeightDType::kFloat16, {kVisionWidth}); + if (!ln_attn_w || !ln_attn_b || !qkv_w || !qkv_b || !o_w || !o_b || + !ln_ffn_w || !ln_ffn_b || !up_w || !up_b || !down_w || !down_b) { + return invalid("Thor vision layer weights are incomplete"); + } + const std::size_t scale = static_cast(layer) * 4; + st = driver_->layer_norm_fp8( + dptr(x), dptr(x_fp8), dptr(ln_attn_w), dptr(ln_attn_b), + static_cast(sequence), kVisionWidth, 1e-5f, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_nn_bias( + dptr(x_fp8), dptr(qkv_w), dptr(qkv), dptr(qkv_b), + static_cast(sequence), 3 * kVisionWidth, kVisionWidth, + weight_scales.vision[scale], stream); + if (!st.ok_status()) return st; + st = driver_->vision_fmha_fp16( + dptr(qkv), offset_bytes(dptr(qkv), kVisionWidth * 2), + offset_bytes(dptr(qkv), 2 * kVisionWidth * 2), dptr(attention), + static_cast(views), 256, 256, kVisionHeads, kVisionHeads, + kVisionHeadDimension, + 3 * kVisionWidth, 3 * kVisionWidth, stream); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(attention), dptr(x_fp8), + static_cast(dptr(unit_scale)), + sequence * kVisionWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_nn_bias_residual( + dptr(x_fp8), dptr(o_w), dptr(x), dptr(o_b), + static_cast(sequence), kVisionWidth, kVisionWidth, + weight_scales.vision[scale + 1], stream); + if (!st.ok_status()) return st; + st = driver_->layer_norm_fp8( + dptr(x), dptr(x_fp8), dptr(ln_ffn_w), dptr(ln_ffn_b), + static_cast(sequence), kVisionWidth, 1e-5f, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_nn_gelu_bias( + dptr(x_fp8), dptr(up_w), dptr(hidden), dptr(up_b), + static_cast(sequence), kVisionHidden, kVisionWidth, + weight_scales.vision[scale + 2], stream); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(hidden), dptr(hidden_fp8), + static_cast(dptr(unit_scale)), + sequence * kVisionHidden, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_nn_bias_residual( + dptr(hidden_fp8), dptr(down_w), dptr(x), dptr(down_b), + static_cast(sequence), kVisionWidth, kVisionHidden, + weight_scales.vision[scale + 3], stream); + if (!st.ok_status()) return st; + } + + const NativeDeviceWeight* final_w = weight( + weights, "vision_final_norm_w", NativeWeightDType::kFloat16, + {kVisionWidth}); + const NativeDeviceWeight* final_b = weight( + weights, "vision_final_norm_b", NativeWeightDType::kFloat16, + {kVisionWidth}); + const NativeDeviceWeight* projector_w = weight( + weights, "encoder_multi_modal_projector_w", + NativeWeightDType::kFloat16, {kVisionWidth, kEncoderWidth}); + const NativeDeviceWeight* projector_b = weight( + weights, "encoder_multi_modal_projector_b", + NativeWeightDType::kFloat16, {kEncoderWidth}); + if (!final_w || !final_b || !projector_w || !projector_b) { + return invalid("Thor vision projection weights are incomplete"); + } + st = driver_->layer_norm_fp16( + dptr(x), dptr(final_w), dptr(final_b), dptr(attention), + static_cast(sequence), kVisionWidth, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->fp16_nn( + dptr(attention), dptr(projector_w), dptr(encoder_x), + static_cast(sequence), kEncoderWidth, kVisionWidth, stream); + if (!st.ok_status()) return st; + return driver_->add_bias_fp16( + dptr(encoder_x), dptr(projector_b), static_cast(sequence), + kEncoderWidth, stream); +} + +modalities::Status NativeThorFp8Forward::encoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const std::vector& activation_weight_alphas, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace || + workspace->flavor() != NativeWorkspaceFlavor::kThorFp8 || + activation_weight_alphas.size() != kLayers * 4) { + return invalid("Thor encoder forward owner is invalid"); + } + const std::uint64_t sequence = workspace->encoder_sequence(); + const std::uint64_t keys = workspace->total_keys(); + const NativeWorkspaceBuffer* x = buffer( + *workspace, "encoder_x", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* x_fp8 = buffer( + *workspace, "encoder_x_fp8", modalities::DType::kUInt8, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* qkv = buffer( + *workspace, "encoder_QKV", modalities::DType::kFloat16, + {sequence, 2560}); + const NativeWorkspaceBuffer* logits = buffer( + *workspace, "encoder_logits", modalities::DType::kFloat16, + {sequence * kHeads, keys}); + const NativeWorkspaceBuffer* attention = buffer( + *workspace, "encoder_attn", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* o_fp8 = buffer( + *workspace, "encoder_o_fp8", modalities::DType::kUInt8, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* gate = buffer( + *workspace, "encoder_gate_merged", modalities::DType::kFloat16, + {sequence, 2 * kEncoderHidden}); + const NativeWorkspaceBuffer* hidden_fp8 = buffer( + *workspace, "encoder_hidden_fp8", modalities::DType::kUInt8, + {sequence, kEncoderHidden}); + const NativeWorkspaceBuffer* fg = buffer( + *workspace, "encoder_fg", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* rope = buffer( + *workspace, "encoder_rope_weights", modalities::DType::kFloat16, + {sequence, kHeadDimension}); + const NativeWorkspaceBuffer* activation_scales = buffer( + *workspace, "encoder_activation_scales", + modalities::DType::kFloat32, {kLayers, 4}); + const NativeWorkspaceBuffer* key_cache = buffer( + *workspace, "encoder_k_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* value_cache = buffer( + *workspace, "encoder_v_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* valid_keys = buffer( + *workspace, "attn_enc_seqused", modalities::DType::kUInt8, + {sizeof(std::int32_t)}); + if (!x || !x_fp8 || !qkv || !logits || !attention || !o_fp8 || !gate || + !hidden_fp8 || !fg || !rope || !activation_scales || !key_cache || + !value_cache || !valid_keys) { + return invalid("Thor encoder workspace is incomplete"); + } + + const float attention_scale = + 1.0f / std::sqrt(static_cast(kHeadDimension)); + const std::size_t cache_layer_elements = keys * kHeadDimension; + for (int layer = 0; layer < kLayers; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* qkv_w = weight( + weights, "encoder_attn_qkv_w_" + suffix, + NativeWeightDType::kFp8E4M3, {2560, kEncoderWidth}); + if (!qkv_w) return invalid("Thor encoder QKV weight is invalid"); + const std::size_t scale = static_cast(layer) * 4; + modalities::Status st = driver_->rms_norm_fp8_noweight( + dptr(x), dptr(x_fp8), static_cast(sequence), kEncoderWidth, + scale_ptr(activation_scales, scale), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(x_fp8), dptr(qkv_w), dptr(qkv), static_cast(sequence), + 2560, kEncoderWidth, activation_weight_alphas[scale], 0.0f, + NativeThorFp8Tactic::kSquare, stream); + if (!st.ok_status()) return st; + st = driver_->qkv_rope_cache_fp16( + dptr(qkv), dptr(rope), dptr(attention), dptr(key_cache), + dptr(value_cache), static_cast(sequence), kEncoderWidth, + kHeadDimension, kHeadDimension, 2560, + static_cast(scale / 4 * cache_layer_elements), + kHeadDimension, stream); + if (!st.ok_status() || layer == kLayers - 1) return st; + + void* layer_key = offset_bytes( + dptr(key_cache), static_cast(layer) * + cache_layer_elements * sizeof(std::uint16_t)); + void* layer_value = offset_bytes( + dptr(value_cache), static_cast(layer) * + cache_layer_elements * sizeof(std::uint16_t)); + st = driver_->attention_seqused_fp16( + dptr(attention), layer_key, layer_value, dptr(logits), + dptr(attention), static_cast(sequence), + static_cast(sequence), kHeads, kHeadDimension, + static_cast(dptr(valid_keys)), attention_scale, stream); + if (!st.ok_status()) return st; + + const NativeDeviceWeight* o_w = weight( + weights, "encoder_attn_o_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kEncoderWidth, kEncoderWidth}); + const NativeDeviceWeight* gate_w = weight( + weights, "encoder_ffn_gate_up_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {2 * kEncoderHidden, kEncoderWidth}); + const NativeDeviceWeight* down_w = weight( + weights, "encoder_ffn_down_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kEncoderWidth, kEncoderHidden}); + if (!o_w || !gate_w || !down_w) { + return invalid("Thor encoder layer weights are incomplete"); + } + st = driver_->quantize_fp8_static( + dptr(attention), dptr(o_fp8), scale_ptr(activation_scales, scale + 1), + sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(o_fp8), dptr(o_w), dptr(fg), static_cast(sequence), + kEncoderWidth, kEncoderWidth, + activation_weight_alphas[scale + 1], 0.0f, + NativeThorFp8Tactic::kSquare, stream); + if (!st.ok_status()) return st; + st = driver_->residual_rms_norm_fp8_noweight( + dptr(x), dptr(fg), dptr(x_fp8), static_cast(sequence), + kEncoderWidth, scale_ptr(activation_scales, scale + 2), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(x_fp8), dptr(gate_w), dptr(gate), static_cast(sequence), + 2 * kEncoderHidden, kEncoderWidth, + activation_weight_alphas[scale + 2], 0.0f, + NativeThorFp8Tactic::kT1, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_fp8( + dptr(gate), dptr(hidden_fp8), static_cast(sequence), + kEncoderHidden, scale_ptr(activation_scales, scale + 3), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(hidden_fp8), dptr(down_w), dptr(fg), + static_cast(sequence), kEncoderWidth, kEncoderHidden, + activation_weight_alphas[scale + 3], 0.0f, + NativeThorFp8Tactic::kWide, stream); + if (!st.ok_status()) return st; + st = driver_->residual_add_fp16( + dptr(x), dptr(fg), sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +modalities::Status NativeThorFp8Forward::calibrate_encoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const NativeThorWeightScales& weight_scales, + std::vector* sample_scales, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace || + !sample_scales || + workspace->flavor() != NativeWorkspaceFlavor::kThorFp8 || + weight_scales.encoder.size() != kLayers * 4) { + return invalid("Thor encoder calibration owner is invalid"); + } + const std::uint64_t sequence = workspace->encoder_sequence(); + const std::uint64_t keys = workspace->total_keys(); + const NativeWorkspaceBuffer* x = buffer( + *workspace, "encoder_x", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* x_fp8 = buffer( + *workspace, "encoder_x_fp8", modalities::DType::kUInt8, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* qkv = buffer( + *workspace, "encoder_QKV", modalities::DType::kFloat16, + {sequence, 2560}); + const NativeWorkspaceBuffer* logits = buffer( + *workspace, "encoder_logits", modalities::DType::kFloat16, + {sequence * kHeads, keys}); + const NativeWorkspaceBuffer* attention = buffer( + *workspace, "encoder_attn", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* o_fp8 = buffer( + *workspace, "encoder_o_fp8", modalities::DType::kUInt8, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* gate = buffer( + *workspace, "encoder_gate_merged", modalities::DType::kFloat16, + {sequence, 2 * kEncoderHidden}); + const NativeWorkspaceBuffer* hidden = buffer( + *workspace, "encoder_hidden", modalities::DType::kFloat16, + {sequence, kEncoderHidden}); + const NativeWorkspaceBuffer* hidden_fp8 = buffer( + *workspace, "encoder_hidden_fp8", modalities::DType::kUInt8, + {sequence, kEncoderHidden}); + const NativeWorkspaceBuffer* fg = buffer( + *workspace, "encoder_fg", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* rope = buffer( + *workspace, "encoder_rope_weights", modalities::DType::kFloat16, + {sequence, kHeadDimension}); + const NativeWorkspaceBuffer* key_cache = buffer( + *workspace, "encoder_k_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* value_cache = buffer( + *workspace, "encoder_v_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* valid_keys = buffer( + *workspace, "attn_enc_seqused", modalities::DType::kUInt8, + {sizeof(std::int32_t)}); + const NativeWorkspaceBuffer* norm_scratch = buffer( + *workspace, "encoder_norm_scratch", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* x_scratch = buffer( + *workspace, "encoder_x_scratch", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* fp8_scratch = buffer( + *workspace, "encoder_fp8_scratch", modalities::DType::kUInt8, + {sequence, kEncoderHidden}); + const NativeWorkspaceBuffer* scales = buffer( + *workspace, "encoder_sample_scales", modalities::DType::kFloat32, + {kLayers, 4}); + const NativeWorkspaceBuffer* dynamic_scale = buffer( + *workspace, "calibration_scale", modalities::DType::kFloat32, {1}); + const NativeWorkspaceBuffer* ones = buffer( + *workspace, "encoder_rms_ones", modalities::DType::kFloat16, + {kEncoderWidth}); + if (!x || !x_fp8 || !qkv || !logits || !attention || !o_fp8 || !gate || + !hidden || !hidden_fp8 || !fg || !rope || !key_cache || !value_cache || + !valid_keys || !norm_scratch || !x_scratch || !fp8_scratch || + !scales || !dynamic_scale || !ones) { + return invalid("Thor encoder calibration workspace is incomplete"); + } + cudaError_t rc = cudaMemsetAsync( + dptr(scales), 0, kLayers * 4 * sizeof(float), + reinterpret_cast(stream)); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + + const float attention_scale = + 1.0f / std::sqrt(static_cast(kHeadDimension)); + const std::size_t cache_layer_elements = keys * kHeadDimension; + for (int layer = 0; layer < kLayers; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* qkv_w = weight( + weights, "encoder_attn_qkv_w_" + suffix, + NativeWeightDType::kFp8E4M3, {2560, kEncoderWidth}); + if (!qkv_w) return invalid("Thor encoder QKV weight is invalid"); + const std::size_t site = static_cast(layer) * 4; + modalities::Status st = driver_->rms_norm_fp16( + dptr(x), dptr(ones), dptr(norm_scratch), + static_cast(sequence), kEncoderWidth, 1e-6f, stream); + if (!st.ok_status()) return st; + float qkv_scale = 0.0f; + st = measure_scale( + *driver_, dptr(norm_scratch), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), scale_ptr(scales, site), + sequence * kEncoderWidth, stream, &qkv_scale); + if (!st.ok_status()) return st; + st = driver_->rms_norm_fp8_noweight( + dptr(x), dptr(x_fp8), static_cast(sequence), kEncoderWidth, + scale_ptr(scales, site), stream); + if (!st.ok_status()) return st; + const float qkv_alpha = qkv_scale * weight_scales.encoder[site]; + st = driver_->fp8_cutlass( + dptr(x_fp8), dptr(qkv_w), dptr(qkv), static_cast(sequence), + 2560, kEncoderWidth, qkv_alpha, 0.0f, + NativeThorFp8Tactic::kSquare, stream); + if (!st.ok_status()) return st; + st = driver_->qkv_rope_cache_fp16( + dptr(qkv), dptr(rope), dptr(attention), dptr(key_cache), + dptr(value_cache), static_cast(sequence), kEncoderWidth, + kHeadDimension, kHeadDimension, 2560, + static_cast(static_cast(layer) * + cache_layer_elements), + kHeadDimension, stream); + if (!st.ok_status()) return st; + if (layer == kLayers - 1) break; + + void* layer_key = offset_bytes( + dptr(key_cache), static_cast(layer) * + cache_layer_elements * 2); + void* layer_value = offset_bytes( + dptr(value_cache), static_cast(layer) * + cache_layer_elements * 2); + st = driver_->attention_seqused_fp16( + dptr(attention), layer_key, layer_value, dptr(logits), + dptr(attention), static_cast(sequence), + static_cast(sequence), kHeads, kHeadDimension, + static_cast(dptr(valid_keys)), attention_scale, stream); + if (!st.ok_status()) return st; + + const NativeDeviceWeight* o_w = weight( + weights, "encoder_attn_o_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kEncoderWidth, kEncoderWidth}); + const NativeDeviceWeight* gate_w = weight( + weights, "encoder_ffn_gate_up_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {2 * kEncoderHidden, kEncoderWidth}); + const NativeDeviceWeight* down_w = weight( + weights, "encoder_ffn_down_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kEncoderWidth, kEncoderHidden}); + if (!o_w || !gate_w || !down_w) { + return invalid("Thor encoder calibration weights are incomplete"); + } + float o_scale = 0.0f; + st = measure_scale( + *driver_, dptr(attention), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(scales, site + 1), sequence * kEncoderWidth, stream, + &o_scale); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(attention), dptr(o_fp8), scale_ptr(scales, site + 1), + sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(o_fp8), dptr(o_w), dptr(fg), static_cast(sequence), + kEncoderWidth, kEncoderWidth, + o_scale * weight_scales.encoder[site + 1], 0.0f, + NativeThorFp8Tactic::kSquare, stream); + if (!st.ok_status()) return st; + + st = copy_device(dptr(x_scratch), dptr(x), + sequence * kEncoderWidth * 2, stream); + if (!st.ok_status()) return st; + st = driver_->residual_add_fp16( + dptr(x_scratch), dptr(fg), sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->rms_norm_fp16( + dptr(x_scratch), dptr(ones), dptr(norm_scratch), + static_cast(sequence), kEncoderWidth, 1e-6f, stream); + if (!st.ok_status()) return st; + float gate_scale = 0.0f; + st = measure_scale( + *driver_, dptr(norm_scratch), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(scales, site + 2), sequence * kEncoderWidth, stream, + &gate_scale); + if (!st.ok_status()) return st; + st = driver_->residual_rms_norm_fp8_noweight( + dptr(x), dptr(fg), dptr(x_fp8), static_cast(sequence), + kEncoderWidth, scale_ptr(scales, site + 2), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(x_fp8), dptr(gate_w), dptr(gate), static_cast(sequence), + 2 * kEncoderHidden, kEncoderWidth, + gate_scale * weight_scales.encoder[site + 2], 0.0f, + NativeThorFp8Tactic::kT1, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_fp16( + dptr(gate), dptr(hidden), static_cast(sequence), + kEncoderHidden, stream); + if (!st.ok_status()) return st; + float down_scale = 0.0f; + st = measure_scale( + *driver_, dptr(hidden), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(scales, site + 3), sequence * kEncoderHidden, stream, + &down_scale); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_fp8( + dptr(gate), dptr(hidden_fp8), static_cast(sequence), + kEncoderHidden, scale_ptr(scales, site + 3), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_cutlass( + dptr(hidden_fp8), dptr(down_w), dptr(fg), + static_cast(sequence), kEncoderWidth, kEncoderHidden, + down_scale * weight_scales.encoder[site + 3], 0.0f, + NativeThorFp8Tactic::kWide, stream); + if (!st.ok_status()) return st; + st = driver_->residual_add_fp16( + dptr(x), dptr(fg), sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + } + + // The last encoder layer only writes Q/K/V. Canonical non-zero values keep + // the artifact valid without advertising measurements for skipped sites. + const float unused_scales[] = {1.0f, 1.0f, 1.0f}; + rc = cudaMemcpyAsync( + scale_ptr(scales, (kLayers - 1) * 4 + 1), unused_scales, + sizeof(unused_scales), cudaMemcpyHostToDevice, + reinterpret_cast(stream)); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + modalities::Status st = synchronize(stream); + if (!st.ok_status()) return st; + sample_scales->resize(kLayers * 4); + rc = cudaMemcpy(sample_scales->data(), dptr(scales), + sample_scales->size() * sizeof(float), + cudaMemcpyDeviceToHost); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +modalities::Status NativeThorFp8Forward::calibrate_decoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::vector* sample_scales, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace || + !sample_scales || + workspace->flavor() != NativeWorkspaceFlavor::kThorFp8) { + return invalid("Thor decoder calibration owner is invalid"); + } + const std::uint64_t sequence = workspace->chunk_size(); + const std::uint64_t steps = workspace->num_steps(); + const std::uint64_t keys = workspace->total_keys(); + const NativeWorkspaceBuffer* noise = buffer( + *workspace, "diffusion_noise", modalities::DType::kFloat16, + {sequence, 32}); + const NativeWorkspaceBuffer* x = buffer( + *workspace, "decoder_x", modalities::DType::kFloat16, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* xn = buffer( + *workspace, "x_normed_buf", modalities::DType::kFloat16, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* gate = buffer( + *workspace, "gate_buf", modalities::DType::kFloat16, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* qkv = buffer( + *workspace, "decoder_QKV", modalities::DType::kFloat16, + {sequence, 2560}); + const NativeWorkspaceBuffer* logits = buffer( + *workspace, "decoder_logits", modalities::DType::kFloat16, + {sequence * kHeads, keys}); + const NativeWorkspaceBuffer* attention = buffer( + *workspace, "decoder_attn", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* hidden = buffer( + *workspace, "decoder_hidden", modalities::DType::kFloat16, + {sequence, 2 * kDecoderHidden}); + const NativeWorkspaceBuffer* fg = buffer( + *workspace, "decoder_fg", modalities::DType::kFloat16, + {sequence, 2 * kDecoderHidden}); + const NativeWorkspaceBuffer* action_f32 = buffer( + *workspace, "decoder_action_f32", modalities::DType::kFloat32, + {sequence, 32}); + const NativeWorkspaceBuffer* xn_fp8 = buffer( + *workspace, "decoder_x_fp8", modalities::DType::kUInt8, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* hidden_fp8 = buffer( + *workspace, "decoder_hidden_fp8", modalities::DType::kUInt8, + {sequence, kDecoderHidden}); + const NativeWorkspaceBuffer* context_fp8 = buffer( + *workspace, "decoder_context_fp8", modalities::DType::kUInt8, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* fp8_scratch = buffer( + *workspace, "decoder_fp8_scratch", modalities::DType::kUInt8, + {sequence, kDecoderHidden}); + const NativeWorkspaceBuffer* sample_scale_buffer = buffer( + *workspace, "decoder_sample_scales", modalities::DType::kFloat32, + {steps, kLayers, 4}); + const NativeWorkspaceBuffer* dynamic_scale = buffer( + *workspace, "calibration_scale", modalities::DType::kFloat32, {1}); + const NativeWorkspaceBuffer* rope = buffer( + *workspace, "decoder_rope_weights", modalities::DType::kFloat16, + {sequence, kHeadDimension}); + const NativeWorkspaceBuffer* style_attn = buffer( + *workspace, "decoder_style_attn", modalities::DType::kFloat16, + {steps, kLayers, sequence, 3 * kDecoderWidth}); + const NativeWorkspaceBuffer* style_ffn = buffer( + *workspace, "decoder_style_ffn", modalities::DType::kFloat16, + {steps, kLayers, sequence, 3 * kDecoderWidth}); + const NativeWorkspaceBuffer* style_final = buffer( + *workspace, "decoder_style_final", modalities::DType::kFloat16, + {steps, sequence, 3 * kDecoderWidth}); + const NativeWorkspaceBuffer* key_cache = buffer( + *workspace, "encoder_k_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* value_cache = buffer( + *workspace, "encoder_v_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* valid_keys = buffer( + *workspace, "attn_dec_seqused", modalities::DType::kUInt8, + {sizeof(std::int32_t)}); + const NativeWorkspaceBuffer* device_position = buffer( + *workspace, "attn_dec_devpos", modalities::DType::kUInt8, + {sizeof(std::int32_t)}); + const NativeDeviceWeight* weight_scales = weight( + weights, "decoder_weight_scales", NativeWeightDType::kFloat32, + {kLayers * 4}); + const NativeDeviceWeight* input_w = weight( + weights, "decoder_action_in_proj_w", NativeWeightDType::kFloat16, + {32, kDecoderWidth}); + const NativeDeviceWeight* input_b = weight( + weights, "decoder_action_in_proj_b", NativeWeightDType::kFloat16, + {kDecoderWidth}); + const NativeDeviceWeight* output_w = weight( + weights, "decoder_action_out_proj_w", NativeWeightDType::kFloat16, + {kDecoderWidth, 32}); + const NativeDeviceWeight* output_b = weight( + weights, "decoder_action_out_proj_b", NativeWeightDType::kFloat16, + {32}); + if (!noise || !x || !xn || !gate || !qkv || !logits || !attention || + !hidden || !fg || !action_f32 || !xn_fp8 || !hidden_fp8 || + !context_fp8 || !fp8_scratch || !sample_scale_buffer || + !dynamic_scale || !rope || !style_attn || !style_ffn || + !style_final || !key_cache || !value_cache || + !device_position || !weight_scales || !input_w || !input_b || + !output_w || !output_b) { + return invalid("Thor decoder calibration workspace is incomplete"); + } + + const std::size_t scale_count = steps * kLayers * 4; + cudaError_t rc = cudaMemsetAsync( + dptr(sample_scale_buffer), 0, scale_count * sizeof(float), + reinterpret_cast(stream)); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + + const float attention_scale = + 1.0f / std::sqrt(static_cast(kHeadDimension)); + const float dt = -1.0f / static_cast(steps); + const std::size_t cache_layer_elements = keys * kHeadDimension; + const std::size_t style_row_elements = sequence * 3 * kDecoderWidth; + for (int step = 0; step < static_cast(steps); ++step) { + modalities::Status st; + st = driver_->gmm_fp16( + dptr(noise), dptr(input_w), dptr(x), static_cast(sequence), + kDecoderWidth, 32, 0.0f, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + dptr(x), dptr(input_b), static_cast(sequence), + kDecoderWidth, stream); + if (!st.ok_status()) return st; + + for (int layer = 0; layer < kLayers; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* qkv_w = weight( + weights, "decoder_attn_qkv_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kDecoderWidth, 2560}); + const NativeDeviceWeight* o_w = weight( + weights, "decoder_attn_o_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kEncoderWidth, kDecoderWidth}); + const NativeDeviceWeight* gate_w = weight( + weights, "decoder_ffn_gate_up_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kDecoderWidth, 2 * kDecoderHidden}); + const NativeDeviceWeight* down_w = weight( + weights, "decoder_ffn_down_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kDecoderHidden, kDecoderWidth}); + if (!qkv_w || !o_w || !gate_w || !down_w) { + return invalid("Thor decoder calibration weights are incomplete"); + } + const std::size_t site = + (static_cast(step) * kLayers + layer) * 4; + const std::size_t style_site = + (static_cast(step) * kLayers + layer) * + style_row_elements; + const void* attn_style = + offset_bytes(dptr(style_attn), style_site * 2); + const void* ffn_style = + offset_bytes(dptr(style_ffn), style_site * 2); + + if (layer == 0) { + st = driver_->adarms_fp16( + dptr(x), attn_style, dptr(xn), dptr(gate), + static_cast(sequence), kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = measure_scale( + *driver_, dptr(xn), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(sample_scale_buffer, site), + sequence * kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fused_adarms_fp8( + dptr(x), attn_style, dptr(xn_fp8), dptr(gate), + static_cast(sequence), kDecoderWidth, + scale_ptr(sample_scale_buffer, site), stream); + if (!st.ok_status()) return st; + } + + st = driver_->fp8_descale( + dptr(xn_fp8), dptr(qkv_w), dptr(qkv), + static_cast(sequence), 2560, kDecoderWidth, + scale_ptr(sample_scale_buffer, site), + scale_ptr(weight_scales, static_cast(layer) * 4), + stream); + if (!st.ok_status()) return st; + st = driver_->qkv_rope_cache_devpos_fp16( + dptr(qkv), dptr(rope), dptr(attention), dptr(key_cache), + dptr(value_cache), static_cast(dptr(device_position)), + static_cast(sequence), kEncoderWidth, kHeadDimension, + kHeadDimension, 2560, + static_cast(static_cast(layer) * + cache_layer_elements), + kHeadDimension, stream); + if (!st.ok_status()) return st; + void* layer_key = offset_bytes( + dptr(key_cache), static_cast(layer) * + cache_layer_elements * 2); + void* layer_value = offset_bytes( + dptr(value_cache), static_cast(layer) * + cache_layer_elements * 2); + st = driver_->attention_seqused_fp16( + dptr(attention), layer_key, layer_value, dptr(logits), + dptr(attention), static_cast(sequence), + static_cast(keys), kHeads, kHeadDimension, + static_cast(dptr(valid_keys)), attention_scale, + stream); + if (!st.ok_status()) return st; + + st = measure_scale( + *driver_, dptr(attention), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(sample_scale_buffer, site + 1), + sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(attention), dptr(context_fp8), + scale_ptr(sample_scale_buffer, site + 1), + sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_descale( + dptr(context_fp8), dptr(o_w), dptr(fg), + static_cast(sequence), kDecoderWidth, kEncoderWidth, + scale_ptr(sample_scale_buffer, site + 1), + scale_ptr(weight_scales, + static_cast(layer) * 4 + 1), + stream); + if (!st.ok_status()) return st; + + st = driver_->gate_res_fp16( + dptr(fg), dptr(gate), dptr(x), + sequence * kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->adarms_fp16( + dptr(x), ffn_style, dptr(xn), dptr(gate), + static_cast(sequence), kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = measure_scale( + *driver_, dptr(xn), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(sample_scale_buffer, site + 2), + sequence * kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(xn), dptr(xn_fp8), + scale_ptr(sample_scale_buffer, site + 2), + sequence * kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_descale( + dptr(xn_fp8), dptr(gate_w), dptr(fg), + static_cast(sequence), 2 * kDecoderHidden, + kDecoderWidth, scale_ptr(sample_scale_buffer, site + 2), + scale_ptr(weight_scales, + static_cast(layer) * 4 + 2), + stream); + if (!st.ok_status()) return st; + + st = driver_->gate_gelu_fp16( + dptr(fg), dptr(hidden), static_cast(sequence), + kDecoderHidden, stream); + if (!st.ok_status()) return st; + st = measure_scale( + *driver_, dptr(hidden), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(sample_scale_buffer, site + 3), + sequence * kDecoderHidden, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_fp8( + dptr(fg), dptr(hidden_fp8), static_cast(sequence), + kDecoderHidden, + scale_ptr(sample_scale_buffer, site + 3), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_descale( + dptr(hidden_fp8), dptr(down_w), dptr(fg), + static_cast(sequence), kDecoderWidth, kDecoderHidden, + scale_ptr(sample_scale_buffer, site + 3), + scale_ptr(weight_scales, + static_cast(layer) * 4 + 3), + stream); + if (!st.ok_status()) return st; + + if (layer + 1 < kLayers) { + const std::size_t next_style_site = + style_site + style_row_elements; + const void* next_attn_style = offset_bytes( + dptr(style_attn), next_style_site * 2); + st = driver_->gate_res_fp16( + dptr(fg), dptr(gate), dptr(x), + sequence * kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->adarms_fp16( + dptr(x), next_attn_style, dptr(xn), dptr(gate), + static_cast(sequence), kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = measure_scale( + *driver_, dptr(xn), dptr(fp8_scratch), + static_cast(dptr(dynamic_scale)), + scale_ptr(sample_scale_buffer, site + 4), + sequence * kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(xn), dptr(xn_fp8), + scale_ptr(sample_scale_buffer, site + 4), + sequence * kDecoderWidth, stream); + } else { + st = driver_->gate_res_fp16( + dptr(fg), dptr(gate), dptr(x), + sequence * kDecoderWidth, stream); + } + if (!st.ok_status()) return st; + } + + const void* final_style = offset_bytes( + dptr(style_final), static_cast(step) * + style_row_elements * 2); + st = driver_->adarms_fp16( + dptr(x), final_style, dptr(xn), dptr(gate), + static_cast(sequence), kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->gmm_fp16_out_fp32( + dptr(xn), dptr(output_w), static_cast(dptr(action_f32)), + static_cast(sequence), 32, kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->action_update_fp16( + static_cast(dptr(action_f32)), dptr(output_b), + dptr(noise), static_cast(sequence), 32, dt, stream); + if (!st.ok_status()) return st; + } + + modalities::Status st = synchronize(stream); + if (!st.ok_status()) return st; + sample_scales->resize(scale_count); + rc = cudaMemcpy(sample_scales->data(), dptr(sample_scale_buffer), + scale_count * sizeof(float), cudaMemcpyDeviceToHost); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +modalities::Status NativeThorFp8Forward::diffusion( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace || + workspace->flavor() != NativeWorkspaceFlavor::kThorFp8) { + return invalid("Thor decoder forward owner is invalid"); + } + const std::uint64_t sequence = workspace->chunk_size(); + const std::uint64_t steps = workspace->num_steps(); + const std::uint64_t keys = workspace->total_keys(); + const NativeWorkspaceBuffer* noise = buffer( + *workspace, "diffusion_noise", modalities::DType::kFloat16, + {sequence, 32}); + const NativeWorkspaceBuffer* x = buffer( + *workspace, "decoder_x", modalities::DType::kFloat16, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* xn = buffer( + *workspace, "x_normed_buf", modalities::DType::kFloat16, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* gate = buffer( + *workspace, "gate_buf", modalities::DType::kFloat16, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* qkv = buffer( + *workspace, "decoder_QKV", modalities::DType::kFloat16, + {sequence, 2560}); + const NativeWorkspaceBuffer* logits = buffer( + *workspace, "decoder_logits", modalities::DType::kFloat16, + {sequence * kHeads, keys}); + const NativeWorkspaceBuffer* attention = buffer( + *workspace, "decoder_attn", modalities::DType::kFloat16, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* hidden = buffer( + *workspace, "decoder_hidden", modalities::DType::kFloat16, + {sequence, 2 * kDecoderHidden}); + const NativeWorkspaceBuffer* fg = buffer( + *workspace, "decoder_fg", modalities::DType::kFloat16, + {sequence, 2 * kDecoderHidden}); + const NativeWorkspaceBuffer* action_f32 = buffer( + *workspace, "decoder_action_f32", modalities::DType::kFloat32, + {sequence, 32}); + const NativeWorkspaceBuffer* xn_fp8 = buffer( + *workspace, "decoder_x_fp8", modalities::DType::kUInt8, + {sequence, kDecoderWidth}); + const NativeWorkspaceBuffer* hidden_fp8 = buffer( + *workspace, "decoder_hidden_fp8", modalities::DType::kUInt8, + {sequence, kDecoderHidden}); + const NativeWorkspaceBuffer* context_fp8 = buffer( + *workspace, "decoder_context_fp8", modalities::DType::kUInt8, + {sequence, kEncoderWidth}); + const NativeWorkspaceBuffer* rope = buffer( + *workspace, "decoder_rope_weights", modalities::DType::kFloat16, + {sequence, kHeadDimension}); + const NativeWorkspaceBuffer* style_attn = buffer( + *workspace, "decoder_style_attn", modalities::DType::kFloat16, + {steps, kLayers, sequence, 3 * kDecoderWidth}); + const NativeWorkspaceBuffer* style_ffn = buffer( + *workspace, "decoder_style_ffn", modalities::DType::kFloat16, + {steps, kLayers, sequence, 3 * kDecoderWidth}); + const NativeWorkspaceBuffer* style_final = buffer( + *workspace, "decoder_style_final", modalities::DType::kFloat16, + {steps, sequence, 3 * kDecoderWidth}); + const NativeWorkspaceBuffer* activation_scales = buffer( + *workspace, "decoder_activation_scales", + modalities::DType::kFloat32, {steps, kLayers, 4}); + const NativeWorkspaceBuffer* key_cache = buffer( + *workspace, "encoder_k_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* value_cache = buffer( + *workspace, "encoder_v_cache", modalities::DType::kFloat16, + {kLayers, keys, kHeadDimension}); + const NativeWorkspaceBuffer* valid_keys = buffer( + *workspace, "attn_dec_seqused", modalities::DType::kUInt8, + {sizeof(std::int32_t)}); + const NativeWorkspaceBuffer* device_position = buffer( + *workspace, "attn_dec_devpos", modalities::DType::kUInt8, + {sizeof(std::int32_t)}); + const NativeDeviceWeight* weight_scales = weight( + weights, "decoder_weight_scales", NativeWeightDType::kFloat32, + {kLayers * 4}); + const NativeDeviceWeight* input_w = weight( + weights, "decoder_action_in_proj_w", NativeWeightDType::kFloat16, + {32, kDecoderWidth}); + const NativeDeviceWeight* input_b = weight( + weights, "decoder_action_in_proj_b", NativeWeightDType::kFloat16, + {kDecoderWidth}); + const NativeDeviceWeight* output_w = weight( + weights, "decoder_action_out_proj_w", NativeWeightDType::kFloat16, + {kDecoderWidth, 32}); + const NativeDeviceWeight* output_b = weight( + weights, "decoder_action_out_proj_b", NativeWeightDType::kFloat16, + {32}); + if (!noise || !x || !xn || !gate || !qkv || !logits || !attention || + !hidden || !fg || !action_f32 || !xn_fp8 || !hidden_fp8 || + !context_fp8 || !rope || !style_attn || !style_ffn || !style_final || + !activation_scales || !key_cache || !value_cache || !valid_keys || + !device_position || !weight_scales || !input_w || !input_b || + !output_w || !output_b) { + return invalid("Thor decoder workspace or global weights are incomplete"); + } + + const float attention_scale = + 1.0f / std::sqrt(static_cast(kHeadDimension)); + const float dt = -1.0f / static_cast(steps); + const std::size_t cache_layer_elements = keys * kHeadDimension; + const std::size_t style_row_elements = sequence * 3 * kDecoderWidth; + for (int step = 0; step < static_cast(steps); ++step) { + modalities::Status st = driver_->gmm_fp16( + dptr(noise), dptr(input_w), dptr(x), static_cast(sequence), + kDecoderWidth, 32, 0.0f, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + dptr(x), dptr(input_b), static_cast(sequence), + kDecoderWidth, stream); + if (!st.ok_status()) return st; + + for (int layer = 0; layer < kLayers; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* qkv_w = weight( + weights, "decoder_attn_qkv_w_" + suffix, + NativeWeightDType::kFp8E4M3, {kDecoderWidth, 2560}); + const NativeDeviceWeight* o_w = weight( + weights, "decoder_attn_o_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kEncoderWidth, kDecoderWidth}); + const NativeDeviceWeight* gate_w = weight( + weights, "decoder_ffn_gate_up_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kDecoderWidth, 2 * kDecoderHidden}); + const NativeDeviceWeight* down_w = weight( + weights, "decoder_ffn_down_w_" + suffix, + NativeWeightDType::kFp8E4M3, + {kDecoderHidden, kDecoderWidth}); + if (!qkv_w || !o_w || !gate_w || !down_w) { + return invalid("Thor decoder layer weights are incomplete"); + } + const std::size_t site = + (static_cast(step) * kLayers + layer) * 4; + const std::size_t style_site = + (static_cast(step) * kLayers + layer) * + style_row_elements; + const void* attn_style = offset_bytes(dptr(style_attn), + style_site * 2); + const void* ffn_style = offset_bytes(dptr(style_ffn), + style_site * 2); + if (layer == 0) { + st = driver_->fused_adarms_fp8( + dptr(x), attn_style, dptr(xn_fp8), dptr(gate), + static_cast(sequence), kDecoderWidth, + scale_ptr(activation_scales, site), stream); + if (!st.ok_status()) return st; + } + st = driver_->fp8_descale( + dptr(xn_fp8), dptr(qkv_w), dptr(qkv), + static_cast(sequence), 2560, kDecoderWidth, + scale_ptr(activation_scales, site), + scale_ptr(weight_scales, static_cast(layer) * 4), + stream); + if (!st.ok_status()) return st; + st = driver_->qkv_rope_cache_devpos_fp16( + dptr(qkv), dptr(rope), dptr(attention), dptr(key_cache), + dptr(value_cache), static_cast(dptr(device_position)), + static_cast(sequence), kEncoderWidth, kHeadDimension, + kHeadDimension, 2560, + static_cast(static_cast(layer) * + cache_layer_elements), + kHeadDimension, stream); + if (!st.ok_status()) return st; + void* layer_key = offset_bytes( + dptr(key_cache), static_cast(layer) * + cache_layer_elements * 2); + void* layer_value = offset_bytes( + dptr(value_cache), static_cast(layer) * + cache_layer_elements * 2); + st = driver_->attention_seqused_fp16( + dptr(attention), layer_key, layer_value, dptr(logits), + dptr(attention), static_cast(sequence), + static_cast(keys), kHeads, kHeadDimension, + static_cast(dptr(valid_keys)), attention_scale, + stream); + if (!st.ok_status()) return st; + st = driver_->quantize_fp8_static( + dptr(attention), dptr(context_fp8), + scale_ptr(activation_scales, site + 1), + sequence * kEncoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->fp8_descale( + dptr(context_fp8), dptr(o_w), dptr(fg), + static_cast(sequence), kDecoderWidth, kEncoderWidth, + scale_ptr(activation_scales, site + 1), + scale_ptr(weight_scales, + static_cast(layer) * 4 + 1), + stream); + if (!st.ok_status()) return st; + st = driver_->gate_res_adarms_fp8( + dptr(fg), dptr(gate), dptr(x), ffn_style, dptr(xn_fp8), + dptr(gate), static_cast(sequence), kDecoderWidth, + scale_ptr(activation_scales, site + 2), stream); + if (!st.ok_status()) return st; + st = driver_->fp8_descale( + dptr(xn_fp8), dptr(gate_w), dptr(fg), + static_cast(sequence), 2 * kDecoderHidden, + kDecoderWidth, scale_ptr(activation_scales, site + 2), + scale_ptr(weight_scales, + static_cast(layer) * 4 + 2), + stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_fp8( + dptr(fg), dptr(hidden_fp8), static_cast(sequence), + kDecoderHidden, scale_ptr(activation_scales, site + 3), + stream); + if (!st.ok_status()) return st; + st = driver_->fp8_descale( + dptr(hidden_fp8), dptr(down_w), dptr(fg), + static_cast(sequence), kDecoderWidth, kDecoderHidden, + scale_ptr(activation_scales, site + 3), + scale_ptr(weight_scales, + static_cast(layer) * 4 + 3), + stream); + if (!st.ok_status()) return st; + if (layer + 1 < kLayers) { + const std::size_t next_style_site = style_site + style_row_elements; + const void* next_attn_style = offset_bytes( + dptr(style_attn), next_style_site * 2); + st = driver_->gate_res_adarms_fp8( + dptr(fg), dptr(gate), dptr(x), next_attn_style, + dptr(xn_fp8), dptr(gate), static_cast(sequence), + kDecoderWidth, scale_ptr(activation_scales, site + 4), + stream); + } else { + st = driver_->gate_res_fp16( + dptr(fg), dptr(gate), dptr(x), + sequence * kDecoderWidth, stream); + } + if (!st.ok_status()) return st; + } + + const void* final_style = offset_bytes( + dptr(style_final), static_cast(step) * + style_row_elements * 2); + st = driver_->adarms_fp16( + dptr(x), final_style, dptr(xn), dptr(gate), + static_cast(sequence), kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->gmm_fp16_out_fp32( + dptr(xn), dptr(output_w), static_cast(dptr(action_f32)), + static_cast(sequence), 32, kDecoderWidth, stream); + if (!st.ok_status()) return st; + st = driver_->action_update_fp16( + static_cast(dptr(action_f32)), dptr(output_b), + dptr(noise), static_cast(sequence), 32, dt, stream); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_graph_owner.cpp b/cpp/models/pi05/src/native_thor_graph_owner.cpp new file mode 100644 index 00000000..abbdcd49 --- /dev/null +++ b/cpp/models/pi05/src/native_thor_graph_owner.cpp @@ -0,0 +1,318 @@ +#include "flashrt/cpp/models/pi05/native_thor_graph_owner.h" + +#include "flashrt/cpp/models/pi05/native_thor_style_precompute.h" +#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" + +#include + +#include +#include +#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); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +modalities::Status copy_scales(const NativeWorkspace& workspace, + const char* name, + const std::vector& values) { + const NativeWorkspaceBuffer* destination = workspace.find(name); + if (!destination || destination->dtype != modalities::DType::kFloat32 || + frt_buffer_bytes(destination->buffer) != + values.size() * sizeof(float)) { + return invalid("Thor activation scale workspace is invalid"); + } + const cudaError_t rc = cudaMemcpy( + frt_buffer_dptr(destination->buffer), values.data(), + values.size() * sizeof(float), cudaMemcpyHostToDevice); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("Thor activation scale upload failed"); +} + +bool calibration_matches(const NativeCalibrationArtifact& artifact, + const NativeGraphConfig& config) { + return artifact.num_views == config.num_views && + artifact.max_prompt_tokens == config.max_prompt_tokens && + artifact.chunk_size == config.chunk_size && + artifact.num_steps == config.num_steps && + artifact.vision_pool_factor == config.vision_pool_factor; +} + +} // namespace + +NativeThorGraphOwner::NativeThorGraphOwner( + frt_ctx ctx, + const NativeGraphConfig& config) + : ctx_(ctx), + config_(config), + weights_(ctx), + workspace_(ctx), + forward_(&driver_), + capture_status_(modalities::Status::ok()) {} + +NativeThorGraphOwner::~NativeThorGraphOwner() { + 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 NativeThorGraphOwner::create( + const std::string& checkpoint_path, + const NativeGraphConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status) { + if (config.num_views < 1 || config.num_views > 3 || + config.max_prompt_tokens < 1 || (config.max_prompt_tokens & 1) || + config.chunk_size < 1 || config.num_steps < 1 || + static_cast(config.max_prompt_tokens) + + static_cast(config.chunk_size) + + static_cast(config.num_views) * 256 > + static_cast(std::numeric_limits::max()) || + config.vision_pool_factor != 1 || + !calibration_matches(calibration, config)) { + if (status) { + *status = invalid("Thor native graph configuration is invalid"); + } + return nullptr; + } + modalities::Status st = + validate_native_calibration_artifact(calibration); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) { + if (status) *status = backend("Thor graph context creation failed"); + return nullptr; + } + std::unique_ptr owner( + new (std::nothrow) NativeThorGraphOwner(ctx, config)); + if (!owner) { + frt_ctx_destroy(ctx); + if (status) *status = backend("Thor graph owner allocation failed"); + return nullptr; + } + st = owner->initialize(checkpoint_path, calibration); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + if (status) *status = modalities::Status::ok(); + return owner; +} + +modalities::Status NativeThorGraphOwner::initialize( + const std::string& checkpoint_path, + const NativeCalibrationArtifact& calibration) { + const bool profile_setup = std::getenv("FLASHRT_PROFILE_NATIVE_SETUP"); + const auto setup_begin = std::chrono::steady_clock::now(); + auto checkpoint = setup_begin; + const auto report = [&](const char* phase) { + const auto now = std::chrono::steady_clock::now(); + if (profile_setup) { + std::fprintf(stderr, "native_thor_setup %s_ms=%.3f\n", phase, + std::chrono::duration( + now - checkpoint).count()); + } + checkpoint = now; + }; + + loader::SafetensorsFile source; + if (!source.open(checkpoint_path + "/model.safetensors")) { + return modalities::Status::error(modalities::StatusCode::kNotFound, + source.error()); + } + report("header"); + NativeThorWeightMaterializer materializer(source, &weights_); + NativeThorMaterializationOptions options; + options.num_steps = config_.num_steps; + options.include_embedding = true; + modalities::Status st = + materializer.materialize_all(options, &weight_scales_); + if (!st.ok_status()) return st; + report("materialize"); + + 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; + workspace_config.flavor = NativeWorkspaceFlavor::kThorFp8; + st = workspace_.allocate(workspace_config); + if (!st.ok_status()) return st; + st = workspace_.expand_vision_position_embedding(weights_); + if (!st.ok_status()) return st; + st = copy_scales(workspace_, "encoder_activation_scales", + calibration.encoder_scales); + if (!st.ok_status()) return st; + st = copy_scales(workspace_, "decoder_activation_scales", + calibration.decoder_scales); + if (!st.ok_status()) return st; + encoder_alphas_.resize(calibration.encoder_scales.size()); + if (encoder_alphas_.size() != weight_scales_.encoder.size()) { + return invalid("Thor encoder scale count is invalid"); + } + for (std::size_t i = 0; i < encoder_alphas_.size(); ++i) { + encoder_alphas_[i] = + calibration.encoder_scales[i] * weight_scales_.encoder[i]; + if (!std::isfinite(encoder_alphas_[i]) || + !(encoder_alphas_[i] > 0.0f)) { + return invalid("Thor encoder alpha is invalid"); + } + } + st = set_prompt_length(0); + if (!st.ok_status()) return st; + NativeThorStylePrecomputer precomputer(&driver_); + st = precomputer.run(weights_, &workspace_, 0); + if (!st.ok_status()) return st; + report("workspace_style"); + + for (const char* name : {"observation_images_normalized", + "prompt_embedding", "diffusion_noise"}) { + const NativeWorkspaceBuffer* input = workspace_.find(name); + if (!input || + cudaMemset(frt_buffer_dptr(input->buffer), 0, + frt_buffer_bytes(input->buffer)) != cudaSuccess) { + return backend("Thor graph input initialization failed"); + } + } + if (cudaDeviceSynchronize() != cudaSuccess) { + return backend("Thor graph setup synchronization failed"); + } + + // CUTLASS, cuBLAS, and FMHA initialize tactics outside CUDA capture. + st = record(nullptr); + if (!st.ok_status()) return st; + if (cudaDeviceSynchronize() != cudaSuccess) { + return backend("Thor graph warmup synchronization failed"); + } + const NativeWorkspaceBuffer* noise = workspace_.find("diffusion_noise"); + if (!noise || + cudaMemset(frt_buffer_dptr(noise->buffer), 0, + frt_buffer_bytes(noise->buffer)) != cudaSuccess) { + return backend("Thor graph warmup reset failed"); + } + report("warmup"); + + 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"); + 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("Thor 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("Thor full graph capture failed") + : capture_status_; + } + if (!capture_status_.ok_status() || + frt_graph_variant_count(infer_graph_) != 1) { + return capture_status_.ok_status() + ? backend("Thor full graph variant is missing") + : capture_status_; + } + report("capture"); + + cudaStream_t stream = nullptr; + if (cudaStreamCreate(&stream) != cudaSuccess) { + return backend("Thor replay stream creation failed"); + } + replay_stream_ = stream; + stream_id_ = frt_ctx_wrap_stream(ctx_, replay_stream_); + if (stream_id_ < 0) return backend("Thor replay stream wrapping failed"); + report("stream"); + if (profile_setup) { + const auto now = std::chrono::steady_clock::now(); + std::fprintf(stderr, "native_thor_setup total_ms=%.3f\n", + std::chrono::duration( + now - setup_begin).count()); + } + return modalities::Status::ok(); +} + +modalities::Status NativeThorGraphOwner::record(void* stream) { + const NativeWorkspaceBuffer* prompt = workspace_.find("prompt_embedding"); + const NativeWorkspaceBuffer* encoder = workspace_.find("encoder_x"); + if (!prompt || !encoder) return invalid("Thor 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("Thor 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("Thor prompt graph copy failed"); + } + const std::uintptr_t stream_id = reinterpret_cast(stream); + modalities::Status st = + forward_.vision(weights_, &workspace_, weight_scales_, stream_id); + if (!st.ok_status()) return st; + st = forward_.encoder(weights_, &workspace_, encoder_alphas_, stream_id); + if (!st.ok_status()) return st; + return forward_.diffusion(weights_, &workspace_, stream_id); +} + +void NativeThorGraphOwner::record_graph(void* user, void* stream) { + auto* owner = static_cast(user); + owner->capture_status_ = owner->record(stream); +} + +modalities::Status NativeThorGraphOwner::set_prompt_length(int prompt_tokens) { + return workspace_.set_fixed_prompt_length(prompt_tokens); +} + +int NativeThorGraphOwner::replay() const { + if (!infer_graph_ || stream_id_ < 0) return FRT_ERR_INVALID; + return frt_graph_replay(infer_graph_, 0, stream_id_); +} + +modalities::Status NativeThorGraphOwner::synchronize() const { + if (!replay_stream_) return invalid("Thor replay stream is missing"); + const cudaError_t rc = + cudaStreamSynchronize(static_cast(replay_stream_)); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("Thor replay synchronization failed"); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_kernel_driver.cu b/cpp/models/pi05/src/native_thor_kernel_driver.cu new file mode 100644 index 00000000..0ab2b298 --- /dev/null +++ b/cpp/models/pi05/src/native_thor_kernel_driver.cu @@ -0,0 +1,622 @@ +#include "flashrt/cpp/models/pi05/native_thor_kernel_driver.h" + +#include "activation.cuh" +#include "attention_cublas.cuh" +#include "decoder_fused.cuh" +#include "elementwise.cuh" +#include "gemm_runner.h" +#include "norm.cuh" +#include "patch_embed.cuh" +#include "quantize.cuh" +#include "rope.cuh" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +extern "C" int cutlass_fp8_sq(void*, void*, void*, int, int, int, float, + float, cudaStream_t); +extern "C" int cutlass_fp8_t1(void*, void*, void*, int, int, int, float, + float, cudaStream_t); +extern "C" int cutlass_fp8_wide(void*, void*, void*, int, int, int, float, + float, cudaStream_t); +extern "C" int cutlass_fp8_plain(void*, void*, void*, int, int, int, float, + float, cudaStream_t); +extern "C" int fmha_fp16_strided( + const void*, const void*, const void*, void*, int, int, int, int, int, + int, int, int, cudaStream_t); +extern "C" cudaError_t frt_pi05_precise_silu_fp16( + __half*, std::size_t, cudaStream_t); +void silu_inplace_fp16(__half*, int, cudaStream_t); + +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)); +} + +bool valid_gemm(void* a, void* b, void* output, int m, int n, int k) { + return a && b && output && m > 0 && n > 0 && k > 0; +} + +} // namespace + +struct NativeThorKernelDriver::Impl { + Impl() { + const cublasStatus_t rc = cublasCreate(&cublas); + if (rc != CUBLAS_STATUS_SUCCESS) { + throw std::runtime_error("Thor cuBLAS handle creation failed"); + } + } + ~Impl() { + if (cublas) cublasDestroy(cublas); + } + + GemmRunner gemm; + cublasHandle_t cublas = nullptr; +}; + +NativeThorKernelDriver::NativeThorKernelDriver() noexcept { + try { + impl_.reset(new Impl()); + } catch (const std::exception& e) { + error_ = e.what(); + } catch (...) { + error_ = "Thor kernel driver initialization failed"; + } +} + +NativeThorKernelDriver::~NativeThorKernelDriver() = default; + +modalities::Status NativeThorKernelDriver::status() const { + return impl_ ? modalities::Status::ok() : backend(error_); +} + +modalities::Status NativeThorKernelDriver::fp16_nn( + void* a, void* b, void* output, int m, int n, int k, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!valid_gemm(a, b, output, m, n, k)) { + return invalid("Thor FP16 GEMM arguments are invalid"); + } + try { + impl_->gemm.fp16_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("Thor FP16 GEMM launch failed"); + } +} + +modalities::Status NativeThorKernelDriver::fp8_nn_bias( + void* a, void* b, void* output, void* bias, int m, int n, int k, + float alpha, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!valid_gemm(a, b, output, m, n, k) || !bias || + !std::isfinite(alpha)) { + return invalid("Thor FP8 bias GEMM arguments are invalid"); + } + try { + impl_->gemm.fp8_nn_bias( + a, b, output, bias, m, n, k, alpha, + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("Thor FP8 bias GEMM launch failed"); + } +} + +modalities::Status NativeThorKernelDriver::fp8_nn_bias_residual( + void* a, void* b, void* output, void* bias, int m, int n, int k, + float alpha, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!valid_gemm(a, b, output, m, n, k) || !bias || + !std::isfinite(alpha)) { + return invalid("Thor FP8 residual GEMM arguments are invalid"); + } + try { + impl_->gemm.fp8_nn_bias_res( + a, b, output, bias, m, n, k, alpha, + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("Thor FP8 residual GEMM launch failed"); + } +} + +modalities::Status NativeThorKernelDriver::fp8_nn_gelu_bias( + void* a, void* b, void* output, void* bias, int m, int n, int k, + float alpha, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!valid_gemm(a, b, output, m, n, k) || !bias || + !std::isfinite(alpha)) { + return invalid("Thor FP8 GELU GEMM arguments are invalid"); + } + try { + impl_->gemm.fp8_nn_gelu_bias( + a, b, output, bias, m, n, k, alpha, + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("Thor FP8 GELU GEMM launch failed"); + } +} + +modalities::Status NativeThorKernelDriver::fp8_cutlass( + void* a, void* b, void* output, int m, int n, int k, + float alpha, float beta, NativeThorFp8Tactic tactic, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!valid_gemm(a, b, output, m, n, k) || !std::isfinite(alpha) || + !std::isfinite(beta)) { + return invalid("Thor CUTLASS FP8 GEMM arguments are invalid"); + } + using Fn = int (*)(void*, void*, void*, int, int, int, float, float, + cudaStream_t); + Fn fn = nullptr; + switch (tactic) { + case NativeThorFp8Tactic::kSquare: fn = cutlass_fp8_sq; break; + case NativeThorFp8Tactic::kT1: fn = cutlass_fp8_t1; break; + case NativeThorFp8Tactic::kWide: fn = cutlass_fp8_wide; break; + case NativeThorFp8Tactic::kPlain: fn = cutlass_fp8_plain; break; + } + if (!fn) return invalid("Thor CUTLASS FP8 tactic is invalid"); + const int rc = fn(a, b, output, m, n, k, alpha, beta, + reinterpret_cast(stream)); + return rc == 0 ? launch_status() + : backend("Thor CUTLASS FP8 GEMM rejected the shape"); +} + +modalities::Status NativeThorKernelDriver::fp8_descale( + void* a, void* b, void* output, int m, int n, int k, + const float* activation_scale, const float* weight_scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!valid_gemm(a, b, output, m, n, k) || !activation_scale || + !weight_scale) { + return invalid("Thor cuBLASLt FP8 GEMM arguments are invalid"); + } + try { + impl_->gemm.fp8_descale_fp16( + a, b, output, m, n, k, const_cast(activation_scale), + const_cast(weight_scale), + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("Thor cuBLASLt FP8 GEMM launch failed"); + } +} + +modalities::Status NativeThorKernelDriver::add_bias_fp16( + void* values, const void* bias, int rows, int columns, + std::uintptr_t stream) const { + if (!values || !bias || rows <= 0 || columns <= 0) { + return invalid("Thor FP16 bias arguments are invalid"); + } + ::add_bias_fp16(static_cast<__half*>(values), + static_cast(bias), rows, columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::layer_norm_fp16( + const void* values, const void* weight, const void* bias, void* output, + int rows, int columns, float epsilon, std::uintptr_t stream) const { + if (!values || !weight || !bias || !output || rows <= 0 || columns <= 0 || + !(epsilon > 0.0f) || !std::isfinite(epsilon)) { + return invalid("Thor FP16 LayerNorm arguments are invalid"); + } + ::layer_norm_fp16( + static_cast(values), static_cast(weight), + static_cast(bias), static_cast<__half*>(output), rows, + columns, epsilon, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::layer_norm_fp8( + const void* values, void* output, const void* weight, const void* bias, + int rows, int columns, float epsilon, std::uintptr_t stream) const { + if (!values || !output || !weight || !bias || rows <= 0 || columns <= 0 || + !(epsilon > 0.0f) || !std::isfinite(epsilon)) { + return invalid("Thor FP8 LayerNorm arguments are invalid"); + } + ::layer_norm_fp8( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), + static_cast(weight), static_cast(bias), + rows, columns, epsilon, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::rms_norm_fp16( + const void* values, const void* weight, void* output, int rows, + int columns, float epsilon, std::uintptr_t stream) const { + if (!values || !weight || !output || rows <= 0 || columns <= 0 || + !(epsilon > 0.0f) || !std::isfinite(epsilon)) { + return invalid("Thor FP16 RMSNorm arguments are invalid"); + } + ::rms_norm_fp16( + static_cast(values), static_cast(weight), + static_cast<__half*>(output), rows, columns, epsilon, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::rms_norm_fp8_noweight( + const void* values, void* output, int rows, int columns, + const float* scale, std::uintptr_t stream) const { + if (!values || !output || !scale || rows <= 0 || columns <= 0) { + return invalid("Thor FP8 RMSNorm arguments are invalid"); + } + ::rms_norm_fp8_noweight_fp16( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), rows, columns, scale, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::residual_rms_norm_fp8_noweight( + void* residual, const void* values, void* output, int rows, int columns, + const float* scale, std::uintptr_t stream) const { + if (!residual || !values || !output || !scale || rows <= 0 || + columns <= 0) { + return invalid("Thor residual FP8 RMSNorm arguments are invalid"); + } + ::residual_add_rms_norm_fp8_noweight_fp16( + static_cast<__half*>(residual), static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), rows, columns, scale, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::quantize_fp8_static( + const void* values, void* output, const float* scale, std::size_t elements, + std::uintptr_t stream) const { + if (!values || !output || !scale || !elements || elements > INT_MAX) { + return invalid("Thor static FP8 quantization arguments are invalid"); + } + ::quantize_fp8_static_fp16( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), scale, + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::quantize_fp8_dynamic( + const void* values, void* output, float* scale, std::size_t elements, + std::uintptr_t stream) const { + if (!values || !output || !scale || !elements || elements > INT_MAX) { + return invalid("Thor dynamic FP8 quantization arguments are invalid"); + } + ::quantize_fp8_device_fp16( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), scale, + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::gelu_fp16( + void* values, std::size_t elements, std::uintptr_t stream) const { + if (!values || !elements || elements > INT_MAX) { + return invalid("Thor FP16 GELU arguments are invalid"); + } + ::gelu_inplace_fp16(static_cast<__half*>(values), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::silu_fp16( + void* values, std::size_t elements, std::uintptr_t stream) const { + if (!values || !elements) { + return invalid("Thor FP16 SiLU arguments are invalid"); + } + if (elements > INT_MAX || elements % 2) { + return invalid("Thor FP16 SiLU element count is invalid"); + } + ::silu_inplace_fp16(static_cast<__half*>(values), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::precise_silu_fp16( + void* values, std::size_t elements, std::uintptr_t stream) const { + if (!values || !elements || elements > INT_MAX) { + return invalid("Thor precise FP16 SiLU arguments are invalid"); + } + const cudaError_t rc = frt_pi05_precise_silu_fp16( + static_cast<__half*>(values), elements, + reinterpret_cast(stream)); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +modalities::Status NativeThorKernelDriver::gate_gelu_fp16( + const void* merged, void* output, int rows, int hidden, + std::uintptr_t stream) const { + if (!merged || !output || rows <= 0 || hidden <= 0) { + return invalid("Thor FP16 gated GELU arguments are invalid"); + } + ::gate_silu_mul_merged_fp16( + static_cast(merged), static_cast<__half*>(output), + rows, hidden, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::gate_gelu_fp8( + const void* merged, void* output, int rows, int hidden, + const float* scale, std::uintptr_t stream) const { + if (!merged || !output || !scale || rows <= 0 || hidden <= 0) { + return invalid("Thor FP8 gated GELU arguments are invalid"); + } + ::gate_silu_mul_merged_fp8_fp16( + static_cast(merged), + static_cast<__nv_fp8_e4m3*>(output), rows, hidden, scale, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::residual_add_fp16( + void* residual, const void* values, std::size_t elements, + std::uintptr_t stream) const { + if (!residual || !values || !elements || elements > INT_MAX) { + return invalid("Thor FP16 residual arguments are invalid"); + } + ::residual_add_fp16( + static_cast<__half*>(residual), static_cast(values), + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::fused_adarms_fp8( + const void* values, const void* style, void* output, void* gate, + int rows, int columns, const float* scale, + std::uintptr_t stream) const { + if (!values || !style || !output || !gate || !scale || rows <= 0 || + columns <= 0) { + return invalid("Thor fused AdaRMSNorm arguments are invalid"); + } + ::fused_adarms_fp8_static_fp16( + static_cast(values), static_cast(style), + static_cast<__nv_fp8_e4m3*>(output), static_cast<__half*>(gate), rows, + columns, scale, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::gate_res_adarms_fp8( + const void* gemm_output, const void* previous_gate, void* residual, + const void* style, void* output, void* gate, int rows, int columns, + const float* scale, std::uintptr_t stream) const { + if (!gemm_output || !previous_gate || !residual || !style || !output || + !gate || !scale || rows <= 0 || columns <= 0) { + return invalid("Thor gated AdaRMSNorm arguments are invalid"); + } + ::gate_res_adarms_fp8_static_fp16( + static_cast(gemm_output), + static_cast(previous_gate), + static_cast<__half*>(residual), static_cast(style), + static_cast<__nv_fp8_e4m3*>(output), static_cast<__half*>(gate), rows, + columns, scale, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::gate_res_fp16( + const void* gemm_output, const void* gate, void* residual, + std::size_t elements, std::uintptr_t stream) const { + if (!gemm_output || !gate || !residual || !elements || elements > INT_MAX) { + return invalid("Thor gated residual arguments are invalid"); + } + ::gate_res_fp16( + static_cast(gemm_output), + static_cast(gate), static_cast<__half*>(residual), + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::adarms_fp16( + const void* values, const void* style, void* output, void* gate, + int rows, int columns, std::uintptr_t stream) const { + if (!values || !style || !output || !gate || rows <= 0 || columns <= 0) { + return invalid("Thor FP16 AdaRMSNorm arguments are invalid"); + } + ::adarms_fp16( + static_cast(values), static_cast(style), + static_cast<__half*>(output), static_cast<__half*>(gate), rows, + columns, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::qkv_rope_cache_fp16( + const void* qkv, const void* rope, void* query, void* key_cache, + void* value_cache, int rows, int query_columns, int key_columns, + int head_dimension, int qkv_stride, int cache_offset, int cache_stride, + std::uintptr_t stream) const { + if (!qkv || !rope || !query || !key_cache || !value_cache || rows <= 0 || + query_columns <= 0 || key_columns <= 0 || head_dimension <= 0 || + qkv_stride <= 0 || cache_offset < 0 || cache_stride <= 0) { + return invalid("Thor FP16 QKV cache arguments are invalid"); + } + ::qkv_split_rope_kvcache_fp16( + static_cast(qkv), static_cast(rope), + static_cast<__half*>(query), static_cast<__half*>(key_cache), + static_cast<__half*>(value_cache), rows, query_columns, key_columns, + head_dimension, qkv_stride, cache_offset, cache_stride, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::qkv_rope_cache_devpos_fp16( + const void* qkv, const void* rope, void* query, void* key_cache, + void* value_cache, const int* device_position, int rows, + int query_columns, int key_columns, int head_dimension, int qkv_stride, + int cache_offset, int cache_stride, std::uintptr_t stream) const { + if (!qkv || !rope || !query || !key_cache || !value_cache || + !device_position || rows <= 0 || query_columns <= 0 || + key_columns <= 0 || head_dimension <= 0 || qkv_stride <= 0 || + cache_offset < 0 || cache_stride <= 0) { + return invalid("Thor FP16 device-position QKV arguments are invalid"); + } + ::qkv_split_rope_kvcache_fp16_devpos( + static_cast(qkv), static_cast(rope), + static_cast<__half*>(query), static_cast<__half*>(key_cache), + static_cast<__half*>(value_cache), device_position, rows, + query_columns, key_columns, head_dimension, qkv_stride, cache_offset, + cache_stride, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::attention_seqused_fp16( + const void* query, const void* key, const void* value, void* logits, + void* output, int query_rows, int max_key_rows, int heads, + int head_dimension, const int* valid_key_rows, float scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!query || !key || !value || !logits || !output || !valid_key_rows || + query_rows <= 0 || max_key_rows <= 0 || heads <= 0 || + head_dimension <= 0 || !(scale > 0.0f) || !std::isfinite(scale)) { + return invalid("Thor FP16 attention arguments are invalid"); + } + ::attention_qkv_fp16_seqused( + impl_->cublas, static_cast(query), + static_cast(key), static_cast(value), + static_cast<__half*>(logits), static_cast<__half*>(output), query_rows, + max_key_rows, heads, head_dimension, valid_key_rows, scale, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::vision_fmha_fp16( + const void* query, const void* key, const void* value, void* output, + int batch, int query_rows, int key_rows, int query_heads, int key_heads, + int head_dimension, int query_stride, int key_stride, + std::uintptr_t stream) const { + if (!query || !key || !value || !output || batch <= 0 || query_rows <= 0 || + key_rows <= 0 || query_heads <= 0 || key_heads <= 0 || + head_dimension <= 0 || query_stride <= 0 || key_stride <= 0) { + return invalid("Thor vision FMHA arguments are invalid"); + } + const int rc = ::fmha_fp16_strided( + query, key, value, output, batch, query_rows, key_rows, query_heads, + key_heads, head_dimension, query_stride, key_stride, + reinterpret_cast(stream)); + return rc == 0 ? launch_status() + : backend("Thor vision FMHA rejected the shape"); +} + +modalities::Status NativeThorKernelDriver::patch_im2col_fp16( + const void* images, void* patches, int num_views, + std::uintptr_t stream) const { + if (!images || !patches || num_views <= 0 || num_views > 3) { + return invalid("Thor patch im2col arguments are invalid"); + } + ::patch_im2col(static_cast(images), + static_cast<__half*>(patches), num_views, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::patch_bias_position_fp16( + void* output, const void* bias, const void* position, int sequence, + int columns, int positions_per_view, std::uintptr_t stream) const { + if (!output || !bias || !position || sequence <= 0 || columns <= 0 || + positions_per_view <= 0) { + return invalid("Thor patch position arguments are invalid"); + } + ::patch_embed_bias_pos( + static_cast<__half*>(output), static_cast(bias), + static_cast(position), sequence, columns, + positions_per_view, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::bias_residual_fp16( + void* residual, const void* values, const void* bias, int rows, + int columns, std::uintptr_t stream) const { + if (!residual || !values || !bias || rows <= 0 || columns <= 0) { + return invalid("Thor FP16 bias residual arguments are invalid"); + } + ::bias_residual_strict_fp16( + static_cast<__half*>(residual), static_cast(values), + static_cast(bias), rows, columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::gmm_fp16( + const void* a, const void* b, void* output, int m, int n, int k, + float beta, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!a || !b || !output || m <= 0 || n <= 0 || k <= 0 || + !std::isfinite(beta)) { + return invalid("Thor decoder FP16 GEMM arguments are invalid"); + } + ::gmm_fp16( + impl_->cublas, static_cast(a), + static_cast(b), static_cast<__half*>(output), + m, n, k, beta, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::gmm_fp16_out_fp32( + const void* a, const void* b, float* 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("Thor FP32-output GEMM arguments are invalid"); + } + ::gmm_fp16_out_fp32( + impl_->cublas, static_cast(a), + static_cast(b), output, m, n, k, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeThorKernelDriver::action_update_fp16( + const float* delta, const void* bias, void* noise, int rows, int columns, + float dt, std::uintptr_t stream) const { + if (!delta || !bias || !noise || rows <= 0 || columns <= 0 || + !std::isfinite(dt)) { + return invalid("Thor action update arguments are invalid"); + } + ::action_update_from_fp32( + delta, static_cast(bias), static_cast<__half*>(noise), + rows, columns, dt, true, reinterpret_cast(stream)); + return launch_status(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_setup_ops.cu b/cpp/models/pi05/src/native_thor_setup_ops.cu new file mode 100644 index 00000000..e1c6452a --- /dev/null +++ b/cpp/models/pi05/src/native_thor_setup_ops.cu @@ -0,0 +1,28 @@ +#include +#include + +#include + +namespace { + +__global__ void precise_silu_fp16_kernel(__half* values, + std::size_t elements) { + const std::size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= elements) return; + const float value = __half2float(values[index]); + values[index] = __float2half_rn( + value / (1.0f + expf(-value))); +} + +} // namespace + +extern "C" cudaError_t frt_pi05_precise_silu_fp16( + __half* values, std::size_t elements, cudaStream_t stream) { + if (!values || !elements) return cudaErrorInvalidValue; + constexpr unsigned int kBlock = 256; + precise_silu_fp16_kernel<<< + static_cast((elements + kBlock - 1) / kBlock), + kBlock, 0, stream>>>(values, elements); + return cudaGetLastError(); +} diff --git a/cpp/models/pi05/src/native_thor_style_precompute.cpp b/cpp/models/pi05/src/native_thor_style_precompute.cpp new file mode 100644 index 00000000..25f996bd --- /dev/null +++ b/cpp/models/pi05/src/native_thor_style_precompute.cpp @@ -0,0 +1,210 @@ +#include "flashrt/cpp/models/pi05/native_thor_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::kFloat16 || + 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 NativeThorStylePrecomputer::run( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace || + workspace->flavor() != NativeWorkspaceFlavor::kThorFp8) { + return invalid("Thor 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->dtype != modalities::DType::kFloat16 || + time_output->shape.size() != 3 || style_attn->shape.size() != 4 || + style_ffn->shape != style_attn->shape || + style_final->shape.size() != 3) { + return invalid("Thor 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("Thor 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("Thor style global weights are incomplete"); + } + + const cudaStream_t cuda_stream = reinterpret_cast(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_->fp16_nn( + time_row, frt_buffer_dptr(time_in_w->buffer), + frt_buffer_dptr(scratch_a->buffer), 1, 1024, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + frt_buffer_dptr(scratch_a->buffer), + frt_buffer_dptr(time_in_b->buffer), 1, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->precise_silu_fp16( + frt_buffer_dptr(scratch_a->buffer), 1024, stream); + if (!st.ok_status()) return st; + st = driver_->fp16_nn( + frt_buffer_dptr(scratch_a->buffer), + frt_buffer_dptr(time_out_w->buffer), + frt_buffer_dptr(scratch_b->buffer), 1, 1024, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + frt_buffer_dptr(scratch_b->buffer), + frt_buffer_dptr(time_out_b->buffer), 1, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->precise_silu_fp16( + frt_buffer_dptr(scratch_b->buffer), 1024, 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("Thor 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("Thor 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_->fp16_nn( + expanded, frt_buffer_dptr(attn_w->buffer), attn_target, + chunk, 3072, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + attn_target, frt_buffer_dptr(attn_b->buffer), chunk, 3072, + stream); + if (!st.ok_status()) return st; + st = driver_->fp16_nn( + expanded, frt_buffer_dptr(ffn_w->buffer), ffn_target, + chunk, 3072, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + ffn_target, frt_buffer_dptr(ffn_b->buffer), chunk, 3072, + stream); + if (!st.ok_status()) return st; + } + void* final_target = offset( + frt_buffer_dptr(style_final->buffer), + static_cast(step) * chunk * 3072); + st = driver_->fp16_nn( + expanded, frt_buffer_dptr(final_w->buffer), final_target, + chunk, 3072, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_fp16( + final_target, frt_buffer_dptr(final_b->buffer), chunk, 3072, + stream); + if (!st.ok_status()) return st; + } + const cudaError_t rc = cudaStreamSynchronize(cuda_stream); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("Thor style precompute synchronization failed"); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_weight_materializer.cpp b/cpp/models/pi05/src/native_thor_weight_materializer.cpp new file mode 100644 index 00000000..163059fa --- /dev/null +++ b/cpp/models/pi05/src/native_thor_weight_materializer.cpp @@ -0,0 +1,545 @@ +#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" + +#include +#include +#include +#include +#include +#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); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +int materialization_workers(int layers) { + int workers = std::min(layers, 8); + const char* setting = std::getenv("FLASHRT_NATIVE_WEIGHT_WORKERS"); + if (!setting || !setting[0]) return workers; + errno = 0; + char* end = nullptr; + const long parsed = std::strtol(setting, &end, 10); + if (errno || !end || *end || parsed < 1 || parsed > 64) return workers; + return std::min(layers, static_cast(parsed)); +} + +template +modalities::Status materialize_layers_parallel( + int layers, + Materialize materialize, + std::vector* scales) { + if (layers <= 0 || !scales) { + return invalid("Thor parallel materialization input is invalid"); + } + std::vector statuses( + static_cast(layers), modalities::Status::ok()); + std::vector> layer_scales( + static_cast(layers)); + std::atomic next{0}; + std::atomic stop{false}; + const int workers = materialization_workers(layers); + std::vector threads; + threads.reserve(static_cast(workers)); + try { + for (int worker = 0; worker < workers; ++worker) { + threads.emplace_back([&] { + while (!stop.load(std::memory_order_relaxed)) { + const int layer = next.fetch_add(1); + if (layer >= layers) break; + try { + statuses[static_cast(layer)] = materialize( + layer, + &layer_scales[static_cast(layer)]); + } catch (const std::exception& error) { + statuses[static_cast(layer)] = + backend(std::string("Thor weight worker failed: ") + + error.what()); + } catch (...) { + statuses[static_cast(layer)] = + backend("Thor weight worker failed"); + } + if (!statuses[static_cast(layer)].ok_status()) { + stop.store(true, std::memory_order_relaxed); + } + } + }); + } + } catch (const std::exception& error) { + stop.store(true, std::memory_order_relaxed); + for (auto& thread : threads) thread.join(); + return backend(std::string("Thor worker creation failed: ") + + error.what()); + } + for (auto& thread : threads) thread.join(); + for (int layer = 0; layer < layers; ++layer) { + const auto& status = statuses[static_cast(layer)]; + if (!status.ok_status()) return status; + const auto& values = layer_scales[static_cast(layer)]; + scales->insert(scales->end(), values.begin(), values.end()); + } + return modalities::Status::ok(); +} + +std::string encoder_prefix(int layer) { + return "paligemma_with_expert.paligemma.model.language_model.layers." + + std::to_string(layer); +} + +std::string decoder_prefix(int layer) { + return "paligemma_with_expert.gemma_expert.model.layers." + + std::to_string(layer); +} + +const std::string& vision_prefix() { + static const std::string prefix = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; + return prefix; +} + +std::string layer_name(const char* stem, int layer) { + return std::string(stem) + std::to_string(layer); +} + +} // namespace + +modalities::Status NativeThorWeightMaterializer::upload_f16( + const std::string& source_key, + const std::string& destination_name, + bool transpose) { + if (!destination_) return invalid("Thor weight destination is null"); + NativeSourceTensorView source; + NativeF16Tensor converted; + modalities::Status st = + load_native_source_tensor(source_, source_key, &source); + if (!st.ok_status()) return st; + st = native_source_to_f16(source, transpose, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); +} + +modalities::Status NativeThorWeightMaterializer::upload_f16( + const std::string& destination_name, + const NativeF16Tensor& tensor) { + if (!destination_) return invalid("Thor weight destination is null"); + return destination_->upload(destination_name, tensor); +} + +modalities::Status NativeThorWeightMaterializer::upload_fp8( + const std::string& destination_name, + const NativeF16Tensor& tensor, + std::vector* scales) { + if (!destination_ || !scales) { + return invalid("Thor FP8 weight destination is invalid"); + } + NativeFp8Tensor quantized; + modalities::Status st = + native_quantize_fp8_e4m3(tensor, false, &quantized); + if (!st.ok_status()) return st; + st = destination_->upload_bytes( + destination_name, quantized.shape, NativeWeightDType::kFp8E4M3, + quantized.values.data(), quantized.values.size()); + if (!st.ok_status()) return st; + scales->push_back(quantized.scale); + return modalities::Status::ok(); +} + +modalities::Status NativeThorWeightMaterializer::upload_scale_vector( + const std::string& name, + const std::vector& values) { + if (!destination_ || values.empty()) { + return invalid("Thor weight scale vector is invalid"); + } + return destination_->upload_bytes( + name, {static_cast(values.size())}, + NativeWeightDType::kFloat32, values.data(), + values.size() * sizeof(float)); +} + +modalities::Status NativeThorWeightMaterializer::materialize_vision_globals() { + if (!destination_) return invalid("Thor weight destination is null"); + const std::string prefix = vision_prefix(); + NativeSourceTensorView patch; + NativeF16Tensor permuted; + modalities::Status st = load_native_source_tensor( + source_, prefix + ".embeddings.patch_embedding.weight", &patch); + if (!st.ok_status()) return st; + st = native_source_patch_oihw_to_hwio_f16(patch, &permuted); + if (!st.ok_status()) return st; + st = upload_f16("vision_patch_embedding_w", permuted); + if (!st.ok_status()) return st; + + const struct { + const char* source; + const char* destination; + bool transpose; + } entries[] = { + {"embeddings.patch_embedding.bias", "vision_patch_embedding_b", false}, + {"embeddings.position_embedding.weight", "vision_position_embedding", false}, + {"post_layernorm.weight", "vision_final_norm_w", false}, + {"post_layernorm.bias", "vision_final_norm_b", false}, + }; + for (const auto& entry : entries) { + st = upload_f16(prefix + "." + entry.source, entry.destination, + entry.transpose); + if (!st.ok_status()) return st; + } + const std::string projector = + "paligemma_with_expert.paligemma.model.multi_modal_projector.linear"; + st = upload_f16(projector + ".weight", + "encoder_multi_modal_projector_w", true); + if (!st.ok_status()) return st; + return upload_f16(projector + ".bias", + "encoder_multi_modal_projector_b", false); +} + +modalities::Status NativeThorWeightMaterializer::materialize_vision_layer( + int layer, + std::vector* scales) { + if (layer < 0 || layer >= 27 || !destination_ || !scales) { + return invalid("Thor vision layer index is invalid"); + } + const std::string prefix = vision_prefix() + ".encoder.layers." + + std::to_string(layer); + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + modalities::Status st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.weight", &q); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.weight", &k); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.weight", &v); + if (!st.ok_status()) return st; + NativeF16Tensor joined; + st = native_source_qkv_to_f16(q, k, v, 0, 0, nullptr, true, &joined); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("vision_attn_qkv_w_", layer), joined, scales); + if (!st.ok_status()) return st; + + st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.bias", &q); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.bias", &k); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.bias", &v); + if (!st.ok_status()) return st; + st = native_source_concat_vectors_to_f16({&q, &k, &v}, &joined); + if (!st.ok_status()) return st; + st = upload_f16(layer_name("vision_attn_qkv_b_", layer), joined); + if (!st.ok_status()) return st; + + const struct { + const char* source; + const char* destination; + } quantized[] = { + {"self_attn.out_proj.weight", "vision_attn_o_w_"}, + {"mlp.fc1.weight", "vision_ffn_up_w_"}, + {"mlp.fc2.weight", "vision_ffn_down_w_"}, + }; + for (const auto& entry : quantized) { + NativeSourceTensorView source; + NativeF16Tensor converted; + st = load_native_source_tensor(source_, prefix + "." + entry.source, + &source); + if (!st.ok_status()) return st; + st = native_source_to_f16(source, true, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name(entry.destination, layer), converted, scales); + if (!st.ok_status()) return st; + } + + const struct { + const char* source; + const char* destination; + } fp16[] = { + {"self_attn.out_proj.bias", "vision_attn_o_b_"}, + {"mlp.fc1.bias", "vision_ffn_up_b_"}, + {"mlp.fc2.bias", "vision_ffn_down_b_"}, + {"layer_norm1.weight", "vision_pre_attn_norm_w_"}, + {"layer_norm1.bias", "vision_pre_attn_norm_b_"}, + {"layer_norm2.weight", "vision_pre_ffn_norm_w_"}, + {"layer_norm2.bias", "vision_pre_ffn_norm_b_"}, + }; + for (const auto& entry : fp16) { + st = upload_f16(prefix + "." + entry.source, + layer_name(entry.destination, layer), false); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +modalities::Status NativeThorWeightMaterializer::materialize_encoder_layer( + int layer, + std::vector* scales) { + if (layer < 0 || layer >= 18 || !destination_ || !scales) { + return invalid("Thor encoder layer index is invalid"); + } + const std::string prefix = encoder_prefix(layer); + NativeFloatTensor norm; + modalities::Status st = load_native_float_tensor( + source_, prefix + ".input_layernorm.weight", &norm); + if (!st.ok_status()) return st; + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.weight", &q); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.weight", &k); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.weight", &v); + if (!st.ok_status()) return st; + NativeF16Tensor converted; + st = native_source_qkv_to_f16( + q, k, v, 8, 1, &norm, false, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("encoder_attn_qkv_w_", layer), + converted, scales); + if (!st.ok_status()) return st; + + NativeSourceTensorView source; + st = load_native_source_tensor( + source_, prefix + ".self_attn.o_proj.weight", &source); + if (!st.ok_status()) return st; + st = native_source_to_f16(source, false, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("encoder_attn_o_w_", layer), converted, scales); + if (!st.ok_status()) return st; + + st = load_native_float_tensor( + source_, prefix + ".post_attention_layernorm.weight", &norm); + if (!st.ok_status()) return st; + NativeSourceTensorView gate; + NativeSourceTensorView up; + st = load_native_source_tensor( + source_, prefix + ".mlp.gate_proj.weight", &gate); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".mlp.up_proj.weight", &up); + if (!st.ok_status()) return st; + st = native_source_pair_to_f16(gate, up, &norm, false, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("encoder_ffn_gate_up_w_", layer), + converted, scales); + if (!st.ok_status()) return st; + + st = load_native_source_tensor( + source_, prefix + ".mlp.down_proj.weight", &source); + if (!st.ok_status()) return st; + st = native_source_to_f16(source, false, &converted); + if (!st.ok_status()) return st; + return upload_fp8(layer_name("encoder_ffn_down_w_", layer), + converted, scales); +} + +modalities::Status NativeThorWeightMaterializer::materialize_decoder_layer( + int layer, + std::vector* scales) { + if (layer < 0 || layer >= 18 || !destination_ || !scales) { + return invalid("Thor decoder layer index is invalid"); + } + const std::string prefix = decoder_prefix(layer); + NativeSourceTensorView q; + NativeSourceTensorView k; + NativeSourceTensorView v; + modalities::Status st = load_native_source_tensor( + source_, prefix + ".self_attn.q_proj.weight", &q); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.k_proj.weight", &k); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".self_attn.v_proj.weight", &v); + if (!st.ok_status()) return st; + NativeF16Tensor converted; + st = native_source_qkv_to_f16( + q, k, v, 8, 1, nullptr, true, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("decoder_attn_qkv_w_", layer), + converted, scales); + if (!st.ok_status()) return st; + + NativeSourceTensorView source; + st = load_native_source_tensor( + source_, prefix + ".self_attn.o_proj.weight", &source); + if (!st.ok_status()) return st; + st = native_source_to_f16(source, true, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("decoder_attn_o_w_", layer), converted, scales); + if (!st.ok_status()) return st; + + NativeSourceTensorView gate; + NativeSourceTensorView up; + st = load_native_source_tensor( + source_, prefix + ".mlp.gate_proj.weight", &gate); + if (!st.ok_status()) return st; + st = load_native_source_tensor( + source_, prefix + ".mlp.up_proj.weight", &up); + if (!st.ok_status()) return st; + st = native_source_pair_to_f16(gate, up, nullptr, true, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("decoder_ffn_gate_up_w_", layer), + converted, scales); + if (!st.ok_status()) return st; + + st = load_native_source_tensor( + source_, prefix + ".mlp.down_proj.weight", &source); + if (!st.ok_status()) return st; + st = native_source_to_f16(source, true, &converted); + if (!st.ok_status()) return st; + st = upload_fp8(layer_name("decoder_ffn_down_w_", layer), + converted, scales); + if (!st.ok_status()) return st; + + const struct { + const char* source; + const char* destination; + bool transpose; + } fp16[] = { + {"input_layernorm.dense.weight", "decoder_pre_attn_norm_mod_w_", true}, + {"input_layernorm.dense.bias", "decoder_pre_attn_norm_mod_b_", false}, + {"post_attention_layernorm.dense.weight", "decoder_pre_ffn_norm_mod_w_", true}, + {"post_attention_layernorm.dense.bias", "decoder_pre_ffn_norm_mod_b_", false}, + }; + for (const auto& entry : fp16) { + st = upload_f16(prefix + "." + entry.source, + layer_name(entry.destination, layer), entry.transpose); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +modalities::Status NativeThorWeightMaterializer::materialize_decoder_globals( + int num_steps) { + if (!destination_ || num_steps <= 0) { + return invalid("Thor 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}, + {"action_out_proj.weight", "decoder_action_out_proj_w", true}, + {"action_out_proj.bias", "decoder_action_out_proj_b", false}, + }; + for (const auto& entry : entries) { + modalities::Status st = upload_f16( + entry.source, entry.destination, entry.transpose); + if (!st.ok_status()) return st; + } + NativeF16Tensor time_embeddings; + modalities::Status st = native_pi05_time_embeddings_f16( + num_steps, 1024, &time_embeddings); + if (!st.ok_status()) return st; + return upload_f16("decoder_time_embeds", time_embeddings); +} + +modalities::Status NativeThorWeightMaterializer::materialize_embedding() { + return upload_f16( + "paligemma_with_expert.paligemma.lm_head.weight", + "embedding_weight", false); +} + +modalities::Status NativeThorWeightMaterializer::materialize_all( + const NativeThorMaterializationOptions& options, + NativeThorWeightScales* scales) { + if (!destination_ || !scales || options.num_steps <= 0) { + return invalid("Thor materialization options are invalid"); + } + scales->vision.clear(); + scales->encoder.clear(); + scales->decoder.clear(); + scales->vision.reserve(27 * 4); + scales->encoder.reserve(18 * 4); + scales->decoder.reserve(18 * 4); + + const bool profile = std::getenv("FLASHRT_PROFILE_NATIVE_SETUP"); + auto checkpoint = std::chrono::steady_clock::now(); + const auto report = [&](const char* phase) { + const auto now = std::chrono::steady_clock::now(); + if (profile) { + std::fprintf(stderr, "native_thor_weights %s_ms=%.3f\n", phase, + std::chrono::duration( + now - checkpoint).count()); + } + checkpoint = now; + }; + + modalities::Status st = materialize_vision_globals(); + if (!st.ok_status()) return st; + st = materialize_layers_parallel( + 27, + [this](int layer, std::vector* layer_scales) { + return materialize_vision_layer(layer, layer_scales); + }, + &scales->vision); + if (!st.ok_status()) return st; + report("vision"); + st = materialize_layers_parallel( + 18, + [this](int layer, std::vector* layer_scales) { + return materialize_encoder_layer(layer, layer_scales); + }, + &scales->encoder); + if (!st.ok_status()) return st; + report("encoder"); + st = materialize_layers_parallel( + 18, + [this](int layer, std::vector* layer_scales) { + return materialize_decoder_layer(layer, layer_scales); + }, + &scales->decoder); + if (!st.ok_status()) return st; + report("decoder"); + st = materialize_decoder_globals(options.num_steps); + if (!st.ok_status()) return st; + if (options.include_embedding) { + st = materialize_embedding(); + if (!st.ok_status()) return st; + } + report("globals"); + + if (scales->vision.size() != 108 || scales->encoder.size() != 72 || + scales->decoder.size() != 72) { + return invalid("Thor materialized weight scale count is invalid"); + } + st = upload_scale_vector("vision_weight_scales", scales->vision); + if (!st.ok_status()) return st; + st = upload_scale_vector("encoder_weight_scales", scales->encoder); + if (!st.ok_status()) return st; + return upload_scale_vector("decoder_weight_scales", scales->decoder); +} + +} // 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 81a0004b..67501029 100644 --- a/cpp/models/pi05/src/native_workspace.cpp +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -1,4 +1,5 @@ #include "flashrt/cpp/models/pi05/native_workspace.h" +#include "flashrt/cpp/models/pi05/native_rope.h" #ifdef FLASHRT_CPP_WITH_CUDA_STAGING #include @@ -98,8 +99,16 @@ modalities::Status NativeWorkspace::initialize_rms_ones() { 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)); + if (target->shape.size() != 1 || + (target->dtype != modalities::DType::kBFloat16 && + target->dtype != modalities::DType::kFloat16)) { + return invalid("native RMS buffer layout is invalid"); + } + const std::uint16_t one = + target->dtype == modalities::DType::kFloat16 + ? modalities::float_to_float16(1.0f) + : modalities::float_to_bfloat16(1.0f); + std::vector ones(target->shape[0], one); const cudaError_t rc = cudaMemcpy( frt_buffer_dptr(target->buffer), ones.data(), ones.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice); @@ -117,6 +126,13 @@ modalities::Status NativeWorkspace::initialize_rope() { #else const int max_positions = encoder_sequence_ + chunk_size_; rope_table_.resize(static_cast(max_positions) * 256); + const NativeWorkspaceBuffer* encoder = find("encoder_rope_weights"); + if (!encoder) return invalid("encoder RoPE buffer was not allocated"); + if (flavor_ == NativeWorkspaceFlavor::kThorFp8) { + modalities::Status st = generate_native_thor_rope_f16( + frt_buffer_dptr(encoder->buffer), 0, encoder_sequence_, 0); + return st.ok_status() ? update_decoder_rope(0) : st; + } for (int position = 0; position < max_positions; ++position) { const std::size_t row = static_cast(position) * 256; for (int i = 0; i < 128; ++i) { @@ -125,16 +141,12 @@ modalities::Status NativeWorkspace::initialize_rope() { 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))); + 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); @@ -146,6 +158,50 @@ modalities::Status NativeWorkspace::initialize_rope() { #endif } +modalities::Status NativeWorkspace::set_fixed_prompt_length( + int prompt_tokens) { + if (flavor_ != NativeWorkspaceFlavor::kThorFp8) { + return update_decoder_rope(prompt_tokens); + } + if (prompt_tokens < 0 || prompt_tokens > max_prompt_tokens_ || + !prompt_embedding_buffer_) { + return invalid("Thor fixed prompt length is invalid"); + } + const int rounded_prompt = prompt_tokens + (prompt_tokens & 1); + if (rounded_prompt > max_prompt_tokens_) { + return invalid("Thor fixed prompt length exceeds its even capacity"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Thor fixed prompt update requires the CUDA build"); +#else + if (rounded_prompt != prompt_tokens && prompt_tokens > 0) { + const std::size_t row_bytes = 2048 * sizeof(std::uint16_t); + auto* base = static_cast( + frt_buffer_dptr(prompt_embedding_buffer_)); + const cudaError_t copy_rc = cudaMemcpy( + base + static_cast(prompt_tokens) * row_bytes, + base + static_cast(prompt_tokens - 1) * row_bytes, + row_bytes, cudaMemcpyDeviceToDevice); + if (copy_rc != cudaSuccess) { + return backend("Thor prompt padding copy failed"); + } + } + const std::int32_t valid = encoder_vision_sequence_ + rounded_prompt; + const std::int32_t values[] = {valid, valid + chunk_size_, valid}; + for (int i = 0; i < 3; ++i) { + if (!prompt_length_buffers_[i] || + cudaMemcpy(frt_buffer_dptr(prompt_length_buffers_[i]), &values[i], + sizeof(values[i]), cudaMemcpyHostToDevice) != + cudaSuccess) { + return backend("Thor fixed prompt control upload failed"); + } + } + return update_decoder_rope(rounded_prompt); +#endif +} + modalities::Status NativeWorkspace::update_decoder_rope(int prompt_tokens) { if (prompt_tokens < 0 || prompt_tokens > max_prompt_tokens_ || rope_table_.empty()) { @@ -158,6 +214,11 @@ modalities::Status NativeWorkspace::update_decoder_rope(int prompt_tokens) { #else if (!decoder_rope_buffer_) return invalid("decoder RoPE buffer was not allocated"); + if (flavor_ == NativeWorkspaceFlavor::kThorFp8) { + return generate_native_thor_rope_f16( + frt_buffer_dptr(decoder_rope_buffer_), + encoder_vision_sequence_ + prompt_tokens, chunk_size_, 0); + } const std::size_t start = static_cast(encoder_vision_sequence_ + prompt_tokens) * 256; @@ -182,8 +243,16 @@ modalities::Status NativeWorkspace::expand_vision_position_embedding( weights.find("vision_position_embedding"); const NativeWorkspaceBuffer* destination = find("vision_pos_embed_expanded"); - if (!source || !destination || - source->dtype != NativeWeightDType::kBf16 || + const NativeWeightDType expected_weight = + flavor_ == NativeWorkspaceFlavor::kThorFp8 + ? NativeWeightDType::kFloat16 + : NativeWeightDType::kBf16; + const modalities::DType expected_buffer = + flavor_ == NativeWorkspaceFlavor::kThorFp8 + ? modalities::DType::kFloat16 + : modalities::DType::kBFloat16; + if (!source || !destination || source->dtype != expected_weight || + destination->dtype != expected_buffer || source->shape != std::vector({256, 1152})) { return invalid("vision position embedding source is invalid"); } @@ -218,9 +287,15 @@ modalities::Status NativeWorkspace::allocate( 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.chunk_size > std::numeric_limits::max() - + config.max_prompt_tokens - + config.num_views * 256 || (config.vision_pool_factor != 1 && config.vision_pool_factor != 2 && - config.vision_pool_factor != 4)) { + config.vision_pool_factor != 4) || + (config.flavor == NativeWorkspaceFlavor::kThorFp8 && + (config.vision_pool_factor != 1 || + (config.max_prompt_tokens & 1)))) { return invalid("Pi0.5 native workspace configuration is invalid"); } const int pool_area = @@ -228,6 +303,8 @@ modalities::Status NativeWorkspace::allocate( num_views_ = config.num_views; max_prompt_tokens_ = config.max_prompt_tokens; chunk_size_ = config.chunk_size; + num_steps_ = config.num_steps; + flavor_ = config.flavor; vision_sequence_ = config.num_views * 256; encoder_vision_sequence_ = vision_sequence_ / pool_area; encoder_sequence_ = @@ -245,6 +322,143 @@ modalities::Status NativeWorkspace::allocate( st = add(__VA_ARGS__); \ if (!st.ok_status()) return st; \ } while (false) + if (flavor_ == NativeWorkspaceFlavor::kThorFp8) { + const std::uint64_t keys = es + ds; + FRT_ADD("observation_images_normalized", {nv, 224, 224, 3}, + modalities::DType::kFloat16); + FRT_ADD("vision_x", {vs, 1152}, modalities::DType::kFloat16); + st = add_alias("vision_x_pooled", "vision_x", {vs, 1152}); + if (!st.ok_status()) return st; + FRT_ADD("vision_x_fp8", {vs, 1152}, modalities::DType::kUInt8); + FRT_ADD("vision_QKV", {vs, 3456}, modalities::DType::kFloat16); + FRT_ADD("vision_attn", {vs, 1152}, modalities::DType::kFloat16); + st = add_alias("vision_postln_scratch", "vision_attn", {vs, 1152}); + if (!st.ok_status()) return st; + FRT_ADD("vision_hidden", {vs, 4304}, modalities::DType::kFloat16); + FRT_ADD("vision_hidden_fp8", {vs, 4304}, + modalities::DType::kUInt8); + FRT_ADD("vision_pos_embed_expanded", {vs, 1152}, + modalities::DType::kFloat16); + FRT_ADD("vision_patches", {vs, 588}, modalities::DType::kFloat16); + FRT_ADD("vision_unit_scale", {1}, modalities::DType::kFloat32); + + FRT_ADD("encoder_rope_weights", {es, 256}, + modalities::DType::kFloat16); + FRT_ADD("prompt_embedding", + {static_cast(max_prompt_tokens_), 2048}, + modalities::DType::kFloat16); + FRT_ADD("encoder_x", {es, 2048}, modalities::DType::kFloat16); + FRT_ADD("encoder_x_fp8", {es, 2048}, modalities::DType::kUInt8); + FRT_ADD("encoder_QKV", {es, 2560}, modalities::DType::kFloat16); + FRT_ADD("encoder_logits", {es * 8, keys}, + modalities::DType::kFloat16); + FRT_ADD("encoder_attn", {es, 2048}, modalities::DType::kFloat16); + FRT_ADD("encoder_o_fp8", {es, 2048}, modalities::DType::kUInt8); + FRT_ADD("encoder_gate_merged", {es, 32768}, + modalities::DType::kFloat16); + FRT_ADD("encoder_hidden", {es, 16384}, + modalities::DType::kFloat16); + FRT_ADD("encoder_hidden_fp8", {es, 16384}, + modalities::DType::kUInt8); + FRT_ADD("encoder_fg", {es, 2048}, modalities::DType::kFloat16); + FRT_ADD("encoder_rms_ones", {2048}, modalities::DType::kFloat16); + FRT_ADD("encoder_activation_scales", {18, 4}, + modalities::DType::kFloat32); + FRT_ADD("encoder_k_cache", {18, keys, 256}, + modalities::DType::kFloat16); + FRT_ADD("encoder_v_cache", {18, keys, 256}, + modalities::DType::kFloat16); + FRT_ADD("attn_enc_seqused", {sizeof(std::int32_t)}, + modalities::DType::kUInt8); + FRT_ADD("attn_dec_seqused", {sizeof(std::int32_t)}, + modalities::DType::kUInt8); + FRT_ADD("attn_dec_devpos", {sizeof(std::int32_t)}, + modalities::DType::kUInt8); + + FRT_ADD("decoder_rope_weights", {ds, 256}, + modalities::DType::kFloat16); + FRT_ADD("decoder_x", {ds, 1024}, modalities::DType::kFloat16); + FRT_ADD("x_normed_buf", {ds, 1024}, modalities::DType::kFloat16); + FRT_ADD("gate_buf", {ds, 1024}, modalities::DType::kFloat16); + FRT_ADD("decoder_QKV", {ds, 2560}, modalities::DType::kFloat16); + FRT_ADD("decoder_logits", {ds * 8, keys}, + modalities::DType::kFloat16); + FRT_ADD("decoder_attn", {ds, 2048}, modalities::DType::kFloat16); + FRT_ADD("decoder_hidden", {ds, 8192}, modalities::DType::kFloat16); + FRT_ADD("decoder_fg", {ds, 8192}, modalities::DType::kFloat16); + FRT_ADD("decoder_action_f32", {ds, 32}, + modalities::DType::kFloat32); + FRT_ADD("decoder_x_fp8", {ds, 1024}, modalities::DType::kUInt8); + FRT_ADD("decoder_hidden_fp8", {ds, 4096}, + modalities::DType::kUInt8); + FRT_ADD("decoder_context_fp8", {ds, 2048}, + modalities::DType::kUInt8); + FRT_ADD("decoder_time_emb", {steps, ds, 1024}, + modalities::DType::kFloat16); + FRT_ADD("decoder_style_attn", {steps, 18, ds, 3072}, + modalities::DType::kFloat16); + FRT_ADD("decoder_style_ffn", {steps, 18, ds, 3072}, + modalities::DType::kFloat16); + FRT_ADD("decoder_style_final", {steps, ds, 3072}, + modalities::DType::kFloat16); + FRT_ADD("decoder_activation_scales", {steps, 18, 4}, + modalities::DType::kFloat32); + FRT_ADD("diffusion_noise", {ds, 32}, modalities::DType::kFloat16); + FRT_ADD("rtc_prev_action_chunk", {ds, 32}, + modalities::DType::kFloat16); + FRT_ADD("rtc_prefix_weights", {ds}, modalities::DType::kFloat32); + FRT_ADD("rtc_guidance_weight", {1}, modalities::DType::kFloat32); + FRT_ADD("decoder_rms_ones", {1024}, modalities::DType::kFloat16); + + if (config.enable_calibration) { + FRT_ADD("encoder_norm_scratch", {es, 2048}, + modalities::DType::kFloat16); + FRT_ADD("encoder_x_scratch", {es, 2048}, + modalities::DType::kFloat16); + FRT_ADD("encoder_fp8_scratch", {es, 16384}, + modalities::DType::kUInt8); + FRT_ADD("encoder_sample_scales", {18, 4}, + modalities::DType::kFloat32); + FRT_ADD("decoder_fp8_scratch", {ds, 4096}, + modalities::DType::kUInt8); + FRT_ADD("decoder_sample_scales", {steps, 18, 4}, + modalities::DType::kFloat32); + FRT_ADD("calibration_scale", {1}, modalities::DType::kFloat32); + } + + const NativeWorkspaceBuffer* decoder = find("decoder_rope_weights"); + const NativeWorkspaceBuffer* prompt = find("prompt_embedding"); + if (!decoder || !prompt) { + return invalid("Thor prompt workspace was not allocated"); + } + decoder_rope_buffer_ = decoder->buffer; + prompt_embedding_buffer_ = prompt->buffer; + const char* controls[] = { + "attn_enc_seqused", "attn_dec_seqused", "attn_dec_devpos"}; + for (int i = 0; i < 3; ++i) { + const NativeWorkspaceBuffer* control = find(controls[i]); + if (!control) return invalid("Thor attention control is missing"); + prompt_length_buffers_[i] = control->buffer; + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Thor workspace initialization requires the CUDA build"); +#else + const float unit_scale = 1.0f; + const NativeWorkspaceBuffer* unit = find("vision_unit_scale"); + if (!unit || cudaMemcpy(frt_buffer_dptr(unit->buffer), &unit_scale, + sizeof(unit_scale), cudaMemcpyHostToDevice) != + cudaSuccess) { + return backend("Thor unit scale upload failed"); + } +#endif + st = initialize_rms_ones(); + if (!st.ok_status()) return st; + st = initialize_rope(); + if (!st.ok_status()) return st; + return set_fixed_prompt_length(0); + } FRT_ADD("observation_images_normalized", {nv, 224, 224, 3}, modalities::DType::kBFloat16); FRT_ADD("vision_x", {vs, 1152}, modalities::DType::kBFloat16); diff --git a/cpp/tests/gate_pi05_native_thor_fp8.py b/cpp/tests/gate_pi05_native_thor_fp8.py new file mode 100644 index 00000000..af0e6cf4 --- /dev/null +++ b/cpp/tests/gate_pi05_native_thor_fp8.py @@ -0,0 +1,385 @@ +#!/usr/bin/env python3 +"""Compare the native Thor FP8 producer with the shipped Torch producer.""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import types +from pathlib import Path + +import numpy as np +import torch +from safetensors import safe_open + +import flash_rt.flash_rt_kernels as fvk +from flash_rt.frontends.torch.pi05_thor import Pi05TorchFrontendThor +from flash_rt.hardware.thor.shared_primitives import encoder_forward_calibrate +from flash_rt.models.pi05.pipeline_thor import decoder_forward_calibrate + + +def _artifact(path: Path) -> tuple[dict[str, str], np.ndarray, np.ndarray]: + with safe_open(path, framework="np") as handle: + metadata = handle.metadata() or {} + encoder = handle.get_tensor("encoder_scales").astype( + np.float32, copy=True) + decoder = handle.get_tensor("decoder_scales").astype( + np.float32, copy=True) + return metadata, encoder, decoder + + +def _sample(index: int, num_views: int + ) -> tuple[str, np.ndarray, list[np.ndarray], np.ndarray]: + pixels = np.arange(224 * 224 * 3, dtype=np.uint64) + image = ((pixels * 3 + index * 17) % 251).astype(np.uint8) + wrist = ((pixels * 7 + index * 29 + 11) % 253).astype(np.uint8) + state = np.asarray( + [((index * 8 + dim) % 17 - 8) / 8.0 for dim in range(8)], + dtype=np.float32, + ) + noise = np.asarray( + [((index * 320 + item) % 31 - 15) / 16.0 for item in range(320)], + dtype=np.float32, + ).reshape(10, 32) + prompt = ( + "move the black bowl to the plate" + if index & 1 + else "pick up the black bowl" + ) + images = [image.reshape(224, 224, 3), wrist.reshape(224, 224, 3)] + return ( + prompt, + state, + images[:num_views], + noise, + ) + + +def _normalized_state(frontend: Pi05TorchFrontendThor, + state: np.ndarray) -> np.ndarray: + q01 = np.asarray(frontend.norm_stats["state"]["q01"], dtype=np.float32) + q99 = np.asarray(frontend.norm_stats["state"]["q99"], dtype=np.float32) + return ( + ((state - q01) / (q99 - q01 + np.float32(1.0e-6))) + * np.float32(2.0) + - np.float32(1.0) + ).astype(np.float32) + + +def _seed_calibration( + encoder: np.ndarray, + decoder: np.ndarray, +): + def seed(self: Pi05TorchFrontendThor, _sequence: int) -> None: + self._enc_calib_scales.copy_(torch.from_numpy(encoder)) + self._ae_calib_scales.copy_(torch.from_numpy(decoder)) + weight_scales = self._enc_w_dev.cpu().numpy() + self._enc_alpha_host = [ + float(np.float32(encoder[i]) * np.float32(weight_scales[i])) + for i in range(encoder.size) + ] + + return seed + + +class CalibrationOracle: + def __init__(self, frontend: Pi05TorchFrontendThor): + self.frontend = frontend + p = frontend + sequence = p.Se + width = p.De + hidden = p.He + action_width = p.Da + action_hidden = p.Ha + self._encoder_norm_scratch = torch.empty( + sequence * width, dtype=torch.float16, device="cuda") + self._encoder_x_scratch = torch.empty( + sequence * width, dtype=torch.float16, device="cuda") + self._encoder_calib_buffer = torch.zeros( + p.Le * 4, dtype=torch.float32, device="cuda") + self._encoder_dynamic_scale = torch.zeros( + 1, dtype=torch.float32, device="cuda") + self._encoder_fp8_scratch = torch.zeros( + sequence * max(width, hidden), dtype=torch.uint8, device="cuda") + self._encoder_ones = torch.ones( + width, dtype=torch.float16, device="cuda") + self._decoder_calib_buffer = torch.zeros( + p.steps * p.La * 4, dtype=torch.float32, device="cuda") + self._decoder_dynamic_scale = torch.zeros( + 1, dtype=torch.float32, device="cuda") + self._decoder_hidden_scratch = torch.empty( + p.Sa * action_hidden, dtype=torch.float16, device="cuda") + self._decoder_fp8_scratch = torch.zeros( + p.Sa * max(action_width, action_hidden), dtype=torch.uint8, + device="cuda") + + self.encoder_buffers = { + "x": p._enc_x.data_ptr(), + "x_fp8": p._enc_x_fp8.data_ptr(), + "qkv": p._enc_qkv_buf.data_ptr(), + "logits": p._enc_logits.data_ptr(), + "attn_out": p._enc_attn.data_ptr(), + "o_fp8": p._enc_o_fp8.data_ptr(), + "gate": p._enc_gate.data_ptr(), + "hidden": p._enc_hidden.data_ptr(), + "hid_fp8": p._enc_hid_fp8.data_ptr(), + "fg": p._enc_fg.data_ptr(), + "ctx": p._ctx, + "norm_scratch": self._encoder_norm_scratch.data_ptr(), + "x_scratch": self._encoder_x_scratch.data_ptr(), + "calib_buf": self._encoder_calib_buffer.data_ptr(), + "d_scale": self._encoder_dynamic_scale.data_ptr(), + "fp8_scratch": self._encoder_fp8_scratch.data_ptr(), + "ones": self._encoder_ones.data_ptr(), + } + self.encoder_weights = { + "qkv_w": [weight.data_ptr() for weight in p._enc_qkv_w], + "o_w": [weight.data_ptr() for weight in p._enc_o_w], + "gate_w": [weight.data_ptr() for weight in p._enc_gu_w], + "down_w": [weight.data_ptr() for weight in p._enc_d_w], + "rope": p._enc_rope.data_ptr(), + "Kc": p._Kc.reshape(-1).data_ptr(), + "Vc": p._Vc.reshape(-1).data_ptr(), + "w_scales": p._enc_w_dev.data_ptr(), + } + self.encoder_dims = { + "Se": sequence, + "D": width, + "H": hidden, + "NH": p.NHe, + "HD": p.HDe, + "L": p.Le, + "total_keys": p.total_keys, + } + + self.decoder_buffers = { + "noise": p._g_noise.data_ptr(), + "x": p._ae_x.data_ptr(), + "xn": p._ae_xn.data_ptr(), + "gate": p._ae_gate.data_ptr(), + "qkv": p._ae_qkv.data_ptr(), + "logits": p._ae_logits.data_ptr(), + "attn_out": p._ae_attn.data_ptr(), + "hid": p._ae_hid.data_ptr(), + "fg": p._ae_fg.data_ptr(), + "action_f32": p._ae_action_f32.data_ptr(), + "xn_fp8": p._ae_xn_fp8.data_ptr(), + "hid_fp8": p._ae_hid_fp8.data_ptr(), + "ctx_fp8": p._ae_ctx_fp8.data_ptr(), + "calib_buf": self._decoder_calib_buffer.data_ptr(), + "d_scale": self._decoder_dynamic_scale.data_ptr(), + "hidden_scratch": self._decoder_hidden_scratch.data_ptr(), + "fp8_scratch": self._decoder_fp8_scratch.data_ptr(), + } + self.decoder_weights = { + "ain_w": p._ain_w.data_ptr(), + "ain_b": p._ain_b.data_ptr(), + "sa": p._sa_all.data_ptr(), + "qw": p._dec_qkv_flat.data_ptr(), + "Kc": p._Kc.reshape(-1).data_ptr(), + "Vc": p._Vc.reshape(-1).data_ptr(), + "dec_devpos": p._attn.dec_devpos.data_ptr(), + "ow": p._dec_o_flat.data_ptr(), + "sf": p._sf_all.data_ptr(), + "gw": p._dec_gu_flat.data_ptr(), + "dw": p._dec_d_flat.data_ptr(), + "aow": p._aow.data_ptr(), + "aob": p._aob.data_ptr(), + "aob_dt": p._aob_dt.data_ptr(), + "dt": p._ae_dt, + "fs": p._fs_all.data_ptr(), + "rope": p._dec_rope.data_ptr(), + "w_scales": p._ae_w_dev.data_ptr(), + } + self.decoder_dims = { + "S": p.Sa, + "D": action_width, + "H": action_hidden, + "NH": 8, + "HD": 256, + "steps": p.steps, + "layers": p.La, + "enc_seq": sequence, + "total_keys": p.total_keys, + "fixed_shape": p._fixed_shape_active, + } + + def observe(self, images: list[np.ndarray], noise: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + p = self.frontend + normalized = np.stack([ + (image.astype(np.float32) / np.float32(127.5) + - np.float32(1.0)).astype(np.float16) + for image in images + ]) + p._img_buf.upload(normalized) + p._siglip_graph.replay() + torch.cuda.synchronize() + + p._enc_calib_scales.zero_() + p._Kc.zero_() + p._Vc.zero_() + encoder_forward_calibrate( + p._gemm, fvk, self.encoder_buffers, self.encoder_weights, + self.encoder_dims, p._enc_calib_scales.data_ptr(), stream=0, + attn=p._attn, + ) + torch.cuda.synchronize() + encoder = p._enc_calib_scales.cpu().numpy().copy() + + p._ae_calib_scales.zero_() + p._g_noise.copy_(torch.from_numpy(noise.astype(np.float16)).cuda()) + decoder_forward_calibrate( + p._ctx, fvk, self.decoder_buffers, self.decoder_weights, + self.decoder_dims, p._ae_calib_scales.data_ptr(), stream=0, + attn=p._attn, + ) + torch.cuda.synchronize() + decoder = p._ae_calib_scales.cpu().numpy().copy() + return encoder, decoder + + +def _assert_equal(label: str, actual: np.ndarray, + expected: np.ndarray) -> None: + if np.array_equal(actual, expected): + print(f"PASS {label}: bit-exact ({actual.size} values)") + return + actual_flat = actual.reshape(-1) + expected_flat = expected.reshape(-1) + difference = np.abs( + actual_flat.astype(np.float64) - expected_flat.astype(np.float64) + ) + index = int(np.argmax(difference)) + first = int(np.flatnonzero(actual_flat != expected_flat)[0]) + mismatches = int(np.count_nonzero(actual_flat != expected_flat)) + raise AssertionError( + f"{label} mismatch: max_abs={difference[index]:.9g} " + f"at {index}, actual={actual_flat[index]!r}, " + f"expected={expected_flat[index]!r}; " + f"first={first}, actual={actual_flat[first]!r}, " + f"expected={expected_flat[first]!r}, " + f"mismatches={mismatches}/{actual.size}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--probe", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--artifact", type=Path, required=True) + parser.add_argument("--samples", type=int, default=3) + parser.add_argument("--views", type=int, choices=(1, 2), default=2) + args = parser.parse_args() + if args.samples < 1 or args.samples > 256: + parser.error("--samples must be in [1, 256]") + + probe = args.probe.resolve() + os.environ["FLASH_RT_PALIGEMMA_TOKENIZER"] = str( + args.tokenizer.resolve() + ) + raw_path = args.artifact.with_suffix(args.artifact.suffix + ".raw") + env = os.environ.copy() + library_paths = [probe.parent, probe.parent / "exec", + probe.parent / "runtime"] + if env.get("LD_LIBRARY_PATH"): + library_paths.append(Path(env["LD_LIBRARY_PATH"])) + env["LD_LIBRARY_PATH"] = os.pathsep.join(map(str, library_paths)) + subprocess.run( + [ + str(probe), + str(args.checkpoint), + str(args.tokenizer), + str(args.artifact), + str(args.samples), + str(args.views), + str(raw_path), + ], + check=True, + env=env, + ) + metadata, expected_encoder, expected_decoder = _artifact(args.artifact) + _, single_encoder, single_decoder = _artifact( + Path(str(args.artifact) + ".single") + ) + if int(metadata.get("sample_count", "0")) != args.samples: + raise AssertionError("native calibration sample_count is incorrect") + if int(metadata.get("num_views", "0")) != args.views: + raise AssertionError("native calibration num_views is incorrect") + + frontend = Pi05TorchFrontendThor( + str(args.checkpoint), num_views=args.views, use_cuda_graph=True, + autotune=0, + use_fp8=True, state_prompt_mode="fixed", + state_prompt_fixed_max_len=200, + ) + frontend._calibrate = types.MethodType( + _seed_calibration(expected_encoder, expected_decoder), frontend + ) + first_prompt, first_state, _, _ = _sample(0, args.views) + frontend.set_prompt( + first_prompt, state=_normalized_state(frontend, first_state) + ) + oracle = CalibrationOracle(frontend) + + encoder_samples = [] + decoder_samples = [] + samples = [] + for index in range(args.samples): + prompt, state, images, noise = _sample(index, args.views) + samples.append((prompt, state, images, noise)) + frontend.set_prompt( + prompt, state=_normalized_state(frontend, state) + ) + encoder, decoder = oracle.observe(images, noise) + # The final encoder layer only emits decoder K/V. Its O/FFN sites are + # skipped; the native artifact uses 1.0 as their valid neutral value. + encoder[-3:] = np.float32(1.0) + encoder_samples.append(encoder) + decoder_samples.append(decoder) + + _assert_equal("single-sample encoder calibration", + encoder_samples[0], single_encoder) + _assert_equal("single-sample decoder calibration", + decoder_samples[0], single_decoder) + reduced_encoder = np.percentile( + np.stack(encoder_samples), 99.9, axis=0 + ).astype(np.float32) + reduced_decoder = np.percentile( + np.stack(decoder_samples), 99.9, axis=0 + ).astype(np.float32) + _assert_equal("dataset encoder calibration", + reduced_encoder, expected_encoder) + _assert_equal("dataset decoder calibration", + reduced_decoder, expected_decoder) + + _, state, images, noise = samples[-1] + frontend.set_prompt( + "pick up the black bowl", state=_normalized_state(frontend, state) + ) + frontend._enc_calib_scales.copy_(torch.from_numpy(expected_encoder)) + frontend._ae_calib_scales.copy_(torch.from_numpy(expected_decoder)) + weight_scales = frontend._enc_w_dev.cpu().numpy() + frontend._enc_alpha_host = [ + float(np.float32(expected_encoder[i]) * np.float32(weight_scales[i])) + for i in range(expected_encoder.size) + ] + frontend._capture_enc_ae_graph() + normalized_images = np.stack([ + (image.astype(np.float32) / np.float32(127.5) + - np.float32(1.0)).astype(np.float16) + for image in images + ]) + frontend._img_buf.upload(normalized_images) + frontend._siglip_graph.replay() + frontend._g_noise.copy_(torch.from_numpy(noise.astype(np.float16)).cuda()) + frontend._enc_ae_graph.replay() + torch.cuda.synchronize() + python_raw = frontend._g_noise.cpu().numpy() + native_raw = np.fromfile(raw_path, dtype=np.float16).reshape(10, 32) + _assert_equal("native raw action", python_raw, native_raw) + + +if __name__ == "__main__": + main() diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index edb0043d..4b37d042 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -20,8 +20,9 @@ extern "C" int frt_model_runtime_open_v1(const char* config_json, 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"; + if (argc != 3 && argc != 4) { + std::cerr << "usage: pi05_native_open_probe CHECKPOINT TOKENIZER " + "[CALIBRATION]\n"; return 2; } int replay_count = 1; @@ -70,7 +71,11 @@ int main(int argc, char** argv) { << "\",\"state_prompt_mode\":\"fixed\"," << "\"max_prompt_tokens\":200,\"state_dim\":8," << "\"num_views\":2,\"chunk\":10,\"num_steps\":10," - << "\"vision_pool_factor\":1}"; + << "\"vision_pool_factor\":1"; + if (argc == 4) { + json << ",\"calibration_path\":\"" << argv[3] << '"'; + } + json << '}'; 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) { @@ -94,6 +99,9 @@ int main(int argc, char** argv) { const std::string hardware_identity = "hardware=" + hardware_id; const std::string hardware_manifest = "\"hardware\":\"" + hardware_id + "\""; + const uint32_t expected_io_dtype = hardware_id == "sm110" + ? FRT_RT_DTYPE_F16 + : FRT_RT_DTYPE_BF16; 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 && @@ -126,7 +134,9 @@ int main(int argc, char** argv) { 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[2].dtype == expected_io_dtype && + model->ports[3].dtype == expected_io_dtype && + model->ports[5].dtype == expected_io_dtype && model->ports[5].update == FRT_RT_PORT_SWAP && model->ports[5].buffer == model->ports[3].buffer && model->ports[5].offset == model->ports[3].offset && @@ -246,8 +256,11 @@ int main(int argc, char** argv) { frt_buffer noise = model->ports[3].buffer; std::vector host_noise(10 * 32); for (std::size_t i = 0; i < host_noise.size(); ++i) { - host_noise[i] = flashrt::modalities::float_to_bfloat16( - static_cast(static_cast(i % 23) - 11) / 12.0f); + const float value = + static_cast(static_cast(i % 23) - 11) / 12.0f; + host_noise[i] = expected_io_dtype == FRT_RT_DTYPE_F16 + ? flashrt::modalities::float_to_float16(value) + : flashrt::modalities::float_to_bfloat16(value); } const bool profile_range = replay_env || std::getenv("FLASHRT_PROFILE_RANGE") != nullptr; diff --git a/cpp/tests/pi05_native_rope_probe.cpp b/cpp/tests/pi05_native_rope_probe.cpp index 16cb0a6f..4e94b581 100644 --- a/cpp/tests/pi05_native_rope_probe.cpp +++ b/cpp/tests/pi05_native_rope_probe.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -35,9 +36,9 @@ std::vector download( } // namespace int main(int argc, char** argv) { - if (argc != 6) { + if (argc < 6 || argc > 8) { std::cerr << "usage: pi05_native_rope_probe VIEWS MAX_PROMPT CHUNK " - "POOL PROMPT\n"; + "POOL PROMPT [thor_fp8 [OUTPUT_PREFIX]]\n"; return 2; } flashrt::models::pi05::NativeWorkspaceConfig config; @@ -45,12 +46,18 @@ int main(int argc, char** argv) { config.max_prompt_tokens = std::stoi(argv[2]); config.chunk_size = std::stoi(argv[3]); config.vision_pool_factor = std::stoi(argv[4]); + if (argc >= 7) { + if (std::string(argv[6]) != "thor_fp8") return 2; + config.flavor = + flashrt::models::pi05::NativeWorkspaceFlavor::kThorFp8; + } 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()) { + !(argc >= 7 ? workspace.set_fixed_prompt_length(prompt) + : workspace.update_decoder_rope(prompt)).ok_status()) { frt_ctx_destroy(ctx); return 1; } @@ -58,6 +65,20 @@ int main(int argc, char** argv) { download(*workspace.find("encoder_rope_weights")); const std::vector decoder = download(*workspace.find("decoder_rope_weights")); + if (argc == 8) { + std::ofstream encoder_file( + std::string(argv[7]) + ".encoder", std::ios::binary); + std::ofstream decoder_file( + std::string(argv[7]) + ".decoder", std::ios::binary); + encoder_file.write(reinterpret_cast(encoder.data()), + encoder.size() * sizeof(encoder[0])); + decoder_file.write(reinterpret_cast(decoder.data()), + decoder.size() * sizeof(decoder[0])); + if (!encoder_file || !decoder_file) { + frt_ctx_destroy(ctx); + return 1; + } + } std::cout << "encoder_shape=" << workspace.encoder_sequence() << ",256" << " encoder_fnv=" << std::hex << std::setw(16) << std::setfill('0') << fnv1a(encoder) diff --git a/cpp/tests/pi05_native_thor_fp8_probe.cpp b/cpp/tests/pi05_native_thor_fp8_probe.cpp new file mode 100644 index 00000000..ac7424ff --- /dev/null +++ b/cpp/tests/pi05_native_thor_fp8_probe.cpp @@ -0,0 +1,357 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/cpp/models/pi05/c_api.h" +#include "flashrt/cpp/models/pi05/native_calibration.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +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; +} + +std::string config_json(const std::string& checkpoint, + const std::string& tokenizer, + int num_views, + const std::string& calibration = "") { + std::ostringstream out; + out << "{\"io\":\"native_v2\",\"checkpoint_path\":" + << json_string(checkpoint) << ",\"tokenizer_model_path\":" + << json_string(tokenizer) + << ",\"state_prompt_mode\":\"fixed\"," + "\"precision\":\"fp8_e4m3fn\"," + "\"max_prompt_tokens\":200,\"state_dim\":8,\"num_views\":" + << num_views << ",\"chunk\":10,\"num_steps\":10," + "\"vision_pool_factor\":1"; + if (!calibration.empty()) { + out << ",\"calibration_path\":" << json_string(calibration); + } + out << '}'; + return out.str(); +} + +int calibration_error(frt_pi05_calibration_session* session, + const char* operation, + int rc) { + std::cerr << operation << " failed: rc=" << rc << " error=" + << (session ? frt_pi05_calibration_last_error_v1(session) + : frt_pi05_calibration_create_last_error_v1()) + << '\n'; + if (session) frt_pi05_calibration_destroy_v1(session); + return 1; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 6 && argc != 7) { + std::cerr << "usage: pi05_native_thor_fp8_probe CHECKPOINT TOKENIZER " + "ARTIFACT SAMPLES VIEWS [RAW_ACTION_OUTPUT]\n"; + return 2; + } + const int samples = std::atoi(argv[4]); + if (samples < 1 || samples > 256) { + std::cerr << "SAMPLES must be in [1, 256]\n"; + return 2; + } + const int num_views = std::atoi(argv[5]); + if (num_views < 1 || num_views > 2) { + std::cerr << "VIEWS must be in [1, 2]\n"; + return 2; + } + const std::string calibration_path = argv[3]; + const std::string single_path = calibration_path + ".single"; + const std::string calibration_config = + config_json(argv[1], argv[2], num_views); + frt_pi05_calibration_session* session = nullptr; + int rc = frt_pi05_calibration_create_v1( + calibration_config.c_str(), 99.9, &session); + if (rc != 0 || !session) { + return calibration_error(nullptr, "calibration create", rc); + } + + std::vector image(224 * 224 * 3); + std::vector wrist(image.size()); + float state[8]{}; + std::vector noise(10 * 32); + std::vector frames(num_views); + const char* names[] = {"image", "wrist_image"}; + std::vector* pixels[] = {&image, &wrist}; + for (int view = 0; view < num_views; ++view) { + frames[view].struct_size = sizeof(frt_pi05_vision_frame); + frames[view].name = names[view]; + frames[view].data = pixels[view]->data(); + frames[view].bytes = pixels[view]->size(); + frames[view].width = 224; + frames[view].height = 224; + frames[view].stride_bytes = 224 * 3; + frames[view].pixel_format = FRT_PI05_PIXEL_RGB8; + } + if (num_views == 2) std::swap(frames[0], frames[1]); + for (int sample_index = 0; sample_index < samples; ++sample_index) { + for (std::size_t i = 0; i < image.size(); ++i) { + image[i] = static_cast( + (i * 3 + sample_index * 17) % 251); + wrist[i] = static_cast( + (i * 7 + sample_index * 29 + 11) % 253); + } + for (int i = 0; i < 8; ++i) { + state[i] = static_cast( + (sample_index * 8 + i) % 17 - 8) / 8.0f; + } + for (std::size_t i = 0; i < noise.size(); ++i) { + const int centered = static_cast( + (static_cast(sample_index) * noise.size() + i) % + 31) - 15; + noise[i] = static_cast(centered) / 16.0f; + } + frt_pi05_calibration_sample_v1 sample{}; + sample.struct_size = sizeof(sample); + sample.prompt = sample_index & 1 + ? "move the black bowl to the plate" + : "pick up the black bowl"; + sample.state = state; + sample.n_state = 8; + sample.frames = frames.data(); + sample.n_frames = frames.size(); + sample.noise = noise.data(); + sample.n_noise = noise.size(); + sample.noise_seed = 1234; + if (sample_index == 0) { + frt_pi05_calibration_sample_v1 incomplete = sample; + incomplete.n_frames = static_cast(num_views - 1); + rc = frt_pi05_calibration_observe_v1(session, &incomplete); + if (rc == 0 || + frt_pi05_calibration_sample_count_v1(session) != 0) { + std::cerr << "incomplete camera set was accepted\n"; + frt_pi05_calibration_destroy_v1(session); + return 1; + } + if (num_views > 1) { + std::vector duplicate = frames; + duplicate[1].name = duplicate[0].name; + frt_pi05_calibration_sample_v1 duplicate_names = sample; + duplicate_names.frames = duplicate.data(); + rc = frt_pi05_calibration_observe_v1( + session, &duplicate_names); + if (rc == 0 || + frt_pi05_calibration_sample_count_v1(session) != 0) { + std::cerr << + "duplicate calibration camera name was accepted\n"; + frt_pi05_calibration_destroy_v1(session); + return 1; + } + } + std::vector unknown = frames; + unknown[0].name = "unknown_camera"; + frt_pi05_calibration_sample_v1 unknown_name = sample; + unknown_name.frames = unknown.data(); + rc = frt_pi05_calibration_observe_v1(session, &unknown_name); + if (rc == 0 || + frt_pi05_calibration_sample_count_v1(session) != 0) { + std::cerr << "unknown calibration camera name was accepted\n"; + frt_pi05_calibration_destroy_v1(session); + return 1; + } + std::vector bgr = frames; + bgr[0].pixel_format = FRT_PI05_PIXEL_BGR8; + frt_pi05_calibration_sample_v1 wrong_format = sample; + wrong_format.frames = bgr.data(); + rc = frt_pi05_calibration_observe_v1(session, &wrong_format); + if (rc == 0 || + frt_pi05_calibration_sample_count_v1(session) != 0) { + std::cerr << "non-RGB calibration frame was accepted\n"; + frt_pi05_calibration_destroy_v1(session); + return 1; + } + frt_pi05_calibration_sample_v1 malformed_noise = sample; + malformed_noise.n_noise--; + rc = frt_pi05_calibration_observe_v1(session, &malformed_noise); + if (rc == 0 || + frt_pi05_calibration_sample_count_v1(session) != 0) { + std::cerr << "malformed calibration noise was accepted\n"; + frt_pi05_calibration_destroy_v1(session); + return 1; + } + } + rc = frt_pi05_calibration_observe_v1(session, &sample); + if (rc != 0) { + return calibration_error(session, "calibration observe", rc); + } + if (sample_index == 0) { + rc = frt_pi05_calibration_finalize_v1( + session, single_path.c_str()); + if (rc != 0) { + return calibration_error(session, "single finalize", rc); + } + } + } + if (frt_pi05_calibration_sample_count_v1(session) != + static_cast(samples)) { + return calibration_error(session, "sample count", -1); + } + rc = frt_pi05_calibration_finalize_v1( + session, calibration_path.c_str()); + if (rc != 0) { + return calibration_error(session, "dataset finalize", rc); + } + frt_pi05_calibration_sample_v1 generated_noise{}; + generated_noise.struct_size = sizeof(generated_noise); + generated_noise.prompt = "pick up the black bowl"; + generated_noise.state = state; + generated_noise.n_state = 8; + generated_noise.frames = frames.data(); + generated_noise.n_frames = frames.size(); + generated_noise.noise_seed = 4321; + rc = frt_pi05_calibration_observe_v1(session, &generated_noise); + if (rc != 0 || frt_pi05_calibration_sample_count_v1(session) != + static_cast(samples + 1)) { + return calibration_error(session, "generated-noise observe", rc); + } + frt_pi05_calibration_destroy_v1(session); + + flashrt::models::pi05::NativeCalibrationArtifact artifact; + auto status = flashrt::models::pi05::load_native_calibration_artifact( + single_path, &artifact); + if (!status.ok_status() || artifact.sample_count != 1) { + std::cerr << "single calibration artifact validation failed\n"; + return 1; + } + status = flashrt::models::pi05::load_native_calibration_artifact( + calibration_path, &artifact); + if (!status.ok_status() || artifact.sample_count != + static_cast(samples) || + artifact.num_views != num_views) { + std::cerr << "dataset calibration artifact validation failed\n"; + return 1; + } + + const std::string runtime_config = + config_json(argv[1], argv[2], num_views, calibration_path); + frt_model_runtime_v1* model = nullptr; + rc = frt_model_runtime_open_v1(runtime_config.c_str(), &model); + if (rc != 0 || !model) { + std::cerr << "native FP8 open failed: rc=" << rc << " error=" + << frt_pi05_native_open_last_error() << '\n'; + return 1; + } + const frt_runtime_export_v1* exp = model->exp; + bool schema_ok = model->n_ports == 6 && exp && exp->n_graphs == 1 && + frt_graph_variant_count(exp->graphs[0].handle) == 1 && + exp->identity && + std::strstr(exp->identity, + "precision=fp8_e4m3fn") && + std::strstr(exp->identity, "hardware=sm110") && + std::strstr(exp->identity, "calibration_sha256=") && + model->ports[2].dtype == FRT_RT_DTYPE_F16 && + model->ports[3].dtype == FRT_RT_DTYPE_F16 && + model->ports[5].dtype == FRT_RT_DTYPE_F16; + if (!schema_ok) { + std::cerr << "native FP8 schema validation failed\n"; + model->release(model->owner); + return 1; + } + + const std::string prompt = "pick up the black bowl"; + 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 FP8 prompt/state staging failed\n"; + model->release(model->owner); + return 1; + } + std::vector views(num_views); + for (int view = 0; view < num_views; ++view) { + views[view].struct_size = sizeof(frt_image_view); + views[view].pixel_format = FRT_RT_PIXEL_RGB8; + views[view].data = pixels[view]->data(); + views[view].bytes = pixels[view]->size(); + views[view].width = 224; + views[view].height = 224; + views[view].stride_bytes = 224 * 3; + } + if (model->verbs.set_input( + model->self, 2, views.data(), + views.size() * sizeof(frt_image_view), -1) != 0) { + std::cerr << "native FP8 image staging failed\n"; + model->release(model->owner); + return 1; + } + std::vector noise_f16(noise.size()); + for (std::size_t i = 0; i < noise.size(); ++i) { + noise_f16[i] = flashrt::modalities::float_to_float16(noise[i]); + } + if (!model->ports[3].buffer || + cudaMemcpy(frt_buffer_dptr(model->ports[3].buffer), noise_f16.data(), + noise_f16.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess || + model->verbs.step(model->self) != 0 || + cudaDeviceSynchronize() != cudaSuccess) { + std::cerr << "native FP8 replay failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + if (argc == 7) { + std::vector raw(noise.size()); + if (!model->ports[5].buffer || + cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), + raw.size() * sizeof(raw[0]), + cudaMemcpyDeviceToHost) != cudaSuccess) { + std::cerr << "native FP8 raw action download failed\n"; + model->release(model->owner); + return 1; + } + std::ofstream output(argv[6], std::ios::binary | std::ios::trunc); + output.write(reinterpret_cast(raw.data()), + static_cast( + raw.size() * sizeof(raw[0]))); + if (!output) { + std::cerr << "native FP8 raw action write failed\n"; + model->release(model->owner); + return 1; + } + } + 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)) { + std::cerr << "native FP8 action read failed\n"; + model->release(model->owner); + return 1; + } + for (float value : actions) { + if (!std::isfinite(value)) { + std::cerr << "native FP8 action is non-finite\n"; + model->release(model->owner); + return 1; + } + } + model->release(model->owner); + std::cout << "PASS native Thor FP8 calibration and runtime lifecycle\n"; + return 0; +} diff --git a/cpp/tests/test_pi05_c_api.cpp b/cpp/tests/test_pi05_c_api.cpp index e7e29360..77429103 100644 --- a/cpp/tests/test_pi05_c_api.cpp +++ b/cpp/tests/test_pi05_c_api.cpp @@ -51,6 +51,11 @@ int main() { return 0; } + assert(frt_pi05_calibration_sample_count_v1(nullptr) == 0); + assert(frt_pi05_calibration_create_v1("{}", 99.9, nullptr) == -1); + assert(std::strstr(frt_pi05_calibration_create_last_error_v1(), "out")); + frt_pi05_calibration_destroy_v1(nullptr); + frt_ctx ctx = frt_ctx_create(); assert(ctx); int sid = frt_ctx_stream(ctx, 0); diff --git a/cpp/tests/test_pi05_native_calibration.cpp b/cpp/tests/test_pi05_native_calibration.cpp new file mode 100644 index 00000000..6759528d --- /dev/null +++ b/cpp/tests/test_pi05_native_calibration.cpp @@ -0,0 +1,91 @@ +#include "flashrt/cpp/models/pi05/native_calibration.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +std::string temp_path() { + char path[] = "/tmp/frt_pi05_calibration_XXXXXX"; + const int fd = ::mkstemp(path); + assert(fd >= 0); + ::close(fd); + assert(::unlink(path) == 0); + return path; +} + +} // namespace + +int main() { + using flashrt::models::pi05::NativeCalibrationArtifact; + using flashrt::models::pi05::load_native_calibration_artifact; + using flashrt::models::pi05::reduce_native_calibration_samples; + using flashrt::models::pi05::save_native_calibration_artifact; + + std::vector reduced; + assert(reduce_native_calibration_samples( + {{1.0f, 10.0f}, {2.0f, 20.0f}, {4.0f, 40.0f}, + {8.0f, 80.0f}}, + 25.0, &reduced) + .ok_status()); + assert(reduced.size() == 2); + assert(std::fabs(reduced[0] - 1.75f) < 1e-6f); + assert(std::fabs(reduced[1] - 17.5f) < 1e-6f); + assert(reduce_native_calibration_samples({{3.0f}}, 99.9, &reduced) + .ok_status()); + assert(reduced == std::vector({3.0f})); + assert(!reduce_native_calibration_samples({}, 99.9, &reduced) + .ok_status()); + assert(!reduce_native_calibration_samples( + {{1.0f}, {1.0f, 2.0f}}, 99.9, &reduced) + .ok_status()); + + NativeCalibrationArtifact expected; + expected.hardware = "sm110"; + expected.weights_sha256 = std::string(64, 'a'); + expected.tokenizer_sha256 = std::string(64, 'b'); + expected.num_views = 2; + expected.max_prompt_tokens = 200; + expected.state_dim = 8; + expected.chunk_size = 10; + expected.num_steps = 10; + expected.vision_pool_factor = 1; + expected.sample_count = 8; + expected.percentile = 99.9; + expected.encoder_scales.resize(18 * 4); + expected.decoder_scales.resize(10 * 18 * 4); + for (std::size_t i = 0; i < expected.encoder_scales.size(); ++i) { + expected.encoder_scales[i] = 0.001f * static_cast(i + 1); + } + for (std::size_t i = 0; i < expected.decoder_scales.size(); ++i) { + expected.decoder_scales[i] = 0.0001f * static_cast(i + 1); + } + + const std::string path = temp_path(); + assert(save_native_calibration_artifact(path, expected).ok_status()); + NativeCalibrationArtifact loaded; + assert(load_native_calibration_artifact(path, &loaded).ok_status()); + assert(loaded.hardware == expected.hardware); + assert(loaded.weights_sha256 == expected.weights_sha256); + assert(loaded.tokenizer_sha256 == expected.tokenizer_sha256); + assert(loaded.num_views == expected.num_views); + assert(loaded.max_prompt_tokens == expected.max_prompt_tokens); + assert(loaded.state_dim == expected.state_dim); + assert(loaded.chunk_size == expected.chunk_size); + assert(loaded.num_steps == expected.num_steps); + assert(loaded.vision_pool_factor == expected.vision_pool_factor); + assert(loaded.sample_count == expected.sample_count); + assert(std::fabs(loaded.percentile - expected.percentile) < 1e-12); + assert(loaded.encoder_scales == expected.encoder_scales); + assert(loaded.decoder_scales == expected.decoder_scales); + assert(::unlink(path.c_str()) == 0); + + expected.weights_sha256 = "short"; + assert(!save_native_calibration_artifact(path, expected).ok_status()); + std::printf("PASS - Pi0.5 native calibration artifact\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 655d94f1..4885fa12 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -209,6 +209,25 @@ 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, ",\"precision\":\"nvfp4\"").c_str(), + &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "precision")); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, + ",\"calibration_path\":\"/missing/calibration.safetensors\"") + .c_str(), + &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), + "calibration_path")); + out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); assert(rc == -2); @@ -234,6 +253,13 @@ int main() { missing_key.c_str())); write_safetensors(root + "/model.safetensors"); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, ",\"max_frame_width\":0").c_str(), &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "max_frame")); + out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); assert(rc == -2); diff --git a/cpp/tests/test_pi05_native_workspace.cpp b/cpp/tests/test_pi05_native_workspace.cpp index 78c59586..459d73b4 100644 --- a/cpp/tests/test_pi05_native_workspace.cpp +++ b/cpp/tests/test_pi05_native_workspace.cpp @@ -4,6 +4,7 @@ #include #include +#include #include namespace { @@ -23,9 +24,11 @@ void check_ones(const flashrt::models::pi05::NativeWorkspaceBuffer& buffer) { 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)); - } + const std::uint16_t expected = + buffer.dtype == flashrt::modalities::DType::kFloat16 + ? flashrt::modalities::float_to_float16(1.0f) + : flashrt::modalities::float_to_bfloat16(1.0f); + for (std::uint16_t value : values) assert(value == expected); } } // namespace @@ -43,6 +46,11 @@ int main() { NativeWorkspaceConfig invalid; invalid.vision_pool_factor = 3; assert(!workspace.allocate(invalid).ok_status()); + invalid.vision_pool_factor = 1; + invalid.num_views = 3; + invalid.max_prompt_tokens = std::numeric_limits::max() - 769; + invalid.chunk_size = 2; + assert(!workspace.allocate(invalid).ok_status()); NativeWorkspaceConfig config; assert(workspace.allocate(config).ok_status()); assert(workspace.logical_size() == 35); @@ -131,6 +139,99 @@ int main() { std::vector({5, 50, 1024})); } frt_ctx_destroy(ctx); + + ctx = frt_ctx_create(); + assert(ctx); + { + NativeWorkspace workspace(ctx); + NativeWorkspaceConfig invalid; + invalid.flavor = NativeWorkspaceFlavor::kThorFp8; + invalid.vision_pool_factor = 2; + assert(!workspace.allocate(invalid).ok_status()); + invalid.vision_pool_factor = 1; + invalid.max_prompt_tokens = 199; + assert(!workspace.allocate(invalid).ok_status()); + + NativeWorkspaceConfig config; + config.flavor = NativeWorkspaceFlavor::kThorFp8; + config.enable_calibration = true; + assert(workspace.allocate(config).ok_status()); + assert(workspace.vision_sequence() == 512); + assert(workspace.encoder_sequence() == 712); + assert(workspace.total_keys() == 722); + assert(workspace.find("observation_images_normalized")->dtype == + flashrt::modalities::DType::kFloat16); + assert(workspace.find("encoder_x_fp8")->dtype == + flashrt::modalities::DType::kUInt8); + assert(workspace.find("encoder_logits")->shape == + std::vector({712 * 8, 722})); + assert(workspace.find("encoder_k_cache")->shape == + std::vector({18, 722, 256})); + assert(workspace.find("decoder_activation_scales")->shape == + std::vector({10, 18, 4})); + assert(workspace.find("encoder_sample_scales")); + assert(workspace.find("decoder_sample_scales")); + check_ones(*workspace.find("encoder_rms_ones")); + check_ones(*workspace.find("decoder_rms_ones")); + std::vector encoder_rope(712 * 256); + assert(cudaMemcpy( + encoder_rope.data(), + frt_buffer_dptr( + workspace.find("encoder_rope_weights")->buffer), + encoder_rope.size() * sizeof(encoder_rope[0]), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(encoder_rope[468 * 256 + 2 * 30 + 1] == 0xb7fdu); + assert(encoder_rope[624 * 256 + 2 * 34 + 1] == 0xb7fdu); + + const auto* prompt = workspace.find("prompt_embedding"); + std::vector row(2048); + for (std::size_t i = 0; i < row.size(); ++i) { + row[i] = flashrt::modalities::float_to_float16( + static_cast(i % 31)); + } + auto* prompt_base = static_cast( + frt_buffer_dptr(prompt->buffer)); + assert(cudaMemcpy(prompt_base + 4 * row.size() * sizeof(row[0]), + row.data(), row.size() * sizeof(row[0]), + cudaMemcpyHostToDevice) == cudaSuccess); + assert(workspace.set_fixed_prompt_length(5).ok_status()); + std::vector padded(row.size()); + assert(cudaMemcpy(padded.data(), + prompt_base + 5 * row.size() * sizeof(row[0]), + padded.size() * sizeof(padded[0]), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(padded == row); + const char* control_names[] = { + "attn_enc_seqused", "attn_dec_seqused", "attn_dec_devpos"}; + const std::int32_t expected_controls[] = {518, 528, 518}; + for (int i = 0; i < 3; ++i) { + std::int32_t value = 0; + assert(cudaMemcpy( + &value, + frt_buffer_dptr(workspace.find(control_names[i])->buffer), + sizeof(value), cudaMemcpyDeviceToHost) == cudaSuccess); + assert(value == expected_controls[i]); + } + const std::size_t allocations = workspace.allocation_count(); + for (int i = 0; i < 1000; ++i) { + assert(workspace.set_fixed_prompt_length(i % 200).ok_status()); + assert(workspace.allocation_count() == allocations); + } + + NativeDeviceWeightStore weights(ctx); + NativeF16Tensor 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_float16( + static_cast(i % 97) / 97.0f); + } + assert(weights.upload("vision_position_embedding", position) + .ok_status()); + assert(workspace.expand_vision_position_embedding(weights) + .ok_status()); + } + frt_ctx_destroy(ctx); std::printf("PASS - Pi0.5 native workspace\n"); return 0; } From 5df8978c43dc1c68d3806f4cbcc137c193e51402 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 11:39:43 -0400 Subject: [PATCH 66/83] docs(pi05): define native FP8 calibration contract --- .github/pull_request_template.md | 3 + CONTRIBUTING.md | 18 ++ USAGE.md | 28 +++ docs/architecture.md | 33 +-- docs/cpp_runtime_design.md | 50 ++-- docs/cpp_runtime_modalities.md | 28 +-- docs/mindon_pi05_integration.md | 38 +-- docs/native_model_runtime_producer.md | 20 ++ docs/pi05_io_contract.md | 115 ++++++--- docs/pi05_thor_native_fp8.md | 323 ++++++++++++++++++++++++++ docs/pr_review_checklist.md | 14 ++ docs/runtime_contract.md | 37 +-- 12 files changed, 593 insertions(+), 114 deletions(-) create mode 100644 docs/pi05_thor_native_fp8.md diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 6c5f513b..897ae578 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -17,9 +17,12 @@ - [ ] Focused tests cover success and rejection paths - [ ] Affected CUDA-off/hardware configurations were checked or disclosed - [ ] Numerical claims use a fixed, justified gate +- [ ] Calibration artifacts and cache invalidation are identity-complete +- [ ] Single/multi-sample calibration and named-input rejection are covered - [ ] 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 +- [ ] Model semantics remain model-local; frozen runtime/exec stay generic diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 328083a5..83d71495 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -247,6 +247,24 @@ calibration, or graph capture behavior, include a precision comparison: - action sanity check for quickstart-only paths - latency before/after for performance-sensitive changes +Native calibration APIs and artifacts have additional producer-boundary rules: + +- Keep model calibration sites, dimensions, camera names, prompt/state + semantics, and artifact schema under `cpp/models//`; do not add them + to the frozen runtime ABI or generic `exec/` mechanism. +- Keep dataset discovery, decoding, synchronization, and sampling policy in the + host. A calibration session accepts complete observations; it does not own a + dataset loader. +- Bind reusable artifacts to observed hardware, model/tokenizer content, + precision, fixed shapes, reducer/schema version, and any other input that can + change scale meaning. Include the artifact digest in producer identity when + changing it changes inference math. +- Test single-observation and repeated-observation reduction. Named multi-input + samples must reject missing/duplicate/unknown names and prove that caller + array order does not change semantic order. +- Use explicit fixed stochastic inputs for reference comparison. A generated + fallback must be deterministic, documented, and separately exercised. + ### Performance Measurement Use the right metric for the claim: diff --git a/USAGE.md b/USAGE.md index 55192325..715870a3 100644 --- a/USAGE.md +++ b/USAGE.md @@ -169,6 +169,34 @@ out the exact command. do I need to migrate?** A: No. The legacy path is still in the search list. +### Pi0.5 Native C++ Runtime + +Pi0.5 can also be loaded without a resident Python producer through +`frt_model_runtime_open_v1`. This is a separate deployment face from +`flash_rt.load_model()` and returns the backend-neutral +`frt_model_runtime_v1` contract. + +| Hardware | Native precision | Required backend | Calibration artifact | +|---|---|---|---| +| SM120 | BF16 | native FA2 + SentencePiece | no | +| Thor SM110 | FP8 E4M3 | Thor FP8/CUTLASS + SentencePiece | yes | + +The SM110 producer includes model-specific C APIs for single-view, +multi-view, and repeated dataset-observation calibration. The resulting +safetensors artifact is bound to hardware, checkpoint, tokenizer, fixed +shapes, sample count, and reducer policy; runtime identity also includes the +artifact SHA-256. Native C++ NVFP4 is not currently supported. Python FP8 and +NVFP4 routes keep their existing behavior. + +Native safetensors setup uses one direct mmap -> transform/quantize -> device +upload path and does not create an implicit weight-cache format. This is +independent of the Python JAX Orbax weight cache documented below. + +See [Pi0.5 Native C++ FP8 on Thor](docs/pi05_thor_native_fp8.md) for build +flags, configuration JSON, C calibration usage, camera-name rules, artifact +invalidation, runtime ports, and validation commands. The complete portable IO +contract is [Pi0.5 Native Model Runtime IO](docs/pi05_io_contract.md). + --- ## Quick Start diff --git a/docs/architecture.md b/docs/architecture.md index aa982dba..4b4ae235 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -9,22 +9,21 @@ ## 1. Runtime layers ``` -serving/ scenario hosts: sessions, schedulers, protocols, robot loops, - OpenAI-compatible HTTP/SSE, streaming audio, rollout policy - │ - ▼ -flash_rt/ Python frontends: weights, calibration, CUDA Graph capture, - model APIs and per-model pipeline composition - │ - ▼ -exec/ C ABI execution contract: Buffer / Graph / Plan, replay-time - mechanism for native hosts - │ - ▼ -csrc/ CUDA/C++ kernels and vendored attention/GEMM building blocks +flash_rt/ Python producers: weights, calibration, graph capture ─┐ +cpp/ native producers: checkpoint/tokenizer IO, capture ────┤ +csrc/ CUDA/C++ kernels and vendor building blocks ───────────┘ + │ + ▼ +runtime/ backend-neutral model hand-off: ports, stages, regions, identity + │ + ▼ +exec/ Buffer / Graph / Plan replay mechanism │ + ▼ +serving/ sessions, protocols, robot loops, and scheduling policy ``` -The serving layer is deliberately above the model runtime. It decides +Python and native producers are alternatives; neither is layered on top of the +other. The serving layer is deliberately above the model runtime. It decides session policy, streaming protocol, episode lifecycle, interrupt/reset behavior, and host language. The core runtime owns the graph and kernel mechanism. See [`../serving/README.md`](../serving/README.md), @@ -88,7 +87,11 @@ mechanism. See [`../serving/README.md`](../serving/README.md), └──────────────────────────────────────────────────────────────┘ ``` -The frontend is the only file you write per model. Everything below it is shared infrastructure. Everything above it is dispatch. +For the Python producer, the frontend is the primary per-model integration. +An optional Python-free producer instead lives under `cpp/models//` and +must present the same backend-neutral model-runtime contract. Model checkpoint +names, dimensions, prompt/state rules, and calibration sites stay model-local +in either producer; `runtime/` and `exec/` do not learn them. --- diff --git a/docs/cpp_runtime_design.md b/docs/cpp_runtime_design.md index 7f4b86fc..625b6b59 100644 --- a/docs/cpp_runtime_design.md +++ b/docs/cpp_runtime_design.md @@ -10,10 +10,9 @@ cannot promise. This document is the structure map; the interface reference is ## One struct, two producers Everything converges on `frt_model_runtime_v1` (the standard face of one -deployed, tickable model). The Python setup bridge produces it today; a native -model-runtime `.so` (`frt_model_runtime_open_v1`) produces the same struct -later. Consumers — FlashRT-Nexus, robot loops, FFI hosts — never change when -the producer does. +deployed, tickable model). Python setup bridges and native model-runtime shared +objects (`frt_model_runtime_open_v1`) produce the same struct. Consumers — +FlashRT-Nexus, robot loops, FFI hosts — never change when the producer does. The clean hybrid path is **verb override**: the setup producer exports the authoritative ports, stage DAG, graph streams, identity and fingerprint; a @@ -67,19 +66,22 @@ GROOT runtime would export its own model factory, and so on. That code owns the model's hot-path transforms: image normalization, state packing, action postprocess, and the names/shapes of public ports it supports. -The **hardware** is selected before the C++ runtime sees the model: the Python -or native setup producer chooses the hardware pipeline, captures the graphs, -allocates live buffers, calibrates precision-specific paths, and writes the -canonical identity/fingerprint. The C++ overlay then inherits those graph, -stream, stage, and buffer declarations with `frt_model_runtime_override_verbs`. +The **hardware** is selected by the producer, not by the ABI consumer. A Python +producer uses its hardware dispatch map. A native producer queries the active +device, resolves an explicitly supported precision/backend pair, captures the +graphs, allocates live buffers, calibrates precision-specific paths when +required, and writes observed hardware into identity. Neither path adds a +backend-kind field to the frozen structs. -So the expected setup shape is: +There are two setup shapes: -1. The hardware-specific pipeline builds a ready model instance. -2. `flash_rt/models//runtime_export.py` exports that instance as the - model family's standard `frt_model_runtime_v1` face. -3. `cpp/models//` overlays native hot verbs on that exact declaration. -4. Nexus or a robot loop consumes only the resulting model-runtime handle. +1. An adopted producer builds a ready model instance, exports its declarations, + and lets `cpp/models//` override only the hot verbs it implements. +2. A native producer under `cpp/models//` loads the checkpoint and + tokenizer, builds/captures its backend, and publishes the complete model + runtime directly through `frt_model_runtime_open_v1`. +3. In both cases Nexus or another host consumes only the resulting + `frt_model_runtime_v1` handle. If two hardware pipelines expose the same logical ports and stage DAG, they can share one native C++ overlay. If their visible contract differs, the difference @@ -112,11 +114,12 @@ Optional cuts are managed outside the C++ runtime under `flash_rt/subgraphs/`. See [`subgraph_stage_plans.md`](subgraph_stage_plans.md) for the customer registration and capture-hook workflow. -The C++ runtime does not parse manifests or hardcode split names. For Pi0.5, -`frt_pi05_model_runtime_create_over` inherits the producer's declarations and -maps only the public ports it implements (`images`, optional `noise`, -`actions`). `step` is convenience only: same-stream stage chains may replay -sequentially; cross-stream dependencies require a host scheduler. +The C++ runtime does not infer stage plans from manifests. For Pi0.5, the +adopted `frt_pi05_model_runtime_create_over` path inherits the producer's graph +and stage declarations. The native `frt_model_runtime_open_v1` producer +currently declares one `infer` stage. `step` is convenience only: same-stream +stage chains may replay sequentially; cross-stream dependencies require a host +scheduler. Pi0.5's default producer plan is: @@ -131,9 +134,12 @@ actions for the same inputs. It also exposes two IO faces over the same captured graphs: - `io="python"`: Python frontend hot loop; normalized tensors are SWAP ports. -- `io="native"`: native C++ hot loop; raw images/actions are STAGED and noise - remains a SWAP port. This is the face consumed by +- `io="native"`: adopted native C++ hot loop; raw images/actions are STAGED and + noise remains a SWAP port. This is the face consumed by `frt_pi05_model_runtime_create_over`. +- `io="native_v2"`: Python-free native producer; checkpoint loading, + tokenizer/prompt/state processing, calibration where required, and graph + capture are all producer-owned setup. The native `actions` port declares the logical output chunk delivered by `get_output`, not necessarily the raw model buffer layout. A Pi0.5 producer may diff --git a/docs/cpp_runtime_modalities.md b/docs/cpp_runtime_modalities.md index f3554e41..dda1c314 100644 --- a/docs/cpp_runtime_modalities.md +++ b/docs/cpp_runtime_modalities.md @@ -60,8 +60,8 @@ Model adapters live in `cpp/models//`: Pi0.5 is the first adapter: -- vision: `image`, `wrist_image`, `wrist_image_right` -> NHWC BF16 224x224, - normalized to `[-1, 1]`; +- vision: `image`, `wrist_image`, `wrist_image_right` -> NHWC 224x224 in the + producer-declared dtype, normalized to `[-1, 1]`; - action: `(chunk, 32)` model output -> first 7 robot dims, unnormalized by deployment stats. - `flashrt::models::pi05::RuntimeIo` binds those specs to concrete tensor @@ -82,23 +82,25 @@ Current Pi0.5 status: resize/normalize/cast directly into export device buffers; - conservative action staging path: device action buffer -> D2H -> CPU reference postprocess; -- native SM120 checkpoint loading, tokenizer/prompt staging, weight - materialization, and graph capture are implemented by the optional - `frt_model_runtime_open_v1` producer. They present the same model-runtime ABI - and remain FlashRT responsibilities, not Nexus features. +- native SM120 BF16 and SM110 FP8 checkpoint loading, tokenizer/prompt staging, + weight materialization, calibration where required, and graph capture are + implemented by the optional `frt_model_runtime_open_v1` producer. They + present the same model-runtime ABI and remain FlashRT responsibilities, not + Nexus features. -## CPU Reference First +## Reference First -The current implementation is a CPU reference path: +The portable semantic references are: - `preprocess_vision_cpu` - `postprocess_action_cpu` -This is intentional. It gives every CUDA/DMA/zero-copy fast path a golden -contract. The current vision device path already uses a CUDA -resize/normalize/cast kernel and is tested against the CPU reference. The -action device path is still conservative D2H staging because the postprocess is -small; it can be moved to CUDA without changing model adapters. +They give every CUDA/DMA/zero-copy fast path a golden contract. The vision +device path uses CUDA resize/normalize/cast and is tested against the CPU +reference at the declared output dtype. The action device path remains +conservative D2H staging because postprocess is small; it can move to CUDA +without changing model adapters. Native SM110 calibration and inference add +separate producer-parity gates on top of these modality references. ## Hot Path Rules diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 38ed51c7..35e41385 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -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: Native SM120 Producer +### Lane C: Native Producer Load a native FlashRT shared object and call: @@ -77,14 +77,25 @@ The returned struct must expose the same public model-runtime contract as the Python setup producer. The host and Nexus adoption code must not change when switching between Lane A and Lane C. -The current C++ shared object implements this symbol as a complete SM120 -native-v2 producer when built with CUDA kernels, native FA2, and SentencePiece. -It validates `io`, checkpoint/tokenizer paths, fixed prompt mode, capacities, -the complete 812-tensor inventory, and OpenPI or LeRobot action/state q01/q99 +The current C++ shared object implements this symbol as a complete native-v2 +producer on SM120 BF16 and SM110 FP8. Both require CUDA kernels and +SentencePiece; SM120 uses native FA2, while SM110 uses the Thor FP8/CUTLASS +backend and an identity-bound calibration artifact. The factory validates `io`, +precision, 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. +and returns the integrated model runtime. Missing backend/SentencePiece support +or unsupported hardware returns unsupported instead of publishing unusable +ports. + +On SM110, create the artifact with the model-specific calibration API before +opening the runtime, then pass it as `calibration_path`. One observation can +contain one, two, or three named camera frames; repeat observations for dataset +calibration. Camera synchronization and dataset policy stay in the Mindon host. +See [`pi05_thor_native_fp8.md`](pi05_thor_native_fp8.md) for exact build flags, +C API usage, artifact invalidation, and validation gates. Native C++ NVFP4 is +not currently advertised; Python precision routes remain independent. Use a Release build for MindOn deployment. Native startup includes full-content checkpoint hashing and CUDA graph capture in addition to safetensors parsing @@ -96,12 +107,13 @@ weight materialization execute concurrently, but the factory publishes no runtime until both have succeeded. The native loader maps the checkpoint read-only and directly emits final BF16 -device layouts from F32, BF16, or F16 source tensors. QKV interleave/concat, -RMS folding, patch permutation, transpose, and output scaling are fused into -that pass; there is no checkpoint-sized Python dictionary or chain of float -intermediates. `FLASHRT_PROFILE_NATIVE_SETUP=1` reports header, materialization, -workspace/style, input initialization, capture, stream, and total setup time. -This diagnostic is setup-only and does not change the runtime contract. +or F16/FP8 device layouts from F32, BF16, or F16 source tensors. QKV +interleave/concat, RMS folding, patch permutation, transpose, and output scaling +are fused into that pass; there is no checkpoint-sized Python dictionary or +chain of float intermediates. `FLASHRT_PROFILE_NATIVE_SETUP=1` reports header, +materialization, workspace/style, input initialization, capture, stream, and +total setup time. This diagnostic is setup-only and does not change the runtime +contract. ## No-HTTP C++ Host Shape diff --git a/docs/native_model_runtime_producer.md b/docs/native_model_runtime_producer.md index 72bfd342..0d9fd5d5 100644 --- a/docs/native_model_runtime_producer.md +++ b/docs/native_model_runtime_producer.md @@ -42,6 +42,26 @@ Manifest fields are discovery metadata, not a substitute for identity. A schema or restore change intentionally produces a new fingerprint and rejects old capsules. +## Calibration artifacts + +Calibration is producer setup, not a new runtime mechanism. Keep model sites, +tensor dimensions, camera names, state/prompt semantics, and artifact format in +`cpp/models//`. Generic loaders may parse a standard container, but +`runtime/`, `exec/`, and consumers must not interpret model calibration data. + +The host owns dataset traversal, decoding, synchronization, and sampling +policy. A model session consumes one complete observation per call and may +reduce repeated observations according to a documented policy. Named inputs +must be canonicalized before model math and reject missing, duplicate, or +unknown names atomically. + +An artifact must bind every fact that changes scale meaning: observed hardware, +model and tokenizer content digests, precision, fixed shapes, schema/reducer +version, and successful sample count. When artifact bytes change inference +math, include the artifact digest in producer identity. Loading incompatible +metadata is a hard setup error, never a warning or fallback recalibration in the +hot process. + ## Multiple producers and backends Python, native CUDA, CPU, llama.cpp, and future producers expose the same diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index e48cdfc2..04f8d965 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -100,15 +100,16 @@ Producer-owned preprocessing: - 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; +- output dtype is the exported tensor dtype: BF16 for the SM120 native backend + and F16 for the SM110 native FP8 backend; - 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). +errors. The native CUDA and CPU reference paths are bit-exact at the exported +rounding boundary 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 @@ -136,7 +137,7 @@ must be read from the port shape; host code must not assume `(10, 32)`. `actions` is the host-visible robot action chunk after producer-owned postprocess. Its declared dtype is F32 because that is the payload returned by -`get_output`; it does not expose the BF16 diffusion window as backing. +`get_output`; it does not expose the backend's diffusion window as backing. The logical output shape is: @@ -218,16 +219,20 @@ There are three supported integration lanes: the hot loop adopts `frt_model_runtime_v1` and runs through C++/Nexus. - Lane B: an adopted setup producer exposes real hot `prompt`/`state` STAGED ports and the C++ overlay owns their transforms. -- Lane C, current on RTX SM120: a C++ shared object implements +- Lane C, current on SM120 and SM110: a C++ shared object implements `frt_model_runtime_open_v1(config_json, &out)` and produces the same public struct without Python setup. The Pi0.5 C++ shared object exports `frt_model_runtime_open_v1` as a complete -native-v2 producer when built with CUDA kernels, native FA2, and SentencePiece. -Execution currently requires RTX SM120. The factory requires `io="native_v2"`, -`checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, -`max_prompt_tokens >= 200`, and a positive `state_dim`; `num_views`, `chunk`, -`num_steps`, and `vision_pool_factor` are optional fixed setup values. It parses +native-v2 producer when built with CUDA kernels, SentencePiece, and the backend +for the executing device: native FA2 for SM120 BF16 or the Thor FP8/CUTLASS +backend for SM110. The factory requires `io="native_v2"`, `checkpoint_path`, +`tokenizer_model_path`, `state_prompt_mode="fixed"`, `max_prompt_tokens >= 200`, +and a positive `state_dim`; `precision`, `num_views`, `chunk`, `num_steps`, +`vision_pool_factor`, and frame capacities are producer setup values. SM110 +FP8 additionally requires an identity-compatible `calibration_path`. See +[`pi05_thor_native_fp8.md`](pi05_thor_native_fp8.md) for the build, calibration, +artifact, and validation contract. 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, @@ -235,9 +240,15 @@ projectors, action projections, and time MLP. It also verifies action/state q01/ dimensions from either openpi `norm_stats.json` or LeRobot policy normalizer/unnormalizer safetensors. Safetensors tensor byte ranges must match dtype/shape, and normalization quantiles must be finite ordered pairs. Builds -without native FA2 or SentencePiece validate the config and return unsupported; -they do not advertise a runtime they cannot execute. The mmap and parsed tensor -views are setup-side assets; they never enter the model-runtime ABI or hot path. +without the selected backend 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 BF16/FA2 implementation details below describe the SM120 backend unless a +paragraph explicitly says otherwise. The SM110 backend uses the same public +ports and lifecycle with F16 tensor windows, producer-equivalent FP8 packing, +FP16 attention/setup math where required, and identity-bound activation scales. The native setup layer also carries CPU reference transforms matching the existing PyTorch producer: source BF16 rounding for vision/decoder weights, @@ -250,12 +261,12 @@ Build native producers with `CMAKE_BUILD_TYPE=Release` for deployment. The native source view accepts F32, BF16, and F16 safetensors without materializing a checkpoint-sized float dictionary. Type dispatch occurs once per tensor; aligned payloads use typed loops and valid unaligned safetensors ranges use -safe scalar reads. Plain copy/transpose preserves source BF16 bits. Conversion, -RMS fold, Q/K interleave, QKV concat/transpose, patch OIHW-to-HWIO layout, and -action scaling write the final BF16 payload directly. They do not create -rounded, permuted, or concatenated F32 intermediates. Tests cover every source -dtype, and real-checkpoint gates byte-compare all fused transform families with -the PyTorch multi-step reference. +alignment-safe scalar reads. Plain copy/transpose preserves source bits. +Conversion, RMS fold, Q/K interleave, QKV concat/transpose, patch OIHW-to-HWIO +layout, and action scaling write the backend's final BF16 or F16 payload +directly. They do not create rounded, permuted, or concatenated checkpoint-sized +F32 intermediates. Tests cover every source dtype, and real-checkpoint gates +byte-compare fused transform families with the matching shipped producer. Checkpoint identity remains a full-file SHA-256, not a path, timestamp, or partial-content surrogate. When OpenSSL is available at configure time the @@ -275,6 +286,14 @@ reference measurements, not a latency ABI. Use deployment system. Compare producer startup using the complete scope; a Python timer that stops after safetensors conversion is not the same metric. +On a representative Thor SM110 Release run with the same checkpoint size, +native setup measured 7.90 seconds and the complete open/infer/output/teardown +lifecycle measured 8.62 seconds. Weight transform, FP8 quantization, and upload +accounted for 7.61 seconds; workspace/style setup, warmup, and capture measured +94 ms, 163 ms, and 26 ms respectively. These measurements do not add a binary +weight cache: native safetensors loading remains one direct, identity-checked +path. + 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 @@ -423,11 +442,13 @@ 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. +includes observed hardware, precision, model/tokenizer SHA-256 values, prompt +mode, fixed shapes, and schedule parameters. SM110 FP8 identity also includes +the calibration artifact SHA-256. 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 @@ -466,9 +487,9 @@ 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. +gates use a SentencePiece-enabled build with either SM120 FA2 or SM110 Thor FP8. +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 @@ -478,13 +499,38 @@ The native factory lifecycle gate is: ``` /pi05_native_open_probe \ - + [calibration.safetensors] ``` 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. +SM110 FP8 calibration and runtime math are gated separately against the shipped +Torch producer: + +``` +python cpp/tests/gate_pi05_native_thor_fp8.py \ + --probe /pi05_native_thor_fp8_probe \ + --checkpoint \ + --tokenizer \ + --artifact \ + --samples 1 --views 1 + +python cpp/tests/gate_pi05_native_thor_fp8.py \ + --probe /pi05_native_thor_fp8_probe \ + --checkpoint \ + --tokenizer \ + --artifact \ + --samples 3 --views 2 +``` + +The two-view gate submits cameras in reverse order to verify name-based +canonicalization. It also rejects incomplete/duplicate camera sets and invalid +noise without committing a sample, exercises deterministic generated noise, +and requires all encoder scales, decoder scales, and raw actions to be +bit-exact. Dataset iteration and camera synchronization remain host policy. + Python and C++ native-v2 producers must also publish identical canonical port/stage/region records (their producer identity and fingerprints remain different): @@ -586,7 +632,7 @@ nsys profile --trace=cuda \ --capture-range=cudaProfilerApi --capture-range-end=stop \ -o \ /pi05_native_open_probe \ - + [calibration.safetensors] nsys stats --report cuda_api_trace --format csv \ .nsys-rep > .csv python cpp/tests/gate_pi05_hot_allocator.py \ @@ -610,13 +656,14 @@ device update): ``` FLASHRT_HOT_STATE_UPDATES=1000 FLASHRT_HOT_STATE_P99_US=1000 \ /pi05_native_open_probe \ - + [calibration.safetensors] ``` 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. +updates, and requires the graph variant count to remain one. A reference RTX +5090 SM120 run measured p50/p99/max of 39.31/41.70/43.44 microseconds. A +reference Thor SM110 FP8 run measured 129.13/265.66/541.68 microseconds. Both +remain 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, diff --git a/docs/pi05_thor_native_fp8.md b/docs/pi05_thor_native_fp8.md new file mode 100644 index 00000000..3ca0f11a --- /dev/null +++ b/docs/pi05_thor_native_fp8.md @@ -0,0 +1,323 @@ +# Pi0.5 Native C++ FP8 On Thor + +This guide covers the Python-free Pi0.5 producer on NVIDIA Thor SM110. It +loads a safetensors checkpoint, calibrates FP8 activation scales, captures one +fixed-shape CUDA Graph, and returns `frt_model_runtime_v1` from +`frt_model_runtime_open_v1`. + +The implementation is model-specific by design. The stable runtime ABI remains +model- and backend-neutral; Pi0.5 checkpoint names, dimensions, prompt/state +semantics, calibration sites, and camera names stay under `cpp/models/pi05/`. +Dataset discovery, decoding, synchronization, and sampling policy remain host +responsibilities. + +## Support Matrix + +| Producer | Hardware | Precision | Calibration | Native IO dtype | +|---|---|---|---|---| +| Native C++ | SM110 | FP8 E4M3 | Required | F16 | +| Native C++ | SM120 | BF16 | Not used | BF16 | +| Python | Backend-specific | Existing FP8/BF16/NVFP4 routes | Existing Python contract | Producer-declared | + +`precision="auto"` selects FP8 E4M3 on SM110 and BF16 on SM120. Production +configuration should normally specify the intended precision explicitly. +Native C++ NVFP4 is not implemented by this producer; `precision="nvfp4"` is +rejected. This does not change the independent Python NVFP4 path. + +## Build + +The SM110 backend requires CUDA kernels, CUDA staging, exec, SentencePiece, +and a compatible CUTLASS checkout. It has been validated with CUDA 13.0. + +```bash +export BUILD_DIR="$PWD/cpp/build-thor" +export CUTLASS_DIR="$PWD/third_party/cutlass" + +cmake -S cpp -B "$BUILD_DIR" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=110 \ + -DFLASHRT_CPP_WITH_EXEC=ON \ + -DFLASHRT_CPP_WITH_CUDA_STAGING=ON \ + -DFLASHRT_CPP_WITH_CUDA_KERNELS=ON \ + -DFLASHRT_CPP_WITH_SENTENCEPIECE=ON \ + -DFLASHRT_CPP_WITH_FA2=OFF \ + -DFLASHRT_CPP_WITH_THOR_FP8=ON \ + -DFLASHRT_CPP_CUTLASS_DIR="$CUTLASS_DIR" +cmake --build "$BUILD_DIR" -j +``` + +Configure fails if the Thor backend is enabled without CUDA kernels, CUDA +staging, exec, or CUTLASS. Build `Release` for deployment. A Debug build is +useful for the same test matrix but is not a latency reference. + +## Native Configuration + +Calibration and runtime open consume the same model configuration. Calibration +does not need `calibration_path`; runtime open requires it on SM110. + +```json +{ + "io": "native_v2", + "checkpoint_path": "/path/to/pi05-checkpoint", + "tokenizer_model_path": "/path/to/paligemma_tokenizer.model", + "state_prompt_mode": "fixed", + "precision": "fp8_e4m3fn", + "max_prompt_tokens": 200, + "state_dim": 8, + "num_views": 2, + "chunk": 10, + "num_steps": 10, + "vision_pool_factor": 1, + "max_frame_width": 1280, + "max_frame_height": 720 +} +``` + +The checkpoint supplies state/action q01 and q99 normalization metadata. All +shape fields are fixed setup values. Increasing a capacity or changing view, +chunk, schedule, tokenizer, checkpoint, precision, or calibration artifact +produces a different deployment identity where applicable. + +## Calibration API + +The model-specific API is declared in +`flashrt/cpp/models/pi05/c_api.h`: + +```c +frt_pi05_calibration_session* session = NULL; +int rc = frt_pi05_calibration_create_v1( + calibration_config_json, 99.9, &session); + +frt_pi05_vision_frame frames[2] = {0}; +frames[0].struct_size = sizeof(frames[0]); +frames[0].name = "image"; +frames[0].data = base_rgb; +frames[0].bytes = base_bytes; +frames[0].width = base_width; +frames[0].height = base_height; +frames[0].stride_bytes = base_stride; +frames[0].pixel_format = FRT_PI05_PIXEL_RGB8; + +frames[1].struct_size = sizeof(frames[1]); +frames[1].name = "wrist_image"; +frames[1].data = wrist_rgb; +frames[1].bytes = wrist_bytes; +frames[1].width = wrist_width; +frames[1].height = wrist_height; +frames[1].stride_bytes = wrist_stride; +frames[1].pixel_format = FRT_PI05_PIXEL_RGB8; + +frt_pi05_calibration_sample_v1 sample = {0}; +sample.struct_size = sizeof(sample); +sample.prompt = task_prompt; +sample.state = state_f32; +sample.n_state = state_dim; +sample.frames = frames; +sample.n_frames = 2; +sample.noise = noise_f32; +sample.n_noise = chunk * 32; + +rc = frt_pi05_calibration_observe_v1(session, &sample); +rc = frt_pi05_calibration_finalize_v1( + session, "/path/to/pi05-sm110-fp8.safetensors"); +frt_pi05_calibration_destroy_v1(session); +``` + +Check every return code. Creation errors are available through +`frt_pi05_calibration_create_last_error_v1`; session errors are available +through `frt_pi05_calibration_last_error_v1`. + +Calls on one session are serialized; the handle does not provide internal +concurrent-observe semantics. Hosts that calibrate independent streams in +parallel use independent sessions and artifacts. + +### Camera Sets + +One `observe` call is one complete observation. Supported camera names are the +prefix of this model-specific order: + +1. `image` +2. `wrist_image` +3. `wrist_image_right` + +A one-view session supplies only `image`. A two-view session supplies `image` +and `wrist_image`. The input array may be in any order; calibration canonicalizes +it by name. Missing, duplicate, or unknown names reject the entire observation +without increasing `sample_count`. + +Calibration image payloads use host `u8/RGB8/HWC` frames. Each frame must fit +the setup-time width/height capacity. `timestamp_ns` is carried by the frame +descriptor but FlashRT does not synchronize cameras; the host must submit a +coherent observation. + +The native model-runtime `images` STAGED payload uses `frt_image_view[]`, which +has no camera-name field and follows the producer-declared positional order. +Do not infer one payload convention from the other. + +### Single And Dataset Calibration + +Call `observe` once for a single-observation artifact. Call it repeatedly for +multi-timestamp or dataset calibration, then call `finalize` once. FlashRT +reduces each activation site independently with NumPy-compatible linear +percentile semantics. `99.9` is the validated dataset setting; `100.0` selects +the observed maximum. + +Dataset iteration is intentionally not part of the runtime. The host chooses +episodes, tasks, timestamps, image decoding, and synchronization. Samples +should represent the deployed prompt, state, camera, and action distribution. +Broadening calibration data can improve coverage while reducing resolution for +common activations, so every artifact still needs an end-to-end action gate. + +`noise` is optional F32 `[chunk, 32]`. Supply fixed noise when comparing with a +reference producer. If it is omitted with `n_noise=0`, FlashRT generates +deterministic normal F16 noise from `noise_seed + successful_sample_index`. +Malformed or non-finite state/noise payloads are rejected before the sample is +committed. + +Calibration materializes the model and runs uncaptured reference forwards. It +is an offline/setup operation, not a control-loop operation. Reuse the produced +artifact until an identity input or calibration policy changes. + +## Artifact Contract + +The calibration file is an atomically published safetensors artifact with two +F32 tensors: + +- `encoder_scales`: 72 values; +- `decoder_scales`: `num_steps * 18 * 4` values. + +Metadata binds the artifact to: + +- schema, model, precision, tensor dtype, and reducer version; +- observed SM architecture; +- full checkpoint and tokenizer SHA-256 digests; +- view count, prompt capacity, state dimension, chunk, denoise steps, and + vision pooling; +- successful sample count and percentile. + +Runtime open rejects incompatible metadata, non-positive/non-finite scales, +shape mismatches, checkpoint/tokenizer digest mismatches, and hardware +mismatches. The runtime identity also includes the calibration file SHA-256, +so changing scale bytes intentionally changes the fingerprint and prevents an +old capsule from restoring into different math. + +Do not edit or merge artifacts manually. Re-run calibration from representative +observations. + +## Runtime Open + +Add the artifact path to the configuration and load the producer through the +standard model-runtime symbol: + +```json +{ + "io": "native_v2", + "checkpoint_path": "/path/to/pi05-checkpoint", + "tokenizer_model_path": "/path/to/paligemma_tokenizer.model", + "state_prompt_mode": "fixed", + "precision": "fp8_e4m3fn", + "calibration_path": "/path/to/pi05-sm110-fp8.safetensors", + "max_prompt_tokens": 200, + "state_dim": 8, + "num_views": 2, + "chunk": 10, + "num_steps": 10, + "vision_pool_factor": 1 +} +``` + +Resolve `FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL` as +`frt_model_runtime_open_v1_fn`, or link the producer library and call +`frt_model_runtime_open_v1` directly. The returned runtime publishes one +`infer` stage and these ordered ports: + +| Port | Update | SM110 dtype | Payload | +|---|---|---|---| +| `prompt` | STAGED | U8 | UTF-8 task text | +| `state` | STAGED | F32 | raw proprioception | +| `images` | STAGED | F16 | host `frt_image_view[]` transformed into the captured window | +| `noise` | SWAP | F16 | device `[chunk, 32]` | +| `actions` | STAGED | F32 | host `[chunk, robot_action_dim]` | +| `actions_raw` | SWAP | F16 | device `[chunk, 32]` alias | + +Prompt formatting, state normalization/discretization, tokenization, embedding, +vision preprocessing, and action postprocessing remain producer-owned. Nexus +or another consumer moves declared payloads and schedules declared stages; it +does not interpret Pi0.5 semantics. + +## Loading Path + +Native setup mmaps `model.safetensors` read-only and accepts F32, BF16, or F16 +source tensors. Independent CPU transforms run in bounded parallel tasks while +typed final F16/FP8/scales are uploaded into context-owned buffers. Full-file +checkpoint hashing runs concurrently. Valid unaligned safetensors payloads use +alignment-safe reads. + +Materialization uses up to eight worker threads by default, bounded by the +layer count. `FLASHRT_NATIVE_WEIGHT_WORKERS=` is a setup-only diagnostic and +tuning override (`1..64`, still bounded by layer count); invalid values keep the +default. It changes scheduling only, not tensor order, bytes, identity, or +runtime math. Validate any deployment override with the same byte/action gates. + +There is no implicit second weight-cache format for safetensors. This avoids +another invalidation and serialization boundary while keeping the mathematical +source path identical to the shipped producer. The OS page cache may improve +repeated file reads but is not part of the FlashRT contract. + +On a representative SM110 run with a 14.47 GB F32 checkpoint, native setup was +7.90 seconds and the complete open/infer/output/teardown lifecycle was 8.62 +seconds. Of setup time, 7.61 seconds was weight transform/quantization/upload, +94 ms was workspace/style setup, 163 ms was warmup, and 26 ms was graph capture. +These are reference measurements, not latency ABI guarantees. + +## Validation + +Run CPU/CUDA-off tests in addition to SM110 Release and Debug builds: + +```bash +ctest --test-dir "$BUILD_DIR" --output-on-failure +``` + +The real-checkpoint gate compares native calibration and inference with the +shipped Torch producer using fixed observations and noise: + +```bash +python cpp/tests/gate_pi05_native_thor_fp8.py \ + --probe "$BUILD_DIR/pi05_native_thor_fp8_probe" \ + --checkpoint "$CHECKPOINT_DIR" \ + --tokenizer "$TOKENIZER_MODEL" \ + --artifact "$CALIBRATION_FILE" \ + --samples 1 --views 1 + +python cpp/tests/gate_pi05_native_thor_fp8.py \ + --probe "$BUILD_DIR/pi05_native_thor_fp8_probe" \ + --checkpoint "$CHECKPOINT_DIR" \ + --tokenizer "$TOKENIZER_MODEL" \ + --artifact "$CALIBRATION_FILE" \ + --samples 3 --views 2 +``` + +The two-view probe deliberately submits reversed camera order and checks +duplicate/incomplete/unknown names, non-RGB input rejection, malformed noise, +deterministic generated noise, artifact loading, runtime identity, one graph +variant, finite logical actions, and teardown. The reference gates require all +72 encoder scales, all 720 decoder scales for ten steps, and all 320 raw action +values to be bit-exact. + +For the complete service loop, profile the CUDA profiler range around 1,000 +iterations of prompt/state/image/noise update, replay, and action output: + +```bash +FLASHRT_PROFILE_REPLAYS=1000 FLASHRT_PROFILE_SERVICE_LOOP=1 \ +nsys profile --trace=cuda \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o "$HOT_REPORT" \ + "$BUILD_DIR/pi05_native_open_probe" \ + "$CHECKPOINT_DIR" "$TOKENIZER_MODEL" "$CALIBRATION_FILE" +``` + +The validated trace contained exactly 1,000 graph launches and no CUDA device +allocation/free, CUDA host allocation/registration, mempool, virtual-memory, +graph instantiation, or capture API in the measured range. A separate +1,000-update prompt/state gate measured 266 microseconds p99 against a 1 ms +limit, with one graph variant throughout. diff --git a/docs/pr_review_checklist.md b/docs/pr_review_checklist.md index 3d9af35e..489ae0bc 100644 --- a/docs/pr_review_checklist.md +++ b/docs/pr_review_checklist.md @@ -454,6 +454,16 @@ Required: - Quantized paths need cosine, token-match, or domain-specific validation against the relevant reference. - Cache reuse must document exactly what is cached and when it is invalidated. +- Native calibration artifacts must bind hardware, model/tokenizer content, + precision, fixed shapes, reducer/schema version, and sample count. If artifact + bytes change inference math, their digest must participate in producer + identity/fingerprint. +- Dataset traversal and camera synchronization stay host policy. Model-local + calibration code consumes complete observations and canonicalizes named + inputs before model math. +- Single-observation and repeated-observation reducers need independent gates; + multi-input gates must cover reordered, missing, duplicate, and unknown + inputs. Reference parity uses explicit fixed stochastic inputs. Blockers: @@ -464,6 +474,10 @@ Blockers: - Speed is reported without correctness. - Low correctness is accepted without comparing to the right reference noise floor. +- A calibration artifact can be reused after checkpoint, tokenizer, hardware, + shape, precision, or reducer semantics change without a hard identity error. +- Generic runtime/exec code learns a model's calibration sites, camera names, + dataset format, or dimensions. Minimum correctness evidence: diff --git a/docs/runtime_contract.md b/docs/runtime_contract.md index f4420b16..4640a4de 100644 --- a/docs/runtime_contract.md +++ b/docs/runtime_contract.md @@ -15,13 +15,12 @@ orchestration is the consumer's job. ``` producer (owns model + capture) consumer (owns loop + state policy) ───────────────────────────────── ────────────────────────────────── -today: e.g. FlashRT-Nexus capsule host, - Python setup/capture a robot loop, a server shell +Python setup/capture e.g. FlashRT-Nexus capsule host, flash_rt/runtime/export.py └─ _flashrt_runtime.Builder ──┐ -later (same struct, host unchanged): │ adopt(export*) +native setup (same struct): │ adopt(export*) native model runtime .so ├──► frt_runtime_export_v1 ◄──┘ - frt_runtime_open_v1(config,&out)─┘ │ ctx, streams[], graphs[], + frt_model_runtime_open_v1(...)───┘ │ ctx, streams[], graphs[], │ buffers[], capsule_regions[], │ fingerprint/identity/manifest, │ owner + retain/release @@ -42,9 +41,9 @@ flash_rt/runtime/export.py Python producer: RuntimeExport / build_export() ## The contract, in five rules -1. **One struct, two producers.** Today Python fills it in-process; later a - native model runtime `.so` exports `frt_runtime_open_v1` (symbol name is in - the header) and fills the *same* struct. Consumers never change. +1. **One struct, two producers.** Python fills it in-process; a native model + runtime `.so` exports `frt_model_runtime_open_v1` (symbol name is in the + model-runtime header) and fills the *same* struct. Consumers never change. 2. **Consumers see handles, never internals.** No Python, torch, model code, or kernel headers cross this boundary — only `frt_*` handles, POD descriptors, and strings owned by the export. @@ -138,16 +137,20 @@ Nexus should not implement or own these rules. It adopts `frt_runtime_export_v1` and drives snapshot/restore/replay; FlashRT model runtimes prepare inputs and decode outputs. -Pi0.5 is the reference C++ model runtime under `cpp/models/pi05/`. The current -implementation is the adopted-export path: setup/capture can still be produced -by Python, while the native runtime owns vision prepare, graph replay dispatch, -action decode, and export lifetime. It includes a CUDA vision -resize/normalize/cast path for device buffers, with CPU reference tests, and a -conservative D2H action staging path. The `flashrt_cpp_pi05_c` target exposes a -small C host ABI around this path so C/Python/serving shells can drive the -adopted export without including C++ classes. A future pure C++ checkpoint -loader/tokenizer/capture path must produce the same `frt_runtime_export_v1`, so -Nexus and serving hosts do not change. +Pi0.5 is the reference C++ model runtime under `cpp/models/pi05/`. It supports +both producer forms. The adopted-export path accepts Python- or native-produced +graphs and overlays native vision/action/prompt/state verbs. The Python-free +path loads safetensors and SentencePiece assets, selects an explicitly built +hardware backend, captures one native graph, and publishes the complete +`frt_model_runtime_v1` through `frt_model_runtime_open_v1`. SM120 currently +uses BF16 plus native FA2; SM110 uses FP8 E4M3 plus an identity-bound native +calibration artifact. Both expose the same backend-neutral contract, so Nexus +and serving hosts do not change. + +The `flashrt_cpp_pi05_c` target also exposes the model-specific host/calibration +C API. These functions own Pi0.5 semantic transforms; they do not extend the +frozen runtime or exec structs. See [`pi05_io_contract.md`](pi05_io_contract.md) +and [`pi05_thor_native_fp8.md`](pi05_thor_native_fp8.md). ## Extending the ABI From 72490a9977e8df5eaead833f37f7b94f2364deea Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 13:37:18 -0400 Subject: [PATCH 67/83] refactor(pi05): make native implementation headers private --- cpp/CMakeLists.txt | 16 ++++++++++++---- .../include/flashrt/cpp/models/pi05/io.h | 0 .../cpp/models/pi05/native_bf16_forward.h | 0 .../flashrt/cpp/models/pi05/native_calibration.h | 0 .../cpp/models/pi05/native_device_weights.h | 0 .../flashrt/cpp/models/pi05/native_graph_owner.h | 0 .../cpp/models/pi05/native_graph_runtime.h | 0 .../cpp/models/pi05/native_kernel_driver.h | 0 .../cpp/models/pi05/native_quantization.h | 0 .../flashrt/cpp/models/pi05/native_rope.h | 0 .../cpp/models/pi05/native_rtx_attention.h | 0 .../models/pi05/native_rtx_attention_driver.h | 0 .../cpp/models/pi05/native_style_precompute.h | 0 .../pi05/native_thor_calibration_session.h | 0 .../cpp/models/pi05/native_thor_fp8_forward.h | 0 .../cpp/models/pi05/native_thor_graph_owner.h | 0 .../cpp/models/pi05/native_thor_kernel_driver.h | 0 .../models/pi05/native_thor_style_precompute.h | 0 .../pi05/native_thor_weight_materializer.h | 0 .../cpp/models/pi05/native_weight_materializer.h | 0 .../flashrt/cpp/models/pi05/native_weight_ops.h | 0 .../cpp/models/pi05/native_weight_packer.h | 0 .../flashrt/cpp/models/pi05/native_weights.h | 0 .../flashrt/cpp/models/pi05/native_workspace.h | 0 .../flashrt/cpp/models/pi05/prompt_embed.h | 0 .../flashrt/cpp/models/pi05/prompt_format.h | 0 .../include/flashrt/cpp/models/pi05/runtime.h | 0 .../include/flashrt/cpp/models/pi05/spec.h | 0 cpp/models/pi05/src/native_model_runtime.cpp | 2 +- 29 files changed, 13 insertions(+), 5 deletions(-) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/io.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_bf16_forward.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_calibration.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_device_weights.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_graph_owner.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_graph_runtime.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_kernel_driver.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_quantization.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_rope.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_rtx_attention.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_style_precompute.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_weight_materializer.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_weight_ops.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_weight_packer.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_weights.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/native_workspace.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/prompt_embed.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/prompt_format.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/runtime.h (100%) rename cpp/models/pi05/{ => internal}/include/flashrt/cpp/models/pi05/spec.h (100%) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index cc8f2da4..d272c996 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -218,7 +218,8 @@ 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) + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include) target_link_libraries(flashrt_cpp_pi05 PUBLIC flashrt_cpp_modalities flashrt_cpp_vla flashrt_cpp_loader) @@ -235,8 +236,9 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) ${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 + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include + ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include + ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels) target_link_libraries(flashrt_cpp_pi05_kernels PUBLIC flashrt_cpp_pi05 CUDA::cublas CUDA::cublasLt CUDA::cudart) @@ -315,13 +317,19 @@ add_library(flashrt_cpp_pi05_c SHARED 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) + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include) if(FLASHRT_CPP_WITH_CUDA_KERNELS) target_link_libraries(flashrt_cpp_pi05_c PRIVATE flashrt_cpp_pi05_kernels) endif() if(BUILD_TESTING) + # PI0.5 implementation headers are build-only. Tests may inspect internal + # invariants without publishing those headers to downstream consumers. + include_directories( + ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include) + # The test executables use assert() as their fail-fast check, including # around setup calls whose side effects are required by later checks. # Keep those checks active for Release test builds as well. diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/io.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/io.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_bf16_forward.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_calibration.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_calibration.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_device_weights.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_owner.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_graph_runtime.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_kernel_driver.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_quantization.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_quantization.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_quantization.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rope.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rope.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rope.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rope.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_style_precompute.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_style_precompute.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_style_precompute.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_style_precompute.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_materializer.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_materializer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_materializer.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_ops.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_ops.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_ops.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_packer.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weight_packer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_packer.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weights.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_weights.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weights.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/native_workspace.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_embed.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_embed.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_embed.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_format.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/prompt_format.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_format.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/runtime.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/runtime.h diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/spec.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h similarity index 100% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/spec.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 706d1bea..94adf830 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -6,12 +6,12 @@ #include "config_map.h" #include "flashrt/cpp/loader/sha256.h" #include "flashrt/cpp/models/pi05/model_runtime.h" +#include "flashrt/cpp/models/pi05/native_calibration.h" #include "flashrt/cpp/models/pi05/native_graph_runtime.h" #if defined(FLASHRT_CPP_WITH_FA2) #include "flashrt/cpp/models/pi05/native_graph_owner.h" #endif #if defined(FLASHRT_CPP_WITH_THOR_FP8) -#include "flashrt/cpp/models/pi05/native_calibration.h" #include "flashrt/cpp/models/pi05/native_thor_graph_owner.h" #endif From 3280fb3275892cfd54a717482738bc146717c457 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 13:53:14 -0400 Subject: [PATCH 68/83] refactor(cpp): move pi05 build ownership into model --- cpp/CMakeLists.txt | 390 +-------------------------- cpp/models/pi05/CMakeLists.txt | 152 +++++++++++ cpp/models/pi05/tests/CMakeLists.txt | 238 ++++++++++++++++ cpp/tests/CMakeLists.txt | 24 ++ 4 files changed, 419 insertions(+), 385 deletions(-) create mode 100644 cpp/models/pi05/CMakeLists.txt create mode 100644 cpp/models/pi05/tests/CMakeLists.txt create mode 100644 cpp/tests/CMakeLists.txt diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d272c996..ac2f928d 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -31,6 +31,7 @@ set(FLASHRT_CPP_FA2_LIBRARY "" CACHE FILEPATH set(FLASHRT_CPP_CUTLASS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../third_party/cutlass" CACHE PATH "Path to CUTLASS for the native Thor SM110 FP8 backend") + if(FLASHRT_CPP_WITH_CUDA_STAGING) find_package(CUDAToolkit REQUIRED) endif() @@ -102,6 +103,7 @@ if(FLASHRT_CPP_WITH_SENTENCEPIECE) 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) @@ -144,10 +146,6 @@ if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") # optimized builds; otherwise the host compiler may contract lerps to FMAs. set_property(SOURCE modalities/src/vision_cpu.cpp APPEND PROPERTY COMPILE_OPTIONS -ffp-contract=off) - # Weight fusion mirrors producer-side add/cast/multiply tensor operations. - # Keep each float operation rounded before the final FP16 conversion. - set_property(SOURCE models/pi05/src/native_weight_ops.cpp APPEND PROPERTY - COMPILE_OPTIONS -ffp-contract=off) endif() target_include_directories(flashrt_cpp_modalities PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/modalities/include) @@ -192,386 +190,8 @@ if(OpenSSL_FOUND) target_link_libraries(flashrt_cpp_loader PRIVATE OpenSSL::Crypto) endif() -set(FLASHRT_CPP_PI05_SRCS - models/pi05/src/spec.cpp - models/pi05/src/native_weights.cpp - models/pi05/src/native_weight_ops.cpp - models/pi05/src/native_device_weights.cpp - models/pi05/src/native_weight_packer.cpp - models/pi05/src/native_weight_materializer.cpp - models/pi05/src/native_calibration.cpp - models/pi05/src/native_workspace.cpp - models/pi05/src/native_rtx_attention.cpp - models/pi05/src/prompt_format.cpp - models/pi05/src/prompt_embed.cpp - models/pi05/src/io.cpp - models/pi05/src/runtime.cpp) -if(FLASHRT_CPP_WITH_CUDA_KERNELS) - list(APPEND FLASHRT_CPP_PI05_SRCS - models/pi05/src/native_quantization.cu - models/pi05/src/native_rope.cu) -else() - list(APPEND FLASHRT_CPP_PI05_SRCS - models/pi05/src/native_quantization_unavailable.cpp - models/pi05/src/native_rope_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 - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include) -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_bf16_forward.cpp - models/pi05/src/native_kernel_driver.cu - models/pi05/src/native_style_precompute.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/activation.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/elementwise.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/norm.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/patch_embed.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/rope.cu) - target_include_directories(flashrt_cpp_pi05_kernels - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include - ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels) - target_link_libraries(flashrt_cpp_pi05_kernels - PUBLIC flashrt_cpp_pi05 CUDA::cublas CUDA::cublasLt CUDA::cudart) - set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES - CUDA_STANDARD 17 - CUDA_STANDARD_REQUIRED ON) - target_compile_options(flashrt_cpp_pi05_kernels PRIVATE - $<$: - --expt-relaxed-constexpr -O3 - --ftz=true --prec-div=false --prec-sqrt=false>) - if(FLASHRT_CPP_WITH_THOR_FP8) - add_library(flashrt_cpp_pi05_thor_precise OBJECT - models/pi05/src/native_thor_setup_ops.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/attention/fmha_fp16_strided.cu) - target_include_directories(flashrt_cpp_pi05_thor_precise PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/attention - ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha - ${FLASHRT_CPP_CUTLASS_DIR}/include - ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) - target_compile_options(flashrt_cpp_pi05_thor_precise PRIVATE - $<$: - --expt-relaxed-constexpr --expt-extended-lambda -O3>) - set_target_properties(flashrt_cpp_pi05_thor_precise PROPERTIES - CUDA_ARCHITECTURES 110a - CUDA_STANDARD 17 - CUDA_STANDARD_REQUIRED ON - POSITION_INDEPENDENT_CODE ON) - target_sources(flashrt_cpp_pi05_kernels PRIVATE - $) - target_sources(flashrt_cpp_pi05_kernels PRIVATE - models/pi05/src/native_thor_kernel_driver.cu - models/pi05/src/native_thor_weight_materializer.cpp - models/pi05/src/native_thor_fp8_forward.cpp - models/pi05/src/native_thor_graph_owner.cpp - models/pi05/src/native_thor_calibration_session.cpp - models/pi05/src/native_thor_style_precompute.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/gemm/cutlass_sm100.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/attention_cublas.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/decoder_fused.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/quantize.cu - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/kernels/softmax.cu) - target_include_directories(flashrt_cpp_pi05_kernels PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc/attention - ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha - ${FLASHRT_CPP_CUTLASS_DIR}/include - ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) - target_compile_definitions(flashrt_cpp_pi05_kernels - PUBLIC FLASHRT_CPP_WITH_THOR_FP8=1) - target_compile_options(flashrt_cpp_pi05_kernels PRIVATE - $<$: - --expt-extended-lambda --use_fast_math - -gencode=arch=compute_110a,code=sm_110a>) - set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES - CUDA_ARCHITECTURES 110a - CUDA_RESOLVE_DEVICE_SYMBOLS ON) - endif() - if(FLASHRT_CPP_WITH_FA2) - target_sources(flashrt_cpp_pi05_kernels PRIVATE - models/pi05/src/native_graph_owner.cpp - models/pi05/src/native_rtx_attention_driver.cu) - target_include_directories(flashrt_cpp_pi05_kernels PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/../csrc) - target_link_libraries(flashrt_cpp_pi05_kernels - PUBLIC flashrt_cpp_fa2_external) - target_compile_definitions(flashrt_cpp_pi05_kernels - PUBLIC FLASHRT_CPP_WITH_FA2=1) - endif() -endif() - -add_library(flashrt_cpp_pi05_c SHARED - models/pi05/src/c_api.cpp - models/pi05/src/model_runtime.cpp - models/pi05/src/native_thor_calibration_c_api.cpp - models/pi05/src/native_model_runtime.cpp - models/pi05/src/native_open.cpp) -target_link_libraries(flashrt_cpp_pi05_c - PUBLIC flashrt_cpp_pi05 flashrt_runtime) -target_include_directories(flashrt_cpp_pi05_c - PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/include - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include) -if(FLASHRT_CPP_WITH_CUDA_KERNELS) - target_link_libraries(flashrt_cpp_pi05_c - PRIVATE flashrt_cpp_pi05_kernels) -endif() - +set(FLASHRT_CPP_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") if(BUILD_TESTING) - # PI0.5 implementation headers are build-only. Tests may inspect internal - # invariants without publishing those headers to downstream consumers. - include_directories( - ${CMAKE_CURRENT_SOURCE_DIR}/models/pi05/internal/include) - - # The test executables use assert() as their fail-fast check, including - # around setup calls whose side effects are required by later checks. - # Keep those checks active for Release test builds as well. - if(MSVC) - add_compile_options($<$:/UNDEBUG>) - else() - add_compile_options($<$:-UNDEBUG>) - endif() - - add_executable(test_safetensors_loader tests/test_safetensors_loader.cpp) - target_link_libraries(test_safetensors_loader PRIVATE flashrt_cpp_loader) - add_test(NAME safetensors_loader COMMAND test_safetensors_loader) - - add_executable(test_sha256 tests/test_sha256.cpp) - target_link_libraries(test_sha256 PRIVATE flashrt_cpp_loader) - add_test(NAME sha256 COMMAND test_sha256) - - add_executable(test_cpp_modalities tests/test_modalities.cpp) - target_link_libraries(test_cpp_modalities - PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) - add_test(NAME cpp_modalities COMMAND test_cpp_modalities) - - add_executable(test_text_modalities tests/test_text_modalities.cpp) - target_link_libraries(test_text_modalities PRIVATE flashrt_cpp_modalities) - add_test(NAME text_modalities COMMAND test_text_modalities) - - add_executable(test_text_tokenizer tests/test_text_tokenizer.cpp) - target_link_libraries(test_text_tokenizer PRIVATE flashrt_cpp_modalities) - add_test(NAME text_tokenizer COMMAND test_text_tokenizer) - - add_executable(test_pi05_runtime tests/test_pi05_runtime.cpp) - target_link_libraries(test_pi05_runtime - PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) - add_test(NAME pi05_runtime COMMAND test_pi05_runtime) - - add_executable(test_pi05_native_weight_ops - tests/test_pi05_native_weight_ops.cpp) - target_link_libraries(test_pi05_native_weight_ops PRIVATE flashrt_cpp_pi05) - add_test(NAME pi05_native_weight_ops COMMAND test_pi05_native_weight_ops) - - add_executable(test_pi05_native_calibration - tests/test_pi05_native_calibration.cpp) - target_link_libraries(test_pi05_native_calibration PRIVATE flashrt_cpp_pi05) - add_test(NAME pi05_native_calibration COMMAND test_pi05_native_calibration) - - add_executable(pi05_native_weight_probe - tests/pi05_native_weight_probe.cpp) - target_link_libraries(pi05_native_weight_probe PRIVATE flashrt_cpp_pi05) - - if(FLASHRT_CPP_WITH_CUDA_KERNELS) - add_executable(test_pi05_native_quantization - tests/test_pi05_native_quantization.cpp) - target_link_libraries(test_pi05_native_quantization PRIVATE flashrt_cpp_pi05) - add_test(NAME pi05_native_quantization - COMMAND test_pi05_native_quantization) - - add_executable(pi05_native_quant_probe - tests/pi05_native_quant_probe.cpp) - target_link_libraries(pi05_native_quant_probe PRIVATE flashrt_cpp_pi05) - - add_executable(test_pi05_native_weight_packer - tests/test_pi05_native_weight_packer.cpp) - target_link_libraries(test_pi05_native_weight_packer - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_weight_packer - COMMAND test_pi05_native_weight_packer) - - add_executable(test_pi05_native_kernel_driver - tests/test_pi05_native_kernel_driver.cpp) - target_link_libraries(test_pi05_native_kernel_driver - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_kernel_driver - COMMAND test_pi05_native_kernel_driver) - - add_executable(test_pi05_native_forward_primitives - tests/test_pi05_native_forward_primitives.cpp) - target_link_libraries(test_pi05_native_forward_primitives - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_forward_primitives - COMMAND test_pi05_native_forward_primitives) - - add_executable(test_pi05_native_bf16_forward - tests/test_pi05_native_bf16_forward.cpp) - target_link_libraries(test_pi05_native_bf16_forward - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_bf16_forward - COMMAND test_pi05_native_bf16_forward) - - add_executable(pi05_native_encoder_qkv_probe - tests/pi05_native_encoder_qkv_probe.cpp) - target_link_libraries(pi05_native_encoder_qkv_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(test_pi05_native_style_precompute - tests/test_pi05_native_style_precompute.cpp) - target_link_libraries(test_pi05_native_style_precompute - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_style_precompute - COMMAND test_pi05_native_style_precompute) - - add_executable(pi05_native_style_probe - tests/pi05_native_style_probe.cpp) - target_link_libraries(pi05_native_style_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(test_pi05_native_workspace - tests/test_pi05_native_workspace.cpp) - target_link_libraries(test_pi05_native_workspace - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_workspace - COMMAND test_pi05_native_workspace) - - add_executable(test_pi05_native_rtx_attention - tests/test_pi05_native_rtx_attention.cpp) - target_link_libraries(test_pi05_native_rtx_attention - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_rtx_attention - COMMAND test_pi05_native_rtx_attention) - - if(FLASHRT_CPP_WITH_FA2) - add_executable(pi05_native_encoder_layer_probe - tests/pi05_native_encoder_layer_probe.cpp) - target_link_libraries(pi05_native_encoder_layer_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_encoder_probe - tests/pi05_native_encoder_probe.cpp) - target_link_libraries(pi05_native_encoder_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_vision_probe - tests/pi05_native_vision_probe.cpp) - target_link_libraries(pi05_native_vision_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_diffusion_probe - tests/pi05_native_diffusion_probe.cpp) - target_link_libraries(pi05_native_diffusion_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_graph_probe - tests/pi05_native_graph_probe.cpp) - target_link_libraries(pi05_native_graph_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - if(FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_e2e_probe - tests/pi05_native_e2e_probe.cpp) - target_link_libraries(pi05_native_e2e_probe - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - - add_executable(pi05_native_dlopen_probe - tests/pi05_native_dlopen_probe.cpp) - target_include_directories(pi05_native_dlopen_probe - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime/include - ${CMAKE_CURRENT_SOURCE_DIR}/../exec/include) - target_link_libraries(pi05_native_dlopen_probe - PRIVATE ${CMAKE_DL_LIBS}) - endif() - - add_executable(test_pi05_native_rtx_attention_driver - tests/test_pi05_native_rtx_attention_driver.cpp) - target_link_libraries(test_pi05_native_rtx_attention_driver - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_rtx_attention_driver - COMMAND test_pi05_native_rtx_attention_driver) - endif() - - if(FLASHRT_CPP_WITH_SENTENCEPIECE AND - (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) - 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() - - 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() - - if(FLASHRT_CPP_WITH_THOR_FP8 AND FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_thor_fp8_probe - tests/pi05_native_thor_fp8_probe.cpp) - target_link_libraries(pi05_native_thor_fp8_probe - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - endif() - - add_executable(pi05_native_rope_probe - tests/pi05_native_rope_probe.cpp) - target_link_libraries(pi05_native_rope_probe - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - endif() - - add_executable(test_pi05_prompt_format tests/test_pi05_prompt_format.cpp) - target_link_libraries(test_pi05_prompt_format PRIVATE flashrt_cpp_pi05) - add_test(NAME pi05_prompt_format COMMAND test_pi05_prompt_format) - - add_executable(test_pi05_prompt_embed tests/test_pi05_prompt_embed.cpp) - target_link_libraries(test_pi05_prompt_embed PRIVATE flashrt_cpp_pi05) - add_test(NAME pi05_prompt_embed COMMAND test_pi05_prompt_embed) - - if(FLASHRT_CPP_WITH_CUDA_STAGING) - add_executable(test_device_staging tests/test_device_staging.cpp) - target_link_libraries(test_device_staging - PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities CUDA::cudart) - add_test(NAME device_staging COMMAND test_device_staging) - - add_executable(test_pi05_native_device_weights - tests/test_pi05_native_device_weights.cpp) - target_link_libraries(test_pi05_native_device_weights - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_device_weights - COMMAND test_pi05_native_device_weights) - - add_executable(test_pi05_native_weight_materializer - tests/test_pi05_native_weight_materializer.cpp) - target_link_libraries(test_pi05_native_weight_materializer - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_weight_materializer - COMMAND test_pi05_native_weight_materializer) - - add_executable(test_pi05_c_api tests/test_pi05_c_api.cpp) - target_link_libraries(test_pi05_c_api - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - add_test(NAME pi05_c_api COMMAND test_pi05_c_api) - - add_executable(test_pi05_model_runtime tests/test_pi05_model_runtime.cpp) - target_link_libraries(test_pi05_model_runtime - PRIVATE flashrt_cpp_pi05_c flashrt_exec flashrt_runtime CUDA::cudart) - add_test(NAME pi05_model_runtime COMMAND test_pi05_model_runtime) - - add_executable(test_pi05_native_open tests/test_pi05_native_open.cpp) - target_link_libraries(test_pi05_native_open - PRIVATE flashrt_cpp_pi05_c flashrt_runtime) - if((FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8) 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() + add_subdirectory(tests) endif() +add_subdirectory(models/pi05) diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt new file mode 100644 index 00000000..8ee8b3fa --- /dev/null +++ b/cpp/models/pi05/CMakeLists.txt @@ -0,0 +1,152 @@ +# PI0.5 owns its model sources, backend variants, public C face, and tests. +# Generic native mechanisms remain in the parent cpp targets. + +set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") + +set(_pi05_public_include "${CMAKE_CURRENT_SOURCE_DIR}/include") +set(_pi05_internal_include "${CMAKE_CURRENT_SOURCE_DIR}/internal/include") + +set(_pi05_sources + src/spec.cpp + src/native_weights.cpp + src/native_weight_ops.cpp + src/native_device_weights.cpp + src/native_weight_packer.cpp + src/native_weight_materializer.cpp + src/native_calibration.cpp + src/native_workspace.cpp + src/native_rtx_attention.cpp + src/prompt_format.cpp + src/prompt_embed.cpp + src/io.cpp + src/runtime.cpp) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + list(APPEND _pi05_sources + src/native_quantization.cu + src/native_rope.cu) +else() + list(APPEND _pi05_sources + src/native_quantization_unavailable.cpp + src/native_rope_unavailable.cpp) +endif() + +add_library(flashrt_cpp_pi05 STATIC ${_pi05_sources}) +target_include_directories(flashrt_cpp_pi05 + PUBLIC ${_pi05_public_include} + PRIVATE ${_pi05_internal_include}) +target_link_libraries(flashrt_cpp_pi05 + PUBLIC flashrt_cpp_modalities flashrt_cpp_vla flashrt_cpp_loader) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + # Preserve producer-side add/cast/multiply operation boundaries. + set_property(SOURCE src/native_weight_ops.cpp APPEND PROPERTY + COMPILE_OPTIONS -ffp-contract=off) +endif() + +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + add_library(flashrt_cpp_pi05_kernels STATIC + src/native_bf16_forward.cpp + src/native_kernel_driver.cu + src/native_style_precompute.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/activation.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/elementwise.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/norm.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/patch_embed.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/rope.cu) + target_include_directories(flashrt_cpp_pi05_kernels + PRIVATE ${_pi05_public_include} + ${_pi05_internal_include} + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels) + target_link_libraries(flashrt_cpp_pi05_kernels + PUBLIC flashrt_cpp_pi05 CUDA::cublas CUDA::cublasLt CUDA::cudart) + set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON) + target_compile_options(flashrt_cpp_pi05_kernels PRIVATE + $<$: + --expt-relaxed-constexpr -O3 + --ftz=true --prec-div=false --prec-sqrt=false>) + + if(FLASHRT_CPP_WITH_THOR_FP8) + add_library(flashrt_cpp_pi05_thor_precise OBJECT + src/native_thor_setup_ops.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention/fmha_fp16_strided.cu) + target_include_directories(flashrt_cpp_pi05_thor_precise PRIVATE + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention + ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha + ${FLASHRT_CPP_CUTLASS_DIR}/include + ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) + target_compile_options(flashrt_cpp_pi05_thor_precise PRIVATE + $<$: + --expt-relaxed-constexpr --expt-extended-lambda -O3>) + set_target_properties(flashrt_cpp_pi05_thor_precise PROPERTIES + CUDA_ARCHITECTURES 110a + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + POSITION_INDEPENDENT_CODE ON) + + target_sources(flashrt_cpp_pi05_kernels PRIVATE + $ + src/native_thor_kernel_driver.cu + src/native_thor_weight_materializer.cpp + src/native_thor_fp8_forward.cpp + src/native_thor_graph_owner.cpp + src/native_thor_calibration_session.cpp + src/native_thor_style_precompute.cpp + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/cutlass_sm100.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/attention_cublas.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/decoder_fused.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/quantize.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/softmax.cu) + target_include_directories(flashrt_cpp_pi05_kernels PRIVATE + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention + ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha + ${FLASHRT_CPP_CUTLASS_DIR}/include + ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) + target_compile_definitions(flashrt_cpp_pi05_kernels + PUBLIC FLASHRT_CPP_WITH_THOR_FP8=1) + target_compile_options(flashrt_cpp_pi05_kernels PRIVATE + $<$: + --expt-extended-lambda --use_fast_math + -gencode=arch=compute_110a,code=sm_110a>) + set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES + CUDA_ARCHITECTURES 110a + CUDA_RESOLVE_DEVICE_SYMBOLS ON) + endif() + + if(FLASHRT_CPP_WITH_FA2) + target_sources(flashrt_cpp_pi05_kernels PRIVATE + src/native_graph_owner.cpp + src/native_rtx_attention_driver.cu) + target_include_directories(flashrt_cpp_pi05_kernels PRIVATE + ${FLASHRT_CPP_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 + src/c_api.cpp + src/model_runtime.cpp + src/native_thor_calibration_c_api.cpp + src/native_model_runtime.cpp + src/native_open.cpp) +target_link_libraries(flashrt_cpp_pi05_c + PUBLIC flashrt_cpp_pi05 flashrt_runtime) +target_include_directories(flashrt_cpp_pi05_c + PUBLIC ${_pi05_public_include} + PRIVATE ${_pi05_internal_include}) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + target_link_libraries(flashrt_cpp_pi05_c + PRIVATE flashrt_cpp_pi05_kernels) +endif() + +if(BUILD_TESTING) + add_subdirectory(tests) +endif() diff --git a/cpp/models/pi05/tests/CMakeLists.txt b/cpp/models/pi05/tests/CMakeLists.txt new file mode 100644 index 00000000..73ae8715 --- /dev/null +++ b/cpp/models/pi05/tests/CMakeLists.txt @@ -0,0 +1,238 @@ +# PI0.5 contract, backend, lifecycle, and opt-in diagnostic targets. + +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") +set(_pi05_test_source_dir "${FLASHRT_CPP_SOURCE_DIR}/tests") +include_directories("${_pi05_internal_include}") +if(MSVC) + add_compile_options($<$:/UNDEBUG>) +else() + add_compile_options($<$:-UNDEBUG>) +endif() + +add_executable(test_cpp_modalities + ${_pi05_test_source_dir}/test_modalities.cpp) +target_link_libraries(test_cpp_modalities + PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) +add_test(NAME cpp_modalities COMMAND test_cpp_modalities) + +add_executable(test_pi05_runtime + ${_pi05_test_source_dir}/test_pi05_runtime.cpp) +target_link_libraries(test_pi05_runtime + PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) +add_test(NAME pi05_runtime COMMAND test_pi05_runtime) + +add_executable(test_pi05_native_weight_ops + ${_pi05_test_source_dir}/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(test_pi05_native_calibration + ${_pi05_test_source_dir}/test_pi05_native_calibration.cpp) +target_link_libraries(test_pi05_native_calibration PRIVATE flashrt_cpp_pi05) +add_test(NAME pi05_native_calibration COMMAND test_pi05_native_calibration) + +add_executable(pi05_native_weight_probe + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/pi05_native_quant_probe.cpp) + target_link_libraries(pi05_native_quant_probe PRIVATE flashrt_cpp_pi05) + + add_executable(test_pi05_native_weight_packer + ${_pi05_test_source_dir}/test_pi05_native_weight_packer.cpp) + target_link_libraries(test_pi05_native_weight_packer + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_weight_packer + COMMAND test_pi05_native_weight_packer) + + add_executable(test_pi05_native_kernel_driver + ${_pi05_test_source_dir}/test_pi05_native_kernel_driver.cpp) + target_link_libraries(test_pi05_native_kernel_driver + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_kernel_driver + COMMAND test_pi05_native_kernel_driver) + + add_executable(test_pi05_native_forward_primitives + ${_pi05_test_source_dir}/test_pi05_native_forward_primitives.cpp) + target_link_libraries(test_pi05_native_forward_primitives + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_forward_primitives + COMMAND test_pi05_native_forward_primitives) + + add_executable(test_pi05_native_bf16_forward + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/test_pi05_native_workspace.cpp) + target_link_libraries(test_pi05_native_workspace + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_workspace + COMMAND test_pi05_native_workspace) + + add_executable(test_pi05_native_rtx_attention + ${_pi05_test_source_dir}/test_pi05_native_rtx_attention.cpp) + target_link_libraries(test_pi05_native_rtx_attention + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_rtx_attention + COMMAND test_pi05_native_rtx_attention) + + if(FLASHRT_CPP_WITH_FA2) + add_executable(pi05_native_encoder_layer_probe + ${_pi05_test_source_dir}/pi05_native_encoder_layer_probe.cpp) + target_link_libraries(pi05_native_encoder_layer_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_encoder_probe + ${_pi05_test_source_dir}/pi05_native_encoder_probe.cpp) + target_link_libraries(pi05_native_encoder_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_vision_probe + ${_pi05_test_source_dir}/pi05_native_vision_probe.cpp) + target_link_libraries(pi05_native_vision_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_diffusion_probe + ${_pi05_test_source_dir}/pi05_native_diffusion_probe.cpp) + target_link_libraries(pi05_native_diffusion_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_graph_probe + ${_pi05_test_source_dir}/pi05_native_graph_probe.cpp) + target_link_libraries(pi05_native_graph_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_native_e2e_probe + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/pi05_native_dlopen_probe.cpp) + target_include_directories(pi05_native_dlopen_probe + PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../runtime/include + ${FLASHRT_CPP_SOURCE_DIR}/../exec/include) + target_link_libraries(pi05_native_dlopen_probe + PRIVATE ${CMAKE_DL_LIBS}) + endif() + + add_executable(test_pi05_native_rtx_attention_driver + ${_pi05_test_source_dir}/test_pi05_native_rtx_attention_driver.cpp) + target_link_libraries(test_pi05_native_rtx_attention_driver + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_rtx_attention_driver + COMMAND test_pi05_native_rtx_attention_driver) + endif() + + if(FLASHRT_CPP_WITH_SENTENCEPIECE AND + (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) + add_executable(pi05_native_open_probe + ${_pi05_test_source_dir}/pi05_native_open_probe.cpp) + target_link_libraries(pi05_native_open_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + endif() + + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_tokenizer_corpus_probe + ${_pi05_test_source_dir}/pi05_tokenizer_corpus_probe.cpp) + target_link_libraries(pi05_tokenizer_corpus_probe + PRIVATE flashrt_cpp_pi05) + endif() + + if(FLASHRT_CPP_WITH_THOR_FP8 AND FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_native_thor_fp8_probe + ${_pi05_test_source_dir}/pi05_native_thor_fp8_probe.cpp) + target_link_libraries(pi05_native_thor_fp8_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + endif() + + add_executable(pi05_native_rope_probe + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/test_pi05_prompt_format.cpp) +target_link_libraries(test_pi05_prompt_format PRIVATE flashrt_cpp_pi05) +add_test(NAME pi05_prompt_format COMMAND test_pi05_prompt_format) + +add_executable(test_pi05_prompt_embed + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/test_device_staging.cpp) + target_link_libraries(test_device_staging + PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities CUDA::cudart) + add_test(NAME device_staging COMMAND test_device_staging) + + add_executable(test_pi05_native_device_weights + ${_pi05_test_source_dir}/test_pi05_native_device_weights.cpp) + target_link_libraries(test_pi05_native_device_weights + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + add_test(NAME pi05_native_device_weights + COMMAND test_pi05_native_device_weights) + + add_executable(test_pi05_native_weight_materializer + ${_pi05_test_source_dir}/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 + ${_pi05_test_source_dir}/test_pi05_c_api.cpp) + target_link_libraries(test_pi05_c_api + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + add_test(NAME pi05_c_api COMMAND test_pi05_c_api) + + add_executable(test_pi05_model_runtime + ${_pi05_test_source_dir}/test_pi05_model_runtime.cpp) + 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 + ${_pi05_test_source_dir}/test_pi05_native_open.cpp) + target_link_libraries(test_pi05_native_open + PRIVATE flashrt_cpp_pi05_c flashrt_runtime) + if((FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8) 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() diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt new file mode 100644 index 00000000..21fc5f7a --- /dev/null +++ b/cpp/tests/CMakeLists.txt @@ -0,0 +1,24 @@ +# Tests for model-independent native loader and modality mechanisms. + +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") +if(MSVC) + add_compile_options($<$:/UNDEBUG>) +else() + add_compile_options($<$:-UNDEBUG>) +endif() + +add_executable(test_safetensors_loader test_safetensors_loader.cpp) +target_link_libraries(test_safetensors_loader PRIVATE flashrt_cpp_loader) +add_test(NAME safetensors_loader COMMAND test_safetensors_loader) + +add_executable(test_sha256 test_sha256.cpp) +target_link_libraries(test_sha256 PRIVATE flashrt_cpp_loader) +add_test(NAME sha256 COMMAND test_sha256) + +add_executable(test_text_modalities test_text_modalities.cpp) +target_link_libraries(test_text_modalities PRIVATE flashrt_cpp_modalities) +add_test(NAME text_modalities COMMAND test_text_modalities) + +add_executable(test_text_tokenizer test_text_tokenizer.cpp) +target_link_libraries(test_text_tokenizer PRIVATE flashrt_cpp_modalities) +add_test(NAME text_tokenizer COMMAND test_text_tokenizer) From d4a0240c52213049152a45f42fc1ba8aa3a0d17c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 15:38:52 -0400 Subject: [PATCH 69/83] feat(cpp): add native context and action graphs --- cpp/models/pi05/CMakeLists.txt | 1 + .../cpp/models/pi05/native_graph_owner.h | 27 +-- .../cpp/models/pi05/native_graph_runtime.h | 54 +++++- .../cpp/models/pi05/native_thor_graph_owner.h | 27 +-- .../include/flashrt/cpp/models/pi05/spec.h | 1 + cpp/models/pi05/src/native_graph_owner.cpp | 125 +++++------- cpp/models/pi05/src/native_graph_runtime.cpp | 170 +++++++++++++++++ cpp/models/pi05/src/native_model_runtime.cpp | 39 +++- cpp/models/pi05/src/native_open.cpp | 15 +- cpp/models/pi05/src/native_open_internal.h | 1 + .../pi05/src/native_thor_graph_owner.cpp | 134 +++++-------- cpp/tests/gate_pi05_native_thor_fp8.py | 17 +- cpp/tests/pi05_native_graph_probe.cpp | 30 ++- cpp/tests/pi05_native_open_probe.cpp | 179 +++++++++++++++--- cpp/tests/pi05_native_thor_fp8_probe.cpp | 86 +++++++-- cpp/tests/test_pi05_native_open.cpp | 15 +- 16 files changed, 670 insertions(+), 251 deletions(-) create mode 100644 cpp/models/pi05/src/native_graph_runtime.cpp diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index 8ee8b3fa..6602e33e 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -47,6 +47,7 @@ endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_library(flashrt_cpp_pi05_kernels STATIC src/native_bf16_forward.cpp + src/native_graph_runtime.cpp src/native_kernel_driver.cu src/native_style_precompute.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h index 2321254e..4275ff62 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h @@ -22,10 +22,12 @@ class NativeGraphOwner final : public NativeGraphRuntime { NativeGraphOwner(const NativeGraphOwner&) = delete; NativeGraphOwner& operator=(const NativeGraphOwner&) = delete; - frt_ctx context() const override { return ctx_; } - frt_graph infer_graph() const override { return infer_graph_; } - int stream_id() const override { return stream_id_; } - void* native_stream() const override { return replay_stream_; } + frt_ctx context() const override { return graphs_.context(); } + frt_graph graph(NativeGraphKind kind) const override { + return graphs_.graph(kind); + } + int stream_id() const override { return graphs_.stream_id(); } + void* native_stream() const override { return graphs_.native_stream(); } const NativeGraphConfig& config() const { return config_; } NativeDeviceWeightStore& weights() override { return weights_; } const NativeDeviceWeightStore& weights() const override { return weights_; } @@ -35,16 +37,19 @@ class NativeGraphOwner final : public NativeGraphRuntime { const NativeRtxAttentionWorkspace& attention() const { return attention_; } modalities::Status set_prompt_length(int prompt_tokens) override; - int replay() const; - modalities::Status synchronize() const; + int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const override; + modalities::Status synchronize() const override; 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); + modalities::Status record(NativeGraphKind kind, void* stream); + modalities::Status record_context(void* stream); + modalities::Status record_action(void* stream); + static modalities::Status record_graph( + void* user, NativeGraphKind kind, void* stream); - frt_ctx ctx_ = nullptr; + NativeGraphCatalog graphs_; NativeGraphConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; @@ -52,10 +57,6 @@ class NativeGraphOwner final : public NativeGraphRuntime { 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 diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h index a2776179..ce2957ba 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h @@ -4,6 +4,9 @@ #include "flashrt/cpp/models/pi05/native_device_weights.h" #include "flashrt/cpp/models/pi05/native_workspace.h" +#include +#include + namespace flashrt { namespace models { namespace pi05 { @@ -16,12 +19,59 @@ struct NativeGraphConfig { int vision_pool_factor = 1; }; +enum class NativeGraphKind : std::size_t { + kInfer = 0, + kDecodeOnly = 1, + kContext = 2, + kCount = 3, +}; + +class NativeGraphCatalog { +public: + using RecordFn = modalities::Status (*)( + void* owner, NativeGraphKind kind, void* stream); + + explicit NativeGraphCatalog(frt_ctx ctx) : ctx_(ctx) {} + ~NativeGraphCatalog(); + + NativeGraphCatalog(const NativeGraphCatalog&) = delete; + NativeGraphCatalog& operator=(const NativeGraphCatalog&) = delete; + + modalities::Status capture( + NativeGraphKind kind, const NativeWorkspace& workspace, + std::initializer_list bindings, + RecordFn record, void* owner); + modalities::Status create_replay_stream(); + + frt_ctx context() const { return ctx_; } + frt_graph graph(NativeGraphKind kind) const; + int stream_id() const { return stream_id_; } + void* native_stream() const { return replay_stream_; } + int replay(NativeGraphKind kind) const; + modalities::Status synchronize() const; + + static const char* name(NativeGraphKind kind); + +private: + struct CaptureCall; + static void record_graph(void* user, void* stream); + + frt_ctx ctx_ = nullptr; + frt_graph graphs_[static_cast(NativeGraphKind::kCount)] = {}; + void* replay_stream_ = nullptr; + int stream_id_ = -1; +}; + +modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, + void* stream); + class NativeGraphRuntime { public: virtual ~NativeGraphRuntime() = default; virtual frt_ctx context() const = 0; - virtual frt_graph infer_graph() const = 0; + virtual frt_graph graph(NativeGraphKind kind) const = 0; + frt_graph infer_graph() const { return graph(NativeGraphKind::kInfer); } virtual int stream_id() const = 0; virtual void* native_stream() const = 0; virtual NativeDeviceWeightStore& weights() = 0; @@ -29,6 +79,8 @@ class NativeGraphRuntime { virtual NativeWorkspace& workspace() = 0; virtual const NativeWorkspace& workspace() const = 0; virtual modalities::Status set_prompt_length(int prompt_tokens) = 0; + virtual int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const = 0; + virtual modalities::Status synchronize() const = 0; }; } // namespace pi05 diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h index b3393739..b57c3ee5 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h @@ -26,28 +26,33 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { NativeThorGraphOwner(const NativeThorGraphOwner&) = delete; NativeThorGraphOwner& operator=(const NativeThorGraphOwner&) = delete; - frt_ctx context() const override { return ctx_; } - frt_graph infer_graph() const override { return infer_graph_; } - int stream_id() const override { return stream_id_; } - void* native_stream() const override { return replay_stream_; } + frt_ctx context() const override { return graphs_.context(); } + frt_graph graph(NativeGraphKind kind) const override { + return graphs_.graph(kind); + } + int stream_id() const override { return graphs_.stream_id(); } + void* native_stream() const override { return graphs_.native_stream(); } NativeDeviceWeightStore& weights() override { return weights_; } const NativeDeviceWeightStore& weights() const override { return weights_; } NativeWorkspace& workspace() override { return workspace_; } const NativeWorkspace& workspace() const override { return workspace_; } modalities::Status set_prompt_length(int prompt_tokens) override; - int replay() const; - modalities::Status synchronize() const; + int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const override; + modalities::Status synchronize() const override; private: NativeThorGraphOwner(frt_ctx ctx, const NativeGraphConfig& config); modalities::Status initialize( const std::string& checkpoint_path, const NativeCalibrationArtifact& calibration); - modalities::Status record(void* stream); - static void record_graph(void* user, void* stream); + modalities::Status record(NativeGraphKind kind, void* stream); + modalities::Status record_context(void* stream); + modalities::Status record_action(void* stream); + static modalities::Status record_graph( + void* user, NativeGraphKind kind, void* stream); - frt_ctx ctx_ = nullptr; + NativeGraphCatalog graphs_; NativeGraphConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; @@ -55,10 +60,6 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { NativeThorFp8Forward forward_; NativeThorWeightScales weight_scales_; std::vector encoder_alphas_; - frt_graph infer_graph_ = nullptr; - void* replay_stream_ = nullptr; - int stream_id_ = -1; - modalities::Status capture_status_; }; } // namespace pi05 diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h index 2866cadd..23913960 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h @@ -14,6 +14,7 @@ static constexpr int kImageSize = 224; static constexpr int kDefaultChunk = 10; static constexpr int kModelActionDim = 32; static constexpr int kLiberoActionDim = 7; +static constexpr int kEncoderWidth = 2048; modalities::VisionPreprocessSpec vision_preprocess_spec(int num_views); diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index f9af08b2..0ea2b9f4 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -32,25 +32,14 @@ modalities::Status backend(const char* message) { NativeGraphOwner::NativeGraphOwner( frt_ctx ctx, const NativeGraphConfig& config) - : ctx_(ctx), + : graphs_(ctx), config_(config), weights_(ctx), workspace_(ctx), attention_(ctx), - forward_(&driver_), - capture_status_(modalities::Status::ok()) {} + forward_(&driver_) {} -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; - } -} +NativeGraphOwner::~NativeGraphOwner() = default; std::unique_ptr NativeGraphOwner::create( const std::string& checkpoint_path, @@ -166,40 +155,28 @@ modalities::Status NativeGraphOwner::initialize( } report("input_init"); - 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_; - } + st = graphs_.capture( + NativeGraphKind::kInfer, workspace_, + {"observation_images_normalized", "prompt_embedding", "encoder_x", + "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", + "rtc_guidance_weight"}, + record_graph, this); + if (!st.ok_status()) return st; + st = graphs_.capture( + NativeGraphKind::kDecodeOnly, workspace_, + {"encoder_x", "diffusion_noise", "rtc_prev_action_chunk", + "rtc_prefix_weights", "rtc_guidance_weight"}, + record_graph, this); + if (!st.ok_status()) return st; + st = graphs_.capture( + NativeGraphKind::kContext, workspace_, + {"observation_images_normalized", "prompt_embedding", "encoder_x"}, + record_graph, this); + if (!st.ok_status()) return st; report("capture"); - 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"); + st = graphs_.create_replay_stream(); + if (!st.ok_status()) return st; report("stream"); if (profile_setup) { const auto now = std::chrono::steady_clock::now(); @@ -210,27 +187,10 @@ modalities::Status NativeGraphOwner::initialize( 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( +modalities::Status NativeGraphOwner::record_context(void* stream) { + modalities::Status st = copy_prompt_to_encoder(&workspace_, stream); + if (!st.ok_status()) return st; + st = forward_.vision( weights_, &workspace_, &attention_, attention_driver_.get(), reinterpret_cast(stream)); if (!st.ok_status()) return st; @@ -238,14 +198,30 @@ modalities::Status NativeGraphOwner::record(void* stream) { attention_driver_.get(), reinterpret_cast(stream)); if (!st.ok_status()) return st; + return st; +} + +modalities::Status NativeGraphOwner::record_action(void* stream) { return forward_.diffusion(weights_, &workspace_, &attention_, attention_driver_.get(), reinterpret_cast(stream)); } -void NativeGraphOwner::record_graph(void* user, void* stream) { +modalities::Status NativeGraphOwner::record(NativeGraphKind kind, + void* stream) { + if (kind == NativeGraphKind::kContext) return record_context(stream); + if (kind == NativeGraphKind::kDecodeOnly) return record_action(stream); + if (kind != NativeGraphKind::kInfer) { + return invalid("native graph kind is invalid"); + } + modalities::Status st = record_context(stream); + return st.ok_status() ? record_action(stream) : st; +} + +modalities::Status NativeGraphOwner::record_graph( + void* user, NativeGraphKind kind, void* stream) { auto* owner = static_cast(user); - owner->capture_status_ = owner->record(stream); + return owner->record(kind, stream); } modalities::Status NativeGraphOwner::set_prompt_length(int prompt_tokens) { @@ -254,17 +230,12 @@ modalities::Status NativeGraphOwner::set_prompt_length(int prompt_tokens) { 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_); +int NativeGraphOwner::replay(NativeGraphKind kind) const { + return graphs_.replay(kind); } 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)); + return graphs_.synchronize(); } } // namespace pi05 diff --git a/cpp/models/pi05/src/native_graph_runtime.cpp b/cpp/models/pi05/src/native_graph_runtime.cpp new file mode 100644 index 00000000..4ee994e3 --- /dev/null +++ b/cpp/models/pi05/src/native_graph_runtime.cpp @@ -0,0 +1,170 @@ +#include "flashrt/cpp/models/pi05/native_graph_runtime.h" + +#include "flashrt/cpp/models/pi05/spec.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); +} + +std::size_t graph_index(NativeGraphKind kind) { + return static_cast(kind); +} + +} // namespace + +struct NativeGraphCatalog::CaptureCall { + RecordFn record = nullptr; + void* owner = nullptr; + NativeGraphKind kind = NativeGraphKind::kInfer; + modalities::Status status = modalities::Status::ok(); +}; + +NativeGraphCatalog::~NativeGraphCatalog() { + if (replay_stream_) { + cudaStreamSynchronize(static_cast(replay_stream_)); + cudaStreamDestroy(static_cast(replay_stream_)); + replay_stream_ = nullptr; + } + if (ctx_) { + frt_ctx_destroy(ctx_); + ctx_ = nullptr; + } +} + +const char* NativeGraphCatalog::name(NativeGraphKind kind) { + switch (kind) { + case NativeGraphKind::kInfer: return "infer"; + case NativeGraphKind::kDecodeOnly: return "decode_only"; + case NativeGraphKind::kContext: return "context"; + case NativeGraphKind::kCount: break; + } + return nullptr; +} + +frt_graph NativeGraphCatalog::graph(NativeGraphKind kind) const { + const std::size_t index = graph_index(kind); + return index < graph_index(NativeGraphKind::kCount) + ? graphs_[index] + : nullptr; +} + +void NativeGraphCatalog::record_graph(void* user, void* stream) { + auto* call = static_cast(user); + if (!call || !call->record) return; + call->status = call->record(call->owner, call->kind, stream); +} + +modalities::Status NativeGraphCatalog::capture( + NativeGraphKind kind, const NativeWorkspace& workspace, + std::initializer_list bindings, + RecordFn record, void* owner) { + const std::size_t index = graph_index(kind); + const char* graph_name = name(kind); + if (!ctx_ || !graph_name || index >= graph_index(NativeGraphKind::kCount) || + graphs_[index] || !record || !owner) { + return invalid("native graph capture request is invalid"); + } + frt_graph captured = frt_graph_create(ctx_, graph_name, 1); + if (!captured) return backend("native graph creation failed"); + graphs_[index] = captured; + for (const char* binding : bindings) { + const NativeWorkspaceBuffer* buffer = workspace.find(binding); + if (!buffer || + frt_graph_bind(captured, binding, buffer->buffer) != FRT_OK) { + frt_graph_destroy(captured); + graphs_[index] = nullptr; + return backend("native graph binding failed"); + } + } + CaptureCall call; + call.record = record; + call.owner = owner; + call.kind = kind; + const int rc = frt_graph_capture(captured, 0, record_graph, &call); + if (!call.status.ok_status() || rc != FRT_OK || + frt_graph_variant_count(captured) != 1) { + frt_graph_destroy(captured); + graphs_[index] = nullptr; + return call.status.ok_status() + ? backend("native graph capture failed") + : call.status; + } + return modalities::Status::ok(); +} + +modalities::Status NativeGraphCatalog::create_replay_stream() { + if (!ctx_ || replay_stream_) { + return invalid("native replay stream request is invalid"); + } + cudaStream_t stream = nullptr; + if (cudaStreamCreate(&stream) != cudaSuccess) { + return backend("native replay stream creation failed"); + } + const int wrapped = frt_ctx_wrap_stream(ctx_, stream); + if (wrapped < 0) { + cudaStreamDestroy(stream); + return backend("native replay stream wrapping failed"); + } + replay_stream_ = stream; + stream_id_ = wrapped; + return modalities::Status::ok(); +} + +int NativeGraphCatalog::replay(NativeGraphKind kind) const { + frt_graph selected = graph(kind); + if (!selected || stream_id_ < 0) return FRT_ERR_INVALID; + return frt_graph_replay(selected, 0, stream_id_); +} + +modalities::Status NativeGraphCatalog::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)); +} + +modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, + void* stream) { + if (!workspace) return invalid("native workspace is missing"); + 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()) * + kEncoderWidth * 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; + const cudaError_t rc = cudaMemcpyAsync( + destination, frt_buffer_dptr(prompt->buffer), prompt_bytes, + cudaMemcpyDeviceToDevice, static_cast(stream)); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("native prompt graph copy failed"); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 94adf830..59b35ad8 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -265,7 +265,16 @@ int build_native_model_runtime(const NativeOpenConfig& config, builder, "main", graph->stream_id(), 0, graph->native_stream()) == 0 && frt_runtime_builder_add_graph( - builder, "infer", graph->infer_graph(), 0, keys, 1, + builder, NativeGraphCatalog::name(NativeGraphKind::kInfer), + graph->graph(NativeGraphKind::kInfer), 0, keys, 1, + graph->stream_id()) == 0 && + frt_runtime_builder_add_graph( + builder, NativeGraphCatalog::name(NativeGraphKind::kDecodeOnly), + graph->graph(NativeGraphKind::kDecodeOnly), 0, keys, 1, + graph->stream_id()) == 0 && + frt_runtime_builder_add_graph( + builder, NativeGraphCatalog::name(NativeGraphKind::kContext), + graph->graph(NativeGraphKind::kContext), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_buffer( builder, "observation_images_normalized", images->buffer, @@ -312,6 +321,7 @@ int build_native_model_runtime(const NativeOpenConfig& config, add_identity(builder, "weights_sha256", weights_sha256.digest) && add_identity(builder, "tokenizer_sha256", tokenizer_sha256) && add_identity(builder, "io", "native_v2") && + add_identity(builder, "stage_plan", config.stage_plan) && add_identity(builder, "state_prompt_mode", "fixed") && add_identity(builder, "num_views", std::to_string(config.num_views)) && add_identity(builder, "max_prompt_len", @@ -334,9 +344,19 @@ int build_native_model_runtime(const NativeOpenConfig& config, << "\"hardware\":\"" << hardware_id << "\",\"precision\":\"" << precision_id << "\",\"io\":\"native_v2\"," - << "\"stage_plan\":{\"name\":\"full\"," - << "\"stages\":[{\"name\":\"infer\"," - << "\"graph\":\"infer\",\"after\":[]}]}}"; + << "\"graphs\":[\"infer\",\"decode_only\",\"context\"]," + << "\"stage_plan\":{\"name\":\"" << config.stage_plan + << "\",\"stages\":"; + if (config.stage_plan == "full") { + manifest << "[{\"name\":\"infer\",\"graph\":\"infer\"," + << "\"after\":[]}]}"; + } else { + manifest << "[{\"name\":\"context\",\"graph\":\"context\"," + << "\"after\":[]},{\"name\":\"action\"," + << "\"graph\":\"decode_only\"," + << "\"after\":[\"context\"]}]}"; + } + manifest << '}'; if (frt_runtime_builder_set_manifest(builder, manifest.str().c_str()) != 0) { return fail_builder(builder, error, "native manifest build failed"); } @@ -378,8 +398,15 @@ int build_native_model_runtime(const NativeOpenConfig& config, builder, "actions_raw", FRT_RT_MOD_TENSOR, io_dtype, 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; + frt_buffer_bytes(noise->buffer)) == 0; + if (ok && config.stage_plan == "full") { + ok = frt_runtime_builder_add_stage(builder, 0, nullptr, 0) == 0; + } else if (ok) { + const uint32_t context_stage[] = {0}; + ok = frt_runtime_builder_add_stage(builder, 2, nullptr, 0) == 0 && + frt_runtime_builder_add_stage( + builder, 1, context_stage, 1) == 0; + } if (!ok) return fail_builder(builder, error, "native port/stage build failed"); NativeGraphRuntime* raw_graph = graph.release(); diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index 6a803f1e..4faea959 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -605,13 +605,15 @@ int validate_config( std::string state_prompt_mode; std::string precision; std::string calibration_path; + std::string stage_plan; 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) || !string_field(obj, "precision", &precision, false) || - !string_field(obj, "calibration_path", &calibration_path, false)) { + !string_field(obj, "calibration_path", &calibration_path, false) || + !string_field(obj, "stage_plan", &stage_plan, false)) { return -1; } if (io != "native_v2") { @@ -630,6 +632,12 @@ int validate_config( "precision must be 'auto', 'bf16', or 'fp8_e4m3fn'"; return -1; } + if (stage_plan.empty()) stage_plan = "full"; + if (stage_plan != "full" && stage_plan != "context_action") { + g_last_error = + "stage_plan must be 'full' or 'context_action'"; + return -1; + } if (!path_exists(checkpoint_path)) { g_last_error = "checkpoint_path does not exist"; return -2; @@ -665,8 +673,8 @@ int validate_config( !integer_field(obj, "max_frame_height", &max_frame_height)) { return -1; } - if (max_prompt_tokens < 200 || max_prompt_tokens > INT_MAX) { - g_last_error = "max_prompt_tokens must be in [200, INT_MAX]"; + if (max_prompt_tokens <= 0 || max_prompt_tokens > INT_MAX) { + g_last_error = "max_prompt_tokens must be in [1, INT_MAX]"; return -1; } if (state_dim <= 0 || state_dim > INT_MAX) { @@ -700,6 +708,7 @@ int validate_config( config.tokenizer_model_path = tokenizer_model_path; config.precision = precision; config.calibration_path = calibration_path; + config.stage_plan = stage_plan; 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); diff --git a/cpp/models/pi05/src/native_open_internal.h b/cpp/models/pi05/src/native_open_internal.h index a32e6747..60fff664 100644 --- a/cpp/models/pi05/src/native_open_internal.h +++ b/cpp/models/pi05/src/native_open_internal.h @@ -15,6 +15,7 @@ struct NativeOpenConfig { std::string tokenizer_model_path; std::string precision = "auto"; std::string calibration_path; + std::string stage_plan = "full"; int max_prompt_tokens = 200; int state_dim = 0; int num_views = 2; diff --git a/cpp/models/pi05/src/native_thor_graph_owner.cpp b/cpp/models/pi05/src/native_thor_graph_owner.cpp index abbdcd49..66fd6ec1 100644 --- a/cpp/models/pi05/src/native_thor_graph_owner.cpp +++ b/cpp/models/pi05/src/native_thor_graph_owner.cpp @@ -59,24 +59,13 @@ bool calibration_matches(const NativeCalibrationArtifact& artifact, NativeThorGraphOwner::NativeThorGraphOwner( frt_ctx ctx, const NativeGraphConfig& config) - : ctx_(ctx), + : graphs_(ctx), config_(config), weights_(ctx), workspace_(ctx), - forward_(&driver_), - capture_status_(modalities::Status::ok()) {} + forward_(&driver_) {} -NativeThorGraphOwner::~NativeThorGraphOwner() { - if (replay_stream_) { - cudaStreamSynchronize(static_cast(replay_stream_)); - cudaStreamDestroy(static_cast(replay_stream_)); - replay_stream_ = nullptr; - } - if (ctx_) { - frt_ctx_destroy(ctx_); - ctx_ = nullptr; - } -} +NativeThorGraphOwner::~NativeThorGraphOwner() = default; std::unique_ptr NativeThorGraphOwner::create( const std::string& checkpoint_path, @@ -205,7 +194,7 @@ modalities::Status NativeThorGraphOwner::initialize( } // CUTLASS, cuBLAS, and FMHA initialize tactics outside CUDA capture. - st = record(nullptr); + st = record(NativeGraphKind::kInfer, nullptr); if (!st.ok_status()) return st; if (cudaDeviceSynchronize() != cudaSuccess) { return backend("Thor graph warmup synchronization failed"); @@ -218,39 +207,28 @@ modalities::Status NativeThorGraphOwner::initialize( } report("warmup"); - 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"); - 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("Thor 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("Thor full graph capture failed") - : capture_status_; - } - if (!capture_status_.ok_status() || - frt_graph_variant_count(infer_graph_) != 1) { - return capture_status_.ok_status() - ? backend("Thor full graph variant is missing") - : capture_status_; - } + st = graphs_.capture( + NativeGraphKind::kInfer, workspace_, + {"observation_images_normalized", "prompt_embedding", "encoder_x", + "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", + "rtc_guidance_weight"}, + record_graph, this); + if (!st.ok_status()) return st; + st = graphs_.capture( + NativeGraphKind::kDecodeOnly, workspace_, + {"encoder_x", "diffusion_noise", "rtc_prev_action_chunk", + "rtc_prefix_weights", "rtc_guidance_weight"}, + record_graph, this); + if (!st.ok_status()) return st; + st = graphs_.capture( + NativeGraphKind::kContext, workspace_, + {"observation_images_normalized", "prompt_embedding", "encoder_x"}, + record_graph, this); + if (!st.ok_status()) return st; report("capture"); - cudaStream_t stream = nullptr; - if (cudaStreamCreate(&stream) != cudaSuccess) { - return backend("Thor replay stream creation failed"); - } - replay_stream_ = stream; - stream_id_ = frt_ctx_wrap_stream(ctx_, replay_stream_); - if (stream_id_ < 0) return backend("Thor replay stream wrapping failed"); + st = graphs_.create_replay_stream(); + if (!st.ok_status()) return st; report("stream"); if (profile_setup) { const auto now = std::chrono::steady_clock::now(); @@ -261,56 +239,48 @@ modalities::Status NativeThorGraphOwner::initialize( return modalities::Status::ok(); } -modalities::Status NativeThorGraphOwner::record(void* stream) { - const NativeWorkspaceBuffer* prompt = workspace_.find("prompt_embedding"); - const NativeWorkspaceBuffer* encoder = workspace_.find("encoder_x"); - if (!prompt || !encoder) return invalid("Thor 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("Thor 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("Thor prompt graph copy failed"); - } - const std::uintptr_t stream_id = reinterpret_cast(stream); - modalities::Status st = - forward_.vision(weights_, &workspace_, weight_scales_, stream_id); +modalities::Status NativeThorGraphOwner::record_context(void* stream) { + modalities::Status st = copy_prompt_to_encoder(&workspace_, stream); if (!st.ok_status()) return st; - st = forward_.encoder(weights_, &workspace_, encoder_alphas_, stream_id); + const std::uintptr_t stream_id = reinterpret_cast(stream); + st = forward_.vision(weights_, &workspace_, weight_scales_, stream_id); if (!st.ok_status()) return st; - return forward_.diffusion(weights_, &workspace_, stream_id); + return forward_.encoder( + weights_, &workspace_, encoder_alphas_, stream_id); +} + +modalities::Status NativeThorGraphOwner::record_action(void* stream) { + return forward_.diffusion( + weights_, &workspace_, reinterpret_cast(stream)); } -void NativeThorGraphOwner::record_graph(void* user, void* stream) { +modalities::Status NativeThorGraphOwner::record(NativeGraphKind kind, + void* stream) { + if (kind == NativeGraphKind::kContext) return record_context(stream); + if (kind == NativeGraphKind::kDecodeOnly) return record_action(stream); + if (kind != NativeGraphKind::kInfer) { + return invalid("Thor graph kind is invalid"); + } + modalities::Status st = record_context(stream); + return st.ok_status() ? record_action(stream) : st; +} + +modalities::Status NativeThorGraphOwner::record_graph( + void* user, NativeGraphKind kind, void* stream) { auto* owner = static_cast(user); - owner->capture_status_ = owner->record(stream); + return owner->record(kind, stream); } modalities::Status NativeThorGraphOwner::set_prompt_length(int prompt_tokens) { return workspace_.set_fixed_prompt_length(prompt_tokens); } -int NativeThorGraphOwner::replay() const { - if (!infer_graph_ || stream_id_ < 0) return FRT_ERR_INVALID; - return frt_graph_replay(infer_graph_, 0, stream_id_); +int NativeThorGraphOwner::replay(NativeGraphKind kind) const { + return graphs_.replay(kind); } modalities::Status NativeThorGraphOwner::synchronize() const { - if (!replay_stream_) return invalid("Thor replay stream is missing"); - const cudaError_t rc = - cudaStreamSynchronize(static_cast(replay_stream_)); - return rc == cudaSuccess - ? modalities::Status::ok() - : backend("Thor replay synchronization failed"); + return graphs_.synchronize(); } } // namespace pi05 diff --git a/cpp/tests/gate_pi05_native_thor_fp8.py b/cpp/tests/gate_pi05_native_thor_fp8.py index af0e6cf4..8fb003aa 100644 --- a/cpp/tests/gate_pi05_native_thor_fp8.py +++ b/cpp/tests/gate_pi05_native_thor_fp8.py @@ -34,6 +34,9 @@ def _sample(index: int, num_views: int pixels = np.arange(224 * 224 * 3, dtype=np.uint64) image = ((pixels * 3 + index * 17) % 251).astype(np.uint8) wrist = ((pixels * 7 + index * 29 + 11) % 253).astype(np.uint8) + right_wrist = ( + (pixels * 11 + index * 37 + 19) % 247 + ).astype(np.uint8) state = np.asarray( [((index * 8 + dim) % 17 - 8) / 8.0 for dim in range(8)], dtype=np.float32, @@ -47,7 +50,11 @@ def _sample(index: int, num_views: int if index & 1 else "pick up the black bowl" ) - images = [image.reshape(224, 224, 3), wrist.reshape(224, 224, 3)] + images = [ + image.reshape(224, 224, 3), + wrist.reshape(224, 224, 3), + right_wrist.reshape(224, 224, 3), + ] return ( prompt, state, @@ -270,10 +277,13 @@ def main() -> None: parser.add_argument("--tokenizer", type=Path, required=True) parser.add_argument("--artifact", type=Path, required=True) parser.add_argument("--samples", type=int, default=3) - parser.add_argument("--views", type=int, choices=(1, 2), default=2) + parser.add_argument("--views", type=int, choices=(1, 2, 3), default=2) + parser.add_argument("--max-prompt-tokens", type=int, default=200) args = parser.parse_args() if args.samples < 1 or args.samples > 256: parser.error("--samples must be in [1, 256]") + if args.max_prompt_tokens < 1: + parser.error("--max-prompt-tokens must be positive") probe = args.probe.resolve() os.environ["FLASH_RT_PALIGEMMA_TOKENIZER"] = str( @@ -286,6 +296,7 @@ def main() -> None: if env.get("LD_LIBRARY_PATH"): library_paths.append(Path(env["LD_LIBRARY_PATH"])) env["LD_LIBRARY_PATH"] = os.pathsep.join(map(str, library_paths)) + env["FLASHRT_MAX_PROMPT_TOKENS"] = str(args.max_prompt_tokens) subprocess.run( [ str(probe), @@ -312,7 +323,7 @@ def main() -> None: str(args.checkpoint), num_views=args.views, use_cuda_graph=True, autotune=0, use_fp8=True, state_prompt_mode="fixed", - state_prompt_fixed_max_len=200, + state_prompt_fixed_max_len=args.max_prompt_tokens, ) frontend._calibrate = types.MethodType( _seed_calibration(expected_encoder, expected_decoder), frontend diff --git a/cpp/tests/pi05_native_graph_probe.cpp b/cpp/tests/pi05_native_graph_probe.cpp index c561bb50..6a5d828f 100644 --- a/cpp/tests/pi05_native_graph_probe.cpp +++ b/cpp/tests/pi05_native_graph_probe.cpp @@ -41,8 +41,13 @@ int main(int argc, char** argv) { 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 || + const frt_graph infer = owner->graph(NativeGraphKind::kInfer); + const frt_graph decode = owner->graph(NativeGraphKind::kDecodeOnly); + const frt_graph context = owner->graph(NativeGraphKind::kContext); + if (!images || !prompt || !noise || !infer || !decode || !context || + frt_graph_variant_count(infer) != 1 || + frt_graph_variant_count(decode) != 1 || + frt_graph_variant_count(context) != 1 || owner->stream_id() < 0 || !owner->native_stream()) { return 1; } @@ -82,22 +87,37 @@ int main(int argc, char** argv) { } const std::vector expected = download(noise->buffer); if (expected.empty()) return 1; + if (cudaMemcpy(frt_buffer_dptr(noise->buffer), host_noise.data(), + frt_buffer_bytes(noise->buffer), + cudaMemcpyHostToDevice) != cudaSuccess || + owner->replay(NativeGraphKind::kContext) != FRT_OK || + owner->replay(NativeGraphKind::kDecodeOnly) != FRT_OK || + !owner->synchronize().ok_status() || + download(noise->buffer) != expected) { + 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) { + owner->replay(replay % 2 == 0 + ? NativeGraphKind::kInfer + : NativeGraphKind::kContext) != FRT_OK || + (replay % 2 != 0 && + owner->replay(NativeGraphKind::kDecodeOnly) != FRT_OK)) { return 1; } } if (!owner->synchronize().ok_status() || - frt_graph_variant_count(owner->infer_graph()) != 1 || + frt_graph_variant_count(infer) != 1 || + frt_graph_variant_count(decode) != 1 || + frt_graph_variant_count(context) != 1 || owner->workspace().allocation_count() != allocation_count || download(noise->buffer) != expected) { return 1; } - std::cout << "PASS native full graph 100 replays\n"; + std::cout << "PASS native full/split graphs 100 replays\n"; return 0; } diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index 4b37d042..6d644225 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -19,6 +19,21 @@ 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 { + +constexpr int kLatencyWarmupReplays = 10; + +bool all_graph_variants_stable(const frt_runtime_export_v1* exp) { + for (std::uint64_t graph = 0; graph < exp->n_graphs; ++graph) { + if (frt_graph_variant_count(exp->graphs[graph].handle) != 1) { + return false; + } + } + return true; +} + +} // namespace + int main(int argc, char** argv) { if (argc != 3 && argc != 4) { std::cerr << "usage: pi05_native_open_probe CHECKPOINT TOKENIZER " @@ -36,13 +51,45 @@ int main(int argc, char** argv) { } replay_count = static_cast(parsed); } - const bool profile_service_loop = + bool profile_service_loop = std::getenv("FLASHRT_PROFILE_SERVICE_LOOP") != nullptr; - if (profile_service_loop && !replay_env) { + int latency_replays = 0; + const char* latency_env = + std::getenv("FLASHRT_E2E_LATENCY_REPLAYS"); + if (latency_env) { + char* end = nullptr; + const long parsed = std::strtol(latency_env, &end, 10); + if (!end || *end != '\0' || parsed <= 0 || parsed > 100000) { + std::cerr << "FLASHRT_E2E_LATENCY_REPLAYS must be in " + "[1, 100000]\n"; + return 2; + } + if (replay_env) { + std::cerr << "FLASHRT_E2E_LATENCY_REPLAYS cannot be combined " + "with FLASHRT_PROFILE_REPLAYS\n"; + return 2; + } + latency_replays = static_cast(parsed); + replay_count = kLatencyWarmupReplays + latency_replays; + profile_service_loop = true; + } + if (profile_service_loop && !replay_env && !latency_env) { std::cerr << "FLASHRT_PROFILE_SERVICE_LOOP requires " "FLASHRT_PROFILE_REPLAYS\n"; return 2; } + double latency_p99_limit_ms = 0.0; + const char* latency_limit_env = std::getenv("FLASHRT_E2E_P99_MS"); + if (latency_limit_env) { + char* end = nullptr; + latency_p99_limit_ms = std::strtod(latency_limit_env, &end); + if (!end || *end != '\0' || !std::isfinite(latency_p99_limit_ms) || + latency_p99_limit_ms <= 0.0 || latency_replays == 0) { + std::cerr << "FLASHRT_E2E_P99_MS must be positive and requires " + "FLASHRT_E2E_LATENCY_REPLAYS\n"; + return 2; + } + } int hot_state_updates = 0; const char* hot_updates_env = std::getenv("FLASHRT_HOT_STATE_UPDATES"); if (hot_updates_env) { @@ -54,6 +101,38 @@ int main(int argc, char** argv) { } hot_state_updates = static_cast(parsed); } + const char* stage_plan_env = std::getenv("FLASHRT_STAGE_PLAN"); + const std::string stage_plan = + stage_plan_env ? stage_plan_env : "full"; + if (stage_plan != "full" && stage_plan != "context_action") { + std::cerr << "FLASHRT_STAGE_PLAN must be full or context_action\n"; + return 2; + } + const bool split_stage_plan = stage_plan == "context_action"; + int num_views = 2; + const char* num_views_env = std::getenv("FLASHRT_NUM_VIEWS"); + if (num_views_env) { + char* end = nullptr; + const long parsed = std::strtol(num_views_env, &end, 10); + if (!end || *end != '\0' || parsed < 1 || parsed > 3) { + std::cerr << "FLASHRT_NUM_VIEWS must be in [1, 3]\n"; + return 2; + } + num_views = static_cast(parsed); + } + int max_prompt_tokens = 200; + const char* max_prompt_env = + std::getenv("FLASHRT_MAX_PROMPT_TOKENS"); + if (max_prompt_env) { + char* end = nullptr; + const long parsed = std::strtol(max_prompt_env, &end, 10); + if (!end || *end != '\0' || parsed < 1 || parsed > 100000) { + std::cerr << "FLASHRT_MAX_PROMPT_TOKENS must be in " + "[1, 100000]\n"; + return 2; + } + max_prompt_tokens = 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) { @@ -69,9 +148,12 @@ int main(int argc, char** argv) { 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"; + << "\"max_prompt_tokens\":" << max_prompt_tokens + << ",\"state_dim\":8," + << "\"num_views\":" << num_views + << ",\"chunk\":10,\"num_steps\":10," + << "\"vision_pool_factor\":1,\"stage_plan\":\"" + << stage_plan << '"'; if (argc == 4) { json << ",\"calibration_path\":\"" << argv[3] << '"'; } @@ -106,8 +188,9 @@ int main(int argc, char** argv) { 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 && + model->n_ports == 6 && + model->n_stages == (split_stage_plan ? 2u : 1u) && + exp->n_graphs == 3 && exp->n_streams == 1 && exp->n_capsule_regions == 1 && exp->n_buffers == 7 && exp->fingerprint != 0 && exp->identity && std::strstr(exp->identity, "producer=native") && @@ -117,8 +200,17 @@ int main(int argc, char** argv) { std::strstr(exp->manifest_json, hardware_manifest.c_str()) && std::strstr(exp->identity, "weights_sha256=") && std::strstr(exp->identity, "tokenizer_sha256=") && - model->stages[0].graph == 0 && - frt_graph_variant_count(exp->graphs[0].handle) == 1; + model->stages[0].graph == (split_stage_plan ? 2u : 0u) && + std::strcmp(exp->graphs[0].name, "infer") == 0 && + std::strcmp(exp->graphs[1].name, "decode_only") == 0 && + std::strcmp(exp->graphs[2].name, "context") == 0; + ok = ok && all_graph_variants_stable(exp); + if (split_stage_plan) { + ok = ok && model->stages[0].n_after == 0 && + model->stages[1].graph == 1 && + model->stages[1].n_after == 1 && + model->stages[1].after[0] == 0; + } for (std::uint64_t i = 0; i < model->n_ports; ++i) { ok = ok && std::strcmp(model->ports[i].name, port_names[i]) == 0; } @@ -170,6 +262,8 @@ int main(int argc, char** argv) { } } if (model->verbs.prepare(model->self, 0, 0) != 0 || + model->verbs.prepare(model->self, 1, 0) != 0 || + model->verbs.prepare(model->self, 2, 0) != 0 || model->verbs.prepare(model->self, 0, 1) != -2) { std::cerr << "native prepare validation failed\n"; model->release(model->owner); @@ -224,30 +318,33 @@ int main(int argc, char** argv) { << " 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) { + !all_graph_variants_stable(exp)) { std::cerr << "native hot state update gate failed\n"; model->release(model->owner); return 1; } } - std::vector rgb0(224 * 224 * 3); - std::vector rgb1(rgb0.size()); - for (std::size_t i = 0; i < rgb0.size(); ++i) { - rgb0[i] = static_cast(i % 251); - rgb1[i] = static_cast((3 * i + 17) % 251); + std::vector> rgb( + num_views, std::vector(224 * 224 * 3)); + for (int view = 0; view < num_views; ++view) { + for (std::size_t i = 0; i < rgb[view].size(); ++i) { + rgb[view][i] = static_cast( + ((2 * view + 1) * i + 17 * view) % 251); + } } - frt_image_view views[2]{}; - for (int i = 0; i < 2; ++i) { + std::vector views(num_views); + for (int i = 0; i < num_views; ++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].data = rgb[i].data(); + views[i].bytes = rgb[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) { + if (model->verbs.set_input( + model->self, 2, views.data(), + views.size() * sizeof(frt_image_view), -1) != 0) { std::cerr << "native image staging failed: " << model->verbs.last_error(model->self) << '\n'; model->release(model->owner); @@ -279,7 +376,11 @@ int main(int argc, char** argv) { } int step_rc = 0; cudaError_t upload_rc = cudaSuccess; + cudaError_t replay_sync_rc = cudaSuccess; + std::vector latency_samples_ms; + latency_samples_ms.reserve(latency_replays); for (int replay = 0; replay < replay_count; ++replay) { + const auto replay_begin = std::chrono::steady_clock::now(); if (profile_service_loop) { for (int dim = 0; dim < 8; ++dim) { state[dim] = std::sin( @@ -294,7 +395,8 @@ int main(int argc, char** argv) { model->verbs.set_input( model->self, 1, state, sizeof(state), -1) != 0 || model->verbs.set_input( - model->self, 2, views, sizeof(views), -1) != 0) { + model->self, 2, views.data(), + views.size() * sizeof(frt_image_view), -1) != 0) { step_rc = -1; break; } @@ -315,18 +417,45 @@ int main(int argc, char** argv) { step_rc = -1; break; } + if (latency_replays != 0) { + replay_sync_rc = cudaDeviceSynchronize(); + if (replay_sync_rc != cudaSuccess) break; + if (replay >= kLatencyWarmupReplays) { + const auto replay_end = std::chrono::steady_clock::now(); + latency_samples_ms.push_back( + std::chrono::duration( + replay_end - replay_begin) + .count()); + } + } } const cudaError_t sync_rc = cudaDeviceSynchronize(); const cudaError_t profiler_rc = profile_range ? cudaProfilerStop() : cudaSuccess; - if (step_rc != 0 || upload_rc != cudaSuccess || sync_rc != cudaSuccess || - profiler_rc != cudaSuccess || - frt_graph_variant_count(exp->graphs[0].handle) != 1) { + if (step_rc != 0 || upload_rc != cudaSuccess || + replay_sync_rc != cudaSuccess || sync_rc != cudaSuccess || + profiler_rc != cudaSuccess || !all_graph_variants_stable(exp)) { std::cerr << "native step failed: " << model->verbs.last_error(model->self) << '\n'; model->release(model->owner); return 1; } + if (latency_replays != 0) { + std::sort(latency_samples_ms.begin(), latency_samples_ms.end()); + const std::size_t p99_index = + (latency_samples_ms.size() * 99 + 99) / 100 - 1; + const double p50 = latency_samples_ms[latency_samples_ms.size() / 2]; + const double p99 = latency_samples_ms[p99_index]; + const double maximum = latency_samples_ms.back(); + std::cout << "native service latency: n=" << latency_samples_ms.size() + << " p50_ms=" << p50 << " p99_ms=" << p99 + << " max_ms=" << maximum << '\n'; + if (latency_p99_limit_ms > 0.0 && p99 > latency_p99_limit_ms) { + std::cerr << "native service latency gate failed\n"; + model->release(model->owner); + return 1; + } + } if (model->verbs.get_output(model->self, 4, actions, sizeof(actions), &written, -1) != 0 || written != sizeof(actions)) { diff --git a/cpp/tests/pi05_native_thor_fp8_probe.cpp b/cpp/tests/pi05_native_thor_fp8_probe.cpp index ac7424ff..9c8ae0a0 100644 --- a/cpp/tests/pi05_native_thor_fp8_probe.cpp +++ b/cpp/tests/pi05_native_thor_fp8_probe.cpp @@ -5,6 +5,7 @@ #include +#include #include #include #include @@ -35,6 +36,7 @@ std::string json_string(const std::string& value) { std::string config_json(const std::string& checkpoint, const std::string& tokenizer, int num_views, + int max_prompt_tokens, const std::string& calibration = "") { std::ostringstream out; out << "{\"io\":\"native_v2\",\"checkpoint_path\":" @@ -42,7 +44,8 @@ std::string config_json(const std::string& checkpoint, << json_string(tokenizer) << ",\"state_prompt_mode\":\"fixed\"," "\"precision\":\"fp8_e4m3fn\"," - "\"max_prompt_tokens\":200,\"state_dim\":8,\"num_views\":" + "\"max_prompt_tokens\":" + << max_prompt_tokens << ",\"state_dim\":8,\"num_views\":" << num_views << ",\"chunk\":10,\"num_steps\":10," "\"vision_pool_factor\":1"; if (!calibration.empty()) { @@ -77,14 +80,25 @@ int main(int argc, char** argv) { return 2; } const int num_views = std::atoi(argv[5]); - if (num_views < 1 || num_views > 2) { - std::cerr << "VIEWS must be in [1, 2]\n"; + if (num_views < 1 || num_views > 3) { + std::cerr << "VIEWS must be in [1, 3]\n"; return 2; } + int max_prompt_tokens = 200; + if (const char* value = std::getenv("FLASHRT_MAX_PROMPT_TOKENS")) { + char* end = nullptr; + const long parsed = std::strtol(value, &end, 10); + if (!end || *end != '\0' || parsed < 1 || parsed > 100000) { + std::cerr << "FLASHRT_MAX_PROMPT_TOKENS must be in " + "[1, 100000]\n"; + return 2; + } + max_prompt_tokens = static_cast(parsed); + } const std::string calibration_path = argv[3]; const std::string single_path = calibration_path + ".single"; const std::string calibration_config = - config_json(argv[1], argv[2], num_views); + config_json(argv[1], argv[2], num_views, max_prompt_tokens); frt_pi05_calibration_session* session = nullptr; int rc = frt_pi05_calibration_create_v1( calibration_config.c_str(), 99.9, &session); @@ -94,11 +108,14 @@ int main(int argc, char** argv) { std::vector image(224 * 224 * 3); std::vector wrist(image.size()); + std::vector right_wrist(image.size()); float state[8]{}; std::vector noise(10 * 32); std::vector frames(num_views); - const char* names[] = {"image", "wrist_image"}; - std::vector* pixels[] = {&image, &wrist}; + const char* names[] = { + "image", "wrist_image", "wrist_image_right"}; + std::vector* pixels[] = { + &image, &wrist, &right_wrist}; for (int view = 0; view < num_views; ++view) { frames[view].struct_size = sizeof(frt_pi05_vision_frame); frames[view].name = names[view]; @@ -109,13 +126,15 @@ int main(int argc, char** argv) { frames[view].stride_bytes = 224 * 3; frames[view].pixel_format = FRT_PI05_PIXEL_RGB8; } - if (num_views == 2) std::swap(frames[0], frames[1]); + if (num_views > 1) std::reverse(frames.begin(), frames.end()); for (int sample_index = 0; sample_index < samples; ++sample_index) { for (std::size_t i = 0; i < image.size(); ++i) { image[i] = static_cast( (i * 3 + sample_index * 17) % 251); wrist[i] = static_cast( (i * 7 + sample_index * 29 + 11) % 253); + right_wrist[i] = static_cast( + (i * 11 + sample_index * 37 + 19) % 247); } for (int i = 0; i < 8; ++i) { state[i] = static_cast( @@ -243,13 +262,15 @@ int main(int argc, char** argv) { calibration_path, &artifact); if (!status.ok_status() || artifact.sample_count != static_cast(samples) || - artifact.num_views != num_views) { + artifact.num_views != num_views || + artifact.max_prompt_tokens != max_prompt_tokens) { std::cerr << "dataset calibration artifact validation failed\n"; return 1; } const std::string runtime_config = - config_json(argv[1], argv[2], num_views, calibration_path); + config_json(argv[1], argv[2], num_views, max_prompt_tokens, + calibration_path); frt_model_runtime_v1* model = nullptr; rc = frt_model_runtime_open_v1(runtime_config.c_str(), &model); if (rc != 0 || !model) { @@ -258,8 +279,11 @@ int main(int argc, char** argv) { return 1; } const frt_runtime_export_v1* exp = model->exp; - bool schema_ok = model->n_ports == 6 && exp && exp->n_graphs == 1 && - frt_graph_variant_count(exp->graphs[0].handle) == 1 && + bool schema_ok = model->n_ports == 6 && model->n_stages == 1 && + exp && exp->n_graphs == 3 && + std::strcmp(exp->graphs[0].name, "infer") == 0 && + std::strcmp(exp->graphs[1].name, "decode_only") == 0 && + std::strcmp(exp->graphs[2].name, "context") == 0 && exp->identity && std::strstr(exp->identity, "precision=fp8_e4m3fn") && @@ -268,6 +292,10 @@ int main(int argc, char** argv) { model->ports[2].dtype == FRT_RT_DTYPE_F16 && model->ports[3].dtype == FRT_RT_DTYPE_F16 && model->ports[5].dtype == FRT_RT_DTYPE_F16; + for (std::uint64_t i = 0; schema_ok && i < exp->n_graphs; ++i) { + schema_ok = + frt_graph_variant_count(exp->graphs[i].handle) == 1; + } if (!schema_ok) { std::cerr << "native FP8 schema validation failed\n"; model->release(model->owner); @@ -314,16 +342,34 @@ int main(int argc, char** argv) { model->release(model->owner); return 1; } + std::vector raw(noise.size()); + if (!model->ports[5].buffer || + cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), + raw.size() * sizeof(raw[0]), + cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(frt_buffer_dptr(model->ports[3].buffer), noise_f16.data(), + noise_f16.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess || + frt_graph_replay(exp->graphs[2].handle, 0, + exp->graphs[2].stream_id) != FRT_OK || + frt_graph_replay(exp->graphs[1].handle, 0, + exp->graphs[1].stream_id) != FRT_OK || + cudaDeviceSynchronize() != cudaSuccess) { + std::cerr << "native FP8 split replay failed\n"; + model->release(model->owner); + return 1; + } + std::vector split_raw(noise.size()); + if (cudaMemcpy(split_raw.data(), + frt_buffer_dptr(model->ports[5].buffer), + split_raw.size() * sizeof(split_raw[0]), + cudaMemcpyDeviceToHost) != cudaSuccess || + split_raw != raw) { + std::cerr << "native FP8 full/split raw actions differ\n"; + model->release(model->owner); + return 1; + } if (argc == 7) { - std::vector raw(noise.size()); - if (!model->ports[5].buffer || - cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), - raw.size() * sizeof(raw[0]), - cudaMemcpyDeviceToHost) != cudaSuccess) { - std::cerr << "native FP8 raw action download failed\n"; - model->release(model->owner); - return 1; - } std::ofstream output(argv[6], std::ios::binary | std::ios::trunc); output.write(reinterpret_cast(raw.data()), static_cast( diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 4885fa12..6c9c6076 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -217,6 +217,14 @@ int main() { assert(out == nullptr); assert(std::strstr(frt_pi05_native_open_last_error(), "precision")); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, ",\"stage_plan\":\"async_magic\"").c_str(), + &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "stage_plan")); + out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1( config(root, tokenizer, @@ -283,15 +291,16 @@ int main() { assert(std::strstr(frt_pi05_native_open_last_error(), "validated")); #endif - const std::string short_prompt = + const std::string invalid_prompt_capacity = std::string("{") + "\"io\":\"native_v2\"," + "\"checkpoint_path\":\"" + root + "\"," + "\"tokenizer_model_path\":\"" + tokenizer + "\"," + "\"state_prompt_mode\":\"fixed\"," + - "\"max_prompt_tokens\":199," + + "\"max_prompt_tokens\":0," + "\"state_dim\":8}"; - rc = frt_model_runtime_open_v1(short_prompt.c_str(), &out); + rc = frt_model_runtime_open_v1( + invalid_prompt_capacity.c_str(), &out); assert(rc == -1); assert(std::strstr(frt_pi05_native_open_last_error(), "max_prompt_tokens")); From cf8456c26265c52ac02263a29ce6912cc8ec8bd9 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 17:50:11 -0400 Subject: [PATCH 70/83] feat(cpp): add native sm120 fp8 runtime --- cpp/models/pi05/CMakeLists.txt | 5 +- .../cpp/models/pi05/native_bf16_forward.h | 10 +- .../cpp/models/pi05/native_calibration.h | 2 + .../cpp/models/pi05/native_device_weights.h | 4 + .../cpp/models/pi05/native_graph_owner.h | 11 +- .../cpp/models/pi05/native_graph_runtime.h | 6 + .../cpp/models/pi05/native_kernel_driver.h | 39 ++ .../cpp/models/pi05/native_rtx_linear.h | 103 +++++ .../models/pi05/native_rtx_weight_packer.h | 38 ++ .../cpp/models/pi05/native_workspace.h | 1 + cpp/models/pi05/src/native_bf16_forward.cpp | 412 +++++++++++++----- cpp/models/pi05/src/native_calibration.cpp | 53 ++- cpp/models/pi05/src/native_device_weights.cpp | 37 ++ cpp/models/pi05/src/native_graph_owner.cpp | 161 ++++++- cpp/models/pi05/src/native_kernel_driver.cu | 244 +++++++++++ cpp/models/pi05/src/native_model_runtime.cpp | 56 ++- .../pi05/src/native_rtx_attention_driver.cu | 11 +- cpp/models/pi05/src/native_rtx_linear.cpp | 212 +++++++++ .../pi05/src/native_rtx_weight_packer.cpp | 149 +++++++ cpp/models/pi05/src/native_workspace.cpp | 13 + cpp/models/pi05/tests/CMakeLists.txt | 3 +- cpp/tests/gate_pi05_native_quantization.py | 41 ++ cpp/tests/pi05_native_open_probe.cpp | 66 +++ cpp/tests/pi05_native_quant_probe.cpp | 102 ++++- cpp/tests/pi05_native_vision_probe.cpp | 305 ++++++++++++- cpp/tests/test_pi05_native_calibration.cpp | 17 + csrc/gemm/gemm_runner.cu | 8 + csrc/kernels/quantize.cu | 48 +- csrc/kernels/quantize.cuh | 5 + 29 files changed, 1983 insertions(+), 179 deletions(-) create mode 100644 cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h create mode 100644 cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h create mode 100644 cpp/models/pi05/src/native_rtx_linear.cpp create mode 100644 cpp/models/pi05/src/native_rtx_weight_packer.cpp diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index 6602e33e..124d91f0 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -49,13 +49,17 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) src/native_bf16_forward.cpp src/native_graph_runtime.cpp src/native_kernel_driver.cu + src/native_rtx_linear.cpp + src/native_rtx_weight_packer.cpp src/native_style_precompute.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/activation.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/elementwise.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/fusion.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/norm.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/patch_embed.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/quantize.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/rope.cu) target_include_directories(flashrt_cpp_pi05_kernels PRIVATE ${_pi05_public_include} @@ -101,7 +105,6 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/cutlass_sm100.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/attention_cublas.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/decoder_fused.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/quantize.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/softmax.cu) target_include_directories(flashrt_cpp_pi05_kernels PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h index 16f4b610..6fb9bb0f 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h @@ -2,6 +2,7 @@ #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_linear.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" @@ -13,7 +14,12 @@ namespace pi05 { class NativeBf16Forward { public: explicit NativeBf16Forward(const NativeKernelDriver* driver) - : driver_(driver) {} + : driver_(driver), fallback_linear_(driver, NativeRtxLinearMode::kBf16), + linear_(&fallback_linear_) {} + NativeBf16Forward(const NativeKernelDriver* driver, + const NativeRtxLinear* linear) + : driver_(driver), fallback_linear_(driver, NativeRtxLinearMode::kBf16), + linear_(linear ? linear : &fallback_linear_) {} modalities::Status encoder_qkv( int layer, const NativeDeviceWeightStore& weights, @@ -59,6 +65,8 @@ class NativeBf16Forward { private: const NativeKernelDriver* driver_ = nullptr; + NativeRtxLinear fallback_linear_; + const NativeRtxLinear* linear_ = nullptr; }; } // namespace pi05 diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h index a1d2e600..06352fcf 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h @@ -12,6 +12,7 @@ namespace models { namespace pi05 { struct NativeCalibrationArtifact { + std::string activation_dtype = "float16"; std::string hardware; std::string weights_sha256; std::string tokenizer_sha256; @@ -23,6 +24,7 @@ struct NativeCalibrationArtifact { int vision_pool_factor = 0; std::uint64_t sample_count = 0; double percentile = 100.0; + std::vector vision_scales; std::vector encoder_scales; std::vector decoder_scales; }; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h index 741cc1e2..401051cf 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h @@ -45,6 +45,10 @@ class NativeDeviceWeightStore { NativeWeightDType dtype, const void* data, std::size_t bytes); + modalities::Status allocate( + const std::string& name, + const std::vector& shape, + NativeWeightDType dtype); modalities::Status download_bf16( const std::string& name, NativeBf16Tensor* out) const; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h index 4275ff62..3e9ff956 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h @@ -2,6 +2,7 @@ #define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H #include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_calibration.h" #include "flashrt/cpp/models/pi05/native_graph_runtime.h" #include @@ -16,6 +17,10 @@ class NativeGraphOwner final : public NativeGraphRuntime { static std::unique_ptr create( const std::string& checkpoint_path, const NativeGraphConfig& config, modalities::Status* status); + static std::unique_ptr create( + const std::string& checkpoint_path, const NativeGraphConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status); ~NativeGraphOwner() override; @@ -42,10 +47,13 @@ class NativeGraphOwner final : public NativeGraphRuntime { private: explicit NativeGraphOwner(frt_ctx ctx, const NativeGraphConfig& config); - modalities::Status initialize(const std::string& checkpoint_path); + modalities::Status initialize( + const std::string& checkpoint_path, + const NativeCalibrationArtifact* calibration); modalities::Status record(NativeGraphKind kind, void* stream); modalities::Status record_context(void* stream); modalities::Status record_action(void* stream); + modalities::Status autotune_fp8(); static modalities::Status record_graph( void* user, NativeGraphKind kind, void* stream); @@ -55,6 +63,7 @@ class NativeGraphOwner final : public NativeGraphRuntime { NativeWorkspace workspace_; NativeRtxAttentionWorkspace attention_; NativeKernelDriver driver_; + NativeRtxLinear linear_; NativeBf16Forward forward_; std::unique_ptr attention_driver_; }; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h index ce2957ba..2f5c3df9 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h @@ -11,12 +11,18 @@ namespace flashrt { namespace models { namespace pi05 { +enum class NativeGraphPrecision { + kBf16, + kFp8E4M3, +}; + struct NativeGraphConfig { int num_views = 2; int max_prompt_tokens = 200; int chunk_size = 10; int num_steps = 10; int vision_pool_factor = 1; + NativeGraphPrecision precision = NativeGraphPrecision::kBf16; }; enum class NativeGraphKind : std::size_t { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h index c6678413..ebde43e1 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h @@ -24,6 +24,22 @@ 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 fp8_nn_bf16( + void* a, void* b, void* output, int m, int n, int k, + const float* activation_scale, const float* weight_scale, + std::uintptr_t stream) const; + modalities::Status autotune_fp8_nn_bf16( + void* a, void* b, void* output, int m, int n, int k, + const float* activation_scale, const float* weight_scale) const; + modalities::Status quantize_fp8_static_bf16( + const void* values, void* output, const float* scale, + std::size_t elements, std::uintptr_t stream) const; + modalities::Status quantize_fp8_dynamic_bf16( + const void* values, void* output, float* scale, + std::size_t elements, std::uintptr_t stream) const; + modalities::Status quantize_fp8_weight_bf16( + const void* values, void* output, float* scale, + std::size_t elements, std::uintptr_t stream) const; modalities::Status add_bias_bf16(void* values, const void* bias, int rows, int columns, std::uintptr_t stream) const; @@ -34,6 +50,12 @@ class NativeKernelDriver { modalities::Status gate_gelu_bf16(const void* gate, const void* up, void* output, std::size_t elements, std::uintptr_t stream) const; + modalities::Status gate_gelu_merged_bf16( + const void* merged, void* output, int rows, int hidden, + std::uintptr_t stream) const; + modalities::Status gate_gelu_merged_fp8_bf16( + const void* merged, void* output, int rows, int hidden, + const float* scale, std::uintptr_t stream) const; modalities::Status residual_add_bf16(void* residual, const void* values, std::size_t elements, std::uintptr_t stream) const; @@ -46,6 +68,14 @@ class NativeKernelDriver { 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 rms_norm_fp8_bf16( + const void* values, const void* weight, void* output, + int rows, int columns, float epsilon, const float* scale, + std::uintptr_t stream) const; + modalities::Status residual_add_rms_norm_fp8_bf16( + void* residual, const void* values, const void* weight, void* output, + int rows, int columns, float epsilon, const float* scale, + 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; @@ -53,6 +83,15 @@ class NativeKernelDriver { 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 ada_rms_norm_style_fp8_bf16( + const void* values, const void* weight, const void* style, + void* output, void* gate_output, int rows, int columns, + float epsilon, const float* scale, std::uintptr_t stream) const; + modalities::Status gate_residual_ada_norm_fp8_bf16( + void* residual, const void* values, const void* gate, + const void* weight, const void* style, void* output, + void* gate_output, int rows, int columns, float epsilon, + const float* scale, 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, diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h new file mode 100644 index 00000000..ec1bc58a --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h @@ -0,0 +1,103 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_LINEAR_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_LINEAR_H + +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/native_workspace.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeRtxLinearMode { + kBf16, + kFp8Dynamic, + kFp8Static, +}; + +enum class NativeRtxScaleDomain { + kVision, + kEncoder, + kDecoder, +}; + +struct NativeRtxScaleSite { + NativeRtxScaleDomain domain = NativeRtxScaleDomain::kVision; + int index = 0; +}; + +class NativeRtxLinear { +public: + NativeRtxLinear(const NativeKernelDriver* driver, + NativeRtxLinearMode mode) + : driver_(driver), mode_(mode) {} + + bool fp8() const { return mode_ != NativeRtxLinearMode::kBf16; } + bool static_fp8() const { + return mode_ == NativeRtxLinearMode::kFp8Static; + } + bool dynamic_fp8() const { + return mode_ == NativeRtxLinearMode::kFp8Dynamic; + } + + const NativeDeviceWeight* find_weight( + const NativeDeviceWeightStore& weights, + const std::string& name) const; + bool weight_shape_is( + const NativeDeviceWeightStore& weights, + const std::string& name, + std::initializer_list shape) const; + + modalities::Status run( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const std::string& weight_name, + NativeRtxScaleSite site, + const void* input, + void* output, + int m, + int n, + int k, + std::uintptr_t stream) const; + modalities::Status autotune( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const std::string& weight_name, + NativeRtxScaleSite site, + void* output, + int m, + int n, + int k) const; + modalities::Status run_prequantized( + const NativeDeviceWeightStore& weights, + const std::string& weight_name, + NativeRtxScaleSite site, + const NativeWorkspace& workspace, + const void* input, + void* output, + int m, + int n, + int k, + std::uintptr_t stream) const; + const float* scale( + const NativeWorkspace& workspace, + NativeRtxScaleSite site) const; + +private: + const NativeWorkspaceBuffer* scale_buffer( + const NativeWorkspace& workspace, + NativeRtxScaleDomain domain) const; + + const NativeKernelDriver* driver_ = nullptr; + NativeRtxLinearMode mode_ = NativeRtxLinearMode::kBf16; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_LINEAR_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h new file mode 100644 index 00000000..3d2be1c8 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h @@ -0,0 +1,38 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_WEIGHT_PACKER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_WEIGHT_PACKER_H + +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeRtxWeightPacker { +public: + NativeRtxWeightPacker(NativeDeviceWeightStore* weights, + const NativeKernelDriver* driver) + : weights_(weights), driver_(driver) {} + + modalities::Status pack_weight( + const std::string& source_name, + const std::string& packed_name = ""); + modalities::Status pack_all(); + +private: + modalities::Status merge_bf16_columns( + const std::string& left_name, + const std::string& right_name, + const std::string& output_name); + + NativeDeviceWeightStore* weights_ = nullptr; + const NativeKernelDriver* driver_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_WEIGHT_PACKER_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h index 2d2aaebb..bc4bc3a0 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h @@ -17,6 +17,7 @@ namespace pi05 { enum class NativeWorkspaceFlavor { kBf16, + kRtxFp8, kThorFp8, }; diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/native_bf16_forward.cpp index 10d0e9a1..b84502aa 100644 --- a/cpp/models/pi05/src/native_bf16_forward.cpp +++ b/cpp/models/pi05/src/native_bf16_forward.cpp @@ -14,6 +14,19 @@ modalities::Status invalid(const char* message) { message); } +NativeRtxScaleSite vision_site(int layer, int slot) { + return {NativeRtxScaleDomain::kVision, layer * 4 + slot}; +} + +NativeRtxScaleSite encoder_site(int layer, int slot) { + return {NativeRtxScaleDomain::kEncoder, layer * 4 + slot}; +} + +NativeRtxScaleSite decoder_site(int step, int layer, int slot) { + return {NativeRtxScaleDomain::kDecoder, + (step * 18 + layer) * 4 + slot}; +} + bool shape_is(const NativeWorkspaceBuffer* buffer, std::initializer_list shape) { return buffer && buffer->dtype == modalities::DType::kBFloat16 && @@ -58,8 +71,8 @@ modalities::Status NativeBf16Forward::encoder_qkv( 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)); + const std::string qkv_name = + "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}) || @@ -69,11 +82,10 @@ modalities::Status NativeBf16Forward::encoder_qkv( !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 || + cache->shape[2] != 1 || cache->shape[3] != 256 || !value_cache || value_cache->dtype != NativeAttentionDType::kBf16 || value_cache->shape != cache->shape || - qkv_weight->dtype != NativeWeightDType::kBf16 || - qkv_weight->shape != std::vector({2048, 2560})) { + !linear_->weight_shape_is(weights, qkv_name, {2048, 2560})) { return invalid("native encoder QKV buffers or weight are invalid"); } void* key = attention->encoder_k_layer_dptr(layer); @@ -81,13 +93,42 @@ modalities::Status NativeBf16Forward::encoder_qkv( 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); + modalities::Status st; + if (linear_->static_fp8()) { + const NativeWorkspaceBuffer* scratch = + workspace->find("rtx_fp8_scratch"); + const float* scale = linear_->scale( + *workspace, encoder_site(layer, 0)); + if (!scratch || !scale) { + return invalid("native encoder FP8 QKV storage is invalid"); + } + st = layer == 0 + ? driver_->rms_norm_fp8_bf16( + frt_buffer_dptr(x->buffer), + frt_buffer_dptr(rms->buffer), + frt_buffer_dptr(scratch->buffer), sequence, 2048, + 1e-6f, scale, stream) + : driver_->residual_add_rms_norm_fp8_bf16( + frt_buffer_dptr(x->buffer), + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(rms->buffer), + frt_buffer_dptr(scratch->buffer), sequence, 2048, + 1e-6f, scale, stream); + if (!st.ok_status()) return st; + st = linear_->run_prequantized( + weights, qkv_name, encoder_site(layer, 0), *workspace, + frt_buffer_dptr(scratch->buffer), frt_buffer_dptr(qkv->buffer), + sequence, 2560, 2048, stream); + } else { + 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 = linear_->run( + weights, workspace, qkv_name, encoder_site(layer, 0), + frt_buffer_dptr(x_norm->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), @@ -117,20 +158,16 @@ modalities::Status NativeBf16Forward::vision_layer( 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 std::string qkv_name = "vision_attn_qkv_w_" + suffix; + const std::string output_name = "vision_attn_o_w_" + suffix; + const std::string up_name = "vision_ffn_up_w_" + suffix; + const std::string down_name = "vision_ffn_down_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 = @@ -147,21 +184,22 @@ modalities::Status NativeBf16Forward::vision_layer( !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}) || + !linear_->weight_shape_is(weights, qkv_name, {1152, 3456}) || !shape_is(qkv_bias, {3456}) || - !shape_is(output_weight, {1152, 1152}) || + !linear_->weight_shape_is(weights, output_name, {1152, 1152}) || !shape_is(output_bias, {1152}) || - !shape_is(up_weight, {1152, 4304}) || + !linear_->weight_shape_is(weights, up_name, {1152, 4304}) || !shape_is(up_bias, {4304}) || - !shape_is(down_weight, {4304, 1152}) || + !linear_->weight_shape_is(weights, down_name, {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); + modalities::Status st = linear_->run( + weights, workspace, qkv_name, vision_site(layer, 0), + frt_buffer_dptr(x_norm->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), @@ -174,9 +212,9 @@ modalities::Status NativeBf16Forward::vision_layer( 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), + st = linear_->run( + weights, workspace, output_name, vision_site(layer, 1), + attention_driver->vision_output(), frt_buffer_dptr(x_norm->buffer), sequence, 1152, 1152, stream); if (!st.ok_status()) return st; st = driver_->bias_residual_bf16( @@ -188,9 +226,10 @@ modalities::Status NativeBf16Forward::vision_layer( 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); + st = linear_->run( + weights, workspace, up_name, vision_site(layer, 2), + frt_buffer_dptr(x_norm->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), @@ -200,9 +239,10 @@ modalities::Status NativeBf16Forward::vision_layer( 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); + st = linear_->run( + weights, workspace, down_name, vision_site(layer, 3), + frt_buffer_dptr(hidden->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), @@ -258,8 +298,7 @@ modalities::Status NativeBf16Forward::vision( 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 std::string projector_name = "encoder_multi_modal_projector_w"; const NativeDeviceWeight* projector_bias = weights.find("encoder_multi_modal_projector_b"); if (sequence <= 0 || sequence % 256 || encoder_sequence <= 0 || @@ -283,7 +322,7 @@ modalities::Status NativeBf16Forward::vision( !shape_is(first_norm_bias, {1152}) || !shape_is(final_norm_weight, {1152}) || !shape_is(final_norm_bias, {1152}) || - !shape_is(projector_weight, {1152, 2048}) || + !linear_->weight_shape_is(weights, projector_name, {1152, 2048}) || !shape_is(projector_bias, {2048})) { return invalid("native vision buffers or weights are invalid"); } @@ -321,11 +360,11 @@ modalities::Status NativeBf16Forward::vision( 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); + st = linear_->run( + weights, workspace, projector_name, + {NativeRtxScaleDomain::kVision, 108}, + frt_buffer_dptr(x_norm->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), @@ -353,32 +392,65 @@ modalities::Status NativeBf16Forward::encoder_layer( 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)); + const std::string suffix = std::to_string(layer); + const std::string output_name = "encoder_attn_o_w_" + suffix; + const std::string gate_name = "encoder_ffn_gate_w_" + suffix; + const std::string up_name = "encoder_ffn_up_w_" + suffix; + const std::string gate_up_name = "encoder_ffn_gate_up_w_" + suffix; + const std::string down_name = "encoder_ffn_down_w_" + suffix; + const bool ffn_weights_valid = + linear_->fp8() + ? linear_->weight_shape_is( + weights, gate_up_name, {2048, 32768}) + : linear_->weight_shape_is(weights, gate_name, {2048, 16384}) && + linear_->weight_shape_is( + weights, up_name, {2048, 16384}); 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})) { + !shape_is(rms, {2048}) || !ffn_weights_valid || + !linear_->weight_shape_is(weights, output_name, {2048, 2048}) || + !linear_->weight_shape_is(weights, down_name, {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), + st = linear_->run( + weights, workspace, output_name, encoder_site(layer, 1), + attention_driver->encoder_output(), frt_buffer_dptr(x_norm->buffer), sequence, 2048, 2048, stream); if (!st.ok_status()) return st; + if (linear_->static_fp8()) { + const NativeWorkspaceBuffer* scratch = + workspace->find("rtx_fp8_scratch"); + const float* gate_up_scale = linear_->scale( + *workspace, encoder_site(layer, 2)); + const float* down_scale = linear_->scale( + *workspace, encoder_site(layer, 3)); + if (!scratch || !gate_up_scale || !down_scale) { + return invalid("native encoder fused FP8 storage is invalid"); + } + st = driver_->residual_add_rms_norm_fp8_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(rms->buffer), frt_buffer_dptr(scratch->buffer), + sequence, 2048, 1e-6f, gate_up_scale, stream); + if (!st.ok_status()) return st; + st = linear_->run_prequantized( + weights, gate_up_name, encoder_site(layer, 2), *workspace, + frt_buffer_dptr(scratch->buffer), frt_buffer_dptr(gate->buffer), + sequence, 32768, 2048, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_merged_fp8_bf16( + frt_buffer_dptr(gate->buffer), + frt_buffer_dptr(scratch->buffer), sequence, 16384, down_scale, + stream); + if (!st.ok_status()) return st; + return linear_->run_prequantized( + weights, down_name, encoder_site(layer, 3), *workspace, + frt_buffer_dptr(scratch->buffer), frt_buffer_dptr(x_norm->buffer), + sequence, 2048, 16384, stream); + } st = driver_->residual_add_bf16( frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), static_cast(sequence) * 2048, stream); @@ -387,22 +459,36 @@ modalities::Status NativeBf16Forward::encoder_layer( 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 (linear_->fp8()) { + st = linear_->run( + weights, workspace, gate_up_name, encoder_site(layer, 2), + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), + sequence, 32768, 2048, stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_merged_bf16( + frt_buffer_dptr(gate->buffer), frt_buffer_dptr(hidden->buffer), + sequence, 16384, stream); + } else { + st = linear_->run( + weights, workspace, gate_name, encoder_site(layer, 2), + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), + sequence, 16384, 2048, stream); + if (!st.ok_status()) return st; + st = linear_->run( + weights, workspace, up_name, encoder_site(layer, 2), + frt_buffer_dptr(x_norm->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); + st = linear_->run( + weights, workspace, down_name, encoder_site(layer, 3), + frt_buffer_dptr(hidden->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), @@ -456,16 +542,18 @@ modalities::Status NativeBf16Forward::decoder_layer( 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); + const std::string qkv_name = "decoder_attn_qkv_w_" + suffix; + const std::string output_name = "decoder_attn_o_w_" + suffix; + const std::string gate_name = "decoder_ffn_gate_w_" + suffix; + const std::string up_name = "decoder_ffn_up_w_" + suffix; + const std::string gate_up_name = "decoder_ffn_gate_up_w_" + suffix; + const std::string down_name = "decoder_ffn_down_w_" + suffix; + const bool ffn_weights_valid = + linear_->fp8() + ? linear_->weight_shape_is( + weights, gate_up_name, {1024, 8192}) + : linear_->weight_shape_is(weights, gate_name, {1024, 4096}) && + linear_->weight_shape_is(weights, up_name, {1024, 4096}); if (sequence <= 0 || step < 0 || !shape_is(x, {static_cast(sequence), 1024}) || !shape_is(x_norm, {static_cast(sequence), 1024}) || @@ -487,11 +575,10 @@ modalities::Status NativeBf16Forward::decoder_layer( !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})) { + !ffn_weights_valid || + !linear_->weight_shape_is(weights, qkv_name, {1024, 2560}) || + !linear_->weight_shape_is(weights, output_name, {2048, 1024}) || + !linear_->weight_shape_is(weights, down_name, {4096, 1024})) { return invalid("native decoder layer buffers or weights are invalid"); } const std::size_t style_offset = @@ -503,14 +590,38 @@ modalities::Status NativeBf16Forward::decoder_layer( 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); + modalities::Status st; + const NativeWorkspaceBuffer* fp8_scratch = + linear_->static_fp8() ? workspace->find("rtx_fp8_scratch") : nullptr; + if (linear_->static_fp8()) { + const float* qkv_scale = linear_->scale( + *workspace, decoder_site(step, layer, 0)); + if (!fp8_scratch || !qkv_scale) { + return invalid("native decoder FP8 QKV storage is invalid"); + } + if (layer == 0) { + st = driver_->ada_rms_norm_style_fp8_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), + attn_style, frt_buffer_dptr(fp8_scratch->buffer), + frt_buffer_dptr(gate->buffer), sequence, 1024, 1e-6f, + qkv_scale, stream); + if (!st.ok_status()) return st; + } + st = linear_->run_prequantized( + weights, qkv_name, decoder_site(step, layer, 0), *workspace, + frt_buffer_dptr(fp8_scratch->buffer), + frt_buffer_dptr(qkv->buffer), sequence, 2560, 1024, stream); + } else { + 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 = linear_->run( + weights, workspace, qkv_name, decoder_site(step, layer, 0), + frt_buffer_dptr(x_norm->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), @@ -521,11 +632,63 @@ modalities::Status NativeBf16Forward::decoder_layer( 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), + st = linear_->run( + weights, workspace, output_name, decoder_site(step, layer, 1), + attention_driver->decoder_output(), frt_buffer_dptr(x_norm->buffer), sequence, 1024, 2048, stream); if (!st.ok_status()) return st; + if (linear_->static_fp8()) { + const float* gate_up_scale = linear_->scale( + *workspace, decoder_site(step, layer, 2)); + const float* down_scale = linear_->scale( + *workspace, decoder_site(step, layer, 3)); + if (!fp8_scratch || !gate_up_scale || !down_scale) { + return invalid("native decoder fused FP8 storage is invalid"); + } + st = driver_->gate_residual_ada_norm_fp8_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(gate->buffer), frt_buffer_dptr(rms->buffer), + ffn_style, frt_buffer_dptr(fp8_scratch->buffer), + frt_buffer_dptr(gate->buffer), sequence, 1024, 1e-6f, + gate_up_scale, stream); + if (!st.ok_status()) return st; + st = linear_->run_prequantized( + weights, gate_up_name, decoder_site(step, layer, 2), *workspace, + frt_buffer_dptr(fp8_scratch->buffer), + frt_buffer_dptr(gate_projection->buffer), sequence, 8192, 1024, + stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_merged_fp8_bf16( + frt_buffer_dptr(gate_projection->buffer), + frt_buffer_dptr(fp8_scratch->buffer), sequence, 4096, + down_scale, stream); + if (!st.ok_status()) return st; + st = linear_->run_prequantized( + weights, down_name, decoder_site(step, layer, 3), *workspace, + frt_buffer_dptr(fp8_scratch->buffer), + frt_buffer_dptr(x_norm->buffer), sequence, 1024, 4096, stream); + if (!st.ok_status()) return st; + if (layer == 17) { + 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); + } + const float* next_qkv_scale = linear_->scale( + *workspace, decoder_site(step, layer + 1, 0)); + if (!next_qkv_scale) { + return invalid("native decoder next-layer FP8 scale is invalid"); + } + const auto* next_attn_style = + attn_style + static_cast(sequence) * 3072 * + sizeof(std::uint16_t); + return driver_->gate_residual_ada_norm_fp8_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(gate->buffer), frt_buffer_dptr(rms->buffer), + next_attn_style, frt_buffer_dptr(fp8_scratch->buffer), + frt_buffer_dptr(gate->buffer), sequence, 1024, 1e-6f, + next_qkv_scale, stream); + } st = driver_->gate_mul_residual_bf16( frt_buffer_dptr(x->buffer), frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), @@ -536,23 +699,38 @@ modalities::Status NativeBf16Forward::decoder_layer( 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 (linear_->fp8()) { + st = linear_->run( + weights, workspace, gate_up_name, decoder_site(step, layer, 2), + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(gate_projection->buffer), sequence, 8192, 1024, + stream); + if (!st.ok_status()) return st; + st = driver_->gate_gelu_merged_bf16( + frt_buffer_dptr(gate_projection->buffer), + frt_buffer_dptr(hidden->buffer), sequence, 4096, stream); + } else { + st = linear_->run( + weights, workspace, gate_name, decoder_site(step, layer, 2), + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(gate_projection->buffer), sequence, 4096, 1024, + stream); + if (!st.ok_status()) return st; + st = linear_->run( + weights, workspace, up_name, decoder_site(step, layer, 2), + frt_buffer_dptr(x_norm->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); + st = linear_->run( + weights, workspace, down_name, decoder_site(step, layer, 3), + frt_buffer_dptr(hidden->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), diff --git a/cpp/models/pi05/src/native_calibration.cpp b/cpp/models/pi05/src/native_calibration.cpp index 6da57b76..b460e360 100644 --- a/cpp/models/pi05/src/native_calibration.cpp +++ b/cpp/models/pi05/src/native_calibration.cpp @@ -21,11 +21,14 @@ namespace models { namespace pi05 { namespace { -constexpr char kSchema[] = "flashrt.pi05.fp8_calibration.v1"; +constexpr char kSchemaV1[] = "flashrt.pi05.fp8_calibration.v1"; +constexpr char kSchemaV2[] = "flashrt.pi05.fp8_calibration.v2"; constexpr char kModel[] = "pi05"; constexpr char kPrecision[] = "fp8_e4m3fn"; -constexpr char kTensorDtype[] = "float16"; +constexpr char kFloat16[] = "float16"; +constexpr char kBfloat16[] = "bfloat16"; constexpr char kReducer[] = "linear_percentile"; +constexpr std::size_t kVisionScaleCount = 27 * 4 + 1; constexpr std::size_t kEncoderScaleCount = 18 * 4; modalities::Status invalid(const std::string& message) { @@ -93,10 +96,12 @@ std::string decimal(double value) { std::map artifact_metadata( const NativeCalibrationArtifact& artifact) { return { - {"schema", kSchema}, + {"schema", artifact.activation_dtype == kBfloat16 + ? kSchemaV2 + : kSchemaV1}, {"model", kModel}, {"precision", kPrecision}, - {"tensor_dtype", kTensorDtype}, + {"tensor_dtype", artifact.activation_dtype}, {"reducer", kReducer}, {"hardware", artifact.hardware}, {"weights_sha256", artifact.weights_sha256}, @@ -114,6 +119,8 @@ std::map artifact_metadata( std::string make_header(const NativeCalibrationArtifact& artifact) { const auto metadata = artifact_metadata(artifact); + const std::uint64_t vision_bytes = + artifact.vision_scales.size() * sizeof(float); const std::uint64_t encoder_bytes = artifact.encoder_scales.size() * sizeof(float); const std::uint64_t decoder_bytes = @@ -126,12 +133,20 @@ std::string make_header(const NativeCalibrationArtifact& artifact) { first = false; out << json_string(entry.first) << ':' << json_string(entry.second); } - out << "},\"encoder_scales\":{\"dtype\":\"F32\",\"shape\":[" + out << '}'; + if (!artifact.vision_scales.empty()) { + out << ",\"vision_scales\":{\"dtype\":\"F32\",\"shape\":[" + << artifact.vision_scales.size() + << "],\"data_offsets\":[0," << vision_bytes << "]}"; + } + out << ",\"encoder_scales\":{\"dtype\":\"F32\",\"shape\":[" << artifact.encoder_scales.size() - << "],\"data_offsets\":[0," << encoder_bytes + << "],\"data_offsets\":[" << vision_bytes << ',' + << vision_bytes + encoder_bytes << "]},\"decoder_scales\":{\"dtype\":\"F32\",\"shape\":[" << artifact.decoder_scales.size() << "],\"data_offsets\":[" - << encoder_bytes << ',' << encoder_bytes + decoder_bytes << "]}}"; + << vision_bytes + encoder_bytes << ',' + << vision_bytes + encoder_bytes + decoder_bytes << "]}}"; std::string header = out.str(); while (header.size() % 8) header.push_back(' '); return header; @@ -240,7 +255,9 @@ bool copy_f32_tensor(const loader::SafetensorsFile& file, modalities::Status validate_native_calibration_artifact( const NativeCalibrationArtifact& artifact) { - if (artifact.hardware.empty() || !valid_digest(artifact.weights_sha256) || + if ((artifact.activation_dtype != kFloat16 && + artifact.activation_dtype != kBfloat16) || + artifact.hardware.empty() || !valid_digest(artifact.weights_sha256) || !valid_digest(artifact.tokenizer_sha256) || artifact.num_views < 1 || artifact.num_views > 3 || artifact.max_prompt_tokens < 1 || artifact.state_dim < 1 || artifact.chunk_size < 1 || @@ -252,7 +269,13 @@ modalities::Status validate_native_calibration_artifact( artifact.percentile < 0.0 || artifact.percentile > 100.0) { return invalid("native calibration metadata is invalid"); } - if (artifact.encoder_scales.size() != kEncoderScaleCount || + const bool vision_valid = + artifact.activation_dtype == kBfloat16 + ? artifact.vision_scales.size() == kVisionScaleCount && + valid_scales(artifact.vision_scales) + : artifact.vision_scales.empty(); + if (!vision_valid || + artifact.encoder_scales.size() != kEncoderScaleCount || artifact.decoder_scales.size() != static_cast(artifact.num_steps) * 18 * 4 || !valid_scales(artifact.encoder_scales) || @@ -291,6 +314,8 @@ modalities::Status save_native_calibration_artifact( } ok = write_all(fd, header_size, sizeof(header_size)) && write_all(fd, header.data(), header.size()) && + write_all(fd, artifact.vision_scales.data(), + artifact.vision_scales.size() * sizeof(float)) && write_all(fd, artifact.encoder_scales.data(), artifact.encoder_scales.size() * sizeof(float)) && write_all(fd, artifact.decoder_scales.data(), @@ -321,14 +346,18 @@ modalities::Status load_native_calibration_artifact( loader::SafetensorsFile file; if (!file.open(path)) return not_found(file.error()); const auto& metadata = file.metadata(); - if (!metadata_is(metadata, "schema", kSchema) || + const bool is_v1 = metadata_is(metadata, "schema", kSchemaV1) && + metadata_is(metadata, "tensor_dtype", kFloat16); + const bool is_v2 = metadata_is(metadata, "schema", kSchemaV2) && + metadata_is(metadata, "tensor_dtype", kBfloat16); + if ((!is_v1 && !is_v2) || !metadata_is(metadata, "model", kModel) || !metadata_is(metadata, "precision", kPrecision) || - !metadata_is(metadata, "tensor_dtype", kTensorDtype) || !metadata_is(metadata, "reducer", kReducer)) { return invalid("native calibration artifact schema is incompatible"); } NativeCalibrationArtifact loaded; + loaded.activation_dtype = is_v2 ? kBfloat16 : kFloat16; const auto hardware = metadata.find("hardware"); const auto weights = metadata.find("weights_sha256"); const auto tokenizer = metadata.find("tokenizer_sha256"); @@ -342,6 +371,8 @@ modalities::Status load_native_calibration_artifact( !parse_int(metadata, "vision_pool_factor", &loaded.vision_pool_factor) || !parse_u64(metadata, "sample_count", &loaded.sample_count) || !parse_double(metadata, "percentile", &loaded.percentile) || + (is_v2 && + !copy_f32_tensor(file, "vision_scales", &loaded.vision_scales)) || !copy_f32_tensor(file, "encoder_scales", &loaded.encoder_scales) || !copy_f32_tensor(file, "decoder_scales", &loaded.decoder_scales)) { return invalid("native calibration artifact metadata is incomplete"); diff --git a/cpp/models/pi05/src/native_device_weights.cpp b/cpp/models/pi05/src/native_device_weights.cpp index 8e70dcc9..a8bddb53 100644 --- a/cpp/models/pi05/src/native_device_weights.cpp +++ b/cpp/models/pi05/src/native_device_weights.cpp @@ -114,6 +114,43 @@ modalities::Status NativeDeviceWeightStore::upload_bytes( #endif } +modalities::Status NativeDeviceWeightStore::allocate( + const std::string& name, + const std::vector& shape, + NativeWeightDType dtype) { + std::lock_guard lock(upload_mutex_); + if (!ctx_ || name.empty() || weights_.find(name) != weights_.end()) { + return invalid("invalid device weight allocation"); + } + std::size_t elements = 0; + const std::size_t width = element_bytes(dtype); + if (!width || !element_count(shape, &elements) || !elements || + elements > std::numeric_limits::max() / width) { + return invalid("device weight allocation shape is invalid"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "device weight allocation requires the CUDA build"); +#else + const std::size_t bytes = elements * width; + frt_buffer buffer = frt_buffer_alloc(ctx_, name.c_str(), bytes); + if (!buffer) { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + cudaMemGetInfo(&free_bytes, &total_bytes); + std::ostringstream message; + message << "device weight allocation failed: " << name + << " requested=" << bytes << " free=" << free_bytes + << " total=" << total_bytes; + return modalities::Status::error(modalities::StatusCode::kBackend, + message.str()); + } + weights_.emplace(name, NativeDeviceWeight{buffer, shape, dtype}); + return modalities::Status::ok(); +#endif +} + const NativeDeviceWeight* NativeDeviceWeightStore::find( const std::string& name) const { const auto it = weights_.find(name); diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index 0ea2b9f4..5e487669 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -2,6 +2,7 @@ #include "flashrt/cpp/models/pi05/native_style_precompute.h" #include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" #include @@ -27,6 +28,26 @@ modalities::Status backend(const char* message) { message); } +modalities::Status upload_scales( + NativeWorkspace* workspace, + const char* name, + const std::vector& values) { + const NativeWorkspaceBuffer* destination = + workspace ? workspace->find(name) : nullptr; + if (!destination || destination->dtype != modalities::DType::kFloat32 || + destination->shape != + std::vector({values.size()}) || + values.empty()) { + return invalid("native RTX FP8 scale payload is invalid"); + } + const cudaError_t rc = cudaMemcpy( + frt_buffer_dptr(destination->buffer), values.data(), + values.size() * sizeof(float), cudaMemcpyHostToDevice); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("native RTX FP8 scale upload failed"); +} + } // namespace NativeGraphOwner::NativeGraphOwner( @@ -37,7 +58,10 @@ NativeGraphOwner::NativeGraphOwner( weights_(ctx), workspace_(ctx), attention_(ctx), - forward_(&driver_) {} + linear_(&driver_, config.precision == NativeGraphPrecision::kFp8E4M3 + ? NativeRtxLinearMode::kFp8Static + : NativeRtxLinearMode::kBf16), + forward_(&driver_, &linear_) {} NativeGraphOwner::~NativeGraphOwner() = default; @@ -69,7 +93,37 @@ std::unique_ptr NativeGraphOwner::create( if (status) *status = backend("native graph owner allocation failed"); return nullptr; } - modalities::Status st = owner->initialize(checkpoint_path); + modalities::Status st = owner->initialize(checkpoint_path, nullptr); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + if (status) *status = modalities::Status::ok(); + return owner; +} + +std::unique_ptr NativeGraphOwner::create( + const std::string& checkpoint_path, + const NativeGraphConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status) { + if (config.precision != NativeGraphPrecision::kFp8E4M3) { + if (status) *status = invalid("native RTX FP8 graph precision 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, &calibration); if (!st.ok_status()) { if (status) *status = st; return nullptr; @@ -79,7 +133,13 @@ std::unique_ptr NativeGraphOwner::create( } modalities::Status NativeGraphOwner::initialize( - const std::string& checkpoint_path) { + const std::string& checkpoint_path, + const NativeCalibrationArtifact* calibration) { + const bool fp8 = config_.precision == NativeGraphPrecision::kFp8E4M3; + if (fp8 != (calibration != nullptr) || + (calibration && calibration->activation_dtype != "bfloat16")) { + return invalid("native RTX FP8 calibration is incompatible"); + } const bool profile_setup = std::getenv("FLASHRT_PROFILE_NATIVE_SETUP"); const auto setup_begin = std::chrono::steady_clock::now(); auto checkpoint = setup_begin; @@ -101,11 +161,17 @@ modalities::Status NativeGraphOwner::initialize( NativeWeightMaterializer materializer(source, &weights_); NativeMaterializationOptions options; options.num_steps = config_.num_steps; - options.merge_decoder_gate_up = false; + options.merge_decoder_gate_up = fp8; options.include_embedding = true; modalities::Status st = materializer.materialize_all(options); if (!st.ok_status()) return st; report("materialize"); + if (fp8) { + NativeRtxWeightPacker packer(&weights_, &driver_); + st = packer.pack_all(); + if (!st.ok_status()) return st; + report("fp8_pack"); + } NativeWorkspaceConfig workspace_config; workspace_config.num_views = config_.num_views; @@ -113,8 +179,24 @@ modalities::Status NativeGraphOwner::initialize( workspace_config.chunk_size = config_.chunk_size; workspace_config.num_steps = config_.num_steps; workspace_config.vision_pool_factor = config_.vision_pool_factor; + workspace_config.flavor = fp8 ? NativeWorkspaceFlavor::kRtxFp8 + : NativeWorkspaceFlavor::kBf16; st = workspace_.allocate(workspace_config); if (!st.ok_status()) return st; + if (fp8) { + st = upload_scales( + &workspace_, "rtx_fp8_vision_scales", + calibration->vision_scales); + if (!st.ok_status()) return st; + st = upload_scales( + &workspace_, "rtx_fp8_encoder_scales", + calibration->encoder_scales); + if (!st.ok_status()) return st; + st = upload_scales( + &workspace_, "rtx_fp8_decoder_scales", + calibration->decoder_scales); + if (!st.ok_status()) return st; + } st = workspace_.expand_vision_position_embedding(weights_); if (!st.ok_status()) return st; @@ -139,10 +221,16 @@ modalities::Status NativeGraphOwner::initialize( } st = attention_driver_->status(); if (!st.ok_status()) return st; + if (fp8) { + st = autotune_fp8(); + if (!st.ok_status()) return st; + report("fp8_autotune"); + } report("workspace_style"); for (const char* name : {"observation_images_normalized", - "prompt_embedding", "diffusion_noise"}) { + "prompt_embedding", "encoder_x", + "diffusion_noise"}) { const NativeWorkspaceBuffer* buffer = workspace_.find(name); if (!buffer || cudaMemset(frt_buffer_dptr(buffer->buffer), 0, @@ -187,6 +275,69 @@ modalities::Status NativeGraphOwner::initialize( return modalities::Status::ok(); } +modalities::Status NativeGraphOwner::autotune_fp8() { + const int vision_sequence = config_.num_views * 256; + const int encoder_vision_sequence = workspace_.encoder_vision_sequence(); + const int encoder_sequence = workspace_.encoder_sequence(); + const int decoder_sequence = config_.chunk_size; + struct Shape { + const char* weight; + NativeRtxScaleSite site; + const char* output; + int m; + int n; + int k; + }; + const Shape shapes[] = { + {"vision_attn_qkv_w_0", {NativeRtxScaleDomain::kVision, 0}, + "vision_QKV", + vision_sequence, 3456, 1152}, + {"vision_attn_o_w_0", {NativeRtxScaleDomain::kVision, 1}, + "vision_x_norm", + vision_sequence, 1152, 1152}, + {"vision_ffn_up_w_0", {NativeRtxScaleDomain::kVision, 2}, + "vision_hidden", + vision_sequence, 4304, 1152}, + {"vision_ffn_down_w_0", {NativeRtxScaleDomain::kVision, 3}, + "vision_x_norm", + vision_sequence, 1152, 4304}, + {"encoder_multi_modal_projector_w", + {NativeRtxScaleDomain::kVision, 108}, "encoder_x", + encoder_vision_sequence, 2048, 1152}, + {"encoder_attn_qkv_w_0", {NativeRtxScaleDomain::kEncoder, 0}, + "encoder_QKV", + encoder_sequence, 2560, 2048}, + {"encoder_attn_o_w_0", {NativeRtxScaleDomain::kEncoder, 1}, + "encoder_x_norm", + encoder_sequence, 2048, 2048}, + {"encoder_ffn_gate_up_w_0", {NativeRtxScaleDomain::kEncoder, 2}, + "encoder_gate_merged", encoder_sequence, 32768, 2048}, + {"encoder_ffn_down_w_0", {NativeRtxScaleDomain::kEncoder, 3}, + "encoder_x_norm", + encoder_sequence, 2048, 16384}, + {"decoder_attn_qkv_w_0", {NativeRtxScaleDomain::kDecoder, 0}, + "decoder_QKV", + decoder_sequence, 2560, 1024}, + {"decoder_attn_o_w_0", {NativeRtxScaleDomain::kDecoder, 1}, + "x_normed_buf", + decoder_sequence, 1024, 2048}, + {"decoder_ffn_gate_up_w_0", {NativeRtxScaleDomain::kDecoder, 2}, + "decoder_gate_merged", decoder_sequence, 8192, 1024}, + {"decoder_ffn_down_w_0", {NativeRtxScaleDomain::kDecoder, 3}, + "x_normed_buf", + decoder_sequence, 1024, 4096}, + }; + for (const Shape& shape : shapes) { + const NativeWorkspaceBuffer* output = workspace_.find(shape.output); + if (!output) return invalid("native FP8 autotune output is missing"); + modalities::Status st = linear_.autotune( + weights_, &workspace_, shape.weight, shape.site, + frt_buffer_dptr(output->buffer), shape.m, shape.n, shape.k); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + modalities::Status NativeGraphOwner::record_context(void* stream) { modalities::Status st = copy_prompt_to_encoder(&workspace_, stream); if (!st.ok_status()) return st; diff --git a/cpp/models/pi05/src/native_kernel_driver.cu b/cpp/models/pi05/src/native_kernel_driver.cu index cff13040..2a4a7964 100644 --- a/cpp/models/pi05/src/native_kernel_driver.cu +++ b/cpp/models/pi05/src/native_kernel_driver.cu @@ -2,9 +2,11 @@ #include "activation.cuh" #include "elementwise.cuh" +#include "fusion.cuh" #include "gemm_runner.h" #include "norm.cuh" #include "patch_embed.cuh" +#include "quantize.cuh" #include "rope.cuh" #include @@ -93,6 +95,112 @@ modalities::Status NativeKernelDriver::bf16_nn( } } +modalities::Status NativeKernelDriver::fp8_nn_bf16( + void* a, + void* b, + void* output, + int m, + int n, + int k, + const float* activation_scale, + const float* weight_scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!a || !b || !output || !activation_scale || !weight_scale || + m <= 0 || n <= 0 || k <= 0) { + return invalid("native FP8 GEMM arguments are invalid"); + } + try { + impl_->gemm.fp8_nn_dev( + a, b, output, m, n, k, const_cast(activation_scale), + const_cast(weight_scale), + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("native FP8 GEMM launch failed"); + } +} + +modalities::Status NativeKernelDriver::autotune_fp8_nn_bf16( + void* a, + void* b, + void* output, + int m, + int n, + int k, + const float* activation_scale, + const float* weight_scale) const { + if (!impl_) return backend(error_); + if (!a || !b || !output || !activation_scale || !weight_scale || + m <= 0 || n <= 0 || k <= 0) { + return invalid("native FP8 GEMM autotune arguments are invalid"); + } + try { + impl_->gemm.autotune_fp8_nn_dev( + a, b, output, m, n, k, + const_cast(activation_scale), + const_cast(weight_scale)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("native FP8 GEMM autotune failed"); + } +} + +modalities::Status NativeKernelDriver::quantize_fp8_static_bf16( + const void* values, + void* output, + const float* scale, + std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !output || !scale || !elements || elements > INT_MAX) { + return invalid("native static FP8 quantization arguments are invalid"); + } + ::quantize_fp8_static( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), scale, + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::quantize_fp8_dynamic_bf16( + const void* values, + void* output, + float* scale, + std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !output || !scale || !elements || elements > INT_MAX) { + return invalid("native dynamic FP8 quantization arguments are invalid"); + } + ::quantize_fp8_device( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), scale, + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::quantize_fp8_weight_bf16( + const void* values, + void* output, + float* scale, + std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !output || !scale || !elements || elements > INT_MAX) { + return invalid("native FP8 weight quantization arguments are invalid"); + } + ::quantize_fp8_weight_device( + static_cast(values), + static_cast<__nv_fp8_e4m3*>(output), scale, + static_cast(elements), reinterpret_cast(stream)); + return launch_status(); +} + modalities::Status NativeKernelDriver::add_bias_bf16( void* values, const void* bias, @@ -152,6 +260,41 @@ modalities::Status NativeKernelDriver::gate_gelu_bf16( return launch_status(); } +modalities::Status NativeKernelDriver::gate_gelu_merged_bf16( + const void* merged, + void* output, + int rows, + int hidden, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!merged || !output || rows <= 0 || hidden <= 0) { + return invalid("native BF16 merged gated GELU arguments are invalid"); + } + ::gate_silu_mul_merged( + static_cast(merged), + static_cast<__nv_bfloat16*>(output), rows, hidden, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gate_gelu_merged_fp8_bf16( + const void* merged, + void* output, + int rows, + int hidden, + const float* scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!merged || !output || !scale || rows <= 0 || hidden <= 0) { + return invalid("native BF16 merged gated GELU FP8 arguments are invalid"); + } + ::gate_silu_mul_merged_fp8( + static_cast(merged), + static_cast<__nv_fp8_e4m3*>(output), rows, hidden, scale, + 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 { @@ -213,6 +356,51 @@ modalities::Status NativeKernelDriver::rms_norm_bf16( return launch_status(); } +modalities::Status NativeKernelDriver::rms_norm_fp8_bf16( + const void* values, + const void* weight, + void* output, + int rows, + int columns, + float epsilon, + const float* scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !weight || !output || !scale || rows <= 0 || columns <= 0) { + return invalid("native BF16 RMS norm FP8 arguments are invalid"); + } + ::rms_norm_fp8( + static_cast(values), + static_cast(weight), + static_cast<__nv_fp8_e4m3*>(output), rows, columns, epsilon, scale, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::residual_add_rms_norm_fp8_bf16( + void* residual, + const void* values, + const void* weight, + void* output, + int rows, + int columns, + float epsilon, + const float* scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!residual || !values || !weight || !output || !scale || rows <= 0 || + columns <= 0) { + return invalid("native BF16 residual RMS norm FP8 arguments are invalid"); + } + ::residual_add_rms_norm_fp8( + static_cast<__nv_bfloat16*>(residual), + static_cast(values), + static_cast(weight), + static_cast<__nv_fp8_e4m3*>(output), rows, columns, epsilon, scale, + 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 { @@ -248,6 +436,62 @@ modalities::Status NativeKernelDriver::ada_rms_norm_style_bf16( return launch_status(); } +modalities::Status NativeKernelDriver::ada_rms_norm_style_fp8_bf16( + const void* values, + const void* weight, + const void* style, + void* output, + void* gate_output, + int rows, + int columns, + float epsilon, + const float* scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !weight || !style || !output || !gate_output || !scale || + rows <= 0 || columns <= 0) { + return invalid("native BF16 Ada RMS norm FP8 arguments are invalid"); + } + ::ada_rms_norm_style_fp8( + static_cast(values), + static_cast(weight), + static_cast(style), + static_cast<__nv_fp8_e4m3*>(output), + static_cast<__nv_bfloat16*>(gate_output), rows, columns, epsilon, + scale, reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gate_residual_ada_norm_fp8_bf16( + void* residual, + const void* values, + const void* gate, + const void* weight, + const void* style, + void* output, + void* gate_output, + int rows, + int columns, + float epsilon, + const float* scale, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!residual || !values || !gate || !weight || !style || !output || + !gate_output || !scale || rows <= 0 || columns <= 0) { + return invalid("native BF16 gated residual Ada norm FP8 arguments are invalid"); + } + ::gate_residual_ada_norm_fp8( + static_cast<__nv_bfloat16*>(residual), + static_cast(values), + static_cast(gate), + static_cast(weight), + static_cast(style), + static_cast<__nv_fp8_e4m3*>(output), + static_cast<__nv_bfloat16*>(gate_output), rows, columns, epsilon, + scale, 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, diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 59b35ad8..3fa32915 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -98,7 +98,9 @@ int build_native_model_runtime(const NativeOpenConfig& config, Precision precision; if (config.precision == "auto") { if (properties.major == 12 && properties.minor == 0) { - precision = Precision::kBf16; + precision = config.calibration_path.empty() + ? Precision::kBf16 + : Precision::kFp8E4M3Fn; } else if (properties.major == 11 && properties.minor == 0) { precision = Precision::kFp8E4M3Fn; } else { @@ -121,8 +123,9 @@ int build_native_model_runtime(const NativeOpenConfig& config, return -3; } if (precision == Precision::kFp8E4M3Fn && - (properties.major != 11 || properties.minor != 0)) { - if (error) *error = "Pi0.5 native FP8 requires SM110"; + !((properties.major == 11 || properties.major == 12) && + properties.minor == 0)) { + if (error) *error = "Pi0.5 native FP8 requires SM110 or SM120"; return -3; } #if !defined(FLASHRT_CPP_WITH_FA2) @@ -132,11 +135,17 @@ int build_native_model_runtime(const NativeOpenConfig& config, } #endif #if !defined(FLASHRT_CPP_WITH_THOR_FP8) - if (precision == Precision::kFp8E4M3Fn) { + if (precision == Precision::kFp8E4M3Fn && properties.major == 11) { if (error) *error = "Pi0.5 native Thor FP8 backend is not built"; return -3; } #endif +#if !defined(FLASHRT_CPP_WITH_FA2) + if (precision == Precision::kFp8E4M3Fn && properties.major == 12) { + if (error) *error = "Pi0.5 native RTX FP8 backend is not built"; + return -3; + } +#endif struct HashResult { bool ok = false; @@ -163,7 +172,6 @@ int build_native_model_runtime(const NativeOpenConfig& config, NativeCalibrationArtifact calibration; std::string calibration_sha256; if (precision == Precision::kFp8E4M3Fn) { -#if defined(FLASHRT_CPP_WITH_THOR_FP8) if (config.calibration_path.empty()) { if (error) *error = "Pi0.5 native FP8 requires calibration_path"; return -1; @@ -176,6 +184,8 @@ int build_native_model_runtime(const NativeOpenConfig& config, return cface::status_code(calibration_status); } if (calibration.hardware != hardware_id || + calibration.activation_dtype != + (properties.major == 11 ? "float16" : "bfloat16") || calibration.tokenizer_sha256 != tokenizer_sha256 || calibration.num_views != config.num_views || calibration.max_prompt_tokens != config.max_prompt_tokens || @@ -191,7 +201,6 @@ int build_native_model_runtime(const NativeOpenConfig& config, if (error) *error = hash_error; return -2; } -#endif } NativeGraphConfig graph_config; @@ -200,14 +209,25 @@ int build_native_model_runtime(const NativeOpenConfig& config, graph_config.chunk_size = config.chunk; graph_config.num_steps = config.num_steps; graph_config.vision_pool_factor = config.vision_pool_factor; + graph_config.precision = precision == Precision::kFp8E4M3Fn + ? NativeGraphPrecision::kFp8E4M3 + : NativeGraphPrecision::kBf16; modalities::Status st; std::unique_ptr graph; - if (precision == Precision::kBf16) { + const bool thor_fp8 = precision == Precision::kFp8E4M3Fn && + properties.major == 11; + const bool rtx_fp8 = precision == Precision::kFp8E4M3Fn && + properties.major == 12; + if (precision == Precision::kBf16 || rtx_fp8) { #if defined(FLASHRT_CPP_WITH_FA2) - graph = NativeGraphOwner::create( - config.checkpoint_path, graph_config, &st); + graph = rtx_fp8 + ? NativeGraphOwner::create( + config.checkpoint_path, graph_config, calibration, + &st) + : NativeGraphOwner::create( + config.checkpoint_path, graph_config, &st); #endif - } else { + } else if (thor_fp8) { #if defined(FLASHRT_CPP_WITH_THOR_FP8) graph = NativeThorGraphOwner::create( config.checkpoint_path, graph_config, calibration, &st); @@ -246,9 +266,8 @@ int build_native_model_runtime(const NativeOpenConfig& config, "embedding_weight"); if (!images || !noise || !encoder || !previous || !prefix_weights || !guidance || !prompt || !embedding || - embedding->dtype != (precision == Precision::kBf16 - ? NativeWeightDType::kBf16 - : NativeWeightDType::kFloat16) || + embedding->dtype != (thor_fp8 ? NativeWeightDType::kFloat16 + : NativeWeightDType::kBf16) || embedding->shape.size() != 2 || embedding->shape[1] != 2048) { if (error) *error = "native graph export buffers are incomplete"; return -6; @@ -308,9 +327,12 @@ int build_native_model_runtime(const NativeOpenConfig& config, FRT_RT_REGION_SNAPSHOT | FRT_RT_REGION_RESTORE) == 0; if (!ok) return fail_builder(builder, error, "native region build failed"); - const bool thor_fp8 = precision == Precision::kFp8E4M3Fn; - const std::string precision_id = thor_fp8 ? "fp8_e4m3fn" : "bf16"; - const std::string pipeline_id = thor_fp8 ? "NativeThorFp8" : "NativeBf16"; + const bool fp8 = precision == Precision::kFp8E4M3Fn; + const std::string precision_id = fp8 ? "fp8_e4m3fn" : "bf16"; + const std::string pipeline_id = thor_fp8 + ? "NativeThorFp8" + : rtx_fp8 ? "NativeRtxFp8" + : "NativeBf16"; const std::string tensor_dtype = thor_fp8 ? "float16" : "bf16"; ok = add_identity(builder, "model", "pi05") && add_identity(builder, "producer", "native") && @@ -334,7 +356,7 @@ int build_native_model_runtime(const NativeOpenConfig& config, add_identity(builder, "model_action_dim", "32") && add_identity(builder, "robot_action_dim", std::to_string(config.action_q01.size())); - if (ok && thor_fp8) { + if (ok && fp8) { ok = add_identity(builder, "calibration_sha256", calibration_sha256); } if (!ok) return fail_builder(builder, error, "native identity build failed"); diff --git a/cpp/models/pi05/src/native_rtx_attention_driver.cu b/cpp/models/pi05/src/native_rtx_attention_driver.cu index c43cce7b..2850b769 100644 --- a/cpp/models/pi05/src/native_rtx_attention_driver.cu +++ b/cpp/models/pi05/src/native_rtx_attention_driver.cu @@ -38,6 +38,11 @@ bool exact_shape(const NativeAttentionBuffer* buffer, return buffer && buffer->shape == std::vector(expected); } +float inverse_sqrt(int dimension) { + // Match the Python producer: evaluate in binary64, then narrow once. + return static_cast(1.0 / std::sqrt(static_cast(dimension))); +} + } // namespace NativeRtxAttentionDriver::NativeRtxAttentionDriver( @@ -118,7 +123,7 @@ modalities::Status NativeRtxAttentionDriver::vision( 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_, + inverse_sqrt(72), num_sms_, reinterpret_cast(stream)); return launch_status(); } @@ -145,7 +150,7 @@ modalities::Status NativeRtxAttentionDriver::encoder( 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_, + inverse_sqrt(256), num_sms_, reinterpret_cast(stream)); return launch_status(); } @@ -183,7 +188,7 @@ modalities::Status NativeRtxAttentionDriver::decoder( 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), + q_batch_stride, q_row_stride, 256, inverse_sqrt(256), num_sms_, reinterpret_cast(stream)); return launch_status(); } diff --git a/cpp/models/pi05/src/native_rtx_linear.cpp b/cpp/models/pi05/src/native_rtx_linear.cpp new file mode 100644 index 00000000..9e018c7a --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_linear.cpp @@ -0,0 +1,212 @@ +#include "flashrt/cpp/models/pi05/native_rtx_linear.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +void* dptr(const NativeWorkspaceBuffer* buffer) { + return buffer ? frt_buffer_dptr(buffer->buffer) : nullptr; +} + +void* dptr(const NativeDeviceWeight* weight) { + return weight ? frt_buffer_dptr(weight->buffer) : nullptr; +} + +} // namespace + +const NativeDeviceWeight* NativeRtxLinear::find_weight( + const NativeDeviceWeightStore& weights, + const std::string& name) const { + if (!fp8()) return weights.find(name); + const std::string packed_name = + name == "encoder_multi_modal_projector_w" + ? "vision_projector_w" + : name; + return weights.find("fp8." + packed_name); +} + +bool NativeRtxLinear::weight_shape_is( + const NativeDeviceWeightStore& weights, + const std::string& name, + std::initializer_list shape) const { + const NativeDeviceWeight* weight = find_weight(weights, name); + return weight && + weight->dtype == (fp8() ? NativeWeightDType::kFp8E4M3 + : NativeWeightDType::kBf16) && + weight->shape == std::vector(shape); +} + +const NativeWorkspaceBuffer* NativeRtxLinear::scale_buffer( + const NativeWorkspace& workspace, + NativeRtxScaleDomain domain) const { + switch (domain) { + case NativeRtxScaleDomain::kVision: + return workspace.find("rtx_fp8_vision_scales"); + case NativeRtxScaleDomain::kEncoder: + return workspace.find("rtx_fp8_encoder_scales"); + case NativeRtxScaleDomain::kDecoder: + return workspace.find("rtx_fp8_decoder_scales"); + } + return nullptr; +} + +const float* NativeRtxLinear::scale( + const NativeWorkspace& workspace, + NativeRtxScaleSite site) const { + if (!fp8() || site.index < 0) return nullptr; + const NativeWorkspaceBuffer* scales = scale_buffer(workspace, site.domain); + if (!scales || scales->dtype != modalities::DType::kFloat32 || + scales->shape.size() != 1 || + static_cast(site.index) >= scales->shape[0]) { + return nullptr; + } + return static_cast(dptr(scales)) + site.index; +} + +modalities::Status NativeRtxLinear::run( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const std::string& weight_name, + NativeRtxScaleSite site, + const void* input, + void* output, + int m, + int n, + int k, + std::uintptr_t stream) const { + if (!driver_ || !workspace || weight_name.empty() || !input || !output || + m <= 0 || n <= 0 || k <= 0) { + return invalid("native RTX linear arguments are invalid"); + } + const NativeDeviceWeight* weight = find_weight(weights, weight_name); + if (!fp8()) { + if (!weight || weight->dtype != NativeWeightDType::kBf16) { + return invalid("native BF16 linear weight is invalid"); + } + return driver_->bf16_nn( + const_cast(input), dptr(weight), output, m, n, k, stream); + } + const std::string packed_name = + weight_name == "encoder_multi_modal_projector_w" + ? "vision_projector_w" + : weight_name; + const NativeDeviceWeight* weight_scale = + weights.find("fp8." + packed_name + ".scale"); + const NativeWorkspaceBuffer* scratch = + workspace->find("rtx_fp8_scratch"); + const NativeWorkspaceBuffer* scales = scale_buffer(*workspace, site.domain); + if (!weight || weight->dtype != NativeWeightDType::kFp8E4M3 || + !weight_scale || weight_scale->dtype != NativeWeightDType::kFloat32 || + weight_scale->shape != std::vector({1}) || + !scratch || scratch->dtype != modalities::DType::kUInt8 || + !scales || scales->dtype != modalities::DType::kFloat32 || + scales->shape.size() != 1 || site.index < 0 || + static_cast(site.index) >= scales->shape[0] || + static_cast(m) * static_cast(k) > + frt_buffer_bytes(scratch->buffer)) { + return invalid("native FP8 linear storage is invalid"); + } + auto* scale = static_cast(dptr(scales)) + site.index; + modalities::Status st = dynamic_fp8() + ? driver_->quantize_fp8_dynamic_bf16( + input, dptr(scratch), scale, + static_cast(m) * k, stream) + : driver_->quantize_fp8_static_bf16( + input, dptr(scratch), scale, + static_cast(m) * k, stream); + if (!st.ok_status()) return st; + return driver_->fp8_nn_bf16( + dptr(scratch), dptr(weight), output, m, n, k, scale, + static_cast(dptr(weight_scale)), stream); +} + +modalities::Status NativeRtxLinear::autotune( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const std::string& weight_name, + NativeRtxScaleSite site, + void* output, + int m, + int n, + int k) const { + if (!fp8() || !driver_ || !workspace || !output || m <= 0 || n <= 0 || + k <= 0) { + return invalid("native FP8 autotune arguments are invalid"); + } + const NativeDeviceWeight* weight = find_weight(weights, weight_name); + const std::string packed_name = + weight_name == "encoder_multi_modal_projector_w" + ? "vision_projector_w" + : weight_name; + const NativeDeviceWeight* weight_scale = + weights.find("fp8." + packed_name + ".scale"); + const NativeWorkspaceBuffer* scratch = + workspace->find("rtx_fp8_scratch"); + const NativeWorkspaceBuffer* scales = scale_buffer(*workspace, site.domain); + if (!weight || weight->dtype != NativeWeightDType::kFp8E4M3 || + weight->shape != std::vector( + {static_cast(k), + static_cast(n)}) || + !weight_scale || weight_scale->dtype != NativeWeightDType::kFloat32 || + weight_scale->shape != std::vector({1}) || !scratch || + !scales || scales->shape.size() != 1 || site.index < 0 || + static_cast(site.index) >= scales->shape[0] || + static_cast(m) * static_cast(k) > + frt_buffer_bytes(scratch->buffer)) { + return invalid("native FP8 autotune storage is invalid"); + } + const auto* scale = static_cast(dptr(scales)) + site.index; + return driver_->autotune_fp8_nn_bf16( + dptr(scratch), dptr(weight), output, m, n, k, scale, + static_cast(dptr(weight_scale))); +} + +modalities::Status NativeRtxLinear::run_prequantized( + const NativeDeviceWeightStore& weights, + const std::string& weight_name, + NativeRtxScaleSite site, + const NativeWorkspace& workspace, + const void* input, + void* output, + int m, + int n, + int k, + std::uintptr_t stream) const { + if (!fp8() || !driver_ || weight_name.empty() || !input || !output || + m <= 0 || n <= 0 || k <= 0) { + return invalid("native prequantized FP8 linear arguments are invalid"); + } + const NativeDeviceWeight* weight = find_weight(weights, weight_name); + const std::string packed_name = + weight_name == "encoder_multi_modal_projector_w" + ? "vision_projector_w" + : weight_name; + const NativeDeviceWeight* weight_scale = + weights.find("fp8." + packed_name + ".scale"); + const float* activation_scale = scale(workspace, site); + if (!weight || weight->dtype != NativeWeightDType::kFp8E4M3 || + weight->shape != std::vector( + {static_cast(k), + static_cast(n)}) || + !weight_scale || weight_scale->dtype != NativeWeightDType::kFloat32 || + weight_scale->shape != std::vector({1}) || + !activation_scale) { + return invalid("native prequantized FP8 linear storage is invalid"); + } + return driver_->fp8_nn_bf16( + const_cast(input), dptr(weight), output, m, n, k, + activation_scale, static_cast(dptr(weight_scale)), + stream); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_rtx_weight_packer.cpp b/cpp/models/pi05/src/native_rtx_weight_packer.cpp new file mode 100644 index 00000000..5511337f --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_weight_packer.cpp @@ -0,0 +1,149 @@ +#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.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 char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +void* dptr(const NativeDeviceWeight* weight) { + return weight ? frt_buffer_dptr(weight->buffer) : nullptr; +} + +} // namespace + +modalities::Status NativeRtxWeightPacker::pack_weight( + const std::string& source_name, + const std::string& packed_name) { + if (!weights_ || !driver_ || source_name.empty()) { + return invalid("native RTX weight packer is invalid"); + } + const NativeDeviceWeight* source = weights_->find(source_name); + if (!source || source->dtype != NativeWeightDType::kBf16 || + source->shape.size() != 2 || !source->shape[0] || !source->shape[1] || + source->shape[0] > std::numeric_limits::max() / + source->shape[1]) { + return invalid("native RTX FP8 source weight is invalid"); + } + const std::string name = packed_name.empty() ? source_name : packed_name; + const std::string prefix = "fp8." + name; + modalities::Status st = weights_->allocate( + prefix, source->shape, NativeWeightDType::kFp8E4M3); + if (!st.ok_status()) return st; + st = weights_->allocate( + prefix + ".scale", {1}, NativeWeightDType::kFloat32); + if (!st.ok_status()) return st; + const NativeDeviceWeight* output = weights_->find(prefix); + const NativeDeviceWeight* scale = weights_->find(prefix + ".scale"); + const std::size_t elements = static_cast(source->shape[0]) * + static_cast(source->shape[1]); + return driver_->quantize_fp8_weight_bf16( + dptr(source), dptr(output), static_cast(dptr(scale)), + elements, 0); +} + +modalities::Status NativeRtxWeightPacker::merge_bf16_columns( + const std::string& left_name, + const std::string& right_name, + const std::string& output_name) { + if (!weights_ || output_name.empty()) { + return invalid("native RTX merged weight arguments are invalid"); + } + const NativeDeviceWeight* left = weights_->find(left_name); + const NativeDeviceWeight* right = weights_->find(right_name); + if (!left || !right || left->dtype != NativeWeightDType::kBf16 || + right->dtype != NativeWeightDType::kBf16 || + left->shape.size() != 2 || right->shape != left->shape || + left->shape[1] > std::numeric_limits::max() / 2) { + return invalid("native RTX merged BF16 weights are invalid"); + } + const std::vector shape = { + left->shape[0], left->shape[1] * 2}; + modalities::Status st = weights_->allocate( + output_name, shape, NativeWeightDType::kBf16); + if (!st.ok_status()) return st; + const NativeDeviceWeight* output = weights_->find(output_name); + const std::size_t rows = static_cast(left->shape[0]); + const std::size_t columns = static_cast(left->shape[1]); + const std::size_t source_pitch = columns * sizeof(std::uint16_t); + const std::size_t output_pitch = source_pitch * 2; + auto* destination = static_cast(dptr(output)); + if (cudaMemcpy2DAsync( + destination, output_pitch, dptr(left), source_pitch, + source_pitch, rows, cudaMemcpyDeviceToDevice, nullptr) != + cudaSuccess || + cudaMemcpy2DAsync( + destination + source_pitch, output_pitch, dptr(right), + source_pitch, source_pitch, rows, cudaMemcpyDeviceToDevice, + nullptr) != cudaSuccess) { + return backend("native RTX merged BF16 copy failed"); + } + return modalities::Status::ok(); +} + +modalities::Status NativeRtxWeightPacker::pack_all() { + if (!weights_ || !driver_) { + return invalid("native RTX 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_weight(std::string(stem) + std::to_string(layer)); + if (!st.ok_status()) return st; + } + } + st = pack_weight( + "encoder_multi_modal_projector_w", "vision_projector_w"); + 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_weight(name); + 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_weight(name); + if (!st.ok_status()) return st; + } + } + return cudaDeviceSynchronize() == cudaSuccess + ? modalities::Status::ok() + : backend("native RTX FP8 weight packing failed"); +} + +} // 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 67501029..01c7ef2e 100644 --- a/cpp/models/pi05/src/native_workspace.cpp +++ b/cpp/models/pi05/src/native_workspace.cpp @@ -5,6 +5,7 @@ #include #endif +#include #include #include @@ -517,6 +518,18 @@ modalities::Status NativeWorkspace::allocate( 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); + if (flavor_ == NativeWorkspaceFlavor::kRtxFp8) { + const std::uint64_t scratch_elements = std::max({ + vs * 4304, es * 16384, ds * 4096}); + FRT_ADD("rtx_fp8_scratch", {scratch_elements}, + modalities::DType::kUInt8); + FRT_ADD("rtx_fp8_vision_scales", {109}, + modalities::DType::kFloat32); + FRT_ADD("rtx_fp8_encoder_scales", {18 * 4}, + modalities::DType::kFloat32); + FRT_ADD("rtx_fp8_decoder_scales", {steps * 18 * 4}, + modalities::DType::kFloat32); + } #undef FRT_ADD const NativeWorkspaceBuffer* decoder = find("decoder_rope_weights"); if (!decoder) return invalid("decoder RoPE buffer was not allocated"); diff --git a/cpp/models/pi05/tests/CMakeLists.txt b/cpp/models/pi05/tests/CMakeLists.txt index 73ae8715..64cf8bcd 100644 --- a/cpp/models/pi05/tests/CMakeLists.txt +++ b/cpp/models/pi05/tests/CMakeLists.txt @@ -44,7 +44,8 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_executable(pi05_native_quant_probe ${_pi05_test_source_dir}/pi05_native_quant_probe.cpp) - target_link_libraries(pi05_native_quant_probe PRIVATE flashrt_cpp_pi05) + target_link_libraries(pi05_native_quant_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) add_executable(test_pi05_native_weight_packer ${_pi05_test_source_dir}/test_pi05_native_weight_packer.cpp) diff --git a/cpp/tests/gate_pi05_native_quantization.py b/cpp/tests/gate_pi05_native_quantization.py index 50d424ad..2640192c 100644 --- a/cpp/tests/gate_pi05_native_quantization.py +++ b/cpp/tests/gate_pi05_native_quantization.py @@ -7,6 +7,14 @@ DECODER = "paligemma_with_expert.gemma_expert.model.layers.0" +VISION = ( + "paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".encoder.layers.4" +) +VISION_ROOT = ( + "paligemma_with_expert.paligemma.model.vision_tower.vision_model" + ".encoder.layers.0" +) def interleave_qk(weight: torch.Tensor, num_heads: int) -> torch.Tensor: @@ -78,6 +86,39 @@ def bf16(key: str) -> torch.Tensor: 1, digest(scale_tensor), ) + if layout == "kn": + expected["decoder_qkv0_fp8_kn_gpu"] = expected[ + "decoder_qkv0_fp8_kn" + ] + + vision_down = bf16(f"{VISION}.mlp.fc2.weight").t().contiguous().cuda() + scale = max(vision_down.float().abs().max().item() / 448.0, 1e-12) + quantized = (vision_down.float() / scale).clamp(-448.0, 448.0).to( + torch.float8_e4m3fn + ) + scale_tensor = torch.tensor([scale], dtype=torch.float32, device="cuda") + expected["vision_ffn_down4_fp8_kn_gpu"] = ( + tuple(quantized.shape), + digest(quantized.view(torch.uint8)), + 1, + digest(scale_tensor), + ) + + vision_qkv = torch.cat([ + bf16(f"{VISION_ROOT}.self_attn.{stem}_proj.weight") + for stem in ("q", "k", "v") + ], dim=0).t().contiguous().cuda() + scale = max(vision_qkv.float().abs().max().item() / 448.0, 1e-12) + quantized = (vision_qkv.float() / scale).clamp(-448.0, 448.0).to( + torch.float8_e4m3fn + ) + scale_tensor = torch.tensor([scale], dtype=torch.float32, device="cuda") + expected["vision_attn_qkv0_fp8_kn_gpu"] = ( + tuple(quantized.shape), + digest(quantized.view(torch.uint8)), + 1, + digest(scale_tensor), + ) transposed = weight.float().transpose(0, 1).contiguous() scales = torch.clamp( diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp index 6d644225..68ceaf50 100644 --- a/cpp/tests/pi05_native_open_probe.cpp +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -32,6 +32,43 @@ bool all_graph_variants_stable(const frt_runtime_export_v1* exp) { return true; } +const frt_runtime_buffer_desc* find_buffer( + const frt_runtime_export_v1* exp, + const char* name) { + if (!exp || !name) return nullptr; + for (std::uint64_t index = 0; index < exp->n_buffers; ++index) { + if (std::strcmp(exp->buffers[index].name, name) == 0) { + return &exp->buffers[index]; + } + } + return nullptr; +} + +bool write_diagnostics( + const frt_runtime_export_v1* exp, + const char* path) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) return false; + for (const char* name : { + "observation_images_normalized", + "prompt_embedding", + "encoder_x"}) { + const frt_runtime_buffer_desc* buffer = find_buffer(exp, name); + if (!buffer || !buffer->handle || !buffer->bytes) return false; + std::vector host( + static_cast(buffer->bytes)); + if (cudaMemcpy(host.data(), frt_buffer_dptr(buffer->handle), + host.size(), cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + output.write( + reinterpret_cast(host.data()), + static_cast(host.size())); + if (!output) return false; + } + return true; +} + } // namespace int main(int argc, char** argv) { @@ -471,6 +508,35 @@ int main(int argc, char** argv) { return 1; } } + const char* raw_output = std::getenv("FLASHRT_RAW_ACTION_OUTPUT"); + if (raw_output && raw_output[0] != '\0') { + std::vector raw(host_noise.size()); + if (!model->ports[5].buffer || + cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), + raw.size() * sizeof(raw[0]), + cudaMemcpyDeviceToHost) != cudaSuccess) { + std::cerr << "native raw action download failed\n"; + model->release(model->owner); + return 1; + } + std::ofstream output(raw_output, std::ios::binary | std::ios::trunc); + output.write(reinterpret_cast(raw.data()), + static_cast( + raw.size() * sizeof(raw[0]))); + if (!output) { + std::cerr << "native raw action output failed\n"; + model->release(model->owner); + return 1; + } + } + const char* diagnostic_output = + std::getenv("FLASHRT_DIAGNOSTIC_OUTPUT"); + if (diagnostic_output && diagnostic_output[0] != '\0' && + !write_diagnostics(exp, diagnostic_output)) { + std::cerr << "native diagnostic output failed\n"; + model->release(model->owner); + return 1; + } model->retain(model->owner); model->release(model->owner); model->release(model->owner); diff --git a/cpp/tests/pi05_native_quant_probe.cpp b/cpp/tests/pi05_native_quant_probe.cpp index b7a59677..da58b579 100644 --- a/cpp/tests/pi05_native_quant_probe.cpp +++ b/cpp/tests/pi05_native_quant_probe.cpp @@ -1,6 +1,14 @@ #include "flashrt/cpp/models/pi05/native_quantization.h" +#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" +#include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include "flashrt/exec.h" + +#include #include +#include #include #include #include @@ -97,13 +105,20 @@ void print_shape(const std::vector& shape) { void print_result(const std::vector& shape, const void* values, std::size_t value_bytes, const std::vector& scales) { + std::uint32_t first_scale_bits = 0; + if (!scales.empty()) { + static_assert(sizeof(first_scale_bits) == sizeof(scales.front())); + std::memcpy(&first_scale_bits, scales.data(), sizeof(first_scale_bits)); + } 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'; + << fnv1a(scales.data(), scales.size() * sizeof(float)) + << " first_scale_bits=" << std::setw(8) + << first_scale_bits << '\n'; } } // namespace @@ -119,6 +134,91 @@ int main(int argc, char** argv) { return 2; } const std::string op = argv[2]; + if (op == "decoder_qkv0_fp8_kn_gpu" || + op == "vision_attn_qkv0_fp8_kn_gpu" || + op == "vision_ffn_down4_fp8_kn_gpu") { + frt_ctx ctx = frt_ctx_create(); + if (!ctx) { + std::cerr << "failed to create FlashRT context\n"; + return 1; + } + std::vector shape; + std::vector values; + std::vector scales; + std::uint64_t input_hash = 0; + std::string error; + { + flashrt::models::pi05::NativeDeviceWeightStore weights(ctx); + flashrt::models::pi05::NativeWeightMaterializer materializer( + file, &weights); + const bool vision_qkv = op.find("vision_attn") == 0; + const bool vision_down = op.find("vision_ffn") == 0; + const std::string name = + vision_qkv ? "vision_attn_qkv_w_0" : + vision_down ? "vision_ffn_down_w_4" : + "decoder_attn_qkv_w_0"; + Status st = vision_qkv + ? materializer.materialize_vision_layer(0) + : vision_down + ? materializer.materialize_vision_layer(4) + : materializer.materialize_decoder_layer(0, true); + flashrt::models::pi05::NativeKernelDriver driver; + if (st.ok_status()) st = driver.status(); + flashrt::models::pi05::NativeRtxWeightPacker packer( + &weights, &driver); + if (st.ok_status()) st = packer.pack_weight(name); + if (!st.ok_status()) { + error = st.message; + } else { + const auto* input = weights.find(name); + const auto* output = weights.find("fp8." + name); + const auto* scale = weights.find("fp8." + name + ".scale"); + if (!input || !output || !scale) { + error = "GPU FP8 weight pack output is missing"; + } else { + std::vector input_values( + frt_buffer_bytes(input->buffer)); + shape = output->shape; + values.resize(frt_buffer_bytes(output->buffer)); + scales.resize(1); + if (cudaMemcpy( + input_values.data(), frt_buffer_dptr(input->buffer), + input_values.size(), cudaMemcpyDeviceToHost) != + cudaSuccess || + cudaMemcpy( + values.data(), frt_buffer_dptr(output->buffer), + values.size(), cudaMemcpyDeviceToHost) != + cudaSuccess || + cudaMemcpy( + scales.data(), frt_buffer_dptr(scale->buffer), + sizeof(float), cudaMemcpyDeviceToHost) != + cudaSuccess) { + error = "GPU FP8 weight download failed"; + } else { + input_hash = fnv1a( + input_values.data(), input_values.size()); + } + } + } + } + frt_ctx_destroy(ctx); + if (!error.empty()) { + std::cerr << error << '\n'; + return 1; + } + std::cout << "input_fnv=" << std::hex << std::setw(16) + << std::setfill('0') << input_hash << ' '; + if (argc == 4) { + std::ofstream output(argv[3], std::ios::binary | std::ios::trunc); + output.write( + reinterpret_cast(values.data()), + static_cast(values.size())); + if (!output) return 1; + } + print_result( + shape, values.data(), values.size(), scales); + return 0; + } if (op == "encoder_gate_up0_fp8") { NativeF16Tensor weight; NativeFp8Tensor output; diff --git a/cpp/tests/pi05_native_vision_probe.cpp b/cpp/tests/pi05_native_vision_probe.cpp index 4ef4771b..83a3b98b 100644 --- a/cpp/tests/pi05_native_vision_probe.cpp +++ b/cpp/tests/pi05_native_vision_probe.cpp @@ -1,8 +1,11 @@ #include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" #include "flashrt/cpp/models/pi05/native_weight_materializer.h" #include +#include #include #include #include @@ -43,14 +46,201 @@ bool write_buffer(std::ofstream* file, const void* device, return file->good(); } +bool download_buffer( + const void* device, + std::size_t elements, + std::vector* host) { + if (!device || !host) return false; + host->resize(elements); + return cudaMemcpy( + host->data(), device, elements * sizeof(std::uint16_t), + cudaMemcpyDeviceToHost) == cudaSuccess; +} + +bool upload_scales( + flashrt::models::pi05::NativeWorkspace* workspace, + const std::vector& scales) { + const auto* output = workspace->find("rtx_fp8_vision_scales"); + return output && output->shape == + std::vector({scales.size()}) && + cudaMemcpy(frt_buffer_dptr(output->buffer), scales.data(), + scales.size() * sizeof(float), + cudaMemcpyHostToDevice) == cudaSuccess; +} + +bool read_images( + const char* path, + std::vector* images) { + if (!path || !images) return false; + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file || + file.tellg() != static_cast( + images->size() * sizeof(std::uint16_t))) { + return false; + } + file.seekg(0); + file.read( + reinterpret_cast(images->data()), + static_cast( + images->size() * sizeof(std::uint16_t))); + return file.good(); +} + +flashrt::modalities::Status write_vision_layer_diagnostics( + const char* path, + const flashrt::models::pi05::NativeKernelDriver& driver, + const flashrt::models::pi05::NativeBf16Forward& forward, + const flashrt::models::pi05::NativeDeviceWeightStore& weights, + flashrt::models::pi05::NativeWorkspace* workspace, + flashrt::models::pi05::NativeRtxAttentionWorkspace* attention, + const flashrt::models::pi05::NativeRtxAttentionDriver& attention_driver, + cudaStream_t stream) { + using flashrt::models::pi05::NativeWorkspaceBuffer; + 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 auto* patch_weight = weights.find("vision_patch_embedding_w"); + const auto* patch_bias = weights.find("vision_patch_embedding_b"); + const auto* norm_weight = weights.find("vision_pre_attn_norm_w_0"); + const auto* norm_bias = weights.find("vision_pre_attn_norm_b_0"); + if (!images || !patches || !position || !x || !x_norm || + !patch_weight || !patch_bias || !norm_weight || !norm_bias) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kInvalidArgument, + "native vision diagnostic buffers are incomplete"); + } + const auto ptr = [](const auto* value) { + return frt_buffer_dptr(value->buffer); + }; + const int sequence = workspace->vision_sequence(); + const std::uintptr_t native_stream = + reinterpret_cast(stream); + flashrt::modalities::Status st = driver.patch_im2col_16bit( + ptr(images), ptr(patches), workspace->num_views(), native_stream); + if (st.ok_status()) { + st = driver.bf16_nn( + ptr(patches), ptr(patch_weight), ptr(x), sequence, 1152, 588, + native_stream); + } + if (st.ok_status()) { + st = driver.bias_residual_bf16( + ptr(x), ptr(position), ptr(patch_bias), sequence, 1152, + native_stream); + } + if (st.ok_status()) { + st = driver.layer_norm_bf16( + ptr(x), ptr(norm_weight), ptr(norm_bias), ptr(x_norm), sequence, + 1152, 1.0e-5f, native_stream); + } + if (!st.ok_status()) return st; + + int diagnostic_layer = 0; + if (const char* value = + std::getenv("FLASHRT_VISION_DIAGNOSTIC_LAYER")) { + char* end = nullptr; + const long parsed = std::strtol(value, &end, 10); + if (!end || *end != '\0' || parsed < 0 || parsed >= 27) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kInvalidArgument, + "native vision diagnostic layer is invalid"); + } + diagnostic_layer = static_cast(parsed); + } + + std::ofstream output(path, std::ios::binary | std::ios::trunc); + const std::size_t elements = + static_cast(sequence) * 1152; + if (!output || cudaStreamSynchronize(stream) != cudaSuccess || + !write_buffer(&output, ptr(x), elements)) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kBackend, + "native vision diagnostic output failed"); + } + std::vector layer_qkv; + std::vector layer_attention; + std::vector layer_hidden; + for (int layer = 0; layer < 27; ++layer) { + st = forward.vision_layer( + layer, weights, workspace, attention, &attention_driver, + native_stream); + if (!st.ok_status()) return st; + if (cudaStreamSynchronize(stream) != cudaSuccess || + !write_buffer(&output, ptr(x), elements)) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kBackend, + "native vision layer diagnostic output failed"); + } + if (layer == diagnostic_layer) { + const NativeWorkspaceBuffer* qkv = + workspace->find("vision_QKV"); + const NativeWorkspaceBuffer* hidden = + workspace->find("vision_hidden"); + const auto* attention_output = + attention->find("attn_vis_O"); + if (!qkv || !hidden || !attention_output || + !download_buffer( + ptr(qkv), static_cast(sequence) * 3456, + &layer_qkv) || + !download_buffer( + frt_buffer_dptr(attention_output->buffer), + static_cast(sequence) * 1152, + &layer_attention) || + !download_buffer( + ptr(hidden), static_cast(sequence) * 4304, + &layer_hidden)) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kBackend, + "native vision sublayer diagnostics failed"); + } + } + } + for (const auto* values : { + &layer_qkv, + &layer_attention, + &layer_hidden}) { + output.write( + reinterpret_cast(values->data()), + static_cast( + values->size() * sizeof(std::uint16_t))); + if (!output) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kBackend, + "native vision diagnostic write failed"); + } + } + return flashrt::modalities::Status::ok(); +} + } // namespace int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_vision_probe CHECKPOINT OUTPUT\n"; + if (argc != 3 && argc != 5) { + std::cerr << "usage: pi05_native_vision_probe CHECKPOINT OUTPUT " + "[CALIBRATION INPUT_BF16]\n"; return 2; } using namespace flashrt::models::pi05; + const bool fp8 = argc == 5; + NativeCalibrationArtifact calibration; + if (fp8) { + const flashrt::modalities::Status calibration_status = + load_native_calibration_artifact(argv[3], &calibration); + if (!calibration_status.ok_status() || + calibration.activation_dtype != "bfloat16" || + calibration.hardware != "sm120") { + std::cerr << (calibration_status.ok_status() + ? "SM120 FP8 calibration is required" + : calibration_status.message) + << '\n'; + return 2; + } + } flashrt::loader::SafetensorsFile source; if (!source.open(std::string(argv[1]) + "/model.safetensors")) { std::cerr << source.error() << '\n'; @@ -64,12 +254,54 @@ int main(int argc, char** argv) { for (int layer = 0; layer < 27 && st.ok_status(); ++layer) { st = materializer.materialize_vision_layer(layer); } + NativeKernelDriver driver; + if (st.ok_status()) st = driver.status(); + if (fp8 && st.ok_status()) { + NativeRtxWeightPacker packer(&weights, &driver); + for (int layer = 0; layer < 27 && st.ok_status(); ++layer) { + for (const char* stem : { + "vision_attn_qkv_w_", "vision_attn_o_w_", + "vision_ffn_up_w_", "vision_ffn_down_w_"}) { + st = packer.pack_weight( + std::string(stem) + std::to_string(layer)); + if (!st.ok_status()) break; + } + } + if (st.ok_status()) { + st = packer.pack_weight( + "encoder_multi_modal_projector_w", "vision_projector_w"); + } + } 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()) { + NativeWorkspaceConfig workspace_config; + if (fp8) { + workspace_config.num_views = calibration.num_views; + workspace_config.max_prompt_tokens = calibration.max_prompt_tokens; + workspace_config.chunk_size = calibration.chunk_size; + workspace_config.num_steps = calibration.num_steps; + workspace_config.vision_pool_factor = + calibration.vision_pool_factor; + workspace_config.flavor = NativeWorkspaceFlavor::kRtxFp8; + } + NativeRtxAttentionConfig attention_config; + if (st.ok_status()) st = workspace.allocate(workspace_config); + if (st.ok_status()) { + st = workspace.expand_vision_position_embedding(weights); + } + attention_config.num_views = workspace.num_views(); + attention_config.encoder_sequence = workspace.encoder_sequence(); + attention_config.encoder_vision_sequence = + workspace.encoder_vision_sequence(); + attention_config.chunk_size = workspace.chunk_size(); + if (st.ok_status()) st = attention.allocate(attention_config); + if (st.ok_status() && fp8 && + !upload_scales(&workspace, calibration.vision_scales)) { + st = flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kBackend, + "native vision scale upload failed"); + } + if (!st.ok_status()) { std::cerr << st.message << '\n'; frt_ctx_destroy(ctx); return 1; @@ -77,11 +309,21 @@ int main(int argc, char** argv) { 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); + std::vector host_images( + static_cast(workspace.num_views()) * 224 * 224 * 3); + if (fp8) { + if (!read_images(argv[4], &host_images)) { + std::cerr << "native vision input is invalid\n"; + frt_ctx_destroy(ctx); + return 1; + } + } else { + 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(), @@ -90,10 +332,32 @@ int main(int argc, char** argv) { 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); + NativeRtxLinear linear( + &driver, fp8 ? NativeRtxLinearMode::kFp8Static + : NativeRtxLinearMode::kBf16); + NativeBf16Forward forward(&driver, &linear); + const int vision_sequence = workspace.vision_sequence(); + if (fp8 && std::getenv("FLASHRT_VISION_LAYER_DIAGNOSTICS")) { + cudaStream_t diagnostic_stream = nullptr; + if (cudaStreamCreate(&diagnostic_stream) != cudaSuccess) { + frt_ctx_destroy(ctx); + return 1; + } + st = write_vision_layer_diagnostics( + argv[2], driver, forward, weights, &workspace, &attention, + attention_driver, diagnostic_stream); + cudaStreamDestroy(diagnostic_stream); + frt_ctx_destroy(ctx); + if (!st.ok_status()) { + std::cerr << st.message << '\n'; + return 1; + } + std::cout << "PASS native vision layer diagnostics\n"; + return 0; + } + frt_graph graph = frt_graph_create( + ctx, "native_vision", vision_sequence); cudaStream_t stream = nullptr; if (!graph || cudaStreamCreate(&stream) != cudaSuccess || frt_graph_bind(graph, "images", images->buffer) != FRT_OK || @@ -105,7 +369,7 @@ int main(int argc, char** argv) { CaptureArgs capture{&forward, &weights, &workspace, &attention, &attention_driver, false, {}}; const int capture_rc = frt_graph_capture( - graph, 512, record_vision, &capture); + graph, vision_sequence, record_vision, &capture); if (capture_rc != FRT_OK || !capture.recorded) { std::cerr << "vision capture failed: rc=" << capture_rc << " status=" << capture.error << '\n'; @@ -119,7 +383,7 @@ int main(int argc, char** argv) { 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_replay(graph, vision_sequence, stream_id) != FRT_OK) { frt_graph_destroy(graph); cudaStreamDestroy(stream); frt_ctx_destroy(ctx); @@ -135,8 +399,13 @@ int main(int argc, char** argv) { } 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); + write_buffer( + &file, frt_buffer_dptr(vision_x->buffer), + static_cast(vision_sequence) * 1152) && + write_buffer( + &file, frt_buffer_dptr(encoder_x->buffer), + static_cast(workspace.encoder_vision_sequence()) * + 2048); frt_graph_destroy(graph); cudaStreamDestroy(stream); frt_ctx_destroy(ctx); diff --git a/cpp/tests/test_pi05_native_calibration.cpp b/cpp/tests/test_pi05_native_calibration.cpp index 6759528d..2069ceb0 100644 --- a/cpp/tests/test_pi05_native_calibration.cpp +++ b/cpp/tests/test_pi05_native_calibration.cpp @@ -79,11 +79,28 @@ int main() { assert(loaded.num_steps == expected.num_steps); assert(loaded.vision_pool_factor == expected.vision_pool_factor); assert(loaded.sample_count == expected.sample_count); + assert(loaded.activation_dtype == expected.activation_dtype); + assert(loaded.vision_scales == expected.vision_scales); assert(std::fabs(loaded.percentile - expected.percentile) < 1e-12); assert(loaded.encoder_scales == expected.encoder_scales); assert(loaded.decoder_scales == expected.decoder_scales); assert(::unlink(path.c_str()) == 0); + expected.activation_dtype = "bfloat16"; + expected.hardware = "sm120"; + expected.vision_scales.resize(27 * 4 + 1); + for (std::size_t i = 0; i < expected.vision_scales.size(); ++i) { + expected.vision_scales[i] = 0.002f * static_cast(i + 1); + } + assert(save_native_calibration_artifact(path, expected).ok_status()); + assert(load_native_calibration_artifact(path, &loaded).ok_status()); + assert(loaded.activation_dtype == expected.activation_dtype); + assert(loaded.hardware == expected.hardware); + assert(loaded.vision_scales == expected.vision_scales); + assert(loaded.encoder_scales == expected.encoder_scales); + assert(loaded.decoder_scales == expected.decoder_scales); + assert(::unlink(path.c_str()) == 0); + expected.weights_sha256 = "short"; assert(!save_native_calibration_artifact(path, expected).ok_status()); std::printf("PASS - Pi0.5 native calibration artifact\n"); diff --git a/csrc/gemm/gemm_runner.cu b/csrc/gemm/gemm_runner.cu index 2cf0bc35..2c732245 100644 --- a/csrc/gemm/gemm_runner.cu +++ b/csrc/gemm/gemm_runner.cu @@ -1,4 +1,7 @@ #include "gemm_runner.h" + +#include +#include #include #include @@ -144,6 +147,11 @@ void GemmRunner::autotune_cached(CachedGemm& entry, void* A, void* B, void* D, CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(entry.matmul_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &d_scale_b, sizeof(d_scale_b))); } + const char* disable_autotune = + std::getenv("FLASHRT_DISABLE_GEMM_AUTOTUNE"); + if (disable_autotune && std::strcmp(disable_autotune, "1") == 0) { + return; + } auto C_layout = entry.has_C_desc ? entry.C_desc : entry.D_desc; diff --git a/csrc/kernels/quantize.cu b/csrc/kernels/quantize.cu index a25a7be7..f696183a 100644 --- a/csrc/kernels/quantize.cu +++ b/csrc/kernels/quantize.cu @@ -49,7 +49,8 @@ __global__ void quantize_fp8_kernel(const __half* in, __nv_fp8_e4m3* out, const __nv_fp8_e4m3 fp8_pack[4]; #pragma unroll for (int j = 0; j < 4; j++) { - fp8_pack[j] = __nv_fp8_e4m3(fminf(fmaxf(fv[j] * inv_scale, -448.f), 448.f)); + fp8_pack[j] = __nv_fp8_e4m3( + fminf(fmaxf(fv[j] * inv_scale, -448.f), 448.f)); } *reinterpret_cast(out + i) = *reinterpret_cast(fp8_pack); } @@ -75,6 +76,28 @@ __global__ void quantize_fp8_kernel_generic(const T* __restrict__ input, template __global__ void quantize_fp8_kernel_generic<__half>(const __half*, __nv_fp8_e4m3*, const float*, int); template __global__ void quantize_fp8_kernel_generic<__nv_bfloat16>(const __nv_bfloat16*, __nv_fp8_e4m3*, const float*, int); +// Weight packing runs during setup. Match the producer's scalar division as a +// correctly-rounded FP32 reciprocal followed by a correctly-rounded multiply. +// Explicit intrinsics keep this stable under --use_fast_math. +__global__ void quantize_fp8_weight_kernel( + const __nv_bfloat16* __restrict__ input, + __nv_fp8_e4m3* __restrict__ output, + const float* scale, int n) { + using T2 = typename packed2<__nv_bfloat16>::type; + const T2* in2 = reinterpret_cast(input); + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < (n >> 1)) { + const float inv_scale = __fdiv_rn(1.0f, *scale); + const T2 value = in2[idx]; + const float v0 = __fmul_rn(to_f32(value.x), inv_scale); + const float v1 = __fmul_rn(to_f32(value.y), inv_scale); + output[2 * idx] = __nv_fp8_e4m3( + fminf(fmaxf(v0, -448.0f), 448.0f)); + output[2 * idx + 1] = __nv_fp8_e4m3( + fminf(fmaxf(v1, -448.0f), 448.0f)); + } +} + float quantize_fp8(const __nv_bfloat16* input, __nv_fp8_e4m3* output, float* d_scale, int n, cudaStream_t stream) { float* d_max; @@ -128,7 +151,9 @@ float quantize_fp8_fp16(const __half* input, __nv_fp8_e4m3* output, // ── FP8 Quantize Device-Only (CUDA Graph compatible) ── __global__ void compute_scale_kernel(const float* d_absmax, float* d_scale) { float amax = *d_absmax; - float scale = amax / 448.0f; + // Producers compute this scalar with IEEE round-to-nearest semantics. + // Keep the result stable even when the translation unit uses fast math. + float scale = __fdiv_rn(amax, 448.0f); if (scale < 1e-12f) scale = 1e-12f; *d_scale = scale; } @@ -308,6 +333,24 @@ void quantize_fp8_device(const __nv_bfloat16* input, __nv_fp8_e4m3* output, quantize_fp8_kernel_generic<__nv_bfloat16><<>>(input, output, d_scale, n); } +void quantize_fp8_weight_device( + const __nv_bfloat16* input, __nv_fp8_e4m3* output, + float* d_scale, int n, cudaStream_t stream) { + cudaMemsetAsync(d_scale, 0, sizeof(float), stream); + + int threads = 256; + int blocks = (n + threads - 1) / threads; + if (blocks > 1024) blocks = 1024; + absmax_kernel<__nv_bfloat16> + <<>>( + input, d_scale, n); + compute_scale_kernel<<<1, 1, 0, stream>>>(d_scale, d_scale); + + blocks = ((n >> 1) + threads - 1) / threads; + quantize_fp8_weight_kernel<<>>( + input, output, d_scale, n); +} + void quantize_fp8_device_fp16(const __half* input, __nv_fp8_e4m3* output, float* d_scale, int n, cudaStream_t stream) { cudaMemsetAsync(d_scale, 0, sizeof(float), stream); @@ -2553,7 +2596,6 @@ __global__ void quantize_bf16_to_mxfp4_cutlass_kernel( float amax = shared[kb]; float scale = amax / 6.0f; if (scale < 1e-12f) scale = 1e-12f; - float inv_scale = 1.0f / scale; uint8_t ue8m0_scale = float_to_ue8m0_ceil(scale); diff --git a/csrc/kernels/quantize.cuh b/csrc/kernels/quantize.cuh index cdc7bded..9473a274 100644 --- a/csrc/kernels/quantize.cuh +++ b/csrc/kernels/quantize.cuh @@ -46,6 +46,11 @@ void dequantize_fp8_static_bf16_6( void quantize_fp8_device(const __nv_bfloat16* input, __nv_fp8_e4m3* output, float* d_scale, int n, cudaStream_t stream = 0); +// Setup-only weight packer with producer-compatible FP8 midpoint rounding. +void quantize_fp8_weight_device( + const __nv_bfloat16* input, __nv_fp8_e4m3* output, + float* d_scale, int n, cudaStream_t stream = 0); + void fp8_accumulate_scale_max(const float* src_scale, float* dst_scale, cudaStream_t stream = 0); From dae89cdcfe5d1b2e7f86ca57929588166f5a3756 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Thu, 16 Jul 2026 18:20:50 -0400 Subject: [PATCH 71/83] feat(cpp): add native sm120 fp8 calibration --- cpp/models/pi05/CMakeLists.txt | 8 +- .../include/flashrt/cpp/models/pi05/c_api.h | 7 +- .../models/pi05/native_calibration_session.h | 68 +++ .../cpp/models/pi05/native_graph_owner.h | 7 +- .../cpp/models/pi05/native_rtx_autotune.h | 21 + .../pi05/native_rtx_calibration_session.h | 47 ++ .../pi05/native_thor_calibration_session.h | 29 +- ...c_api.cpp => native_calibration_c_api.cpp} | 73 ++- .../pi05/src/native_calibration_session.cpp | 133 ++++++ cpp/models/pi05/src/native_graph_owner.cpp | 137 +++--- cpp/models/pi05/src/native_rtx_autotune.cpp | 80 ++++ .../src/native_rtx_calibration_session.cpp | 449 ++++++++++++++++++ .../src/native_thor_calibration_session.cpp | 105 +--- cpp/models/pi05/tests/CMakeLists.txt | 9 +- ... => pi05_native_fp8_calibration_probe.cpp} | 54 ++- cpp/tests/test_pi05_native_calibration.cpp | 52 ++ 16 files changed, 1050 insertions(+), 229 deletions(-) create mode 100644 cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration_session.h create mode 100644 cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h create mode 100644 cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h rename cpp/models/pi05/src/{native_thor_calibration_c_api.cpp => native_calibration_c_api.cpp} (77%) create mode 100644 cpp/models/pi05/src/native_calibration_session.cpp create mode 100644 cpp/models/pi05/src/native_rtx_autotune.cpp create mode 100644 cpp/models/pi05/src/native_rtx_calibration_session.cpp rename cpp/tests/{pi05_native_thor_fp8_probe.cpp => pi05_native_fp8_calibration_probe.cpp} (88%) diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index 124d91f0..59e04d9b 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -16,6 +16,7 @@ set(_pi05_sources src/native_weight_packer.cpp src/native_weight_materializer.cpp src/native_calibration.cpp + src/native_calibration_session.cpp src/native_workspace.cpp src/native_rtx_attention.cpp src/prompt_format.cpp @@ -125,7 +126,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) if(FLASHRT_CPP_WITH_FA2) target_sources(flashrt_cpp_pi05_kernels PRIVATE src/native_graph_owner.cpp + src/native_rtx_autotune.cpp src/native_rtx_attention_driver.cu) + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + target_sources(flashrt_cpp_pi05_kernels PRIVATE + src/native_rtx_calibration_session.cpp) + endif() target_include_directories(flashrt_cpp_pi05_kernels PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../csrc) target_link_libraries(flashrt_cpp_pi05_kernels @@ -138,7 +144,7 @@ endif() add_library(flashrt_cpp_pi05_c SHARED src/c_api.cpp src/model_runtime.cpp - src/native_thor_calibration_c_api.cpp + src/native_calibration_c_api.cpp src/native_model_runtime.cpp src/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c 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 d4011c9a..d2740ddc 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 @@ -141,10 +141,11 @@ int frt_pi05_runtime_read_actions(frt_pi05_runtime*, const frt_runtime_export_v1* frt_pi05_runtime_export(frt_pi05_runtime*); const char* frt_pi05_runtime_last_error(frt_pi05_runtime*); -/* Native SM110 FP8 calibration. `config_json` uses the same checkpoint, +/* Native FP8 calibration. `config_json` uses the same checkpoint, * tokenizer, shape, and normalization fields as frt_model_runtime_open_v1. - * The precision must resolve to fp8_e4m3fn. A session accepts one or many - * observations and publishes one identity-bound safetensors artifact. */ + * The precision must resolve to fp8_e4m3fn on a compiled backend. A session + * accepts one or many observations and publishes one identity-bound + * safetensors artifact. */ int frt_pi05_calibration_create_v1( const char* config_json, double percentile, diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration_session.h new file mode 100644 index 00000000..6d5cfdbe --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration_session.h @@ -0,0 +1,68 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_CALIBRATION_SESSION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_CALIBRATION_SESSION_H + +#include "flashrt/cpp/modalities/vision.h" +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeCalibrationConfig { + std::string checkpoint_path; + std::string tokenizer_model_path; + int max_prompt_tokens = 200; + int state_dim = 0; + int num_views = 2; + int chunk_size = 10; + int num_steps = 10; + int vision_pool_factor = 1; + int max_frame_width = 1280; + int max_frame_height = 720; + std::vector state_q01; + std::vector state_q99; +}; + +class NativeCalibrationSession { +public: + virtual ~NativeCalibrationSession() = default; + + virtual modalities::Status observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed) = 0; + virtual modalities::Status finalize( + const std::string& artifact_path) const = 0; + virtual std::uint64_t sample_count() const = 0; +}; + +bool valid_native_calibration_config(const NativeCalibrationConfig& config); + +modalities::Status normalize_native_calibration_state( + const NativeCalibrationConfig& config, + const float* state, + std::uint64_t n_state, + std::vector* output); + +modalities::Status prepare_native_calibration_noise( + const float* noise, + std::uint64_t n_noise, + std::uint64_t seed, + std::size_t elements, + modalities::DType dtype, + std::vector* output); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_CALIBRATION_SESSION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h index 3e9ff956..65141be7 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h @@ -21,6 +21,9 @@ class NativeGraphOwner final : public NativeGraphRuntime { const std::string& checkpoint_path, const NativeGraphConfig& config, const NativeCalibrationArtifact& calibration, modalities::Status* status); + static std::unique_ptr create_calibration( + const std::string& checkpoint_path, const NativeGraphConfig& config, + modalities::Status* status); ~NativeGraphOwner() override; @@ -46,14 +49,14 @@ class NativeGraphOwner final : public NativeGraphRuntime { modalities::Status synchronize() const override; private: - explicit NativeGraphOwner(frt_ctx ctx, const NativeGraphConfig& config); + NativeGraphOwner(frt_ctx ctx, const NativeGraphConfig& config, + NativeRtxLinearMode linear_mode); modalities::Status initialize( const std::string& checkpoint_path, const NativeCalibrationArtifact* calibration); modalities::Status record(NativeGraphKind kind, void* stream); modalities::Status record_context(void* stream); modalities::Status record_action(void* stream); - modalities::Status autotune_fp8(); static modalities::Status record_graph( void* user, NativeGraphKind kind, void* stream); diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h new file mode 100644 index 00000000..eba22efb --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h @@ -0,0 +1,21 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_AUTOTUNE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_AUTOTUNE_H + +#include "flashrt/cpp/models/pi05/native_rtx_linear.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +modalities::Status autotune_native_rtx_fp8( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const NativeRtxLinear& linear, + int num_views, + int chunk_size); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_AUTOTUNE_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h new file mode 100644 index 00000000..fb2c24ae --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h @@ -0,0 +1,47 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_CALIBRATION_SESSION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_CALIBRATION_SESSION_H + +#include "flashrt/cpp/models/pi05/native_calibration_session.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeRtxCalibrationSession final : public NativeCalibrationSession { +public: + static std::unique_ptr create( + const NativeCalibrationConfig& config, + double percentile, + modalities::Status* status); + + ~NativeRtxCalibrationSession() override; + + NativeRtxCalibrationSession(const NativeRtxCalibrationSession&) = delete; + NativeRtxCalibrationSession& operator=( + const NativeRtxCalibrationSession&) = delete; + + modalities::Status observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed) override; + modalities::Status finalize( + const std::string& artifact_path) const override; + std::uint64_t sample_count() const override; + +private: + struct Impl; + explicit NativeRtxCalibrationSession(std::unique_ptr impl); + std::unique_ptr impl_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_CALIBRATION_SESSION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h index 352fb8ff..6c7d6fdf 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h @@ -1,8 +1,7 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H -#include "flashrt/cpp/modalities/vision.h" -#include "flashrt/cpp/modalities/types.h" +#include "flashrt/cpp/models/pi05/native_calibration_session.h" #include #include @@ -13,29 +12,16 @@ namespace flashrt { namespace models { namespace pi05 { -struct NativeThorCalibrationConfig { - std::string checkpoint_path; - std::string tokenizer_model_path; - int max_prompt_tokens = 200; - int state_dim = 0; - int num_views = 2; - int chunk_size = 10; - int num_steps = 10; - int vision_pool_factor = 1; - int max_frame_width = 1280; - int max_frame_height = 720; - std::vector state_q01; - std::vector state_q99; -}; +using NativeThorCalibrationConfig = NativeCalibrationConfig; -class NativeThorCalibrationSession { +class NativeThorCalibrationSession final : public NativeCalibrationSession { public: static std::unique_ptr create( const NativeThorCalibrationConfig& config, double percentile, modalities::Status* status); - ~NativeThorCalibrationSession(); + ~NativeThorCalibrationSession() override; NativeThorCalibrationSession(const NativeThorCalibrationSession&) = delete; NativeThorCalibrationSession& operator=( @@ -48,9 +34,10 @@ class NativeThorCalibrationSession { const std::vector& frames, const float* noise, std::uint64_t n_noise, - std::uint64_t noise_seed); - modalities::Status finalize(const std::string& artifact_path) const; - std::uint64_t sample_count() const; + std::uint64_t noise_seed) override; + modalities::Status finalize( + const std::string& artifact_path) const override; + std::uint64_t sample_count() const override; private: struct Impl; diff --git a/cpp/models/pi05/src/native_thor_calibration_c_api.cpp b/cpp/models/pi05/src/native_calibration_c_api.cpp similarity index 77% rename from cpp/models/pi05/src/native_thor_calibration_c_api.cpp rename to cpp/models/pi05/src/native_calibration_c_api.cpp index e55d8440..28373b67 100644 --- a/cpp/models/pi05/src/native_thor_calibration_c_api.cpp +++ b/cpp/models/pi05/src/native_calibration_c_api.cpp @@ -4,10 +4,19 @@ #include "flashrt/cpp/models/pi05/spec.h" #include "native_open_internal.h" -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_SENTENCEPIECE) && \ + (defined(FLASHRT_CPP_WITH_FA2) || \ + defined(FLASHRT_CPP_WITH_THOR_FP8)) +#define FLASHRT_CPP_HAS_NATIVE_CALIBRATION 1 +#include "flashrt/cpp/models/pi05/native_calibration_session.h" +#if defined(FLASHRT_CPP_WITH_FA2) +#include "flashrt/cpp/models/pi05/native_rtx_calibration_session.h" +#endif +#if defined(FLASHRT_CPP_WITH_THOR_FP8) #include "flashrt/cpp/models/pi05/native_thor_calibration_session.h" #endif +#include +#endif #include #include @@ -26,9 +35,8 @@ using flashrt::modalities::VisionFrame; thread_local std::string g_calibration_create_error; struct frt_pi05_calibration_session_s { -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) - std::unique_ptr impl; +#if defined(FLASHRT_CPP_HAS_NATIVE_CALIBRATION) + std::unique_ptr impl; #endif std::string last_error; std::vector view_names; @@ -37,8 +45,7 @@ struct frt_pi05_calibration_session_s { namespace { -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_NATIVE_CALIBRATION) int set_error(frt_pi05_calibration_session* session, const flashrt::modalities::Status& status) { if (session) session->last_error = status.message; @@ -107,8 +114,7 @@ extern "C" int frt_pi05_calibration_create_v1( return -1; } *out = nullptr; -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_NATIVE_CALIBRATION) flashrt::models::pi05::NativeOpenConfig open_config; int rc = flashrt::models::pi05::parse_native_open_config( config_json, &open_config, &g_calibration_create_error); @@ -118,7 +124,7 @@ extern "C" int frt_pi05_calibration_create_v1( "Pi0.5 calibration precision must resolve to fp8_e4m3fn"; return -1; } - flashrt::models::pi05::NativeThorCalibrationConfig config; + flashrt::models::pi05::NativeCalibrationConfig config; config.checkpoint_path = open_config.checkpoint_path; config.tokenizer_model_path = open_config.tokenizer_model_path; config.max_prompt_tokens = open_config.max_prompt_tokens; @@ -131,10 +137,42 @@ extern "C" int frt_pi05_calibration_create_v1( config.max_frame_height = open_config.max_frame_height; config.state_q01 = std::move(open_config.state_q01); config.state_q99 = std::move(open_config.state_q99); + + int device = 0; + cudaDeviceProp properties{}; + cudaError_t cuda_rc = cudaGetDevice(&device); + if (cuda_rc == cudaSuccess) { + cuda_rc = cudaGetDeviceProperties(&properties, device); + } + if (cuda_rc != cudaSuccess) { + g_calibration_create_error = cudaGetErrorString(cuda_rc); + return -6; + } flashrt::modalities::Status status; - auto impl = - flashrt::models::pi05::NativeThorCalibrationSession::create( + std::unique_ptr impl; + if (properties.major == 12 && properties.minor == 0) { +#if defined(FLASHRT_CPP_WITH_FA2) + impl = flashrt::models::pi05::NativeRtxCalibrationSession::create( config, percentile, &status); +#else + status = flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kUnsupported, + "Pi0.5 SM120 calibration backend is not built"); +#endif + } else if (properties.major == 11 && properties.minor == 0) { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) + impl = flashrt::models::pi05::NativeThorCalibrationSession::create( + config, percentile, &status); +#else + status = flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kUnsupported, + "Pi0.5 SM110 calibration backend is not built"); +#endif + } else { + status = flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kUnsupported, + "Pi0.5 native FP8 calibration has no backend for this device"); + } if (!impl) { g_calibration_create_error = status.message; return flashrt::models::pi05::cface::status_code(status); @@ -155,7 +193,7 @@ extern "C" int frt_pi05_calibration_create_v1( (void)config_json; (void)percentile; g_calibration_create_error = - "Pi0.5 calibration requires Thor FP8 and SentencePiece"; + "Pi0.5 calibration requires a native FP8 backend and SentencePiece"; return -3; #endif } catch (const std::exception& error) { @@ -171,8 +209,7 @@ extern "C" int frt_pi05_calibration_create_v1( extern "C" int frt_pi05_calibration_observe_v1( frt_pi05_calibration_session* session, const frt_pi05_calibration_sample_v1* sample) try { -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_NATIVE_CALIBRATION) if (!session || !session->impl || !sample || sample->struct_size < sizeof(frt_pi05_calibration_sample_v1) || !sample->prompt || (!sample->state && sample->n_state) || @@ -203,8 +240,7 @@ extern "C" int frt_pi05_calibration_observe_v1( extern "C" int frt_pi05_calibration_finalize_v1( frt_pi05_calibration_session* session, const char* artifact_path) try { -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_NATIVE_CALIBRATION) if (!session || !session->impl || !artifact_path || !artifact_path[0]) { if (session) session->last_error = "calibration output path is invalid"; return -1; @@ -229,8 +265,7 @@ extern "C" int frt_pi05_calibration_finalize_v1( extern "C" uint64_t frt_pi05_calibration_sample_count_v1( const frt_pi05_calibration_session* session) { -#if defined(FLASHRT_CPP_WITH_THOR_FP8) && \ - defined(FLASHRT_CPP_HAS_SENTENCEPIECE) +#if defined(FLASHRT_CPP_HAS_NATIVE_CALIBRATION) return session && session->impl ? session->impl->sample_count() : 0; #else (void)session; diff --git a/cpp/models/pi05/src/native_calibration_session.cpp b/cpp/models/pi05/src/native_calibration_session.cpp new file mode 100644 index 00000000..07df2c5d --- /dev/null +++ b/cpp/models/pi05/src/native_calibration_session.cpp @@ -0,0 +1,133 @@ +#include "flashrt/cpp/models/pi05/native_calibration_session.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +std::uint64_t splitmix64(std::uint64_t* state) { + std::uint64_t value = (*state += 0x9e3779b97f4a7c15ull); + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ull; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebull; + return value ^ (value >> 31); +} + +double uniform_open(std::uint64_t* state) { + constexpr double kDenominator = 9007199254740993.0; + return (static_cast(splitmix64(state) >> 11) + 1.0) / + kDenominator; +} + +std::uint16_t encode(float value, modalities::DType dtype) { + return dtype == modalities::DType::kFloat16 + ? modalities::float_to_float16(value) + : modalities::float_to_bfloat16(value); +} + +} // namespace + +bool valid_native_calibration_config(const NativeCalibrationConfig& config) { + const std::uint64_t width = + static_cast(config.max_frame_width); + const std::uint64_t height = + static_cast(config.max_frame_height); + bool valid_quantiles = + config.state_q01.size() == + static_cast(config.state_dim) && + config.state_q99.size() == config.state_q01.size(); + for (std::size_t i = 0; + valid_quantiles && i < config.state_q01.size(); ++i) { + valid_quantiles = std::isfinite(config.state_q01[i]) && + std::isfinite(config.state_q99[i]) && + config.state_q99[i] > config.state_q01[i]; + } + return !config.checkpoint_path.empty() && + !config.tokenizer_model_path.empty() && config.state_dim > 0 && + config.num_views >= 1 && config.num_views <= 3 && + config.max_prompt_tokens >= 1 && config.chunk_size > 0 && + config.num_steps > 0 && + (config.vision_pool_factor == 1 || + config.vision_pool_factor == 2 || + config.vision_pool_factor == 4) && + static_cast(config.max_prompt_tokens) + + static_cast(config.chunk_size) + + static_cast(config.num_views) * 256 <= + static_cast( + std::numeric_limits::max()) && + config.max_frame_width > 0 && config.max_frame_height > 0 && + width <= std::numeric_limits::max() / height / 4 && + valid_quantiles; +} + +modalities::Status normalize_native_calibration_state( + const NativeCalibrationConfig& config, + const float* state, + std::uint64_t n_state, + std::vector* output) { + if (!state || !output || + n_state != static_cast(config.state_dim)) { + return invalid("native calibration state shape is invalid"); + } + output->resize(static_cast(config.state_dim)); + for (std::size_t i = 0; i < output->size(); ++i) { + if (!std::isfinite(state[i])) { + return invalid("native calibration state contains non-finite data"); + } + const float lo = config.state_q01[i]; + const float hi = config.state_q99[i]; + (*output)[i] = ((state[i] - lo) / (hi - lo + 1e-6f)) * 2.0f - 1.0f; + } + return modalities::Status::ok(); +} + +modalities::Status prepare_native_calibration_noise( + const float* noise, + std::uint64_t n_noise, + std::uint64_t seed, + std::size_t elements, + modalities::DType dtype, + std::vector* output) { + if (!output || !elements || + (dtype != modalities::DType::kFloat16 && + dtype != modalities::DType::kBFloat16) || + (noise && n_noise != elements) || (!noise && n_noise != 0)) { + return invalid("native calibration noise shape is invalid"); + } + output->resize(elements); + if (noise) { + for (std::size_t i = 0; i < elements; ++i) { + if (!std::isfinite(noise[i])) { + return invalid( + "native calibration noise contains non-finite data"); + } + (*output)[i] = encode(noise[i], dtype); + } + return modalities::Status::ok(); + } + + constexpr double kTwoPi = 6.283185307179586476925286766559; + std::uint64_t state = seed ^ 0x243f6a8885a308d3ull; + for (std::size_t i = 0; i < elements; i += 2) { + const double radius = std::sqrt(-2.0 * std::log(uniform_open(&state))); + const double angle = kTwoPi * uniform_open(&state); + (*output)[i] = encode( + static_cast(radius * std::cos(angle)), dtype); + if (i + 1 < elements) { + (*output)[i + 1] = encode( + static_cast(radius * std::sin(angle)), dtype); + } + } + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index 5e487669..a1399348 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -2,6 +2,7 @@ #include "flashrt/cpp/models/pi05/native_style_precompute.h" #include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include "flashrt/cpp/models/pi05/native_rtx_autotune.h" #include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" #include @@ -52,15 +53,14 @@ modalities::Status upload_scales( NativeGraphOwner::NativeGraphOwner( frt_ctx ctx, - const NativeGraphConfig& config) + const NativeGraphConfig& config, + NativeRtxLinearMode linear_mode) : graphs_(ctx), config_(config), weights_(ctx), workspace_(ctx), attention_(ctx), - linear_(&driver_, config.precision == NativeGraphPrecision::kFp8E4M3 - ? NativeRtxLinearMode::kFp8Static - : NativeRtxLinearMode::kBf16), + linear_(&driver_, linear_mode), forward_(&driver_, &linear_) {} NativeGraphOwner::~NativeGraphOwner() = default; @@ -69,7 +69,8 @@ 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 || + if (config.precision != NativeGraphPrecision::kBf16 || + config.num_views < 1 || config.num_views > 3 || config.max_prompt_tokens < 1 || config.chunk_size < 1 || config.num_steps < 1 || static_cast(config.max_prompt_tokens) + @@ -87,7 +88,8 @@ std::unique_ptr NativeGraphOwner::create( return nullptr; } std::unique_ptr owner( - new (std::nothrow) NativeGraphOwner(ctx, config)); + new (std::nothrow) NativeGraphOwner( + ctx, config, NativeRtxLinearMode::kBf16)); if (!owner) { frt_ctx_destroy(ctx); if (status) *status = backend("native graph owner allocation failed"); @@ -117,7 +119,8 @@ std::unique_ptr NativeGraphOwner::create( return nullptr; } std::unique_ptr owner( - new (std::nothrow) NativeGraphOwner(ctx, config)); + new (std::nothrow) NativeGraphOwner( + ctx, config, NativeRtxLinearMode::kFp8Static)); if (!owner) { frt_ctx_destroy(ctx); if (status) *status = backend("native graph owner allocation failed"); @@ -132,11 +135,43 @@ std::unique_ptr NativeGraphOwner::create( return owner; } +std::unique_ptr NativeGraphOwner::create_calibration( + const std::string& checkpoint_path, + const NativeGraphConfig& config, + modalities::Status* status) { + if (config.precision != NativeGraphPrecision::kFp8E4M3) { + if (status) { + *status = invalid("native RTX calibration precision 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, NativeRtxLinearMode::kFp8Dynamic)); + if (!owner) { + frt_ctx_destroy(ctx); + if (status) *status = backend("native graph owner allocation failed"); + return nullptr; + } + modalities::Status st = owner->initialize(checkpoint_path, nullptr); + 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, const NativeCalibrationArtifact* calibration) { - const bool fp8 = config_.precision == NativeGraphPrecision::kFp8E4M3; - if (fp8 != (calibration != nullptr) || + const bool fp8 = linear_.fp8(); + if ((calibration != nullptr) != linear_.static_fp8() || (calibration && calibration->activation_dtype != "bfloat16")) { return invalid("native RTX FP8 calibration is incompatible"); } @@ -183,7 +218,7 @@ modalities::Status NativeGraphOwner::initialize( : NativeWorkspaceFlavor::kBf16; st = workspace_.allocate(workspace_config); if (!st.ok_status()) return st; - if (fp8) { + if (linear_.static_fp8()) { st = upload_scales( &workspace_, "rtx_fp8_vision_scales", calibration->vision_scales); @@ -196,6 +231,21 @@ modalities::Status NativeGraphOwner::initialize( &workspace_, "rtx_fp8_decoder_scales", calibration->decoder_scales); if (!st.ok_status()) return st; + } else if (linear_.dynamic_fp8()) { + st = upload_scales( + &workspace_, "rtx_fp8_vision_scales", + std::vector(109, 1.0f)); + if (!st.ok_status()) return st; + st = upload_scales( + &workspace_, "rtx_fp8_encoder_scales", + std::vector(18 * 4, 1.0f)); + if (!st.ok_status()) return st; + st = upload_scales( + &workspace_, "rtx_fp8_decoder_scales", + std::vector( + static_cast(config_.num_steps) * 18 * 4, + 1.0f)); + if (!st.ok_status()) return st; } st = workspace_.expand_vision_position_embedding(weights_); if (!st.ok_status()) return st; @@ -222,7 +272,9 @@ modalities::Status NativeGraphOwner::initialize( st = attention_driver_->status(); if (!st.ok_status()) return st; if (fp8) { - st = autotune_fp8(); + st = autotune_native_rtx_fp8( + weights_, &workspace_, linear_, config_.num_views, + config_.chunk_size); if (!st.ok_status()) return st; report("fp8_autotune"); } @@ -275,69 +327,6 @@ modalities::Status NativeGraphOwner::initialize( return modalities::Status::ok(); } -modalities::Status NativeGraphOwner::autotune_fp8() { - const int vision_sequence = config_.num_views * 256; - const int encoder_vision_sequence = workspace_.encoder_vision_sequence(); - const int encoder_sequence = workspace_.encoder_sequence(); - const int decoder_sequence = config_.chunk_size; - struct Shape { - const char* weight; - NativeRtxScaleSite site; - const char* output; - int m; - int n; - int k; - }; - const Shape shapes[] = { - {"vision_attn_qkv_w_0", {NativeRtxScaleDomain::kVision, 0}, - "vision_QKV", - vision_sequence, 3456, 1152}, - {"vision_attn_o_w_0", {NativeRtxScaleDomain::kVision, 1}, - "vision_x_norm", - vision_sequence, 1152, 1152}, - {"vision_ffn_up_w_0", {NativeRtxScaleDomain::kVision, 2}, - "vision_hidden", - vision_sequence, 4304, 1152}, - {"vision_ffn_down_w_0", {NativeRtxScaleDomain::kVision, 3}, - "vision_x_norm", - vision_sequence, 1152, 4304}, - {"encoder_multi_modal_projector_w", - {NativeRtxScaleDomain::kVision, 108}, "encoder_x", - encoder_vision_sequence, 2048, 1152}, - {"encoder_attn_qkv_w_0", {NativeRtxScaleDomain::kEncoder, 0}, - "encoder_QKV", - encoder_sequence, 2560, 2048}, - {"encoder_attn_o_w_0", {NativeRtxScaleDomain::kEncoder, 1}, - "encoder_x_norm", - encoder_sequence, 2048, 2048}, - {"encoder_ffn_gate_up_w_0", {NativeRtxScaleDomain::kEncoder, 2}, - "encoder_gate_merged", encoder_sequence, 32768, 2048}, - {"encoder_ffn_down_w_0", {NativeRtxScaleDomain::kEncoder, 3}, - "encoder_x_norm", - encoder_sequence, 2048, 16384}, - {"decoder_attn_qkv_w_0", {NativeRtxScaleDomain::kDecoder, 0}, - "decoder_QKV", - decoder_sequence, 2560, 1024}, - {"decoder_attn_o_w_0", {NativeRtxScaleDomain::kDecoder, 1}, - "x_normed_buf", - decoder_sequence, 1024, 2048}, - {"decoder_ffn_gate_up_w_0", {NativeRtxScaleDomain::kDecoder, 2}, - "decoder_gate_merged", decoder_sequence, 8192, 1024}, - {"decoder_ffn_down_w_0", {NativeRtxScaleDomain::kDecoder, 3}, - "x_normed_buf", - decoder_sequence, 1024, 4096}, - }; - for (const Shape& shape : shapes) { - const NativeWorkspaceBuffer* output = workspace_.find(shape.output); - if (!output) return invalid("native FP8 autotune output is missing"); - modalities::Status st = linear_.autotune( - weights_, &workspace_, shape.weight, shape.site, - frt_buffer_dptr(output->buffer), shape.m, shape.n, shape.k); - if (!st.ok_status()) return st; - } - return modalities::Status::ok(); -} - modalities::Status NativeGraphOwner::record_context(void* stream) { modalities::Status st = copy_prompt_to_encoder(&workspace_, stream); if (!st.ok_status()) return st; diff --git a/cpp/models/pi05/src/native_rtx_autotune.cpp b/cpp/models/pi05/src/native_rtx_autotune.cpp new file mode 100644 index 00000000..31cbd319 --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_autotune.cpp @@ -0,0 +1,80 @@ +#include "flashrt/cpp/models/pi05/native_rtx_autotune.h" + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +struct Shape { + const char* weight; + NativeRtxScaleSite site; + const char* output; + int m; + int n; + int k; +}; + +} // namespace + +modalities::Status autotune_native_rtx_fp8( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + const NativeRtxLinear& linear, + int num_views, + int chunk_size) { + if (!workspace || !linear.fp8() || num_views < 1 || num_views > 3 || + chunk_size <= 0 || workspace->num_views() != num_views || + workspace->chunk_size() != chunk_size) { + return invalid("native RTX FP8 autotune configuration is invalid"); + } + const int vision_sequence = num_views * 256; + const int encoder_vision_sequence = workspace->encoder_vision_sequence(); + const int encoder_sequence = workspace->encoder_sequence(); + const Shape shapes[] = { + {"vision_attn_qkv_w_0", {NativeRtxScaleDomain::kVision, 0}, + "vision_QKV", vision_sequence, 3456, 1152}, + {"vision_attn_o_w_0", {NativeRtxScaleDomain::kVision, 1}, + "vision_x_norm", vision_sequence, 1152, 1152}, + {"vision_ffn_up_w_0", {NativeRtxScaleDomain::kVision, 2}, + "vision_hidden", vision_sequence, 4304, 1152}, + {"vision_ffn_down_w_0", {NativeRtxScaleDomain::kVision, 3}, + "vision_x_norm", vision_sequence, 1152, 4304}, + {"encoder_multi_modal_projector_w", + {NativeRtxScaleDomain::kVision, 108}, "encoder_x", + encoder_vision_sequence, 2048, 1152}, + {"encoder_attn_qkv_w_0", {NativeRtxScaleDomain::kEncoder, 0}, + "encoder_QKV", encoder_sequence, 2560, 2048}, + {"encoder_attn_o_w_0", {NativeRtxScaleDomain::kEncoder, 1}, + "encoder_x_norm", encoder_sequence, 2048, 2048}, + {"encoder_ffn_gate_up_w_0", {NativeRtxScaleDomain::kEncoder, 2}, + "encoder_gate_merged", encoder_sequence, 32768, 2048}, + {"encoder_ffn_down_w_0", {NativeRtxScaleDomain::kEncoder, 3}, + "encoder_x_norm", encoder_sequence, 2048, 16384}, + {"decoder_attn_qkv_w_0", {NativeRtxScaleDomain::kDecoder, 0}, + "decoder_QKV", chunk_size, 2560, 1024}, + {"decoder_attn_o_w_0", {NativeRtxScaleDomain::kDecoder, 1}, + "x_normed_buf", chunk_size, 1024, 2048}, + {"decoder_ffn_gate_up_w_0", {NativeRtxScaleDomain::kDecoder, 2}, + "decoder_gate_merged", chunk_size, 8192, 1024}, + {"decoder_ffn_down_w_0", {NativeRtxScaleDomain::kDecoder, 3}, + "x_normed_buf", chunk_size, 1024, 4096}, + }; + for (const Shape& shape : shapes) { + const NativeWorkspaceBuffer* output = workspace->find(shape.output); + if (!output) return invalid("native FP8 autotune output is missing"); + modalities::Status st = linear.autotune( + weights, workspace, shape.weight, shape.site, + frt_buffer_dptr(output->buffer), shape.m, shape.n, shape.k); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_rtx_calibration_session.cpp b/cpp/models/pi05/src/native_rtx_calibration_session.cpp new file mode 100644 index 00000000..53f4b223 --- /dev/null +++ b/cpp/models/pi05/src/native_rtx_calibration_session.cpp @@ -0,0 +1,449 @@ +#include "flashrt/cpp/models/pi05/native_rtx_calibration_session.h" + +#include "flashrt/cpp/loader/sha256.h" +#include "flashrt/cpp/models/pi05/io.h" +#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/prompt_embed.h" + +#include + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +constexpr std::size_t kVisionScales = 27 * 4 + 1; +constexpr std::size_t kEncoderScales = 18 * 4; + +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); +} + +struct HashResult { + bool ok = false; + std::string digest; + std::string error; +}; + +modalities::TensorView device_view(const NativeWorkspaceBuffer* buffer, + modalities::DType dtype, + modalities::Layout layout, + modalities::Shape shape) { + modalities::TensorView view; + if (!buffer) return view; + view.data = frt_buffer_dptr(buffer->buffer); + view.bytes = frt_buffer_bytes(buffer->buffer); + view.dtype = dtype; + view.place = modalities::MemoryPlace::kDevice; + view.layout = layout; + view.shape = shape; + return view; +} + +} // namespace + +struct NativeRtxCalibrationSession::Impl { + Impl(NativeCalibrationConfig value, double requested_percentile) + : config(std::move(value)), percentile(requested_percentile) {} + + ~Impl() { + modalities::text_embedding_staging_destroy(&text_staging); + modalities::vision_staging_destroy(&vision_staging); + } + + modalities::Status initialize() { + int device = 0; + cudaDeviceProp properties{}; + cudaError_t rc = cudaGetDevice(&device); + if (rc == cudaSuccess) { + rc = cudaGetDeviceProperties(&properties, device); + } + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + if (properties.major != 12 || properties.minor != 0) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Pi0.5 RTX FP8 calibration requires SM120"); + } + hardware = + "sm" + std::to_string(properties.major * 10 + properties.minor); + + const std::string weights_path = + config.checkpoint_path + "/model.safetensors"; + std::future weights_hash = std::async( + std::launch::async, [weights_path] { + HashResult result; + result.ok = loader::sha256_file( + weights_path, &result.digest, &result.error); + return result; + }); + const std::string tokenizer_path = config.tokenizer_model_path; + std::future tokenizer_hash = std::async( + std::launch::async, [tokenizer_path] { + HashResult result; + result.ok = loader::sha256_file( + tokenizer_path, &result.digest, &result.error); + return result; + }); + + 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_size; + graph_config.num_steps = config.num_steps; + graph_config.vision_pool_factor = config.vision_pool_factor; + graph_config.precision = NativeGraphPrecision::kFp8E4M3; + modalities::Status st; + graph = NativeGraphOwner::create_calibration( + config.checkpoint_path, graph_config, &st); + if (!graph) return st; + + HashResult weights_digest = weights_hash.get(); + if (!weights_digest.ok) { + return modalities::Status::error( + modalities::StatusCode::kNotFound, weights_digest.error); + } + weights_sha256 = std::move(weights_digest.digest); + HashResult tokenizer_digest = tokenizer_hash.get(); + if (!tokenizer_digest.ok) { + return modalities::Status::error( + modalities::StatusCode::kNotFound, tokenizer_digest.error); + } + tokenizer_sha256 = std::move(tokenizer_digest.digest); + + const std::uint64_t frame_capacity = + static_cast(config.max_frame_width) * + static_cast(config.max_frame_height) * 4; + st = modalities::vision_staging_create( + &vision_staging, static_cast(config.num_views), + frame_capacity); + if (!st.ok_status()) return st; + st = modalities::text_embedding_staging_create( + &text_staging, config.max_prompt_tokens); + if (!st.ok_status()) return st; + st = tokenizer.load_model(config.tokenizer_model_path); + if (!st.ok_status()) return st; + tokenizer.reserve(config.max_prompt_tokens); + + const NativeWorkspaceBuffer* images = + graph->workspace().find("observation_images_normalized"); + const NativeWorkspaceBuffer* noise = + graph->workspace().find("diffusion_noise"); + image_output = device_view( + images, modalities::DType::kBFloat16, + modalities::Layout::kNHWC, + {static_cast(config.num_views), 224, 224, 3}); + action_output = device_view( + noise, modalities::DType::kBFloat16, + modalities::Layout::kFlat, + {static_cast(config.chunk_size), 32}); + if (!image_output.data || !action_output.data) { + return invalid("RTX calibration IO buffers are incomplete"); + } + io.reset(new (std::nothrow) RuntimeIo( + config.num_views, image_output, action_output, {}, {}, + graph->native_stream(), config.chunk_size, 32, 32, + modalities::DType::kBFloat16, &vision_staging, nullptr, true)); + if (!io) return backend("RTX calibration IO allocation failed"); + + const NativeDeviceWeight* embedding = + graph->weights().find("embedding_weight"); + const NativeWorkspaceBuffer* prompt = + graph->workspace().find("prompt_embedding"); + if (!embedding || !prompt || + embedding->dtype != NativeWeightDType::kBf16 || + embedding->shape.size() != 2 || embedding->shape[1] != 2048) { + return invalid("RTX calibration embedding buffers are invalid"); + } + embedding_table.data = frt_buffer_dptr(embedding->buffer); + embedding_table.bytes = frt_buffer_bytes(embedding->buffer); + embedding_table.dtype = modalities::DType::kBFloat16; + embedding_table.place = modalities::MemoryPlace::kDevice; + embedding_table.layout = modalities::Layout::kFlat; + embedding_table.shape = {embedding->shape[0], embedding->shape[1]}; + prompt_output = device_view( + prompt, modalities::DType::kBFloat16, + modalities::Layout::kFlat, + {static_cast(config.max_prompt_tokens), 2048}); + prompt_spec.vocab_size = embedding->shape[0]; + prompt_spec.hidden_dim = 2048; + prompt_spec.max_tokens = config.max_prompt_tokens; + prompt_spec.scale = std::sqrt(2048.0f); + token_ids.reserve(static_cast(config.max_prompt_tokens) + 1); + const std::size_t max_prompt_bytes = + static_cast(config.max_prompt_tokens) * 8; + formatted_prompt.reserve( + max_prompt_bytes + static_cast(config.state_dim) * 5 + + 32); + vision_scale_ones.assign(kVisionScales, 1.0f); + encoder_scale_ones.assign(kEncoderScales, 1.0f); + decoder_scale_ones.assign( + static_cast(config.num_steps) * 18 * 4, 1.0f); + return modalities::Status::ok(); + } + + const NativeWorkspaceBuffer* scale_buffer( + const char* name, + std::size_t elements) const { + const NativeWorkspaceBuffer* buffer = + graph ? graph->workspace().find(name) : nullptr; + if (!buffer || buffer->dtype != modalities::DType::kFloat32 || + buffer->shape != std::vector({elements}) || + frt_buffer_bytes(buffer->buffer) != elements * sizeof(float)) { + return nullptr; + } + return buffer; + } + + modalities::Status upload_scales( + const char* name, + const std::vector& values) const { + const NativeWorkspaceBuffer* buffer = + scale_buffer(name, values.size()); + if (!buffer) return invalid("RTX calibration scale buffer is invalid"); + const cudaError_t rc = cudaMemcpyAsync( + frt_buffer_dptr(buffer->buffer), values.data(), + values.size() * sizeof(float), cudaMemcpyHostToDevice, + static_cast(graph->native_stream())); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); + } + + modalities::Status download_scales( + const char* name, + std::size_t elements, + std::vector* output) const { + if (!output) return invalid("RTX calibration scale output is null"); + const NativeWorkspaceBuffer* buffer = scale_buffer(name, elements); + if (!buffer) return invalid("RTX calibration scale buffer is invalid"); + output->resize(elements); + const cudaError_t rc = cudaMemcpyAsync( + output->data(), frt_buffer_dptr(buffer->buffer), + elements * sizeof(float), cudaMemcpyDeviceToHost, + static_cast(graph->native_stream())); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); + } + + modalities::Status reset_scales() const { + modalities::Status st = upload_scales( + "rtx_fp8_vision_scales", vision_scale_ones); + if (!st.ok_status()) return st; + st = upload_scales( + "rtx_fp8_encoder_scales", encoder_scale_ones); + if (!st.ok_status()) return st; + return upload_scales( + "rtx_fp8_decoder_scales", decoder_scale_ones); + } + + modalities::Status collect_scales( + std::vector* vision, + std::vector* encoder, + std::vector* decoder) const { + if (!vision || !encoder || !decoder) { + return invalid("RTX calibration scale outputs are null"); + } + modalities::Status st = download_scales( + "rtx_fp8_vision_scales", vision_scale_ones.size(), vision); + if (!st.ok_status()) return st; + st = download_scales( + "rtx_fp8_encoder_scales", encoder_scale_ones.size(), encoder); + if (!st.ok_status()) return st; + return download_scales( + "rtx_fp8_decoder_scales", decoder_scale_ones.size(), decoder); + } + + modalities::Status observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed) { + modalities::Status st = normalize_native_calibration_state( + config, state, n_state, &normalized_state); + if (!st.ok_status()) return st; + st = prepare_native_calibration_noise( + noise, n_noise, + noise_seed + static_cast(vision_samples.size()), + static_cast(config.chunk_size) * 32, + modalities::DType::kBFloat16, &noise_bf16); + if (!st.ok_status()) return st; + + std::uint64_t prompt_len = 0; + st = embed_prompt( + tokenizer, prompt_spec, prompt, normalized_state.data(), + normalized_state.size(), embedding_table, prompt_output, + &token_ids, &prompt_len, graph->native_stream(), &text_staging, + &formatted_prompt); + if (!st.ok_status()) return st; + st = graph->set_prompt_length(static_cast(prompt_len)); + if (!st.ok_status()) return st; + st = io->prepare_vision(frames); + if (!st.ok_status()) return st; + st = reset_scales(); + if (!st.ok_status()) return st; + + cudaError_t rc = cudaMemcpyAsync( + action_output.data, noise_bf16.data(), + noise_bf16.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, + static_cast(graph->native_stream())); + if (rc != cudaSuccess) return backend(cudaGetErrorString(rc)); + if (graph->replay() != FRT_OK) { + return backend("RTX calibration graph replay failed"); + } + + std::vector vision_scale; + std::vector encoder_scale; + std::vector decoder_scale; + st = collect_scales( + &vision_scale, &encoder_scale, &decoder_scale); + if (!st.ok_status()) return st; + st = graph->synchronize(); + if (!st.ok_status()) return st; + vision_samples.push_back(std::move(vision_scale)); + encoder_samples.push_back(std::move(encoder_scale)); + decoder_samples.push_back(std::move(decoder_scale)); + return modalities::Status::ok(); + } + + modalities::Status finalize(const std::string& artifact_path) const { + if (vision_samples.empty() || + vision_samples.size() != encoder_samples.size() || + vision_samples.size() != decoder_samples.size()) { + return invalid("RTX calibration has no complete samples"); + } + NativeCalibrationArtifact artifact; + artifact.activation_dtype = "bfloat16"; + artifact.hardware = hardware; + artifact.weights_sha256 = weights_sha256; + artifact.tokenizer_sha256 = tokenizer_sha256; + artifact.num_views = config.num_views; + artifact.max_prompt_tokens = config.max_prompt_tokens; + artifact.state_dim = config.state_dim; + artifact.chunk_size = config.chunk_size; + artifact.num_steps = config.num_steps; + artifact.vision_pool_factor = config.vision_pool_factor; + artifact.sample_count = vision_samples.size(); + artifact.percentile = percentile; + modalities::Status st = reduce_native_calibration_samples( + vision_samples, percentile, &artifact.vision_scales); + if (!st.ok_status()) return st; + st = reduce_native_calibration_samples( + encoder_samples, percentile, &artifact.encoder_scales); + if (!st.ok_status()) return st; + st = reduce_native_calibration_samples( + decoder_samples, percentile, &artifact.decoder_scales); + if (!st.ok_status()) return st; + return save_native_calibration_artifact(artifact_path, artifact); + } + + NativeCalibrationConfig config; + double percentile = 99.9; + std::string hardware; + std::string weights_sha256; + std::string tokenizer_sha256; + std::unique_ptr graph; + modalities::VisionStaging vision_staging; + modalities::TextEmbeddingStaging text_staging; + modalities::SentencePieceTokenizer tokenizer; + std::unique_ptr io; + modalities::TensorView image_output; + modalities::TensorView action_output; + modalities::TensorView embedding_table; + modalities::TensorView prompt_output; + PromptEmbeddingSpec prompt_spec; + std::vector token_ids; + std::vector normalized_state; + std::string formatted_prompt; + std::vector noise_bf16; + std::vector vision_scale_ones; + std::vector encoder_scale_ones; + std::vector decoder_scale_ones; + std::vector> vision_samples; + std::vector> encoder_samples; + std::vector> decoder_samples; +}; + +NativeRtxCalibrationSession::NativeRtxCalibrationSession( + std::unique_ptr impl) + : impl_(std::move(impl)) {} + +NativeRtxCalibrationSession::~NativeRtxCalibrationSession() = default; + +std::unique_ptr +NativeRtxCalibrationSession::create( + const NativeCalibrationConfig& config, + double percentile, + modalities::Status* status) { + if (!valid_native_calibration_config(config) || + !std::isfinite(percentile) || percentile < 0.0 || + percentile > 100.0) { + if (status) *status = invalid("RTX calibration config is invalid"); + return nullptr; + } + std::unique_ptr impl( + new (std::nothrow) Impl(config, percentile)); + if (!impl) { + if (status) *status = backend("RTX calibration allocation failed"); + return nullptr; + } + modalities::Status st = impl->initialize(); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + std::unique_ptr session( + new (std::nothrow) NativeRtxCalibrationSession(std::move(impl))); + if (!session) { + if (status) { + *status = backend("RTX calibration session allocation failed"); + } + return nullptr; + } + if (status) *status = modalities::Status::ok(); + return session; +} + +modalities::Status NativeRtxCalibrationSession::observe( + const std::string& prompt, + const float* state, + std::uint64_t n_state, + const std::vector& frames, + const float* noise, + std::uint64_t n_noise, + std::uint64_t noise_seed) { + return impl_ ? impl_->observe( + prompt, state, n_state, frames, noise, n_noise, + noise_seed) + : invalid("RTX calibration session is invalid"); +} + +modalities::Status NativeRtxCalibrationSession::finalize( + const std::string& artifact_path) const { + return impl_ ? impl_->finalize(artifact_path) + : invalid("RTX calibration session is invalid"); +} + +std::uint64_t NativeRtxCalibrationSession::sample_count() const { + return impl_ ? impl_->vision_samples.size() : 0; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/native_thor_calibration_session.cpp b/cpp/models/pi05/src/native_thor_calibration_session.cpp index 93e56241..36696a28 100644 --- a/cpp/models/pi05/src/native_thor_calibration_session.cpp +++ b/cpp/models/pi05/src/native_thor_calibration_session.cpp @@ -14,7 +14,6 @@ #include #include #include -#include #include namespace flashrt { @@ -38,66 +37,6 @@ struct HashResult { std::string error; }; -std::uint64_t splitmix64(std::uint64_t* state) { - std::uint64_t value = (*state += 0x9e3779b97f4a7c15ull); - value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ull; - value = (value ^ (value >> 27)) * 0x94d049bb133111ebull; - return value ^ (value >> 31); -} - -double uniform_open(std::uint64_t* state) { - constexpr double kDenominator = 9007199254740993.0; - return (static_cast(splitmix64(state) >> 11) + 1.0) / - kDenominator; -} - -void deterministic_normal_f16(std::uint64_t seed, - std::vector* output) { - constexpr double kTwoPi = 6.283185307179586476925286766559; - std::uint64_t state = seed ^ 0x243f6a8885a308d3ull; - for (std::size_t i = 0; i < output->size(); i += 2) { - const double radius = std::sqrt(-2.0 * std::log(uniform_open(&state))); - const double angle = kTwoPi * uniform_open(&state); - (*output)[i] = modalities::float_to_float16( - static_cast(radius * std::cos(angle))); - if (i + 1 < output->size()) { - (*output)[i + 1] = modalities::float_to_float16( - static_cast(radius * std::sin(angle))); - } - } -} - -bool valid_config(const NativeThorCalibrationConfig& config) { - const std::uint64_t width = - static_cast(config.max_frame_width); - const std::uint64_t height = - static_cast(config.max_frame_height); - bool valid_quantiles = - config.state_q01.size() == - static_cast(config.state_dim) && - config.state_q99.size() == config.state_q01.size(); - for (std::size_t i = 0; - valid_quantiles && i < config.state_q01.size(); ++i) { - valid_quantiles = std::isfinite(config.state_q01[i]) && - std::isfinite(config.state_q99[i]) && - config.state_q99[i] > config.state_q01[i]; - } - return !config.checkpoint_path.empty() && - !config.tokenizer_model_path.empty() && config.state_dim > 0 && - config.num_views >= 1 && config.num_views <= 3 && - config.max_prompt_tokens >= 1 && - !(config.max_prompt_tokens & 1) && config.chunk_size > 0 && - config.num_steps > 0 && config.vision_pool_factor == 1 && - static_cast(config.max_prompt_tokens) + - static_cast(config.chunk_size) + - static_cast(config.num_views) * 256 <= - static_cast( - std::numeric_limits::max()) && - config.max_frame_width > 0 && config.max_frame_height > 0 && - width <= std::numeric_limits::max() / height / 4 && - valid_quantiles; -} - modalities::TensorView device_view(const NativeWorkspaceBuffer* buffer, modalities::DType dtype, modalities::Layout layout, @@ -289,24 +228,11 @@ struct NativeThorCalibrationSession::Impl { const float* noise, std::uint64_t n_noise, std::uint64_t noise_seed) { - if (!state || n_state != static_cast(config.state_dim)) { - return invalid("Thor calibration state shape is invalid"); - } - if ((noise && n_noise != noise_f16.size()) || - (!noise && n_noise != 0)) { - return invalid("Thor calibration noise shape is invalid"); - } - for (std::size_t i = 0; i < normalized_state.size(); ++i) { - if (!std::isfinite(state[i])) { - return invalid("Thor calibration state contains non-finite data"); - } - 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; - } + modalities::Status st = normalize_native_calibration_state( + config, state, n_state, &normalized_state); + if (!st.ok_status()) return st; std::uint64_t prompt_len = 0; - modalities::Status st = embed_prompt( + st = embed_prompt( tokenizer, prompt_spec, prompt, normalized_state.data(), normalized_state.size(), embedding_table, prompt_output, &token_ids, &prompt_len, stream, &text_staging, @@ -336,19 +262,12 @@ struct NativeThorCalibrationSession::Impl { stream); if (rc_check != cudaSuccess) return backend(cudaGetErrorString(rc_check)); - if (noise) { - for (std::size_t i = 0; i < noise_f16.size(); ++i) { - if (!std::isfinite(noise[i])) { - return invalid( - "Thor calibration noise contains non-finite data"); - } - noise_f16[i] = modalities::float_to_float16(noise[i]); - } - } else { - deterministic_normal_f16( - noise_seed + static_cast(encoder_samples.size()), - &noise_f16); - } + st = prepare_native_calibration_noise( + noise, n_noise, + noise_seed + static_cast(encoder_samples.size()), + static_cast(config.chunk_size) * 32, + modalities::DType::kFloat16, &noise_f16); + if (!st.ok_status()) return st; rc_check = cudaMemcpyAsync( action_output.data, noise_f16.data(), noise_f16.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice, @@ -439,7 +358,9 @@ NativeThorCalibrationSession::create( const NativeThorCalibrationConfig& config, double percentile, modalities::Status* status) { - if (!valid_config(config) || !std::isfinite(percentile) || + if (!valid_native_calibration_config(config) || + config.vision_pool_factor != 1 || (config.max_prompt_tokens & 1) || + !std::isfinite(percentile) || percentile < 0.0 || percentile > 100.0) { if (status) *status = invalid("Thor calibration config is invalid"); return nullptr; diff --git a/cpp/models/pi05/tests/CMakeLists.txt b/cpp/models/pi05/tests/CMakeLists.txt index 64cf8bcd..7dbaa0ec 100644 --- a/cpp/models/pi05/tests/CMakeLists.txt +++ b/cpp/models/pi05/tests/CMakeLists.txt @@ -170,10 +170,11 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) PRIVATE flashrt_cpp_pi05) endif() - if(FLASHRT_CPP_WITH_THOR_FP8 AND FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_thor_fp8_probe - ${_pi05_test_source_dir}/pi05_native_thor_fp8_probe.cpp) - target_link_libraries(pi05_native_thor_fp8_probe + if(FLASHRT_CPP_WITH_SENTENCEPIECE AND + (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) + add_executable(pi05_native_fp8_calibration_probe + ${_pi05_test_source_dir}/pi05_native_fp8_calibration_probe.cpp) + target_link_libraries(pi05_native_fp8_calibration_probe PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) endif() diff --git a/cpp/tests/pi05_native_thor_fp8_probe.cpp b/cpp/tests/pi05_native_fp8_calibration_probe.cpp similarity index 88% rename from cpp/tests/pi05_native_thor_fp8_probe.cpp rename to cpp/tests/pi05_native_fp8_calibration_probe.cpp index 9c8ae0a0..7502fb1a 100644 --- a/cpp/tests/pi05_native_thor_fp8_probe.cpp +++ b/cpp/tests/pi05_native_fp8_calibration_probe.cpp @@ -70,10 +70,27 @@ int calibration_error(frt_pi05_calibration_session* session, int main(int argc, char** argv) { if (argc != 6 && argc != 7) { - std::cerr << "usage: pi05_native_thor_fp8_probe CHECKPOINT TOKENIZER " + std::cerr << "usage: pi05_native_fp8_calibration_probe CHECKPOINT " + "TOKENIZER " "ARTIFACT SAMPLES VIEWS [RAW_ACTION_OUTPUT]\n"; return 2; } + int device = 0; + cudaDeviceProp properties{}; + cudaError_t cuda_rc = cudaGetDevice(&device); + if (cuda_rc == cudaSuccess) { + cuda_rc = cudaGetDeviceProperties(&properties, device); + } + if (cuda_rc != cudaSuccess || + !((properties.major == 11 || properties.major == 12) && + properties.minor == 0)) { + std::cerr << "native FP8 calibration requires SM110 or SM120\n"; + return 1; + } + const bool rtx = properties.major == 12; + const std::string hardware = + "sm" + std::to_string(properties.major * 10 + properties.minor); + const std::string hardware_identity = "hardware=" + hardware; const int samples = std::atoi(argv[4]); if (samples < 1 || samples > 256) { std::cerr << "SAMPLES must be in [1, 256]\n"; @@ -254,7 +271,9 @@ int main(int argc, char** argv) { flashrt::models::pi05::NativeCalibrationArtifact artifact; auto status = flashrt::models::pi05::load_native_calibration_artifact( single_path, &artifact); - if (!status.ok_status() || artifact.sample_count != 1) { + if (!status.ok_status() || artifact.sample_count != 1 || + artifact.hardware != hardware || + artifact.activation_dtype != (rtx ? "bfloat16" : "float16")) { std::cerr << "single calibration artifact validation failed\n"; return 1; } @@ -287,11 +306,15 @@ int main(int argc, char** argv) { exp->identity && std::strstr(exp->identity, "precision=fp8_e4m3fn") && - std::strstr(exp->identity, "hardware=sm110") && + std::strstr(exp->identity, + hardware_identity.c_str()) && std::strstr(exp->identity, "calibration_sha256=") && - model->ports[2].dtype == FRT_RT_DTYPE_F16 && - model->ports[3].dtype == FRT_RT_DTYPE_F16 && - model->ports[5].dtype == FRT_RT_DTYPE_F16; + model->ports[2].dtype == + (rtx ? FRT_RT_DTYPE_BF16 : FRT_RT_DTYPE_F16) && + model->ports[3].dtype == + (rtx ? FRT_RT_DTYPE_BF16 : FRT_RT_DTYPE_F16) && + model->ports[5].dtype == + (rtx ? FRT_RT_DTYPE_BF16 : FRT_RT_DTYPE_F16); for (std::uint64_t i = 0; schema_ok && i < exp->n_graphs; ++i) { schema_ok = frt_graph_variant_count(exp->graphs[i].handle) == 1; @@ -327,13 +350,16 @@ int main(int argc, char** argv) { model->release(model->owner); return 1; } - std::vector noise_f16(noise.size()); + std::vector noise_payload(noise.size()); for (std::size_t i = 0; i < noise.size(); ++i) { - noise_f16[i] = flashrt::modalities::float_to_float16(noise[i]); + noise_payload[i] = + rtx ? flashrt::modalities::float_to_bfloat16(noise[i]) + : flashrt::modalities::float_to_float16(noise[i]); } if (!model->ports[3].buffer || - cudaMemcpy(frt_buffer_dptr(model->ports[3].buffer), noise_f16.data(), - noise_f16.size() * sizeof(std::uint16_t), + cudaMemcpy(frt_buffer_dptr(model->ports[3].buffer), + noise_payload.data(), + noise_payload.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice) != cudaSuccess || model->verbs.step(model->self) != 0 || cudaDeviceSynchronize() != cudaSuccess) { @@ -347,8 +373,9 @@ int main(int argc, char** argv) { cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), raw.size() * sizeof(raw[0]), cudaMemcpyDeviceToHost) != cudaSuccess || - cudaMemcpy(frt_buffer_dptr(model->ports[3].buffer), noise_f16.data(), - noise_f16.size() * sizeof(std::uint16_t), + cudaMemcpy(frt_buffer_dptr(model->ports[3].buffer), + noise_payload.data(), + noise_payload.size() * sizeof(std::uint16_t), cudaMemcpyHostToDevice) != cudaSuccess || frt_graph_replay(exp->graphs[2].handle, 0, exp->graphs[2].stream_id) != FRT_OK || @@ -398,6 +425,7 @@ int main(int argc, char** argv) { } } model->release(model->owner); - std::cout << "PASS native Thor FP8 calibration and runtime lifecycle\n"; + std::cout << "PASS native " << hardware + << " FP8 calibration and runtime lifecycle\n"; return 0; } diff --git a/cpp/tests/test_pi05_native_calibration.cpp b/cpp/tests/test_pi05_native_calibration.cpp index 2069ceb0..65448417 100644 --- a/cpp/tests/test_pi05_native_calibration.cpp +++ b/cpp/tests/test_pi05_native_calibration.cpp @@ -1,4 +1,5 @@ #include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/native_calibration_session.h" #include #include @@ -22,9 +23,60 @@ std::string temp_path() { int main() { using flashrt::models::pi05::NativeCalibrationArtifact; + using flashrt::models::pi05::NativeCalibrationConfig; using flashrt::models::pi05::load_native_calibration_artifact; + using flashrt::models::pi05::normalize_native_calibration_state; + using flashrt::models::pi05::prepare_native_calibration_noise; using flashrt::models::pi05::reduce_native_calibration_samples; using flashrt::models::pi05::save_native_calibration_artifact; + using flashrt::models::pi05::valid_native_calibration_config; + + NativeCalibrationConfig config; + config.checkpoint_path = "checkpoint"; + config.tokenizer_model_path = "tokenizer.model"; + config.state_dim = 2; + config.num_views = 3; + config.state_q01 = {-2.0f, 0.0f}; + config.state_q99 = {2.0f, 4.0f}; + assert(valid_native_calibration_config(config)); + std::vector normalized; + const float state[] = {0.0f, 1.0f}; + assert(normalize_native_calibration_state( + config, state, 2, &normalized) + .ok_status()); + assert(normalized.size() == 2); + assert(std::fabs(normalized[0]) < 1e-6f); + assert(std::fabs(normalized[1] + 0.5f) < 1e-6f); + const float invalid_state[] = {0.0f, INFINITY}; + assert(!normalize_native_calibration_state( + config, invalid_state, 2, &normalized) + .ok_status()); + + const float explicit_noise[] = {-1.0f, 0.0f, 1.0f}; + std::vector encoded; + assert(prepare_native_calibration_noise( + explicit_noise, 3, 0, 3, + flashrt::modalities::DType::kBFloat16, &encoded) + .ok_status()); + assert(encoded == std::vector({ + flashrt::modalities::float_to_bfloat16(-1.0f), + flashrt::modalities::float_to_bfloat16(0.0f), + flashrt::modalities::float_to_bfloat16(1.0f)})); + std::vector generated; + assert(prepare_native_calibration_noise( + nullptr, 0, 1234, 7, + flashrt::modalities::DType::kFloat16, &generated) + .ok_status()); + std::vector repeated; + assert(prepare_native_calibration_noise( + nullptr, 0, 1234, 7, + flashrt::modalities::DType::kFloat16, &repeated) + .ok_status()); + assert(generated == repeated); + assert(!prepare_native_calibration_noise( + explicit_noise, 2, 0, 3, + flashrt::modalities::DType::kFloat16, &encoded) + .ok_status()); std::vector reduced; assert(reduce_native_calibration_samples( From b871c47e23de4d35a954baf0eeeb781d6b603915 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 04:14:16 -0400 Subject: [PATCH 72/83] test(cpp): separate pi05 diagnostic oracles --- .github/pull_request_template.md | 3 + CONTRIBUTING.md | 22 ++ cpp/CMakeLists.txt | 2 + cpp/models/pi05/tests/CMakeLists.txt | 113 +++--- cpp/tests/gate_pi05_native_diffusion.py | 273 -------------- cpp/tests/gate_pi05_native_e2e.py | 275 -------------- cpp/tests/gate_pi05_native_encoder.py | 184 --------- cpp/tests/gate_pi05_native_encoder_layer.py | 125 ------ cpp/tests/gate_pi05_native_encoder_qkv.py | 100 ----- cpp/tests/gate_pi05_native_quantization.py | 153 -------- cpp/tests/gate_pi05_native_rope.py | 70 ---- cpp/tests/gate_pi05_native_style.py | 132 ------- cpp/tests/gate_pi05_native_thor_fp8.py | 396 -------------------- cpp/tests/gate_pi05_native_vision.py | 190 ---------- cpp/tests/gate_pi05_native_weight_ops.py | 141 ------- cpp/tests/gate_pi05_tokenizer_corpus.py | 140 ------- cpp/tests/profile_pi05_python_replay.py | 102 ----- docs/cpp_runtime_modalities.md | 12 +- docs/mindon_pi05_integration.md | 31 +- docs/pi05_cpp_runtime_migration.md | 21 +- docs/pi05_io_contract.md | 67 ++-- docs/pi05_thor_native_fp8.md | 146 +++++--- docs/pr_review_checklist.md | 20 + docs/runtime_contract.md | 12 +- 24 files changed, 300 insertions(+), 2430 deletions(-) delete mode 100644 cpp/tests/gate_pi05_native_diffusion.py delete mode 100644 cpp/tests/gate_pi05_native_e2e.py delete mode 100644 cpp/tests/gate_pi05_native_encoder.py delete mode 100644 cpp/tests/gate_pi05_native_encoder_layer.py delete mode 100644 cpp/tests/gate_pi05_native_encoder_qkv.py delete mode 100644 cpp/tests/gate_pi05_native_quantization.py delete mode 100644 cpp/tests/gate_pi05_native_rope.py delete mode 100644 cpp/tests/gate_pi05_native_style.py delete mode 100644 cpp/tests/gate_pi05_native_thor_fp8.py delete mode 100644 cpp/tests/gate_pi05_native_vision.py delete mode 100644 cpp/tests/gate_pi05_native_weight_ops.py delete mode 100644 cpp/tests/gate_pi05_tokenizer_corpus.py delete mode 100644 cpp/tests/profile_pi05_python_replay.py diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 897ae578..13d0f3be 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -17,6 +17,8 @@ - [ ] Focused tests cover success and rejection paths - [ ] Affected CUDA-off/hardware configurations were checked or disclosed - [ ] Numerical claims use a fixed, justified gate +- [ ] Low-precision claims include identity, artifact, and kernel-dispatch evidence +- [ ] Compared producer binaries were built from the same source revision - [ ] Calibration artifacts and cache invalidation are identity-complete - [ ] Single/multi-sample calibration and named-input rejection are covered - [ ] STAGED ports have real matching verbs @@ -26,3 +28,4 @@ - [ ] Diff contains no private paths, hosts, containers, credentials, or logs - [ ] Shared kernel/CMake ownership and packaging were reviewed - [ ] Model semantics remain model-local; frozen runtime/exec stay generic +- [ ] External test migration preserves core contract/result coverage diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 83d71495..8dae9c99 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -264,6 +264,28 @@ Native calibration APIs and artifacts have additional producer-boundary rules: array order does not change semantic order. - Use explicit fixed stochastic inputs for reference comparison. A generated fallback must be deterministic, documented, and separately exercised. +- Do not infer compute precision from a public staging dtype. Native low- + precision evidence must agree across producer identity, calibration artifact + metadata, and captured kernel dispatch; mixed-precision boundary kernels are + documented separately from the main GEMM route. +- Build every producer used in a parity comparison from the same source + revision. A stale pybind extension or shared library is an invalid numerical + oracle even when its Python sources match the checkout. + +### Native Producer Test Ownership + +Keep frozen ABI/schema, lifecycle/ownership, artifact validation, final-result, +subgraph-equivalence, and hot-path invariant tests in FlashRT. Large-checkpoint +official-framework oracles, layer-by-layer diagnostics, profiler reports, +deployment shadow comparisons, and soak workloads may live in an integration +validation repository. Detailed native probes must be opt-in, non-installed +targets; do not expose intermediate model state through the production ABI to +make an external test convenient. + +Moving a test requires a checked-in coverage map, byte-for-byte migration (or +an explicitly reviewed rewrite), and one overlapping successful run before the +old path is removed. Default builds must still exercise failure paths and final +observable behavior without private checkpoints. ### Performance Measurement diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ac2f928d..5a9a3594 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -26,6 +26,8 @@ option(FLASHRT_CPP_WITH_SENTENCEPIECE "Enable native SentencePiece text tokeniza option(FLASHRT_CPP_WITH_FA2 "Enable the Python-free vendored FA2 driver" OFF) option(FLASHRT_CPP_WITH_THOR_FP8 "Enable the native Thor SM110 FP8 kernel backend" OFF) +option(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS + "Build detailed PI0.5 real-checkpoint diagnostic probes" OFF) set(FLASHRT_CPP_FA2_LIBRARY "" CACHE FILEPATH "Path to the Python-free libflashrt_fa2_raw shared library") set(FLASHRT_CPP_CUTLASS_DIR diff --git a/cpp/models/pi05/tests/CMakeLists.txt b/cpp/models/pi05/tests/CMakeLists.txt index 7dbaa0ec..876ce6a6 100644 --- a/cpp/models/pi05/tests/CMakeLists.txt +++ b/cpp/models/pi05/tests/CMakeLists.txt @@ -31,9 +31,11 @@ add_executable(test_pi05_native_calibration target_link_libraries(test_pi05_native_calibration PRIVATE flashrt_cpp_pi05) add_test(NAME pi05_native_calibration COMMAND test_pi05_native_calibration) -add_executable(pi05_native_weight_probe - ${_pi05_test_source_dir}/pi05_native_weight_probe.cpp) -target_link_libraries(pi05_native_weight_probe PRIVATE flashrt_cpp_pi05) +if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) + add_executable(pi05_native_weight_probe + ${_pi05_test_source_dir}/pi05_native_weight_probe.cpp) + target_link_libraries(pi05_native_weight_probe PRIVATE flashrt_cpp_pi05) +endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_executable(test_pi05_native_quantization @@ -42,10 +44,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_test(NAME pi05_native_quantization COMMAND test_pi05_native_quantization) - add_executable(pi05_native_quant_probe - ${_pi05_test_source_dir}/pi05_native_quant_probe.cpp) - target_link_libraries(pi05_native_quant_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) + add_executable(pi05_native_quant_probe + ${_pi05_test_source_dir}/pi05_native_quant_probe.cpp) + target_link_libraries(pi05_native_quant_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + endif() add_executable(test_pi05_native_weight_packer ${_pi05_test_source_dir}/test_pi05_native_weight_packer.cpp) @@ -75,10 +79,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_test(NAME pi05_native_bf16_forward COMMAND test_pi05_native_bf16_forward) - add_executable(pi05_native_encoder_qkv_probe - ${_pi05_test_source_dir}/pi05_native_encoder_qkv_probe.cpp) - target_link_libraries(pi05_native_encoder_qkv_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) + add_executable(pi05_native_encoder_qkv_probe + ${_pi05_test_source_dir}/pi05_native_encoder_qkv_probe.cpp) + target_link_libraries(pi05_native_encoder_qkv_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + endif() add_executable(test_pi05_native_style_precompute ${_pi05_test_source_dir}/test_pi05_native_style_precompute.cpp) @@ -87,10 +93,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_test(NAME pi05_native_style_precompute COMMAND test_pi05_native_style_precompute) - add_executable(pi05_native_style_probe - ${_pi05_test_source_dir}/pi05_native_style_probe.cpp) - target_link_libraries(pi05_native_style_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) + add_executable(pi05_native_style_probe + ${_pi05_test_source_dir}/pi05_native_style_probe.cpp) + target_link_libraries(pi05_native_style_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + endif() add_executable(test_pi05_native_workspace ${_pi05_test_source_dir}/test_pi05_native_workspace.cpp) @@ -107,37 +115,41 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) COMMAND test_pi05_native_rtx_attention) if(FLASHRT_CPP_WITH_FA2) - add_executable(pi05_native_encoder_layer_probe - ${_pi05_test_source_dir}/pi05_native_encoder_layer_probe.cpp) - target_link_libraries(pi05_native_encoder_layer_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_encoder_probe - ${_pi05_test_source_dir}/pi05_native_encoder_probe.cpp) - target_link_libraries(pi05_native_encoder_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_vision_probe - ${_pi05_test_source_dir}/pi05_native_vision_probe.cpp) - target_link_libraries(pi05_native_vision_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_diffusion_probe - ${_pi05_test_source_dir}/pi05_native_diffusion_probe.cpp) - target_link_libraries(pi05_native_diffusion_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_graph_probe - ${_pi05_test_source_dir}/pi05_native_graph_probe.cpp) - target_link_libraries(pi05_native_graph_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) + add_executable(pi05_native_encoder_layer_probe + ${_pi05_test_source_dir}/pi05_native_encoder_layer_probe.cpp) + target_link_libraries(pi05_native_encoder_layer_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_encoder_probe + ${_pi05_test_source_dir}/pi05_native_encoder_probe.cpp) + target_link_libraries(pi05_native_encoder_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_vision_probe + ${_pi05_test_source_dir}/pi05_native_vision_probe.cpp) + target_link_libraries(pi05_native_vision_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_diffusion_probe + ${_pi05_test_source_dir}/pi05_native_diffusion_probe.cpp) + target_link_libraries(pi05_native_diffusion_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + add_executable(pi05_native_graph_probe + ${_pi05_test_source_dir}/pi05_native_graph_probe.cpp) + target_link_libraries(pi05_native_graph_probe + PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) + + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + add_executable(pi05_native_e2e_probe + ${_pi05_test_source_dir}/pi05_native_e2e_probe.cpp) + target_link_libraries(pi05_native_e2e_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) + endif() + endif() if(FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_e2e_probe - ${_pi05_test_source_dir}/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 ${_pi05_test_source_dir}/pi05_native_dlopen_probe.cpp) target_include_directories(pi05_native_dlopen_probe @@ -163,7 +175,8 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) endif() - if(FLASHRT_CPP_WITH_SENTENCEPIECE) + if(FLASHRT_CPP_WITH_SENTENCEPIECE AND + FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) add_executable(pi05_tokenizer_corpus_probe ${_pi05_test_source_dir}/pi05_tokenizer_corpus_probe.cpp) target_link_libraries(pi05_tokenizer_corpus_probe @@ -178,10 +191,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) endif() - add_executable(pi05_native_rope_probe - ${_pi05_test_source_dir}/pi05_native_rope_probe.cpp) - target_link_libraries(pi05_native_rope_probe - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) + add_executable(pi05_native_rope_probe + ${_pi05_test_source_dir}/pi05_native_rope_probe.cpp) + target_link_libraries(pi05_native_rope_probe + PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) + endif() endif() add_executable(test_pi05_prompt_format diff --git a/cpp/tests/gate_pi05_native_diffusion.py b/cpp/tests/gate_pi05_native_diffusion.py deleted file mode 100644 index ebd4be25..00000000 --- a/cpp/tests/gate_pi05_native_diffusion.py +++ /dev/null @@ -1,273 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import gc -import math -import pathlib -import subprocess -import tempfile - -import numpy as np -import torch -import torch.nn.functional as F -from safetensors import safe_open - - -CHUNK = 10 -PREFIX = 712 - - -def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: - output, inputs = weight.shape - head_dim = output // heads - return ( - weight.reshape(heads, head_dim, inputs) - .reshape(heads, 2, head_dim // 2, inputs) - .permute(0, 2, 1, 3) - .reshape(output, inputs) - ) - - -def ada_rms(values: torch.Tensor, style: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - source = values.float() - normalized = source * torch.rsqrt(source.square().mean(-1, keepdim=True) + 1e-6) - scale, shift, gate = style.float().chunk(3, dim=-1) - output = (normalized * (1.0 + scale) + shift).to(torch.bfloat16) - return output, gate.to(torch.bfloat16) - - -def rotate(tensor: torch.Tensor, heads: int, rope: torch.Tensor) -> torch.Tensor: - pairs = tensor.reshape(CHUNK, heads, 128, 2).float() - cosine = rope[:, None, :, 0].float() - sine = rope[:, None, :, 1].float() - return torch.stack( - [ - pairs[..., 0] * cosine - pairs[..., 1] * sine, - pairs[..., 1] * cosine + pairs[..., 0] * sine, - ], - -1, - ).to(torch.bfloat16).reshape(CHUNK, heads, 256) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - parser.add_argument("--steps", type=int, default=1) - parser.add_argument("--start-step", type=int, default=0) - args = parser.parse_args() - if args.steps < 1 or args.start_step < 0 or args.start_step + args.steps > 10: - raise ValueError("start-step and steps must select a subset of [0, 10)") - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - root = "model." if "model.action_in_proj.weight" in keys else "" - - def raw(name: str) -> torch.Tensor: - return file.get_tensor(root + name) - - def bf16(name: str) -> torch.Tensor: - return raw(name).to(device="cuda", dtype=torch.bfloat16) - - decoder = "paligemma_with_expert.gemma_expert.model.layers" - time_in_w = bf16("time_mlp_in.weight").t().contiguous() - time_in_b = bf16("time_mlp_in.bias") - time_out_w = bf16("time_mlp_out.weight").t().contiguous() - time_out_b = bf16("time_mlp_out.bias") - attn_mod_w = [ - bf16(f"{decoder}.{i}.input_layernorm.dense.weight").t().contiguous() - for i in range(18) - ] - attn_mod_b = [ - bf16(f"{decoder}.{i}.input_layernorm.dense.bias") for i in range(18) - ] - ffn_mod_w = [ - bf16(f"{decoder}.{i}.post_attention_layernorm.dense.weight") - .t() - .contiguous() - for i in range(18) - ] - ffn_mod_b = [ - bf16(f"{decoder}.{i}.post_attention_layernorm.dense.bias") - for i in range(18) - ] - final_mod_w = bf16( - "paligemma_with_expert.gemma_expert.model.norm.dense.weight" - ).t().contiguous() - final_mod_b = bf16( - "paligemma_with_expert.gemma_expert.model.norm.dense.bias" - ) - fraction = torch.linspace(0.0, 1.0, 512) - period = 4e-3 * (4.0 / 4e-3) ** fraction - current = torch.tensor(1.0, dtype=torch.float32) - schedule = [] - for _ in range(10): - angle = current * (1.0 / period) * 2 * math.pi - schedule.append( - torch.cat([torch.sin(angle), torch.cos(angle)]).to( - device="cuda", dtype=torch.bfloat16 - ) - ) - current = current - 0.1 - styles_attn = torch.empty( - 10, 18, CHUNK, 3072, device="cuda", dtype=torch.bfloat16 - ) - styles_ffn = torch.empty_like(styles_attn) - styles_final = torch.empty( - 10, CHUNK, 3072, device="cuda", dtype=torch.bfloat16 - ) - step_range = range(args.start_step, args.start_step + args.steps) - for step in step_range: - value = schedule[step][None, :] @ time_in_w - value = (value.float() + time_in_b.float()).to(torch.bfloat16) - value_float = value.float() - value = (value_float * torch.sigmoid(value_float)).to(torch.bfloat16) - value = value @ time_out_w - value = (value.float() + time_out_b.float()).to(torch.bfloat16) - value_float = value.float() - value = (value_float * torch.sigmoid(value_float)).to(torch.bfloat16) - expanded = value.expand(CHUNK, -1).contiguous() - for layer in range(18): - styles_attn[step, layer] = ( - (expanded @ attn_mod_w[layer]).float() - + attn_mod_b[layer].float() - ).to(torch.bfloat16) - styles_ffn[step, layer] = ( - (expanded @ ffn_mod_w[layer]).float() - + ffn_mod_b[layer].float() - ).to(torch.bfloat16) - styles_final[step] = ( - (expanded @ final_mod_w).float() + final_mod_b.float() - ).to(torch.bfloat16) - - layers = [] - for index in range(18): - prefix = f"{decoder}.{index}" - q = interleave_qk(raw(f"{prefix}.self_attn.q_proj.weight"), 8) - k = interleave_qk(raw(f"{prefix}.self_attn.k_proj.weight"), 1) - v = raw(f"{prefix}.self_attn.v_proj.weight") - layers.append( - { - "qkv": torch.cat([q, k, v], dim=0) - .t() - .to(device="cuda", dtype=torch.bfloat16) - .contiguous(), - "output": bf16(f"{prefix}.self_attn.o_proj.weight") - .t() - .contiguous(), - "gate": bf16(f"{prefix}.mlp.gate_proj.weight").t().contiguous(), - "up": bf16(f"{prefix}.mlp.up_proj.weight").t().contiguous(), - "down": bf16(f"{prefix}.mlp.down_proj.weight").t().contiguous(), - } - ) - - positions = torch.arange(PREFIX, PREFIX + CHUNK, dtype=torch.float64)[:, None] - pair = torch.arange(128, dtype=torch.float64)[None, :] - phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) - rope = torch.stack([torch.cos(phase), torch.sin(phase)], -1).to( - device="cuda", dtype=torch.bfloat16 - ) - layer_index = torch.arange(18, device="cuda")[:, None, None] - row_index = torch.arange(722, device="cuda")[None, :, None] - column_index = torch.arange(256, device="cuda")[None, None, :] - cache_k = ( - ((layer_index + row_index + column_index) % 17 - 8).float() / 16.0 - ).to(torch.bfloat16) - cache_v = ( - ((2 * layer_index + row_index + 3 * column_index) % 19 - 9).float() - / 16.0 - ).to(torch.bfloat16) - flat = torch.arange(CHUNK * 32, device="cuda") - noise = (((flat % 23) - 11).float() / 12.0).to(torch.bfloat16).reshape( - CHUNK, 32 - ) - input_weight = bf16("action_in_proj.weight").t().contiguous() - input_bias = bf16("action_in_proj.bias") - output_weight = ( - bf16("action_out_proj.weight").float() * -0.1 - ).to(torch.bfloat16).t().contiguous() - output_bias = ( - bf16("action_out_proj.bias").float() * -0.1 - ).to(torch.bfloat16) - - for step in step_range: - x = noise @ input_weight - x = (x.float() + input_bias.float()).to(torch.bfloat16) - for index, weights in enumerate(layers): - x_norm, residual_gate = ada_rms(x, styles_attn[step, index]) - qkv = x_norm @ weights["qkv"] - query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) - query = rotate(query, 8, rope) - key = rotate(key, 1, rope).reshape(CHUNK, 256) - value = value.reshape(CHUNK, 256) - cache_k[index, PREFIX : PREFIX + CHUNK] = key - cache_v[index, PREFIX : PREFIX + CHUNK] = value - attended = F.scaled_dot_product_attention( - query.transpose(0, 1).unsqueeze(0), - cache_k[index].reshape(722, 1, 256).transpose(0, 1).unsqueeze(0), - cache_v[index].reshape(722, 1, 256).transpose(0, 1).unsqueeze(0), - scale=1.0 / 16.0, - enable_gqa=True, - ).squeeze(0).transpose(0, 1).reshape(CHUNK, 2048) - projected = attended @ weights["output"] - x = ( - x.float() + projected.float() * residual_gate.float() - ).to(torch.bfloat16) - x_norm, residual_gate = ada_rms(x, styles_ffn[step, index]) - gate = x_norm @ weights["gate"] - up = x_norm @ weights["up"] - gate_float = gate.float() - activated = gate_float / ( - 1.0 - + torch.exp( - -1.5957691216057308 - * gate_float - * (1.0 + 0.044715 * gate_float.square()) - ) - ) - hidden = (activated * up.float()).to(torch.bfloat16) - down = hidden @ weights["down"] - x = (x.float() + down.float() * residual_gate.float()).to( - torch.bfloat16 - ) - x_norm, _ = ada_rms(x, styles_final[step]) - action = x_norm @ output_weight - action = (action.float() + output_bias.float()).to(torch.bfloat16) - noise = (noise.float() + action.float()).to(torch.bfloat16) - - expected = noise.cpu().float() - del cache_k, cache_v, noise, x, x_norm, action - gc.collect() - torch.cuda.empty_cache() - with tempfile.TemporaryDirectory() as directory: - output = str(pathlib.Path(directory) / "diffusion.bin") - subprocess.check_call( - [ - args.probe, - args.checkpoint, - output, - str(args.steps), - str(args.start_step), - ] - ) - bits = np.fromfile(output, dtype=np.uint16) - if bits.size != CHUNK * 32: - raise AssertionError(f"diffusion probe output elements={bits.size}") - actual = torch.from_numpy(bits.copy()).view(torch.bfloat16).float().reshape( - CHUNK, 32 - ) - cosine = float( - F.cosine_similarity( - actual.flatten().double(), expected.flatten().double(), dim=0 - ) - ) - maximum = float((actual - expected).abs().max()) - if cosine < 0.9999: - raise AssertionError(f"cosine={cosine:.8f} max={maximum:.6f}") - print( - f"PASS diffusion steps {args.start_step}.." - f"{args.start_step + args.steps - 1} " - f"cosine={cosine:.8f} max={maximum:.6f}" - ) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_e2e.py b/cpp/tests/gate_pi05_native_e2e.py deleted file mode 100644 index 5a012c6f..00000000 --- a/cpp/tests/gate_pi05_native_e2e.py +++ /dev/null @@ -1,275 +0,0 @@ -"""Compare native SM120 Pi0.5 against the official OpenPI PyTorch policy.""" - -from __future__ import annotations - -import argparse -import io -import json -import os -from pathlib import Path -import subprocess -import sys -import tempfile - -import ml_dtypes -import numpy as np -from PIL import Image -import pyarrow.compute as pc -import pyarrow.parquet as pq - - -def _cosine(a: np.ndarray, b: np.ndarray) -> float: - x = np.asarray(a, dtype=np.float64).reshape(-1) - y = np.asarray(b, dtype=np.float64).reshape(-1) - nx = np.linalg.norm(x) - ny = np.linalg.norm(y) - return float(x @ y / (nx * ny)) if nx and ny else float("nan") - - -def _decode_image(cell) -> np.ndarray: - raw = cell["bytes"] if isinstance(cell, dict) else cell - image = Image.open(io.BytesIO(raw)).convert("RGB") - image = image.resize((224, 224), Image.Resampling.BILINEAR) - return np.ascontiguousarray(image, dtype=np.uint8) - - -def _task(root: Path, task_index: int) -> str: - with (root / "meta" / "tasks.jsonl").open(encoding="utf-8") as stream: - for line in stream: - item = json.loads(line) - if int(item["task_index"]) == task_index: - return str(item["task"]) - raise KeyError(f"task_index={task_index} is missing") - - -def _make_fixture(args, fixture: Path) -> None: - info = json.loads((args.dataset / "meta" / "info.json").read_text()) - chunk = args.episode // int(info.get("chunks_size", 1000)) - relative = info.get( - "data_path", - "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", - ).format(episode_chunk=chunk, episode_index=args.episode) - table = pq.read_table(args.dataset / relative) - rows = table.filter(pc.equal(table["frame_index"], args.frame)).to_pylist() - if len(rows) != 1: - raise RuntimeError( - f"expected one episode={args.episode} frame={args.frame} row" - ) - row = rows[0] - _decode_image(row["image"]).tofile(fixture / "image_0.rgb") - _decode_image(row["wrist_image"]).tofile(fixture / "image_1.rgb") - np.asarray(row["state"], dtype=np.float32).tofile(fixture / "state.f32") - (fixture / "prompt.txt").write_text( - _task(args.dataset, int(row["task_index"])), encoding="utf-8" - ) - values = np.random.default_rng(args.seed).standard_normal(10 * 32) - np.asarray(values, dtype=np.float32).astype(ml_dtypes.bfloat16).tofile( - fixture / "noise.bf16" - ) - - -def _official_baseline(checkpoint: Path, fixture: Path, output: Path) -> None: - import torch - - from openpi.models import model as model_api - from openpi.models import tokenizer as tokenizer_api - from openpi.training import config as training_config - - def image(name: str) -> np.ndarray: - return np.fromfile(fixture / name, dtype=np.uint8).reshape(224, 224, 3) - - state = np.fromfile(fixture / "state.f32", dtype=np.float32) - prompt = (fixture / "prompt.txt").read_text(encoding="utf-8") - noise = np.fromfile(fixture / "noise.bf16", dtype=ml_dtypes.bfloat16) - noise = noise.astype(np.float32).reshape(10, 32) - stats = json.loads( - (checkpoint / "assets" / "physical-intelligence" / "libero" / - "norm_stats.json").read_text() - )["norm_stats"] - state_q01 = np.asarray(stats["state"]["q01"], dtype=np.float32) - state_q99 = np.asarray(stats["state"]["q99"], dtype=np.float32) - normalized_state = ( - (state - state_q01) / (state_q99 - state_q01 + 1e-6) * 2.0 - 1.0 - ) - tokens, token_mask = tokenizer_api.PaligemmaTokenizer(200).tokenize( - prompt, normalized_state - ) - padded_state = np.zeros(32, dtype=np.float32) - padded_state[:state.size] = normalized_state - base = image("image_0.rgb") - wrist = image("image_1.rgb") - device_inputs = { - "image": { - "base_0_rgb": torch.from_numpy(base).to("cuda")[None, ...], - "left_wrist_0_rgb": torch.from_numpy(wrist).to("cuda")[None, ...], - "right_wrist_0_rgb": torch.zeros_like( - torch.from_numpy(base).to("cuda")[None, ...] - ), - }, - "image_mask": { - "base_0_rgb": torch.ones(1, dtype=torch.bool, device="cuda"), - "left_wrist_0_rgb": torch.ones(1, dtype=torch.bool, device="cuda"), - "right_wrist_0_rgb": torch.zeros(1, dtype=torch.bool, device="cuda"), - }, - "state": torch.from_numpy(padded_state).to("cuda")[None, ...], - "tokenized_prompt": torch.from_numpy(tokens).to("cuda")[None, ...], - "tokenized_prompt_mask": torch.from_numpy(token_mask).to("cuda")[None, ...], - } - model_observation = model_api.Observation.from_dict(device_inputs) - train_config = training_config.get_config("pi05_libero") - model = train_config.model.load_pytorch( - train_config, str(checkpoint / "model.safetensors") - ) - model.paligemma_with_expert.to_bfloat16_for_selected_params("bfloat16") - model.to("cuda").eval() - noise_tensor = torch.from_numpy(noise).to("cuda")[None, ...] - with torch.inference_mode(): - raw = model.sample_actions( - "cuda", model_observation, noise=noise_tensor, num_steps=10 - )[0].float().cpu().numpy() - action_q01 = np.asarray(stats["actions"]["q01"], dtype=np.float32) - action_q99 = np.asarray(stats["actions"]["q99"], dtype=np.float32) - clipped = np.clip(raw[:, :action_q01.size], -1.0, 1.0) - actions = ((clipped + 1.0) * 0.5 * - (action_q99 - action_q01 + 1e-6) + action_q01) - np.asarray(raw, dtype=np.float32).tofile(output / "openpi_raw.f32") - np.asarray(actions, dtype=np.float32).tofile( - output / "openpi_actions.f32" - ) - - -def _run(args) -> None: - with tempfile.TemporaryDirectory(prefix="pi05_native_e2e_") as temp: - root = Path(temp) - _make_fixture(args, root) - env = dict(os.environ) - env["TORCH_COMPILE_DISABLE"] = "1" - baseline_prefix = env.get("OPENPI_BASELINE_SITE_PACKAGES") - if baseline_prefix: - baseline_packages = Path(baseline_prefix) - if not baseline_packages.is_dir(): - raise FileNotFoundError(baseline_packages) - existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = str(baseline_packages) + ( - os.pathsep + existing if existing else "" - ) - subprocess.run( - [ - sys.executable, - __file__, - "--baseline-fixture", - str(root), - "--checkpoint", - str(args.checkpoint), - ], - check=True, - env=env, - ) - subprocess.run( - [ - str(args.probe), - str(args.checkpoint), - str(args.tokenizer), - str(root), - str(root), - ], - check=True, - ) - openpi_raw = np.fromfile(root / "openpi_raw.f32", dtype=np.float32) - native_raw = np.fromfile( - root / "native_raw.bf16", dtype=ml_dtypes.bfloat16 - ).astype(np.float32) - openpi_actions = np.fromfile( - root / "openpi_actions.f32", dtype=np.float32 - ) - native_actions = np.fromfile( - root / "native_actions.f32", dtype=np.float32 - ) - sizes = { - "openpi_raw": openpi_raw.size, - "native_raw": native_raw.size, - "openpi_actions": openpi_actions.size, - "native_actions": native_actions.size, - } - expected_sizes = { - "openpi_raw": 10 * 32, - "native_raw": 10 * 32, - "openpi_actions": 10 * 7, - "native_actions": 10 * 7, - } - if sizes != expected_sizes: - raise RuntimeError(f"unexpected E2E output sizes: {sizes}") - raw_cos = _cosine(openpi_raw, native_raw) - action_cos = _cosine(openpi_actions, native_actions) - raw_max = float(np.max(np.abs(openpi_raw - native_raw))) - action_max = float(np.max(np.abs(openpi_actions - native_actions))) - stats = json.loads( - (args.checkpoint / "assets" / "physical-intelligence" / - "libero" / "norm_stats.json").read_text() - )["norm_stats"]["actions"] - q01 = np.asarray(stats["q01"], dtype=np.float32) - q99 = np.asarray(stats["q99"], dtype=np.float32) - native_model = native_raw.reshape(10, 32)[:, :q01.size] - native_contract_actions = ( - (np.clip(native_model, -1.0, 1.0) + 1.0) * 0.5 * - (q99 - q01 + 1e-6) + q01 - ) - contract_max = float(np.max(np.abs( - native_contract_actions.reshape(-1) - native_actions - ))) - contract_close = bool( - np.allclose(native_contract_actions.reshape(-1), native_actions, - rtol=1e-6, atol=1e-6) - ) - print("\n===== PI0.5 NATIVE VS OFFICIAL OPENPI =====") - print(f"episode/frame : {args.episode}/{args.frame}") - print(f"raw action cosine : {raw_cos:.8f} max_abs={raw_max:.6g}") - print( - f"robot action : cos={action_cos:.8f} " - f"max_abs_vs_fp32={action_max:.6g}" - ) - print( - f"action postprocess: allclose={contract_close} " - f"max_abs={contract_max:.6g}" - ) - if raw_cos < 0.9999: - raise AssertionError(f"raw action cosine {raw_cos:.8f} < 0.9999") - if action_cos < 0.9999: - raise AssertionError( - f"robot action cosine {action_cos:.8f} < 0.9999" - ) - if not contract_close: - raise AssertionError( - f"native action postprocess differs; max_abs={contract_max:.6g}" - ) - print("PASS native Pi0.5 real-episode E2E") - - -def _parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", type=Path, required=True) - parser.add_argument("--tokenizer", type=Path) - parser.add_argument("--dataset", type=Path) - parser.add_argument("--probe", type=Path) - parser.add_argument("--episode", type=int, default=0) - parser.add_argument("--frame", type=int, default=0) - parser.add_argument("--seed", type=int, default=20260709) - parser.add_argument("--baseline-fixture", type=Path, help=argparse.SUPPRESS) - args = parser.parse_args() - if args.baseline_fixture is None: - for name in ("tokenizer", "dataset", "probe"): - if getattr(args, name) is None: - parser.error(f"--{name} is required") - return args - - -if __name__ == "__main__": - parsed = _parse_args() - if parsed.baseline_fixture is not None: - _official_baseline( - parsed.checkpoint, - parsed.baseline_fixture, - parsed.baseline_fixture, - ) - else: - _run(parsed) diff --git a/cpp/tests/gate_pi05_native_encoder.py b/cpp/tests/gate_pi05_native_encoder.py deleted file mode 100644 index c7d04ba6..00000000 --- a/cpp/tests/gate_pi05_native_encoder.py +++ /dev/null @@ -1,184 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import gc -import pathlib -import subprocess -import tempfile - -import numpy as np -import torch -import torch.nn.functional as F -from safetensors import safe_open - - -SEQUENCE = 712 -WIDTH = 2048 - - -def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: - output, inputs = weight.shape - head_dim = output // heads - return ( - weight.reshape(heads, head_dim, inputs) - .reshape(heads, 2, head_dim // 2, inputs) - .permute(0, 2, 1, 3) - .reshape(output, inputs) - ) - - -def rms(values: torch.Tensor) -> torch.Tensor: - source = values.float() - return (source * torch.rsqrt(source.square().mean(-1, keepdim=True) + 1e-6)).to( - torch.bfloat16 - ) - - -def rotate(tensor: torch.Tensor, heads: int, rope: torch.Tensor) -> torch.Tensor: - pairs = tensor.reshape(SEQUENCE, heads, 128, 2).float() - cosine = rope[:, None, :, 0].float() - sine = rope[:, None, :, 1].float() - return torch.stack( - [ - pairs[..., 0] * cosine - pairs[..., 1] * sine, - pairs[..., 1] * cosine + pairs[..., 0] * sine, - ], - -1, - ).to(torch.bfloat16).reshape(SEQUENCE, heads, 256) - - -def compare(name: str, actual: torch.Tensor, expected: torch.Tensor) -> str: - cosine = float( - F.cosine_similarity( - actual.flatten().double(), expected.flatten().double(), dim=0 - ) - ) - maximum = float((actual - expected).abs().max()) - if cosine < 0.9999: - raise AssertionError(f"{name}: cosine={cosine:.8f} max={maximum:.6f}") - return f"{name} cosine={cosine:.8f} max={maximum:.6f}" - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - args = parser.parse_args() - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - root = "model." if "model.action_in_proj.weight" in keys else "" - - def raw(name: str) -> torch.Tensor: - return file.get_tensor(root + name) - - x = torch.zeros((SEQUENCE, WIDTH), device="cuda", dtype=torch.bfloat16) - rows = torch.arange(SEQUENCE, device="cuda")[:, None] - columns = torch.arange(512, device="cuda")[None, :] - x[:, :512] = (((rows + columns) % 15 - 7).float() / 8).to(torch.bfloat16) - positions = torch.arange(SEQUENCE, dtype=torch.float64)[:, None] - pair = torch.arange(128, dtype=torch.float64)[None, :] - phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) - rope = torch.stack([torch.cos(phase), torch.sin(phase)], -1).to( - device="cuda", dtype=torch.bfloat16 - ) - - final_q = final_k = final_v = None - for index in range(18): - layer = ( - "paligemma_with_expert.paligemma.model.language_model.layers." - f"{index}" - ) - input_norm = 1.0 + raw(f"{layer}.input_layernorm.weight").float() - q = interleave_qk(raw(f"{layer}.self_attn.q_proj.weight").float(), 8) - k = interleave_qk(raw(f"{layer}.self_attn.k_proj.weight").float(), 1) - v = raw(f"{layer}.self_attn.v_proj.weight").float() - qkv_weight = torch.cat( - [ - q * input_norm[None, :], - k * input_norm[None, :], - v * input_norm[None, :], - ], - dim=0, - ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() - qkv = rms(x) @ qkv_weight - query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) - query = rotate(query, 8, rope) - key = rotate(key, 1, rope) - value = value.reshape(SEQUENCE, 1, 256) - if index == 17: - final_q = query.reshape(SEQUENCE, 2048).cpu().float() - final_k = key.reshape(SEQUENCE, 256).cpu().float() - final_v = value.reshape(SEQUENCE, 256).cpu().float() - break - - attended = F.scaled_dot_product_attention( - query.transpose(0, 1).unsqueeze(0), - key.transpose(0, 1).unsqueeze(0), - value.transpose(0, 1).unsqueeze(0), - scale=1.0 / 16.0, - enable_gqa=True, - ).squeeze(0).transpose(0, 1).reshape(SEQUENCE, 2048) - output_weight = raw(f"{layer}.self_attn.o_proj.weight").to( - device="cuda", dtype=torch.bfloat16 - ).t().contiguous() - x = (x.float() + (attended @ output_weight).float()).to(torch.bfloat16) - post_norm = 1.0 + raw( - f"{layer}.post_attention_layernorm.weight" - ).float() - gate_weight = ( - raw(f"{layer}.mlp.gate_proj.weight").float() * post_norm[None, :] - ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() - up_weight = ( - raw(f"{layer}.mlp.up_proj.weight").float() * post_norm[None, :] - ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() - down_weight = raw(f"{layer}.mlp.down_proj.weight").to( - device="cuda", dtype=torch.bfloat16 - ).t().contiguous() - normalized = rms(x) - gate = normalized @ gate_weight - up = normalized @ up_weight - gate_float = gate.float() - activated = gate_float / ( - 1.0 - + torch.exp( - -1.5957691216057308 - * gate_float - * (1.0 + 0.044715 * gate_float.square()) - ) - ) - hidden = (activated * up.float()).to(torch.bfloat16) - x = (x.float() + (hidden @ down_weight).float()).to(torch.bfloat16) - del q, k, v, qkv_weight, qkv, query, key, value - del attended, output_weight, gate_weight, up_weight, down_weight - del normalized, gate, up, gate_float, activated, hidden - gc.collect() - - expected_x = x.cpu().float() - del x, rope - torch.cuda.empty_cache() - with tempfile.TemporaryDirectory() as directory: - output = str(pathlib.Path(directory) / "encoder.bin") - subprocess.check_call([args.probe, args.checkpoint, output]) - bits = np.fromfile(output, dtype=np.uint16) - sizes = [SEQUENCE * 2048, SEQUENCE * 2048, SEQUENCE * 256, SEQUENCE * 256] - if bits.size != sum(sizes): - raise AssertionError(f"encoder probe output elements={bits.size}") - tensors = [] - offset = 0 - for size in sizes: - tensors.append( - torch.from_numpy(bits[offset : offset + size].copy()) - .view(torch.bfloat16) - .float() - ) - offset += size - messages = [ - compare("x", tensors[0].reshape(SEQUENCE, 2048), expected_x), - compare("q", tensors[1].reshape(SEQUENCE, 2048), final_q), - compare("k", tensors[2].reshape(SEQUENCE, 256), final_k), - compare("v", tensors[3].reshape(SEQUENCE, 256), final_v), - ] - print("PASS encoder 18 layers " + "; ".join(messages)) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_encoder_layer.py b/cpp/tests/gate_pi05_native_encoder_layer.py deleted file mode 100644 index cfabab74..00000000 --- a/cpp/tests/gate_pi05_native_encoder_layer.py +++ /dev/null @@ -1,125 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import pathlib -import subprocess -import tempfile - -import numpy as np -import torch -import torch.nn.functional as F -from safetensors import safe_open - - -def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: - output, inputs = weight.shape - head_dim = output // heads - return ( - weight.reshape(heads, head_dim, inputs) - .reshape(heads, 2, head_dim // 2, inputs) - .permute(0, 2, 1, 3) - .reshape(output, inputs) - ) - - -def rms(values: torch.Tensor) -> torch.Tensor: - source = values.float() - return (source * torch.rsqrt(source.square().mean(-1, keepdim=True) + 1e-6)).to( - torch.bfloat16 - ) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - args = parser.parse_args() - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - root = "model." if "model.action_in_proj.weight" in keys else "" - layer = "paligemma_with_expert.paligemma.model.language_model.layers.0" - - def raw(name: str) -> torch.Tensor: - return file.get_tensor(root + name) - - input_norm = 1.0 + raw(f"{layer}.input_layernorm.weight").float() - q = interleave_qk(raw(f"{layer}.self_attn.q_proj.weight").float(), 8) - k = interleave_qk(raw(f"{layer}.self_attn.k_proj.weight").float(), 1) - v = raw(f"{layer}.self_attn.v_proj.weight").float() - qkv_weight = torch.cat( - [q * input_norm[None, :], k * input_norm[None, :], - v * input_norm[None, :]], dim=0 - ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() - output_weight = raw(f"{layer}.self_attn.o_proj.weight").to( - device="cuda", dtype=torch.bfloat16 - ).t().contiguous() - post_norm = 1.0 + raw(f"{layer}.post_attention_layernorm.weight").float() - gate_weight = (raw(f"{layer}.mlp.gate_proj.weight").float() * - post_norm[None, :]).t().to( - device="cuda", dtype=torch.bfloat16 - ).contiguous() - up_weight = (raw(f"{layer}.mlp.up_proj.weight").float() * - post_norm[None, :]).t().to( - device="cuda", dtype=torch.bfloat16 - ).contiguous() - down_weight = raw(f"{layer}.mlp.down_proj.weight").to( - device="cuda", dtype=torch.bfloat16 - ).t().contiguous() - - x = torch.zeros((712, 2048), device="cuda", dtype=torch.bfloat16) - rows = torch.arange(712, device="cuda")[:, None] - columns = torch.arange(512, device="cuda")[None, :] - x[:, :512] = (((rows + columns) % 15 - 7).float() / 8).to(torch.bfloat16) - qkv = rms(x) @ qkv_weight - query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) - positions = torch.arange(712, dtype=torch.float64)[:, None] - pair = torch.arange(128, dtype=torch.float64)[None, :] - phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) - rope = torch.stack([torch.cos(phase), torch.sin(phase)], -1).to( - device="cuda", dtype=torch.bfloat16 - ) - - def rotate(tensor: torch.Tensor, heads: int) -> torch.Tensor: - pairs = tensor.reshape(712, heads, 128, 2).float() - cosine = rope[:, None, :, 0].float() - sine = rope[:, None, :, 1].float() - return torch.stack( - [pairs[..., 0] * cosine - pairs[..., 1] * sine, - pairs[..., 1] * cosine + pairs[..., 0] * sine], -1 - ).to(torch.bfloat16).reshape(712, heads, 256) - - query = rotate(query, 8).transpose(0, 1).unsqueeze(0) - key = rotate(key, 1).transpose(0, 1).unsqueeze(0) - value = value.reshape(712, 1, 256).transpose(0, 1).unsqueeze(0) - attended = F.scaled_dot_product_attention( - query, key, value, scale=1.0 / 16.0, enable_gqa=True - ).squeeze(0).transpose(0, 1).reshape(712, 2048) - projected = attended @ output_weight - x = (x.float() + projected.float()).to(torch.bfloat16) - normalized = rms(x) - gate = normalized @ gate_weight - up = normalized @ up_weight - gate_float = gate.float() - activated = gate_float / ( - 1.0 + torch.exp(-1.5957691216057308 * gate_float * - (1.0 + 0.044715 * gate_float.square())) - ) - hidden = (activated * up.float()).to(torch.bfloat16) - down = hidden @ down_weight - expected = (x.float() + down.float()).to(torch.bfloat16).cpu().float() - - with tempfile.TemporaryDirectory() as directory: - output = str(pathlib.Path(directory) / "encoder.bin") - subprocess.check_call([args.probe, args.checkpoint, output]) - bits = np.fromfile(output, dtype=np.uint16).reshape(712, 2048) - actual = torch.from_numpy(bits.copy()).view(torch.bfloat16).float() - cosine = float(F.cosine_similarity( - actual.flatten().double(), expected.flatten().double(), dim=0 - )) - maximum = float((actual - expected).abs().max()) - if cosine < 0.9999: - raise AssertionError(f"cosine={cosine:.8f} max={maximum:.6f}") - print(f"PASS encoder layer0 cosine={cosine:.8f} max={maximum:.6f}") - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_encoder_qkv.py b/cpp/tests/gate_pi05_native_encoder_qkv.py deleted file mode 100644 index 4bf57dca..00000000 --- a/cpp/tests/gate_pi05_native_encoder_qkv.py +++ /dev/null @@ -1,100 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import pathlib -import subprocess -import tempfile - -import numpy as np -import torch -from safetensors import safe_open - - -def interleave_qk(weight: torch.Tensor, heads: int) -> torch.Tensor: - output, inputs = weight.shape - head_dim = output // heads - return ( - weight.reshape(heads, head_dim, inputs) - .reshape(heads, 2, head_dim // 2, inputs) - .permute(0, 2, 1, 3) - .reshape(output, inputs) - ) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - args = parser.parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for the encoder QKV gate") - - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - prefix = "model." if "model.action_in_proj.weight" in keys else "" - layer = "paligemma_with_expert.paligemma.model.language_model.layers.17" - - def raw(name: str) -> torch.Tensor: - return file.get_tensor(prefix + name) - - q = interleave_qk(raw(f"{layer}.self_attn.q_proj.weight").float(), 8) - k = interleave_qk(raw(f"{layer}.self_attn.k_proj.weight").float(), 1) - v = raw(f"{layer}.self_attn.v_proj.weight").float() - norm = 1.0 + raw(f"{layer}.input_layernorm.weight").float() - weight = torch.cat( - [q * norm[None, :], k * norm[None, :], v * norm[None, :]], dim=0 - ).t().to(device="cuda", dtype=torch.bfloat16).contiguous() - - x = torch.zeros((712, 2048), device="cuda", dtype=torch.bfloat16) - rows = torch.arange(712, device="cuda")[:, None] - columns = torch.arange(512, device="cuda")[None, :] - x[:, :512] = (((rows + columns) % 15 - 7).float() / 8).to(torch.bfloat16) - x_float = x.float() - normalized = ( - x_float * torch.rsqrt(x_float.square().mean(dim=-1, keepdim=True) + 1e-6) - ).to(torch.bfloat16) - qkv = normalized @ weight - query, key, value = torch.split(qkv, [2048, 256, 256], dim=-1) - - positions = torch.arange(712, dtype=torch.float64)[:, None] - pair = torch.arange(128, dtype=torch.float64)[None, :] - phase = positions / torch.pow(10000.0, (2 * pair) / 256.0) - rope = torch.stack([torch.cos(phase), torch.sin(phase)], dim=-1).to( - device="cuda", dtype=torch.bfloat16 - ) - - def apply_rope(tensor: torch.Tensor, heads: int) -> torch.Tensor: - pairs = tensor.reshape(712, heads, 128, 2).float() - cosine = rope[:, None, :, 0].float() - sine = rope[:, None, :, 1].float() - even = pairs[..., 0] * cosine - pairs[..., 1] * sine - odd = pairs[..., 1] * cosine + pairs[..., 0] * sine - return torch.stack([even, odd], dim=-1).to(torch.bfloat16).reshape( - 712, heads * 256 - ) - - expected = { - "q": apply_rope(query, 8), - "k": apply_rope(key, 1), - "v": value.contiguous(), - } - with tempfile.TemporaryDirectory() as directory: - output = str(pathlib.Path(directory) / "encoder") - subprocess.check_call([args.probe, args.checkpoint, output]) - for name, reference in expected.items(): - actual_bits = np.fromfile(f"{output}.{name}.bin", dtype=np.uint16) - actual = torch.from_numpy(actual_bits.copy()).view(torch.bfloat16) - actual = actual.reshape(reference.shape).float() - target = reference.cpu().float() - cosine = float(torch.nn.functional.cosine_similarity( - actual.flatten().double(), target.flatten().double(), dim=0 - )) - maximum = float((actual - target).abs().max()) - if cosine < 0.9999: - raise AssertionError( - f"{name}: cosine={cosine:.8f} max={maximum:.6f}" - ) - print(f"PASS encoder17 {name} cosine={cosine:.8f} max={maximum:.6f}") - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_quantization.py b/cpp/tests/gate_pi05_native_quantization.py deleted file mode 100644 index 2640192c..00000000 --- a/cpp/tests/gate_pi05_native_quantization.py +++ /dev/null @@ -1,153 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import subprocess - -import torch -from safetensors import safe_open - - -DECODER = "paligemma_with_expert.gemma_expert.model.layers.0" -VISION = ( - "paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".encoder.layers.4" -) -VISION_ROOT = ( - "paligemma_with_expert.paligemma.model.vision_tower.vision_model" - ".encoder.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), - ) - if layout == "kn": - expected["decoder_qkv0_fp8_kn_gpu"] = expected[ - "decoder_qkv0_fp8_kn" - ] - - vision_down = bf16(f"{VISION}.mlp.fc2.weight").t().contiguous().cuda() - scale = max(vision_down.float().abs().max().item() / 448.0, 1e-12) - quantized = (vision_down.float() / scale).clamp(-448.0, 448.0).to( - torch.float8_e4m3fn - ) - scale_tensor = torch.tensor([scale], dtype=torch.float32, device="cuda") - expected["vision_ffn_down4_fp8_kn_gpu"] = ( - tuple(quantized.shape), - digest(quantized.view(torch.uint8)), - 1, - digest(scale_tensor), - ) - - vision_qkv = torch.cat([ - bf16(f"{VISION_ROOT}.self_attn.{stem}_proj.weight") - for stem in ("q", "k", "v") - ], dim=0).t().contiguous().cuda() - scale = max(vision_qkv.float().abs().max().item() / 448.0, 1e-12) - quantized = (vision_qkv.float() / scale).clamp(-448.0, 448.0).to( - torch.float8_e4m3fn - ) - scale_tensor = torch.tensor([scale], dtype=torch.float32, device="cuda") - expected["vision_attn_qkv0_fp8_kn_gpu"] = ( - tuple(quantized.shape), - digest(quantized.view(torch.uint8)), - 1, - digest(scale_tensor), - ) - - transposed = weight.float().transpose(0, 1).contiguous() - scales = torch.clamp( - transposed.abs().amax(dim=1) / 127.0, min=1e-12 - ).to(dtype=torch.float32).contiguous() - quantized = torch.clamp( - torch.round(transposed / scales[:, None]), -127, 127 - ).to(torch.int8).contiguous() - expected["decoder_qkv0_int8"] = ( - tuple(quantized.shape), - digest(quantized), - scales.numel(), - digest(scales), - ) - - for operation, reference in expected.items(): - output = subprocess.check_output( - [args.probe, args.checkpoint, operation], text=True - ) - actual = parse_probe(output) - if actual != reference: - raise AssertionError( - f"{operation}: C++ {actual} != PyTorch {reference}" - ) - print( - f"PASS {operation} shape={actual[0]} " - f"values_fnv={actual[1]:016x} scales_fnv={actual[3]:016x}" - ) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_rope.py b/cpp/tests/gate_pi05_native_rope.py deleted file mode 100644 index 8afbd013..00000000 --- a/cpp/tests/gate_pi05_native_rope.py +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import subprocess - -import ml_dtypes -import numpy as np - - -def fnv1a(data: bytes) -> int: - value = 14695981039346656037 - for byte in data: - value ^= byte - value = (value * 1099511628211) & 0xFFFFFFFFFFFFFFFF - return value - - -def parse_probe(text: str) -> dict[str, str]: - return dict(field.split("=", 1) for field in text.strip().split()) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--probe", required=True) - args = parser.parse_args() - cases = [(2, 200, 10, 1, 37), (3, 256, 50, 2, 256)] - for views, max_prompt, chunk, pool, prompt in cases: - vision = views * 256 // (pool * pool) - encoder_length = vision + max_prompt - max_positions = encoder_length + chunk - inverse_frequency = 1.0 / ( - 10000 ** (np.arange(0, 256, 2, dtype=np.float64) / 256) - ) - positions = np.arange(max_positions, dtype=np.float64) - phase = positions[:, None] * inverse_frequency[None, :] - cosine = np.cos(phase).astype(ml_dtypes.bfloat16) - sine = np.sin(phase).astype(ml_dtypes.bfloat16) - table = np.stack([cosine, sine], axis=-1).reshape(max_positions, 256) - encoder = np.ascontiguousarray(table[:encoder_length]) - decoder = np.ascontiguousarray( - table[vision + prompt : vision + prompt + chunk] - ) - output = subprocess.check_output( - [ - args.probe, - str(views), - str(max_prompt), - str(chunk), - str(pool), - str(prompt), - ], - text=True, - ) - actual = parse_probe(output) - expected = { - "encoder_shape": f"{encoder_length},256", - "encoder_fnv": f"{fnv1a(encoder.tobytes()):016x}", - "decoder_shape": f"{chunk},256", - "decoder_fnv": f"{fnv1a(decoder.tobytes()):016x}", - } - if actual != expected: - raise AssertionError(f"C++ {actual} != NumPy {expected}") - print( - f"PASS views={views} pool={pool} prompt={prompt} " - f"encoder_fnv={actual['encoder_fnv']} " - f"decoder_fnv={actual['decoder_fnv']}" - ) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_style.py b/cpp/tests/gate_pi05_native_style.py deleted file mode 100644 index 87bef640..00000000 --- a/cpp/tests/gate_pi05_native_style.py +++ /dev/null @@ -1,132 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import math -import pathlib -import subprocess -import tempfile - -import numpy as np -import torch -from safetensors import safe_open - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - args = parser.parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for the style precompute gate") - - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - prefix = "model." if "model.action_in_proj.weight" in keys else "" - - def bf16(key: str) -> torch.Tensor: - return file.get_tensor(prefix + key).to( - device="cuda", dtype=torch.bfloat16 - ) - - decoder = "paligemma_with_expert.gemma_expert.model.layers" - time_in_w = bf16("time_mlp_in.weight").t().contiguous() - time_in_b = bf16("time_mlp_in.bias") - time_out_w = bf16("time_mlp_out.weight").t().contiguous() - time_out_b = bf16("time_mlp_out.bias") - attn_w = torch.stack( - [bf16(f"{decoder}.{i}.input_layernorm.dense.weight").t() for i in range(18)] - ) - attn_b = torch.stack( - [bf16(f"{decoder}.{i}.input_layernorm.dense.bias") for i in range(18)] - ) - ffn_w = torch.stack( - [ - bf16(f"{decoder}.{i}.post_attention_layernorm.dense.weight").t() - for i in range(18) - ] - ) - ffn_b = torch.stack( - [ - bf16(f"{decoder}.{i}.post_attention_layernorm.dense.bias") - for i in range(18) - ] - ) - final_w = bf16( - "paligemma_with_expert.gemma_expert.model.norm.dense.weight" - ).t() - final_b = bf16("paligemma_with_expert.gemma_expert.model.norm.dense.bias") - - fraction = torch.linspace(0.0, 1.0, 512) - period = 4e-3 * (4.0 / 4e-3) ** fraction - t = torch.tensor(1.0, dtype=torch.float32) - rows = [] - for _ in range(10): - angle = t * (1.0 / period) * 2 * math.pi - rows.append( - torch.cat([torch.sin(angle), torch.cos(angle)]).to( - device="cuda", dtype=torch.bfloat16 - ) - ) - t = t - 0.1 - schedule = torch.stack(rows) - expected = { - "decoder_time_emb": torch.empty( - 10, 10, 1024, dtype=torch.bfloat16, device="cuda" - ), - "decoder_style_attn": torch.empty( - 10, 18, 10, 3072, dtype=torch.bfloat16, device="cuda" - ), - "decoder_style_ffn": torch.empty_like( - torch.empty(10, 18, 10, 3072, dtype=torch.bfloat16, device="cuda") - ), - "decoder_style_final": torch.empty( - 10, 10, 3072, dtype=torch.bfloat16, device="cuda" - ), - } - for step in range(10): - value = schedule[step : step + 1] - value = (value @ time_in_w + time_in_b[None, :]).float() - value = (value * torch.sigmoid(value)).to(torch.bfloat16) - value = (value @ time_out_w + time_out_b[None, :]).float() - value = (value * torch.sigmoid(value)).to(torch.bfloat16) - expanded = value.expand(10, -1).contiguous() - expected["decoder_time_emb"][step] = expanded - for layer in range(18): - expected["decoder_style_attn"][step, layer] = ( - expanded @ attn_w[layer] + attn_b[layer][None, :] - ) - expected["decoder_style_ffn"][step, layer] = ( - expanded @ ffn_w[layer] + ffn_b[layer][None, :] - ) - expected["decoder_style_final"][step] = ( - expanded @ final_w + final_b[None, :] - ) - - with tempfile.TemporaryDirectory() as directory: - output_prefix = str(pathlib.Path(directory) / "styles") - subprocess.check_call([args.probe, args.checkpoint, output_prefix]) - for name, reference in expected.items(): - actual_bits = np.fromfile( - f"{output_prefix}.{name}.bin", dtype=np.uint16 - ).reshape(tuple(reference.shape)) - reference_bits = reference.contiguous().view(torch.uint16).cpu().numpy() - exact = float(np.mean(actual_bits == reference_bits)) - actual = torch.from_numpy(actual_bits.copy()).view(torch.bfloat16).float() - target = reference.cpu().float() - maximum = float((actual - target).abs().max()) - cosine = float( - torch.nn.functional.cosine_similarity( - actual.flatten().double(), target.flatten().double(), dim=0 - ) - ) - if cosine < 0.9999: - raise AssertionError( - f"{name}: exact={exact} max={maximum} cosine={cosine}" - ) - print( - f"PASS {name} exact={exact:.6f} " - f"max={maximum:.6f} cosine={cosine:.8f}" - ) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_thor_fp8.py b/cpp/tests/gate_pi05_native_thor_fp8.py deleted file mode 100644 index 8fb003aa..00000000 --- a/cpp/tests/gate_pi05_native_thor_fp8.py +++ /dev/null @@ -1,396 +0,0 @@ -#!/usr/bin/env python3 -"""Compare the native Thor FP8 producer with the shipped Torch producer.""" - -from __future__ import annotations - -import argparse -import os -import subprocess -import types -from pathlib import Path - -import numpy as np -import torch -from safetensors import safe_open - -import flash_rt.flash_rt_kernels as fvk -from flash_rt.frontends.torch.pi05_thor import Pi05TorchFrontendThor -from flash_rt.hardware.thor.shared_primitives import encoder_forward_calibrate -from flash_rt.models.pi05.pipeline_thor import decoder_forward_calibrate - - -def _artifact(path: Path) -> tuple[dict[str, str], np.ndarray, np.ndarray]: - with safe_open(path, framework="np") as handle: - metadata = handle.metadata() or {} - encoder = handle.get_tensor("encoder_scales").astype( - np.float32, copy=True) - decoder = handle.get_tensor("decoder_scales").astype( - np.float32, copy=True) - return metadata, encoder, decoder - - -def _sample(index: int, num_views: int - ) -> tuple[str, np.ndarray, list[np.ndarray], np.ndarray]: - pixels = np.arange(224 * 224 * 3, dtype=np.uint64) - image = ((pixels * 3 + index * 17) % 251).astype(np.uint8) - wrist = ((pixels * 7 + index * 29 + 11) % 253).astype(np.uint8) - right_wrist = ( - (pixels * 11 + index * 37 + 19) % 247 - ).astype(np.uint8) - state = np.asarray( - [((index * 8 + dim) % 17 - 8) / 8.0 for dim in range(8)], - dtype=np.float32, - ) - noise = np.asarray( - [((index * 320 + item) % 31 - 15) / 16.0 for item in range(320)], - dtype=np.float32, - ).reshape(10, 32) - prompt = ( - "move the black bowl to the plate" - if index & 1 - else "pick up the black bowl" - ) - images = [ - image.reshape(224, 224, 3), - wrist.reshape(224, 224, 3), - right_wrist.reshape(224, 224, 3), - ] - return ( - prompt, - state, - images[:num_views], - noise, - ) - - -def _normalized_state(frontend: Pi05TorchFrontendThor, - state: np.ndarray) -> np.ndarray: - q01 = np.asarray(frontend.norm_stats["state"]["q01"], dtype=np.float32) - q99 = np.asarray(frontend.norm_stats["state"]["q99"], dtype=np.float32) - return ( - ((state - q01) / (q99 - q01 + np.float32(1.0e-6))) - * np.float32(2.0) - - np.float32(1.0) - ).astype(np.float32) - - -def _seed_calibration( - encoder: np.ndarray, - decoder: np.ndarray, -): - def seed(self: Pi05TorchFrontendThor, _sequence: int) -> None: - self._enc_calib_scales.copy_(torch.from_numpy(encoder)) - self._ae_calib_scales.copy_(torch.from_numpy(decoder)) - weight_scales = self._enc_w_dev.cpu().numpy() - self._enc_alpha_host = [ - float(np.float32(encoder[i]) * np.float32(weight_scales[i])) - for i in range(encoder.size) - ] - - return seed - - -class CalibrationOracle: - def __init__(self, frontend: Pi05TorchFrontendThor): - self.frontend = frontend - p = frontend - sequence = p.Se - width = p.De - hidden = p.He - action_width = p.Da - action_hidden = p.Ha - self._encoder_norm_scratch = torch.empty( - sequence * width, dtype=torch.float16, device="cuda") - self._encoder_x_scratch = torch.empty( - sequence * width, dtype=torch.float16, device="cuda") - self._encoder_calib_buffer = torch.zeros( - p.Le * 4, dtype=torch.float32, device="cuda") - self._encoder_dynamic_scale = torch.zeros( - 1, dtype=torch.float32, device="cuda") - self._encoder_fp8_scratch = torch.zeros( - sequence * max(width, hidden), dtype=torch.uint8, device="cuda") - self._encoder_ones = torch.ones( - width, dtype=torch.float16, device="cuda") - self._decoder_calib_buffer = torch.zeros( - p.steps * p.La * 4, dtype=torch.float32, device="cuda") - self._decoder_dynamic_scale = torch.zeros( - 1, dtype=torch.float32, device="cuda") - self._decoder_hidden_scratch = torch.empty( - p.Sa * action_hidden, dtype=torch.float16, device="cuda") - self._decoder_fp8_scratch = torch.zeros( - p.Sa * max(action_width, action_hidden), dtype=torch.uint8, - device="cuda") - - self.encoder_buffers = { - "x": p._enc_x.data_ptr(), - "x_fp8": p._enc_x_fp8.data_ptr(), - "qkv": p._enc_qkv_buf.data_ptr(), - "logits": p._enc_logits.data_ptr(), - "attn_out": p._enc_attn.data_ptr(), - "o_fp8": p._enc_o_fp8.data_ptr(), - "gate": p._enc_gate.data_ptr(), - "hidden": p._enc_hidden.data_ptr(), - "hid_fp8": p._enc_hid_fp8.data_ptr(), - "fg": p._enc_fg.data_ptr(), - "ctx": p._ctx, - "norm_scratch": self._encoder_norm_scratch.data_ptr(), - "x_scratch": self._encoder_x_scratch.data_ptr(), - "calib_buf": self._encoder_calib_buffer.data_ptr(), - "d_scale": self._encoder_dynamic_scale.data_ptr(), - "fp8_scratch": self._encoder_fp8_scratch.data_ptr(), - "ones": self._encoder_ones.data_ptr(), - } - self.encoder_weights = { - "qkv_w": [weight.data_ptr() for weight in p._enc_qkv_w], - "o_w": [weight.data_ptr() for weight in p._enc_o_w], - "gate_w": [weight.data_ptr() for weight in p._enc_gu_w], - "down_w": [weight.data_ptr() for weight in p._enc_d_w], - "rope": p._enc_rope.data_ptr(), - "Kc": p._Kc.reshape(-1).data_ptr(), - "Vc": p._Vc.reshape(-1).data_ptr(), - "w_scales": p._enc_w_dev.data_ptr(), - } - self.encoder_dims = { - "Se": sequence, - "D": width, - "H": hidden, - "NH": p.NHe, - "HD": p.HDe, - "L": p.Le, - "total_keys": p.total_keys, - } - - self.decoder_buffers = { - "noise": p._g_noise.data_ptr(), - "x": p._ae_x.data_ptr(), - "xn": p._ae_xn.data_ptr(), - "gate": p._ae_gate.data_ptr(), - "qkv": p._ae_qkv.data_ptr(), - "logits": p._ae_logits.data_ptr(), - "attn_out": p._ae_attn.data_ptr(), - "hid": p._ae_hid.data_ptr(), - "fg": p._ae_fg.data_ptr(), - "action_f32": p._ae_action_f32.data_ptr(), - "xn_fp8": p._ae_xn_fp8.data_ptr(), - "hid_fp8": p._ae_hid_fp8.data_ptr(), - "ctx_fp8": p._ae_ctx_fp8.data_ptr(), - "calib_buf": self._decoder_calib_buffer.data_ptr(), - "d_scale": self._decoder_dynamic_scale.data_ptr(), - "hidden_scratch": self._decoder_hidden_scratch.data_ptr(), - "fp8_scratch": self._decoder_fp8_scratch.data_ptr(), - } - self.decoder_weights = { - "ain_w": p._ain_w.data_ptr(), - "ain_b": p._ain_b.data_ptr(), - "sa": p._sa_all.data_ptr(), - "qw": p._dec_qkv_flat.data_ptr(), - "Kc": p._Kc.reshape(-1).data_ptr(), - "Vc": p._Vc.reshape(-1).data_ptr(), - "dec_devpos": p._attn.dec_devpos.data_ptr(), - "ow": p._dec_o_flat.data_ptr(), - "sf": p._sf_all.data_ptr(), - "gw": p._dec_gu_flat.data_ptr(), - "dw": p._dec_d_flat.data_ptr(), - "aow": p._aow.data_ptr(), - "aob": p._aob.data_ptr(), - "aob_dt": p._aob_dt.data_ptr(), - "dt": p._ae_dt, - "fs": p._fs_all.data_ptr(), - "rope": p._dec_rope.data_ptr(), - "w_scales": p._ae_w_dev.data_ptr(), - } - self.decoder_dims = { - "S": p.Sa, - "D": action_width, - "H": action_hidden, - "NH": 8, - "HD": 256, - "steps": p.steps, - "layers": p.La, - "enc_seq": sequence, - "total_keys": p.total_keys, - "fixed_shape": p._fixed_shape_active, - } - - def observe(self, images: list[np.ndarray], noise: np.ndarray - ) -> tuple[np.ndarray, np.ndarray]: - p = self.frontend - normalized = np.stack([ - (image.astype(np.float32) / np.float32(127.5) - - np.float32(1.0)).astype(np.float16) - for image in images - ]) - p._img_buf.upload(normalized) - p._siglip_graph.replay() - torch.cuda.synchronize() - - p._enc_calib_scales.zero_() - p._Kc.zero_() - p._Vc.zero_() - encoder_forward_calibrate( - p._gemm, fvk, self.encoder_buffers, self.encoder_weights, - self.encoder_dims, p._enc_calib_scales.data_ptr(), stream=0, - attn=p._attn, - ) - torch.cuda.synchronize() - encoder = p._enc_calib_scales.cpu().numpy().copy() - - p._ae_calib_scales.zero_() - p._g_noise.copy_(torch.from_numpy(noise.astype(np.float16)).cuda()) - decoder_forward_calibrate( - p._ctx, fvk, self.decoder_buffers, self.decoder_weights, - self.decoder_dims, p._ae_calib_scales.data_ptr(), stream=0, - attn=p._attn, - ) - torch.cuda.synchronize() - decoder = p._ae_calib_scales.cpu().numpy().copy() - return encoder, decoder - - -def _assert_equal(label: str, actual: np.ndarray, - expected: np.ndarray) -> None: - if np.array_equal(actual, expected): - print(f"PASS {label}: bit-exact ({actual.size} values)") - return - actual_flat = actual.reshape(-1) - expected_flat = expected.reshape(-1) - difference = np.abs( - actual_flat.astype(np.float64) - expected_flat.astype(np.float64) - ) - index = int(np.argmax(difference)) - first = int(np.flatnonzero(actual_flat != expected_flat)[0]) - mismatches = int(np.count_nonzero(actual_flat != expected_flat)) - raise AssertionError( - f"{label} mismatch: max_abs={difference[index]:.9g} " - f"at {index}, actual={actual_flat[index]!r}, " - f"expected={expected_flat[index]!r}; " - f"first={first}, actual={actual_flat[first]!r}, " - f"expected={expected_flat[first]!r}, " - f"mismatches={mismatches}/{actual.size}" - ) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--probe", type=Path, required=True) - parser.add_argument("--checkpoint", type=Path, required=True) - parser.add_argument("--tokenizer", type=Path, required=True) - parser.add_argument("--artifact", type=Path, required=True) - parser.add_argument("--samples", type=int, default=3) - parser.add_argument("--views", type=int, choices=(1, 2, 3), default=2) - parser.add_argument("--max-prompt-tokens", type=int, default=200) - args = parser.parse_args() - if args.samples < 1 or args.samples > 256: - parser.error("--samples must be in [1, 256]") - if args.max_prompt_tokens < 1: - parser.error("--max-prompt-tokens must be positive") - - probe = args.probe.resolve() - os.environ["FLASH_RT_PALIGEMMA_TOKENIZER"] = str( - args.tokenizer.resolve() - ) - raw_path = args.artifact.with_suffix(args.artifact.suffix + ".raw") - env = os.environ.copy() - library_paths = [probe.parent, probe.parent / "exec", - probe.parent / "runtime"] - if env.get("LD_LIBRARY_PATH"): - library_paths.append(Path(env["LD_LIBRARY_PATH"])) - env["LD_LIBRARY_PATH"] = os.pathsep.join(map(str, library_paths)) - env["FLASHRT_MAX_PROMPT_TOKENS"] = str(args.max_prompt_tokens) - subprocess.run( - [ - str(probe), - str(args.checkpoint), - str(args.tokenizer), - str(args.artifact), - str(args.samples), - str(args.views), - str(raw_path), - ], - check=True, - env=env, - ) - metadata, expected_encoder, expected_decoder = _artifact(args.artifact) - _, single_encoder, single_decoder = _artifact( - Path(str(args.artifact) + ".single") - ) - if int(metadata.get("sample_count", "0")) != args.samples: - raise AssertionError("native calibration sample_count is incorrect") - if int(metadata.get("num_views", "0")) != args.views: - raise AssertionError("native calibration num_views is incorrect") - - frontend = Pi05TorchFrontendThor( - str(args.checkpoint), num_views=args.views, use_cuda_graph=True, - autotune=0, - use_fp8=True, state_prompt_mode="fixed", - state_prompt_fixed_max_len=args.max_prompt_tokens, - ) - frontend._calibrate = types.MethodType( - _seed_calibration(expected_encoder, expected_decoder), frontend - ) - first_prompt, first_state, _, _ = _sample(0, args.views) - frontend.set_prompt( - first_prompt, state=_normalized_state(frontend, first_state) - ) - oracle = CalibrationOracle(frontend) - - encoder_samples = [] - decoder_samples = [] - samples = [] - for index in range(args.samples): - prompt, state, images, noise = _sample(index, args.views) - samples.append((prompt, state, images, noise)) - frontend.set_prompt( - prompt, state=_normalized_state(frontend, state) - ) - encoder, decoder = oracle.observe(images, noise) - # The final encoder layer only emits decoder K/V. Its O/FFN sites are - # skipped; the native artifact uses 1.0 as their valid neutral value. - encoder[-3:] = np.float32(1.0) - encoder_samples.append(encoder) - decoder_samples.append(decoder) - - _assert_equal("single-sample encoder calibration", - encoder_samples[0], single_encoder) - _assert_equal("single-sample decoder calibration", - decoder_samples[0], single_decoder) - reduced_encoder = np.percentile( - np.stack(encoder_samples), 99.9, axis=0 - ).astype(np.float32) - reduced_decoder = np.percentile( - np.stack(decoder_samples), 99.9, axis=0 - ).astype(np.float32) - _assert_equal("dataset encoder calibration", - reduced_encoder, expected_encoder) - _assert_equal("dataset decoder calibration", - reduced_decoder, expected_decoder) - - _, state, images, noise = samples[-1] - frontend.set_prompt( - "pick up the black bowl", state=_normalized_state(frontend, state) - ) - frontend._enc_calib_scales.copy_(torch.from_numpy(expected_encoder)) - frontend._ae_calib_scales.copy_(torch.from_numpy(expected_decoder)) - weight_scales = frontend._enc_w_dev.cpu().numpy() - frontend._enc_alpha_host = [ - float(np.float32(expected_encoder[i]) * np.float32(weight_scales[i])) - for i in range(expected_encoder.size) - ] - frontend._capture_enc_ae_graph() - normalized_images = np.stack([ - (image.astype(np.float32) / np.float32(127.5) - - np.float32(1.0)).astype(np.float16) - for image in images - ]) - frontend._img_buf.upload(normalized_images) - frontend._siglip_graph.replay() - frontend._g_noise.copy_(torch.from_numpy(noise.astype(np.float16)).cuda()) - frontend._enc_ae_graph.replay() - torch.cuda.synchronize() - python_raw = frontend._g_noise.cpu().numpy() - native_raw = np.fromfile(raw_path, dtype=np.float16).reshape(10, 32) - _assert_equal("native raw action", python_raw, native_raw) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_vision.py b/cpp/tests/gate_pi05_native_vision.py deleted file mode 100644 index c40eb17b..00000000 --- a/cpp/tests/gate_pi05_native_vision.py +++ /dev/null @@ -1,190 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import gc -import pathlib -import subprocess -import tempfile - -import numpy as np -import torch -import torch.nn.functional as F -from safetensors import safe_open - - -NUM_VIEWS = 2 -SEQUENCE = NUM_VIEWS * 256 -WIDTH = 1152 -HIDDEN = 4304 - - -def layer_norm( - values: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor -) -> torch.Tensor: - source = values.float() - mean = source.mean(-1, keepdim=True) - variance = (source - mean).square().mean(-1, keepdim=True) - return ( - (source - mean) - * torch.rsqrt(variance + 1e-5) - * weight.float() - + bias.float() - ).to(torch.bfloat16) - - -def compare(name: str, actual: torch.Tensor, expected: torch.Tensor) -> str: - cosine = float( - F.cosine_similarity( - actual.flatten().double(), expected.flatten().double(), dim=0 - ) - ) - maximum = float((actual - expected).abs().max()) - if cosine < 0.9999: - raise AssertionError(f"{name}: cosine={cosine:.8f} max={maximum:.6f}") - return f"{name} cosine={cosine:.8f} max={maximum:.6f}" - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - args = parser.parse_args() - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - root = "model." if "model.action_in_proj.weight" in keys else "" - vision = "paligemma_with_expert.paligemma.model.vision_tower.vision_model" - - def raw(name: str) -> torch.Tensor: - return file.get_tensor(root + name) - - def bf16(name: str) -> torch.Tensor: - return raw(name).to(device="cuda", dtype=torch.bfloat16) - - flat = torch.arange(NUM_VIEWS * 224 * 224 * 3, device="cuda") - images = (((flat % 257) - 128).float() / 128.0).to(torch.bfloat16).reshape( - NUM_VIEWS, 224, 224, 3 - ) - patches = ( - images.reshape(NUM_VIEWS, 16, 14, 16, 14, 3) - .permute(0, 1, 3, 2, 4, 5) - .reshape(SEQUENCE, 588) - ) - patch_weight = bf16(f"{vision}.embeddings.patch_embedding.weight") - patch_weight = patch_weight.permute(2, 3, 1, 0).reshape(588, WIDTH) - patch_bias = bf16(f"{vision}.embeddings.patch_embedding.bias") - position = bf16(f"{vision}.embeddings.position_embedding.weight").repeat( - NUM_VIEWS, 1 - ) - x = (patches @ patch_weight).to(torch.bfloat16) - x = (x.float() + position.float() + patch_bias.float()).to(torch.bfloat16) - first = f"{vision}.encoder.layers.0" - x_norm = layer_norm( - x, bf16(f"{first}.layer_norm1.weight"), bf16(f"{first}.layer_norm1.bias") - ) - - for index in range(27): - layer = f"{vision}.encoder.layers.{index}" - q_weight = bf16(f"{layer}.self_attn.q_proj.weight") - k_weight = bf16(f"{layer}.self_attn.k_proj.weight") - v_weight = bf16(f"{layer}.self_attn.v_proj.weight") - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0).t().contiguous() - qkv_bias = torch.cat( - [ - bf16(f"{layer}.self_attn.q_proj.bias"), - bf16(f"{layer}.self_attn.k_proj.bias"), - bf16(f"{layer}.self_attn.v_proj.bias"), - ] - ) - qkv = x_norm @ qkv_weight - qkv = (qkv.float() + qkv_bias.float()).to(torch.bfloat16) - query, key, value = qkv.reshape(NUM_VIEWS, 256, 3, 16, 72).unbind(2) - attended = F.scaled_dot_product_attention( - query.transpose(1, 2), - key.transpose(1, 2), - value.transpose(1, 2), - scale=1.0 / np.sqrt(72.0), - ).transpose(1, 2).reshape(SEQUENCE, WIDTH) - output_weight = bf16(f"{layer}.self_attn.out_proj.weight").t().contiguous() - output_bias = bf16(f"{layer}.self_attn.out_proj.bias") - projected = attended @ output_weight - x = (x.float() + projected.float() + output_bias.float()).to(torch.bfloat16) - x_norm = layer_norm( - x, - bf16(f"{layer}.layer_norm2.weight"), - bf16(f"{layer}.layer_norm2.bias"), - ) - up_weight = bf16(f"{layer}.mlp.fc1.weight").t().contiguous() - up_bias = bf16(f"{layer}.mlp.fc1.bias") - hidden = x_norm @ up_weight - hidden = (hidden.float() + up_bias.float()).to(torch.bfloat16) - hidden_float = hidden.float() - hidden = ( - hidden_float - * 0.5 - * ( - 1.0 - + torch.tanh( - 0.7978845608 - * (hidden_float + 0.044715 * hidden_float.pow(3)) - ) - ) - ).to(torch.bfloat16) - down_weight = bf16(f"{layer}.mlp.fc2.weight").t().contiguous() - down_bias = bf16(f"{layer}.mlp.fc2.bias") - down = hidden @ down_weight - x = (x.float() + down.float() + down_bias.float()).to(torch.bfloat16) - if index != 26: - next_layer = f"{vision}.encoder.layers.{index + 1}" - x_norm = layer_norm( - x, - bf16(f"{next_layer}.layer_norm1.weight"), - bf16(f"{next_layer}.layer_norm1.bias"), - ) - del q_weight, k_weight, v_weight, qkv_weight, qkv_bias, qkv - del query, key, value, attended, output_weight, output_bias, projected - del up_weight, up_bias, hidden, hidden_float, down_weight, down_bias, down - gc.collect() - - expected_vision = x.cpu().float() - final_norm = layer_norm( - x, - bf16(f"{vision}.post_layernorm.weight"), - bf16(f"{vision}.post_layernorm.bias"), - ) - projector = ( - "paligemma_with_expert.paligemma.model.multi_modal_projector.linear" - ) - projected = final_norm @ bf16(f"{projector}.weight").t().contiguous() - expected_encoder = ( - projected.float() + bf16(f"{projector}.bias").float() - ).to(torch.bfloat16).cpu().float() - del x, x_norm, final_norm, projected, images, patches - torch.cuda.empty_cache() - - with tempfile.TemporaryDirectory() as directory: - output = str(pathlib.Path(directory) / "vision.bin") - subprocess.check_call([args.probe, args.checkpoint, output]) - bits = np.fromfile(output, dtype=np.uint16) - sizes = [SEQUENCE * WIDTH, SEQUENCE * 2048] - if bits.size != sum(sizes): - raise AssertionError(f"vision probe output elements={bits.size}") - vision_bits = bits[: sizes[0]].copy() - encoder_bits = bits[sizes[0] :].copy() - actual_vision = ( - torch.from_numpy(vision_bits).view(torch.bfloat16).float().reshape(SEQUENCE, WIDTH) - ) - actual_encoder = ( - torch.from_numpy(encoder_bits) - .view(torch.bfloat16) - .float() - .reshape(SEQUENCE, 2048) - ) - print( - "PASS vision 27 layers " - + compare("vision", actual_vision, expected_vision) - + "; " - + compare("encoder", actual_encoder, expected_encoder) - ) - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_native_weight_ops.py b/cpp/tests/gate_pi05_native_weight_ops.py deleted file mode 100644 index 2276606d..00000000 --- a/cpp/tests/gate_pi05_native_weight_ops.py +++ /dev/null @@ -1,141 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import math -import subprocess - -import torch -from safetensors import safe_open - - -VISION = "paligemma_with_expert.paligemma.model.vision_tower.vision_model" -ENCODER = "paligemma_with_expert.paligemma.model.language_model.layers.0" -DECODER = "paligemma_with_expert.gemma_expert.model.layers.0" - - -def interleave_qk(weight: torch.Tensor, num_heads: int) -> torch.Tensor: - out_dim, in_dim = weight.shape - head_dim = out_dim // num_heads - return ( - weight.reshape(num_heads, head_dim, in_dim) - .reshape(num_heads, 2, head_dim // 2, in_dim) - .permute(0, 2, 1, 3) - .reshape(out_dim, in_dim) - ) - - -def fnv1a(data: bytes) -> int: - value = 14695981039346656037 - for byte in data: - value ^= byte - value = (value * 1099511628211) & 0xFFFFFFFFFFFFFFFF - return value - - -def digest(tensor: torch.Tensor) -> tuple[tuple[int, ...], int]: - tensor = tensor.contiguous().view(torch.uint16).cpu() - return tuple(tensor.shape), fnv1a(tensor.numpy().tobytes()) - - -def parse_probe(text: str) -> tuple[tuple[int, ...], int]: - fields = dict(field.split("=", 1) for field in text.strip().split()) - shape = tuple(int(dim) for dim in fields["shape"].split(",")) - return shape, int(fields["fnv"], 16) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--probe", required=True) - args = parser.parse_args() - - file = safe_open(f"{args.checkpoint}/model.safetensors", framework="pt") - keys = set(file.keys()) - prefix = "model." if "model.action_in_proj.weight" in keys else "" - - def raw(key: str) -> torch.Tensor: - return file.get_tensor(prefix + key) - - def bf16(key: str) -> torch.Tensor: - return raw(key).to(torch.bfloat16) - - patch = bf16(f"{VISION}.embeddings.patch_embedding.weight") - expected = { - "patch": patch.permute(2, 3, 1, 0).contiguous(), - } - - q = interleave_qk(raw(f"{ENCODER}.self_attn.q_proj.weight").float(), 8) - k = interleave_qk(raw(f"{ENCODER}.self_attn.k_proj.weight").float(), 1) - v = raw(f"{ENCODER}.self_attn.v_proj.weight").float() - norm = 1.0 + raw(f"{ENCODER}.input_layernorm.weight").float() - expected["encoder_qkv0"] = torch.cat( - [q * norm.unsqueeze(0), k * norm.unsqueeze(0), v * norm.unsqueeze(0)], - dim=0, - ).t().to(torch.bfloat16).contiguous() - - q = interleave_qk(bf16(f"{DECODER}.self_attn.q_proj.weight").float(), 8) - k = interleave_qk(bf16(f"{DECODER}.self_attn.k_proj.weight").float(), 1) - v = bf16(f"{DECODER}.self_attn.v_proj.weight") - expected["decoder_qkv0"] = torch.cat([q, k, v], dim=0).t().to( - torch.bfloat16 - ).contiguous() - - gate = bf16(f"{DECODER}.mlp.gate_proj.weight").t() - up = bf16(f"{DECODER}.mlp.up_proj.weight").t() - expected["decoder_gate_up0"] = torch.cat([gate, up], dim=1).contiguous() - - expected["encoder_o0_fast"] = bf16( - f"{ENCODER}.self_attn.o_proj.weight" - ).t().contiguous() - ffn_norm = 1.0 + raw( - f"{ENCODER}.post_attention_layernorm.weight" - ).float() - expected["encoder_gate0_fast"] = ( - raw(f"{ENCODER}.mlp.gate_proj.weight").float() - * ffn_norm.unsqueeze(0) - ).t().to(torch.bfloat16).contiguous() - expected["decoder_mod_bias0_fast"] = bf16( - f"{DECODER}.input_layernorm.dense.bias" - ).contiguous() - - def time_embeds(num_steps: int) -> torch.Tensor: - fraction = torch.linspace(0.0, 1.0, 512) - period = 4e-3 * (4.0 / 4e-3) ** fraction - t = torch.tensor(1.0, dtype=torch.float32) - rows = [] - for _ in range(num_steps): - angle = ( - t.unsqueeze(-1) - * (1.0 / period).unsqueeze(0) - * 2 - * math.pi - ) - rows.append( - torch.cat([torch.sin(angle), torch.cos(angle)], dim=-1).to( - torch.bfloat16 - ) - ) - t = t - 1.0 / num_steps - return torch.cat(rows, dim=0).contiguous() - - action_out = bf16("action_out_proj.weight").t().to(torch.bfloat16) - for num_steps in (10, 5): - expected[f"action_out{num_steps}"] = ( - action_out * (-1.0 / num_steps) - ).contiguous() - expected[f"time_embeds{num_steps}"] = time_embeds(num_steps) - - for operation, tensor in expected.items(): - output = subprocess.check_output( - [args.probe, args.checkpoint, operation], text=True - ) - actual = parse_probe(output) - reference = digest(tensor) - if actual != reference: - raise AssertionError( - f"{operation}: C++ {actual} != PyTorch {reference}" - ) - print(f"PASS {operation} shape={actual[0]} fnv={actual[1]:016x}") - - -if __name__ == "__main__": - main() diff --git a/cpp/tests/gate_pi05_tokenizer_corpus.py b/cpp/tests/gate_pi05_tokenizer_corpus.py deleted file mode 100644 index 9f9c3c6c..00000000 --- a/cpp/tests/gate_pi05_tokenizer_corpus.py +++ /dev/null @@ -1,140 +0,0 @@ -"""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(" 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/cpp_runtime_modalities.md b/docs/cpp_runtime_modalities.md index dda1c314..e91399ad 100644 --- a/docs/cpp_runtime_modalities.md +++ b/docs/cpp_runtime_modalities.md @@ -82,9 +82,10 @@ Current Pi0.5 status: resize/normalize/cast directly into export device buffers; - conservative action staging path: device action buffer -> D2H -> CPU reference postprocess; -- native SM120 BF16 and SM110 FP8 checkpoint loading, tokenizer/prompt staging, - weight materialization, calibration where required, and graph capture are - implemented by the optional `frt_model_runtime_open_v1` producer. They +- native SM120 BF16/FP8 and SM110 FP8 checkpoint loading, tokenizer/prompt + staging, weight materialization, calibration where required, and graph + capture are implemented by the optional `frt_model_runtime_open_v1` + producer. They present the same model-runtime ABI and remain FlashRT responsibilities, not Nexus features. @@ -99,8 +100,9 @@ They give every CUDA/DMA/zero-copy fast path a golden contract. The vision device path uses CUDA resize/normalize/cast and is tested against the CPU reference at the declared output dtype. The action device path remains conservative D2H staging because postprocess is small; it can move to CUDA -without changing model adapters. Native SM110 calibration and inference add -separate producer-parity gates on top of these modality references. +without changing model adapters. Native SM110/SM120 FP8 calibration and +inference add separate producer-parity gates on top of these modality +references. ## Hot Path Rules diff --git a/docs/mindon_pi05_integration.md b/docs/mindon_pi05_integration.md index 35e41385..f96e84c7 100644 --- a/docs/mindon_pi05_integration.md +++ b/docs/mindon_pi05_integration.md @@ -78,21 +78,24 @@ Python setup producer. The host and Nexus adoption code must not change when switching between Lane A and Lane C. The current C++ shared object implements this symbol as a complete native-v2 -producer on SM120 BF16 and SM110 FP8. Both require CUDA kernels and -SentencePiece; SM120 uses native FA2, while SM110 uses the Thor FP8/CUTLASS -backend and an identity-bound calibration artifact. The factory validates `io`, -precision, checkpoint/tokenizer paths, fixed prompt mode, capacities, the -complete 812-tensor inventory, and OpenPI or LeRobot action/state q01/q99 +producer for SM120 BF16, SM120 FP8, and SM110 FP8. All routes require CUDA +kernels and SentencePiece. SM120 uses native FA2 with either BF16 GEMMs or +calibrated static FP8 GEMMs; SM110 uses the Thor FP8/CUTLASS backend. Every +FP8 route requires an identity-bound calibration artifact. The factory +validates `io`, precision, 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 backend/SentencePiece support -or unsupported hardware returns unsupported instead of publishing unusable -ports. - -On SM110, create the artifact with the model-specific calibration API before -opening the runtime, then pass it as `calibration_path`. One observation can -contain one, two, or three named camera frames; repeat observations for dataset -calibration. Camera synchronization and dataset policy stay in the Mindon host. +materializes context-owned weights/workspace, captures the `infer`, `context`, +and `decode_only` graph catalog, and returns the integrated model runtime. +Missing backend/SentencePiece support or unsupported hardware returns +unsupported instead of publishing unusable ports. + +For either FP8 backend, create the artifact with the model-specific calibration +API before opening the runtime, then pass it as `calibration_path`. One +observation can contain one, two, or three named camera frames; repeat +observations for dataset calibration. Camera synchronization and dataset +policy stay in the Mindon host. See [`pi05_thor_native_fp8.md`](pi05_thor_native_fp8.md) for exact build flags, C API usage, artifact invalidation, and validation gates. Native C++ NVFP4 is not currently advertised; Python precision routes remain independent. diff --git a/docs/pi05_cpp_runtime_migration.md b/docs/pi05_cpp_runtime_migration.md index a8f5e728..5a0ff9ea 100644 --- a/docs/pi05_cpp_runtime_migration.md +++ b/docs/pi05_cpp_runtime_migration.md @@ -22,9 +22,24 @@ 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. +postprocessing. `actions_raw` is the producer-declared SWAP alias for consumers +that need the model-space result: BF16 on SM120 and F16 on SM110. Consumers +must select the declared port rather than infer dtype or normalization from a +model name. + +## Native precision routes + +The native producer supports SM120 BF16, SM120 static FP8 E4M3, and SM110 +static FP8 E4M3. Every FP8 route requires a compatible calibration artifact; +SM120 uses the v2 artifact with vision, encoder, and decoder scales, while +SM110 uses the v1 artifact with encoder and decoder scales. `precision="auto"` +selects SM120 FP8 only when `calibration_path` is present, otherwise SM120 BF16; +SM110 auto-selects FP8 and therefore still requires the artifact. + +Public BF16 windows on SM120 describe staging and attention-boundary storage, +not the GEMM precision. An FP8 claim must be verified from producer identity, +artifact metadata, and captured kernel dispatch. Native C++ NVFP4 is not +implemented; Python precision routes remain independent. ## Runtime adoption diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 04f8d965..42c43603 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -225,12 +225,13 @@ There are three supported integration lanes: The Pi0.5 C++ shared object exports `frt_model_runtime_open_v1` as a complete native-v2 producer when built with CUDA kernels, SentencePiece, and the backend -for the executing device: native FA2 for SM120 BF16 or the Thor FP8/CUTLASS -backend for SM110. The factory requires `io="native_v2"`, `checkpoint_path`, +for the executing device: native FA2 for SM120 BF16/static FP8 or the Thor +FP8/CUTLASS backend for SM110. The factory requires `io="native_v2"`, +`checkpoint_path`, `tokenizer_model_path`, `state_prompt_mode="fixed"`, `max_prompt_tokens >= 200`, and a positive `state_dim`; `precision`, `num_views`, `chunk`, `num_steps`, -`vision_pool_factor`, and frame capacities are producer setup values. SM110 -FP8 additionally requires an identity-compatible `calibration_path`. See +`vision_pool_factor`, and frame capacities are producer setup values. Every +FP8 route additionally requires an identity-compatible `calibration_path`. See [`pi05_thor_native_fp8.md`](pi05_thor_native_fp8.md) for the build, calibration, artifact, and validation contract. It parses `checkpoint_path/model.safetensors` through the native read-only mmap loader to @@ -440,11 +441,13 @@ capture, replays 100 times with one variant, and verifies the new device-side valid length is observed. `flash_rt_fa2` remains a thin Python adapter over the same `libflashrt_fa2_raw` kernel owner. -The native builder publishes one `infer` graph/stage and the ordered ports -`prompt`, `state`, `images`, `noise`, `actions`, and `actions_raw`. Identity -includes observed hardware, precision, model/tokenizer SHA-256 values, prompt -mode, fixed shapes, and schedule parameters. SM110 FP8 identity also includes -the calibration artifact SHA-256. The only capsule region is +The native builder publishes `infer`, `context`, and `decode_only` graphs and +selects either one `infer` stage or the `context -> decode_only` stage DAG. It +also publishes the ordered ports `prompt`, `state`, `images`, `noise`, +`actions`, and `actions_raw`. Identity includes observed hardware, precision, +model/tokenizer SHA-256 values, prompt mode, fixed shapes, stage plan, and +schedule parameters. FP8 identity also includes the calibration artifact +SHA-256. 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 @@ -470,10 +473,16 @@ PYTHONPATH=.:./exec/build:./runtime/build python runtime/tests/test_model_runtim ctest --test-dir cpp/build --output-on-failure ``` -Real-checkpoint gates: +Detailed real-checkpoint oracles are maintained in the +[MindOn validation suite](https://github.com/LiangSu8899/MindOn-dev/tree/main/validation/pi05_cpp). +Build their stage probes with +`FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS=ON`; default FlashRT builds retain contract, +lifecycle, calibration, final-result, and hot-path tests. + +Real-checkpoint oracle example: ``` -python cpp/tests/gate_pi05_native_weight_ops.py \ +python /oracles/gate_pi05_native_weight_ops.py \ --checkpoint \ --probe cpp/build/pi05_native_weight_probe ``` @@ -488,8 +497,9 @@ 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 a SentencePiece-enabled build with either SM120 FA2 or SM110 Thor FP8. -In either lane, pybind modules and the producer library must come from the same -build directory; graph and buffer handles cannot cross exec builds. +In either lane, pybind modules and the producer library must be compiled from +the same source revision; graph and buffer handles cannot cross exec builds. +A stale extension is not a numerical reference for a newly built producer. Prompt/state STAGED ports require token-exact, formatter string-exact, embedding bit-exact, fixed-vs-exact E2E cosine, and hot-contract coverage; a @@ -503,33 +513,34 @@ The native factory lifecycle gate is: ``` 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. +public schema, one variant per captured graph, prompt/state/image staging, +direct SWAP noise input, finite action output, and retain/release teardown. -SM110 FP8 calibration and runtime math are gated separately against the shipped -Torch producer: +SM110 and SM120 FP8 calibration/runtime math are gated separately against a +Torch producer built from the same source revision: ``` -python cpp/tests/gate_pi05_native_thor_fp8.py \ - --probe /pi05_native_thor_fp8_probe \ +python /oracles/gate_pi05_native_thor_fp8.py \ + --probe /pi05_native_fp8_calibration_probe \ --checkpoint \ --tokenizer \ --artifact \ --samples 1 --views 1 -python cpp/tests/gate_pi05_native_thor_fp8.py \ - --probe /pi05_native_thor_fp8_probe \ +python /oracles/gate_pi05_native_calibration.py \ + --probe /pi05_native_fp8_calibration_probe \ --checkpoint \ --tokenizer \ --artifact \ - --samples 3 --views 2 + --samples 3 --views 3 ``` -The two-view gate submits cameras in reverse order to verify name-based +The multi-view gate submits cameras in reverse order to verify name-based canonicalization. It also rejects incomplete/duplicate camera sets and invalid noise without committing a sample, exercises deterministic generated noise, -and requires all encoder scales, decoder scales, and raw actions to be -bit-exact. Dataset iteration and camera synchronization remain host policy. +and requires encoder scales, decoder scales, and raw actions to be bit-exact. +The SM120 gate additionally requires all 109 vision scales to be bit-exact. +Dataset iteration and camera synchronization remain host policy. Python and C++ native-v2 producers must also publish identical canonical port/stage/region records (their producer identity and fingerprints remain @@ -552,7 +563,7 @@ The native formatter and tokenizer must also remain token-exact over real prompt/state traffic: ``` -python cpp/tests/gate_pi05_tokenizer_corpus.py \ +python /oracles/gate_pi05_tokenizer_corpus.py \ --dataset \ --checkpoint \ --tokenizer \ @@ -572,7 +583,7 @@ 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 \ +python /oracles/gate_pi05_native_e2e.py \ --checkpoint \ --tokenizer \ --dataset \ @@ -603,7 +614,7 @@ FLASHRT_PROFILE_RANGE=1 nsys profile --trace=cuda \ nsys profile --trace=cuda --cuda-graph-trace=node \ --capture-range=cudaProfilerApi --capture-range-end=stop \ -o \ - python cpp/tests/profile_pi05_python_replay.py \ + python /perf/profile_pi05_python_replay.py \ --checkpoint --num-views 2 --steps 10 nsys stats --report cuda_gpu_trace --format csv \ diff --git a/docs/pi05_thor_native_fp8.md b/docs/pi05_thor_native_fp8.md index 3ca0f11a..190bb8ad 100644 --- a/docs/pi05_thor_native_fp8.md +++ b/docs/pi05_thor_native_fp8.md @@ -1,8 +1,8 @@ -# Pi0.5 Native C++ FP8 On Thor +# Pi0.5 Native C++ FP8 -This guide covers the Python-free Pi0.5 producer on NVIDIA Thor SM110. It -loads a safetensors checkpoint, calibrates FP8 activation scales, captures one -fixed-shape CUDA Graph, and returns `frt_model_runtime_v1` from +This guide covers the Python-free Pi0.5 FP8 producer on NVIDIA SM110 and +SM120. It loads a safetensors checkpoint, calibrates FP8 activation scales, +captures fixed-shape CUDA Graphs, and returns `frt_model_runtime_v1` from `frt_model_runtime_open_v1`. The implementation is model-specific by design. The stable runtime ABI remains @@ -17,17 +17,24 @@ responsibilities. |---|---|---|---|---| | Native C++ | SM110 | FP8 E4M3 | Required | F16 | | Native C++ | SM120 | BF16 | Not used | BF16 | +| Native C++ | SM120 | FP8 E4M3 | Required | BF16 | | Python | Backend-specific | Existing FP8/BF16/NVFP4 routes | Existing Python contract | Producer-declared | -`precision="auto"` selects FP8 E4M3 on SM110 and BF16 on SM120. Production -configuration should normally specify the intended precision explicitly. +`precision="auto"` selects FP8 E4M3 on SM110. On SM120 it selects static FP8 +when `calibration_path` is present and BF16 otherwise. Production +configuration should specify the intended precision explicitly. BF16 native +windows on SM120 describe activation storage and attention boundaries; they +do not mean that an FP8 runtime fell back to BF16 GEMMs. + Native C++ NVFP4 is not implemented by this producer; `precision="nvfp4"` is rejected. This does not change the independent Python NVFP4 path. ## Build -The SM110 backend requires CUDA kernels, CUDA staging, exec, SentencePiece, -and a compatible CUTLASS checkout. It has been validated with CUDA 13.0. +Both backends require CUDA kernels, CUDA staging, exec, and SentencePiece. +SM110 additionally requires a compatible CUTLASS checkout. SM120 requires the +Python-free FA2 raw library from the same source revision; see +[`pi05_cpp_runtime_migration.md`](pi05_cpp_runtime_migration.md). ```bash export BUILD_DIR="$PWD/cpp/build-thor" @@ -50,10 +57,17 @@ Configure fails if the Thor backend is enabled without CUDA kernels, CUDA staging, exec, or CUTLASS. Build `Release` for deployment. A Debug build is useful for the same test matrix but is not a latency reference. +For SM120, build the shared FA2 raw library first, then configure the C++ tree +with `CMAKE_CUDA_ARCHITECTURES=120`, `FLASHRT_CPP_WITH_FA2=ON`, and +`FLASHRT_CPP_FA2_LIBRARY=`. Detailed real-checkpoint +diagnostic probes are opt-in through +`FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS=ON`; default builds retain the contract, +lifecycle, calibration, final-result, and hot-path tests. + ## Native Configuration Calibration and runtime open consume the same model configuration. Calibration -does not need `calibration_path`; runtime open requires it on SM110. +does not need `calibration_path`; every FP8 runtime open requires it. ```json { @@ -119,7 +133,7 @@ sample.n_noise = chunk * 32; rc = frt_pi05_calibration_observe_v1(session, &sample); rc = frt_pi05_calibration_finalize_v1( - session, "/path/to/pi05-sm110-fp8.safetensors"); + session, "/path/to/pi05-fp8.safetensors"); frt_pi05_calibration_destroy_v1(session); ``` @@ -170,9 +184,9 @@ common activations, so every artifact still needs an end-to-end action gate. `noise` is optional F32 `[chunk, 32]`. Supply fixed noise when comparing with a reference producer. If it is omitted with `n_noise=0`, FlashRT generates -deterministic normal F16 noise from `noise_seed + successful_sample_index`. -Malformed or non-finite state/noise payloads are rejected before the sample is -committed. +deterministic normal noise from `noise_seed + successful_sample_index`, rounded +to the backend activation dtype: F16 on SM110 and BF16 on SM120. Malformed or +non-finite state/noise payloads are rejected before the sample is committed. Calibration materializes the model and runs uncaptured reference forwards. It is an offline/setup operation, not a control-loop operation. Reuse the produced @@ -180,12 +194,21 @@ artifact until an identity input or calibration policy changes. ## Artifact Contract -The calibration file is an atomically published safetensors artifact with two -F32 tensors: +The calibration file is an atomically published safetensors artifact. SM110 +uses schema `flashrt.pi05.fp8_calibration.v1` and two F32 tensors: - `encoder_scales`: 72 values; - `decoder_scales`: `num_steps * 18 * 4` values. +SM120 uses schema `flashrt.pi05.fp8_calibration.v2`, BF16 activation metadata, +and one additional F32 tensor: + +- `vision_scales`: 109 values. + +The same C API dispatches by the active CUDA device. SM120 calibration reuses +the production graph pipeline in dynamic-scale mode; static runtime open +reuses it with the artifact scales. There is no duplicate calibration forward. + Metadata binds the artifact to: - schema, model, precision, tensor dtype, and reducer version; @@ -216,7 +239,7 @@ standard model-runtime symbol: "tokenizer_model_path": "/path/to/paligemma_tokenizer.model", "state_prompt_mode": "fixed", "precision": "fp8_e4m3fn", - "calibration_path": "/path/to/pi05-sm110-fp8.safetensors", + "calibration_path": "/path/to/pi05-fp8.safetensors", "max_prompt_tokens": 200, "state_dim": 8, "num_views": 2, @@ -228,17 +251,19 @@ standard model-runtime symbol: Resolve `FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL` as `frt_model_runtime_open_v1_fn`, or link the producer library and call -`frt_model_runtime_open_v1` directly. The returned runtime publishes one -`infer` stage and these ordered ports: - -| Port | Update | SM110 dtype | Payload | -|---|---|---|---| -| `prompt` | STAGED | U8 | UTF-8 task text | -| `state` | STAGED | F32 | raw proprioception | -| `images` | STAGED | F16 | host `frt_image_view[]` transformed into the captured window | -| `noise` | SWAP | F16 | device `[chunk, 32]` | -| `actions` | STAGED | F32 | host `[chunk, robot_action_dim]` | -| `actions_raw` | SWAP | F16 | device `[chunk, 32]` alias | +`frt_model_runtime_open_v1` directly. The returned runtime publishes `infer`, +`decode_only`, and `context` graphs. The selected stage plan is either one +`infer` stage or the ordered `context -> decode_only` DAG. Both plans use these +ordered ports: + +| Port | Update | SM110 FP8 | SM120 BF16/FP8 | Payload | +|---|---|---|---|---| +| `prompt` | STAGED | U8 | U8 | UTF-8 task text | +| `state` | STAGED | F32 | F32 | raw proprioception | +| `images` | STAGED | F16 | BF16 | host `frt_image_view[]` transformed into the captured window | +| `noise` | SWAP | F16 | BF16 | device `[chunk, 32]` | +| `actions` | STAGED | F32 | F32 | host `[chunk, robot_action_dim]` | +| `actions_raw` | SWAP | F16 | BF16 | device `[chunk, 32]` alias | Prompt formatting, state normalization/discretization, tokenization, embedding, vision preprocessing, and action postprocessing remain producer-owned. Nexus @@ -264,11 +289,12 @@ another invalidation and serialization boundary while keeping the mathematical source path identical to the shipped producer. The OS page cache may improve repeated file reads but is not part of the FlashRT contract. -On a representative SM110 run with a 14.47 GB F32 checkpoint, native setup was -7.90 seconds and the complete open/infer/output/teardown lifecycle was 8.62 -seconds. Of setup time, 7.61 seconds was weight transform/quantization/upload, -94 ms was workspace/style setup, 163 ms was warmup, and 26 ms was graph capture. -These are reference measurements, not latency ABI guarantees. +On a representative current SM110 run, model-owner setup was 7.65 seconds and +the complete `open_v1` publication path was 10.3 seconds. Of owner setup, +7.38 seconds was weight transform/quantization/upload, 85 ms was +workspace/style setup, 142 ms was warmup, and 44 ms was three-graph capture. +These are reference measurements, not latency ABI guarantees, and should not +be compared with a timer that stops after Python safetensors loading. ## Validation @@ -278,31 +304,37 @@ Run CPU/CUDA-off tests in addition to SM110 Release and Debug builds: ctest --test-dir "$BUILD_DIR" --output-on-failure ``` -The real-checkpoint gate compares native calibration and inference with the -shipped Torch producer using fixed observations and noise: +Real-checkpoint producer oracles and profiling tools are maintained in the +[MindOn validation suite](https://github.com/LiangSu8899/MindOn-dev/tree/main/validation/pi05_cpp). +Configure diagnostic probes when an oracle needs a stage-level binary, then +run the matching script. The public paths below are placeholders: ```bash -python cpp/tests/gate_pi05_native_thor_fp8.py \ - --probe "$BUILD_DIR/pi05_native_thor_fp8_probe" \ +python /oracles/gate_pi05_native_thor_fp8.py \ + --probe "$BUILD_DIR/pi05_native_fp8_calibration_probe" \ --checkpoint "$CHECKPOINT_DIR" \ --tokenizer "$TOKENIZER_MODEL" \ --artifact "$CALIBRATION_FILE" \ --samples 1 --views 1 -python cpp/tests/gate_pi05_native_thor_fp8.py \ - --probe "$BUILD_DIR/pi05_native_thor_fp8_probe" \ +python /oracles/gate_pi05_native_calibration.py \ + --probe "$BUILD_DIR/pi05_native_fp8_calibration_probe" \ --checkpoint "$CHECKPOINT_DIR" \ --tokenizer "$TOKENIZER_MODEL" \ --artifact "$CALIBRATION_FILE" \ - --samples 3 --views 2 + --samples 3 --views 3 ``` -The two-view probe deliberately submits reversed camera order and checks +The multi-view probe deliberately submits reversed camera order and checks duplicate/incomplete/unknown names, non-RGB input rejection, malformed noise, -deterministic generated noise, artifact loading, runtime identity, one graph -variant, finite logical actions, and teardown. The reference gates require all -72 encoder scales, all 720 decoder scales for ten steps, and all 320 raw action -values to be bit-exact. +deterministic generated noise, artifact loading, runtime identity, one variant +per captured graph, full/split equivalence, finite logical actions, and +teardown. SM110 +requires all 72 encoder scales, all 720 decoder scales for ten steps, and all +320 raw action values to be bit-exact. SM120 additionally requires all 109 +vision scales to be bit-exact. Build the Python extension and C++ producer from +the same source revision before producer parity testing; stale compiled kernels +are not a valid reference. For the complete service loop, profile the CUDA profiler range around 1,000 iterations of prompt/state/image/noise update, replay, and action output: @@ -320,4 +352,28 @@ The validated trace contained exactly 1,000 graph launches and no CUDA device allocation/free, CUDA host allocation/registration, mempool, virtual-memory, graph instantiation, or capture API in the measured range. A separate 1,000-update prompt/state gate measured 266 microseconds p99 against a 1 ms -limit, with one graph variant throughout. +limit, with the selected graph's variant count fixed at one. + +## Reference Latency And Precision Evidence + +The complete service-loop timer includes prompt/state/image staging, noise +upload, graph replay, synchronization, and logical action read. With a fixed +64-token prompt window, 10 warmups, and 100 measured ticks, representative +results were: + +| Hardware | Views | Precision | p50 | p99 | +|---|---:|---|---:|---:| +| RTX 5090 SM120 | 1 | static FP8 E4M3 | 15.62 ms | 15.65 ms | +| RTX 5090 SM120 | 2 | static FP8 E4M3 | 18.47 ms | 18.51 ms | +| RTX 5090 SM120 | 3 | static FP8 E4M3 | 21.22 ms | 21.25 ms | +| Thor SM110 | 1 | static FP8 E4M3 | 38.82 ms | 39.01 ms | +| Thor SM110 | 2 | static FP8 E4M3 | 46.55 ms | 46.87 ms | +| Thor SM110 | 3 | static FP8 E4M3 | 56.10 ms | 56.24 ms | + +The SM120 label is backed by runtime identity, v2 calibration metadata, +producer bit-exact raw actions, and a graph-node trace. One one-view replay +contained 898 `nvjet_sm120_qqtst_mma` FP8 GEMMs, 306 BF16-to-E4M3 quantization +kernels, and the fused E4M3 norm/gating kernels. BF16 FA2 attention kernels in +the same trace are intentional boundaries, not evidence of BF16 GEMM fallback. +Latency varies with clocks, prompt capacity, build flags, and system load; it +is validation evidence rather than a public performance guarantee. diff --git a/docs/pr_review_checklist.md b/docs/pr_review_checklist.md index 489ae0bc..cd5e81ad 100644 --- a/docs/pr_review_checklist.md +++ b/docs/pr_review_checklist.md @@ -464,6 +464,11 @@ Required: - Single-observation and repeated-observation reducers need independent gates; multi-input gates must cover reordered, missing, duplicate, and unknown inputs. Reference parity uses explicit fixed stochastic inputs. +- Public tensor-window dtype is not proof of compute precision. Native low- + precision claims must agree across producer identity, artifact metadata, and + captured kernel dispatch. Document expected mixed-precision boundaries. +- Producer binaries used for parity must be built from the same source + revision. A stale extension or shared library invalidates the comparison. Blockers: @@ -511,6 +516,17 @@ Tests must not: - Depend on test ordering or global CUDA state from another test. - Use network downloads in default tests. +Native producer test ownership: + +- FlashRT retains ABI/schema, lifecycle/ownership, artifact, final-result, + subgraph-equivalence, and hot-path invariant coverage. +- Layer diagnostics, official-framework/real-dataset oracles, profiler traces, + deployment shadow comparison, and soak tests may be maintained externally. +- Detailed probe targets are opt-in and non-installed. Never add production + intermediate-output ports only to satisfy an external oracle. +- A migrated gate needs a coverage map, byte-identical transfer or reviewed + rewrite, and one overlapping successful run before removal. + Commands: ```bash @@ -658,6 +674,10 @@ Before merge, verify: - [ ] Hardware-specific behavior is isolated. - [ ] Unsupported hardware/platform combinations fail clearly. - [ ] Precision/cache/graph changes have correctness evidence. +- [ ] Precision claims include identity/artifact/kernel-dispatch evidence and + same-revision producer binaries. +- [ ] Test migrations preserve core in-repository coverage and have an audited + coverage map. - [ ] Docs match actual API, flags, modules, and behavior. - [ ] Performance claims include reproducible commands and correctness. diff --git a/docs/runtime_contract.md b/docs/runtime_contract.md index 4640a4de..b8bdccf0 100644 --- a/docs/runtime_contract.md +++ b/docs/runtime_contract.md @@ -141,11 +141,13 @@ Pi0.5 is the reference C++ model runtime under `cpp/models/pi05/`. It supports both producer forms. The adopted-export path accepts Python- or native-produced graphs and overlays native vision/action/prompt/state verbs. The Python-free path loads safetensors and SentencePiece assets, selects an explicitly built -hardware backend, captures one native graph, and publishes the complete -`frt_model_runtime_v1` through `frt_model_runtime_open_v1`. SM120 currently -uses BF16 plus native FA2; SM110 uses FP8 E4M3 plus an identity-bound native -calibration artifact. Both expose the same backend-neutral contract, so Nexus -and serving hosts do not change. +hardware backend, captures the model-owned graph catalog, selects a stage plan, +and publishes the complete `frt_model_runtime_v1` through +`frt_model_runtime_open_v1`. SM120 supports BF16 and calibrated static FP8 with +native FA2; SM110 supports calibrated FP8 E4M3 through the Thor backend. Every +FP8 route requires an identity-bound native calibration artifact. All routes +expose the same backend-neutral contract, so Nexus and serving hosts do not +change. The `flashrt_cpp_pi05_c` target also exposes the model-specific host/calibration C API. These functions own Pi0.5 semantic transforms; they do not extend the From 5bbe3643545fe2dfac91b5f786a8607685538585 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 04:23:35 -0400 Subject: [PATCH 73/83] docs(pi05): align native precision support --- USAGE.md | 18 +++++++++++------- docs/pi05_io_contract.md | 12 ++++++------ 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/USAGE.md b/USAGE.md index 715870a3..6bb28ef1 100644 --- a/USAGE.md +++ b/USAGE.md @@ -179,20 +179,24 @@ Pi0.5 can also be loaded without a resident Python producer through | Hardware | Native precision | Required backend | Calibration artifact | |---|---|---|---| | SM120 | BF16 | native FA2 + SentencePiece | no | +| SM120 | static FP8 E4M3 | native FA2 + SentencePiece | yes (schema v2) | | Thor SM110 | FP8 E4M3 | Thor FP8/CUTLASS + SentencePiece | yes | -The SM110 producer includes model-specific C APIs for single-view, -multi-view, and repeated dataset-observation calibration. The resulting -safetensors artifact is bound to hardware, checkpoint, tokenizer, fixed -shapes, sample count, and reducer policy; runtime identity also includes the -artifact SHA-256. Native C++ NVFP4 is not currently supported. Python FP8 and -NVFP4 routes keep their existing behavior. +The native producer includes model-specific C APIs for single-view, +multi-view, and repeated dataset-observation calibration on SM110 and SM120. +SM110 artifacts contain encoder and decoder scales (schema v1); SM120 also +contains vision scales (schema v2). Every artifact is bound to hardware, +checkpoint, tokenizer, fixed shapes, sample count, and reducer policy; runtime +identity also includes its SHA-256. With `precision="auto"`, SM120 selects +static FP8 when `calibration_path` is present and BF16 otherwise; SM110 selects +FP8 and therefore still requires an artifact. Native C++ NVFP4 is not currently +supported. Python FP8 and NVFP4 routes keep their existing behavior. Native safetensors setup uses one direct mmap -> transform/quantize -> device upload path and does not create an implicit weight-cache format. This is independent of the Python JAX Orbax weight cache documented below. -See [Pi0.5 Native C++ FP8 on Thor](docs/pi05_thor_native_fp8.md) for build +See [Pi0.5 Native C++ FP8](docs/pi05_thor_native_fp8.md) for build flags, configuration JSON, C calibration usage, camera-name rules, artifact invalidation, runtime ports, and validation commands. The complete portable IO contract is [Pi0.5 Native Model Runtime IO](docs/pi05_io_contract.md). diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 42c43603..b6ed7511 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -288,12 +288,12 @@ deployment system. Compare producer startup using the complete scope; a Python timer that stops after safetensors conversion is not the same metric. On a representative Thor SM110 Release run with the same checkpoint size, -native setup measured 7.90 seconds and the complete open/infer/output/teardown -lifecycle measured 8.62 seconds. Weight transform, FP8 quantization, and upload -accounted for 7.61 seconds; workspace/style setup, warmup, and capture measured -94 ms, 163 ms, and 26 ms respectively. These measurements do not add a binary -weight cache: native safetensors loading remains one direct, identity-checked -path. +model-owner setup measured 7.65 seconds and complete `open_v1` publication +measured 10.3 seconds. Weight transform, FP8 quantization, and upload accounted +for 7.38 seconds; workspace/style setup, warmup, and three-graph capture +measured 85 ms, 142 ms, and 44 ms respectively. These measurements do not add +a binary weight cache: native safetensors loading remains one direct, +identity-checked path. 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 From c93624d400482bc6d45831b58ce309226b5c5e7f Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 04:32:54 -0400 Subject: [PATCH 74/83] docs(pi05): sanitize validation guidance --- cpp/tests/gate_pi05_model_runtime_export.py | 9 +++------ docs/pi05_io_contract.md | 4 ++-- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/cpp/tests/gate_pi05_model_runtime_export.py b/cpp/tests/gate_pi05_model_runtime_export.py index d0fd8335..410f7138 100644 --- a/cpp/tests/gate_pi05_model_runtime_export.py +++ b/cpp/tests/gate_pi05_model_runtime_export.py @@ -1,11 +1,8 @@ """Gate a real Pi0.5 export through the generic frt_model_runtime_v1 face. -Run inside the CUDA container from the repo root: - - 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-sm120-debug/libflashrt_cpp_pi05_c.so +Run from the repository root with the matching producer build, checkpoint, +and optional ``--fp8`` flag. The producer library and Python extension used by +the comparison must come from the same source revision. The gate compares three surfaces: 1. Python frontend staging/replay/postprocess. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index b6ed7511..3ba57610 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -596,8 +596,8 @@ 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. +independent acceptance checks. Run the gate against the official OpenPI +reference; the external validation suite documents reference dependency setup. 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 From f36a03d19970a86a6b4a50fa6736d3b8b84d3d68 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 06:13:30 -0400 Subject: [PATCH 75/83] fix(pi05): tighten native runtime contracts --- .../include/flashrt/cpp/models/pi05/c_api.h | 6 +- .../flashrt/cpp/models/pi05/model_runtime.h | 43 +++++------ cpp/models/pi05/src/native_open.cpp | 12 +-- .../pi05_native_fp8_calibration_probe.cpp | 18 ++++- cpp/tests/test_pi05_native_open.cpp | 14 ++++ docs/pi05_io_contract.md | 35 ++++++--- docs/pi05_thor_native_fp8.md | 74 +++++++++++++++++-- 7 files changed, 149 insertions(+), 53 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 d2740ddc..b0873087 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 @@ -45,9 +45,9 @@ typedef struct frt_pi05_runtime_config { const char* image_buffer_name; const char* action_buffer_name; - /* Optional ABI extension. Zero keeps the v1 default: BF16 buffers, which - * is the production FP8 Pi0.5 path. FP16 reference exports set both to - * FRT_PI05_DTYPE_FLOAT16. */ + /* Optional ABI extension. Zero keeps the legacy v1 BF16 default. Complete + * producers set these from their declared tensor windows: BF16 on SM120, + * F16 on the SM110 FP8 producer. */ int image_dtype; int action_dtype; diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/model_runtime.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/model_runtime.h index e7c0b9dc..23382f99 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/model_runtime.h +++ b/cpp/models/pi05/include/flashrt/cpp/models/pi05/model_runtime.h @@ -1,27 +1,20 @@ -/* Pi0.5 as an frt_model_runtime_v1 producer (the generic model-runtime face). +/* Pi0.5 adapters for the generic frt_model_runtime_v1 face. * - * The standard hand-off for hosts: instead of the model-specific - * frt_pi05_runtime_* verbs, a host receives one frt_model_runtime_v1 and - * drives it through the generic port/stage/verb contract - * (flashrt/model_runtime.h). Pi0.5 semantics map onto it as: + * The complete native_v2 producer returned by frt_model_runtime_open_v1 + * declares prompt, state, images, noise, actions, and actions_raw ports. Its + * graph catalog is infer, decode_only, and context. The selected stage plan is + * either one infer stage or context followed by an action stage backed by the + * decode_only graph. Hosts discover those declarations from the returned + * model; this header does not define a second execution contract. * - * port "images" IN STAGED IMAGE set_input <- frt_image_view[] in the - * declared camera-view order - * no "prompt" port adopted-export path: prompt embedding - * is prepared by the producer before - * capture. A native tokenizer producer - * adds a real STAGED TEXT port later. - * port "noise" IN SWAP TENSOR the diffusion seed window — the host - * writes raw bytes directly - * port "actions" OUT STAGED ACTION get_output -> unnormalized f32 robot - * actions (capacity/written in bytes) - * stage 0 the configured infer graph - * - * Two construction paths are exposed: - * - create(exp, ...): legacy adapter path for an export that did not already - * carry a model-runtime declaration; it declares the single infer stage. - * - create_over(model, ...): production path. The producer owns ports, - * stage DAG, identity and fingerprint; Pi0.5 C++ only replaces verbs. + * Two lower-level construction paths remain for producer integration: + * - create(exp, ...): legacy adapter for an export without a model-runtime + * declaration. It creates images/noise/actions ports and one infer stage; + * prompt embedding remains producer setup state. + * - create_over(model, ...): verb overlay for a producer-owned declaration. + * Ports, stage DAG, identity, and fingerprint are inherited exactly. The + * complete native_v2 factory uses this path after installing real prompt, + * state, image, inference, and action behavior. */ #ifndef FLASHRT_CPP_MODELS_PI05_MODEL_RUNTIME_H #define FLASHRT_CPP_MODELS_PI05_MODEL_RUNTIME_H @@ -42,11 +35,13 @@ int frt_pi05_model_runtime_create(const frt_runtime_export_v1* exp, const frt_pi05_runtime_config* config, frt_model_runtime_v1** out); -/* Build a retained Pi0.5 native verb overlay over an existing model-runtime +/* Build a retained Pi0.5 verb overlay over an existing model-runtime * declaration. Ports/stages/identity/fingerprint are inherited exactly from * `model`; the returned object replaces only set_input/get_output/prepare/step. * Required ports by name: "images" (IMAGE IN STAGED) and "actions" (ACTION OUT - * STAGED). Optional "noise" must be TENSOR IN SWAP if present. */ + * STAGED). Optional "noise"/"actions_raw" must be matching TENSOR SWAP ports; + * optional "prompt"/"state" must be TEXT/STATE IN STAGED and require a fully + * configured tokenizer/embedding/state-normalization implementation. */ int frt_pi05_model_runtime_create_over(const frt_model_runtime_v1* model, const frt_pi05_runtime_config* config, frt_model_runtime_v1** out); diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index 4faea959..0f48a0cc 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -657,8 +657,8 @@ int validate_config( int64_t max_prompt_tokens = 0; int64_t state_dim = 0; - int64_t num_views = 0; - int64_t chunk = 0; + int64_t num_views = 2; + int64_t chunk = 10; int64_t num_steps = 10; int64_t vision_pool_factor = 1; int64_t max_frame_width = 1280; @@ -681,11 +681,11 @@ int validate_config( g_last_error = "state_dim must be in [1, INT_MAX]"; return -1; } - if (num_views && (num_views < 1 || num_views > 3)) { + if (num_views < 1 || num_views > 3) { g_last_error = "num_views must be in [1, 3]"; return -1; } - if (chunk && (chunk <= 0 || chunk > INT_MAX)) { + if (chunk <= 0 || chunk > INT_MAX) { g_last_error = "chunk must be in [1, INT_MAX]"; return -1; } @@ -711,8 +711,8 @@ int validate_config( config.stage_plan = stage_plan; 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_views = static_cast(num_views); + config.chunk = static_cast(chunk); config.num_steps = static_cast(num_steps); config.vision_pool_factor = static_cast(vision_pool_factor); config.max_frame_width = static_cast(max_frame_width); diff --git a/cpp/tests/pi05_native_fp8_calibration_probe.cpp b/cpp/tests/pi05_native_fp8_calibration_probe.cpp index 7502fb1a..f0089adb 100644 --- a/cpp/tests/pi05_native_fp8_calibration_probe.cpp +++ b/cpp/tests/pi05_native_fp8_calibration_probe.cpp @@ -69,10 +69,11 @@ int calibration_error(frt_pi05_calibration_session* session, } // namespace int main(int argc, char** argv) { - if (argc != 6 && argc != 7) { + if (argc != 6 && argc != 7 && argc != 8) { std::cerr << "usage: pi05_native_fp8_calibration_probe CHECKPOINT " "TOKENIZER " - "ARTIFACT SAMPLES VIEWS [RAW_ACTION_OUTPUT]\n"; + "ARTIFACT SAMPLES VIEWS [RAW_ACTION_OUTPUT " + "ACTION_OUTPUT]\n"; return 2; } int device = 0; @@ -396,7 +397,7 @@ int main(int argc, char** argv) { model->release(model->owner); return 1; } - if (argc == 7) { + if (argc >= 7) { std::ofstream output(argv[6], std::ios::binary | std::ios::trunc); output.write(reinterpret_cast(raw.data()), static_cast( @@ -424,6 +425,17 @@ int main(int argc, char** argv) { return 1; } } + if (argc == 8) { + std::ofstream output(argv[7], std::ios::binary | std::ios::trunc); + output.write(reinterpret_cast(actions.data()), + static_cast( + actions.size() * sizeof(actions[0]))); + if (!output) { + std::cerr << "native FP8 logical action write failed\n"; + model->release(model->owner); + return 1; + } + } model->release(model->owner); std::cout << "PASS native " << hardware << " FP8 calibration and runtime lifecycle\n"; diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index 6c9c6076..bc6005d7 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -261,6 +261,20 @@ int main() { missing_key.c_str())); write_safetensors(root + "/model.safetensors"); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, ",\"num_views\":0").c_str(), &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "num_views")); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, ",\"chunk\":0").c_str(), &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "chunk")); + out = reinterpret_cast(0x1); rc = frt_model_runtime_open_v1( config(root, tokenizer, ",\"max_frame_width\":0").c_str(), &out); diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 3ba57610..a43e360c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -13,10 +13,11 @@ 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 +## Adopted-export Legacy 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`. +The Python-produced Pi0.5 `io="native"` export declares three host-visible +ports. This is the legacy adopted-export lane, not the complete native C++ +factory returned by `frt_model_runtime_open_v1`. | port | modality/update | direction | dtype/layout/shape | backing | |---|---|---|---|---| @@ -37,17 +38,17 @@ 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 +There is deliberately no `prompt` port on this adopted-export path. The prompt embedding is prepared by the producer before graph capture/export. A producer must not declare a `TEXT/STAGED` or `STATE/STAGED` port until the native verb can really update that input on the hot path. -## Native V2 Face +## Complete Native V2 Face -The `io="native_v2"` export adds prompt/state staging and a raw action alias. -Adding these ports changes the model-runtime identity and therefore the -fingerprint. Existing capsules from the old face must refuse restore into this -face. +The `io="native_v2"` face returned by `frt_model_runtime_open_v1` adds working +prompt/state staging and a raw action alias. Adding these ports changes the +model-runtime identity and therefore the fingerprint. Existing capsules from +the old face must refuse restore into this face. | port | modality/update | direction | dtype/layout/shape | backing | |---|---|---|---|---| @@ -64,6 +65,12 @@ 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. +`num_views`, `chunk`, `num_steps`, `max_prompt_tokens`, `state_dim`, precision, +and stage plan are fixed before graph capture. They determine declared shapes, +captured workspaces, calibration compatibility, and deployment identity; none +is a per-tick setter. A host discovers action dimensions from the `actions` +descriptor and never assumes `(10, 7)` or the model-space width of 32. + 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 @@ -442,7 +449,10 @@ valid length is observed. `flash_rt_fa2` remains a thin Python adapter over the same `libflashrt_fa2_raw` kernel owner. The native builder publishes `infer`, `context`, and `decode_only` graphs and -selects either one `infer` stage or the `context -> decode_only` stage DAG. It +selects either one `infer` stage or a logical `context -> action` stage DAG, +where the action stage is backed by the `decode_only` graph. Graph names and +stage roles are distinct: the frozen stage descriptor references graph indices, +while the producer manifest preserves the plan and logical stage names. It also publishes the ordered ports `prompt`, `state`, `images`, `noise`, `actions`, and `actions_raw`. Identity includes observed hardware, precision, model/tokenizer SHA-256 values, prompt mode, fixed shapes, stage plan, and @@ -453,6 +463,11 @@ 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. +Both plans use one logical workspace/buffer set. `context_action` enables +stage-level scheduling and host-asynchronous action production, but it does not +claim safe cross-tick context/action overlap. That requires a producer to +publish double-buffered hand-off state and a new compatible declaration. + 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 diff --git a/docs/pi05_thor_native_fp8.md b/docs/pi05_thor_native_fp8.md index 190bb8ad..c9c9c214 100644 --- a/docs/pi05_thor_native_fp8.md +++ b/docs/pi05_thor_native_fp8.md @@ -92,6 +92,27 @@ shape fields are fixed setup values. Increasing a capacity or changing view, chunk, schedule, tokenizer, checkpoint, precision, or calibration artifact produces a different deployment identity where applicable. +| Field | Requirement/default | Accepted value | Contract effect | +|---|---|---|---| +| `io` | required | `native_v2` | selects the complete six-port face | +| `checkpoint_path` | required | Pi0.5 safetensors directory | content hash enters identity | +| `tokenizer_model_path` | required | SentencePiece model file | content hash enters identity | +| `state_prompt_mode` | required | `fixed` | graph-safe prompt/state staging; enters identity | +| `precision` | default `auto` | `auto`, `bf16`, `fp8_e4m3fn` | resolves against hardware; resolved precision enters identity | +| `calibration_path` | required for FP8 | compatible artifact file | artifact hash enters FP8 identity | +| `stage_plan` | default `full` | `full`, `context_action` | changes stage DAG and fingerprint | +| `max_prompt_tokens` | required | positive integer | fixed text window; enters identity | +| `state_dim` | required | positive integer matching checkpoint stats | state port shape; enters identity | +| `num_views` | default `2` | `1`, `2`, or `3` | image shape/workspace/calibration; enters identity | +| `chunk` | default `10` | positive integer | noise/action shapes and workspace; enters identity | +| `num_steps` | default `10` | positive integer | denoise graph/calibration; enters identity | +| `vision_pool_factor` | default `1` | `1`, `2`, or `4` | vision graph/workspace/calibration; enters identity | +| `max_frame_width`, `max_frame_height` | defaults `1280`, `720` | positive integers | persistent host staging capacity; not model math | + +Omitting an optional numeric field selects its documented default. Supplying +zero does not mean default and is rejected. Configuration is setup-only; there +is no API to mutate these fields after publication. + ## Calibration API The model-specific API is declared in @@ -253,8 +274,8 @@ Resolve `FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL` as `frt_model_runtime_open_v1_fn`, or link the producer library and call `frt_model_runtime_open_v1` directly. The returned runtime publishes `infer`, `decode_only`, and `context` graphs. The selected stage plan is either one -`infer` stage or the ordered `context -> decode_only` DAG. Both plans use these -ordered ports: +`infer` stage or the ordered logical `context -> action` DAG, where `action` +references the `decode_only` graph. Both plans use these ordered ports: | Port | Update | SM110 FP8 | SM120 BF16/FP8 | Payload | |---|---|---|---|---| @@ -270,6 +291,28 @@ vision preprocessing, and action postprocessing remain producer-owned. Nexus or another consumer moves declared payloads and schedules declared stages; it does not interpret Pi0.5 semantics. +Initialize `state` before `prompt` when both first become available so the +first embedding update already contains the complete pair. During a normal +control loop the task prompt is stable and only raw F32 `state`, images, and +noise change each tick. If both prompt and state change, stage both before the +next replay; an embedded step does this before firing the DAG. Runtime image +views are positional in the declared camera order, unlike the named frames used +by calibration. + +The complete native producer captures one graph catalog and selects a stage +plan over it: + +| `stage_plan` | Published stages | Intended use | +|---|---|---| +| `full` | stage 0 -> graph `infer` | one-call serialized inference | +| `context_action` | stage 0 -> graph `context`; stage 1 -> graph `decode_only`, after stage 0 | explicit context/action scheduling | + +`context_action` does not duplicate hand-off buffers. It supports independent +stage fire/query/sync and asynchronous execution relative to the host control +loop, but context from the next tick must not overwrite buffers while the prior +action stage is still reading them. Safe cross-tick GPU overlap requires a +future producer-owned double-buffered declaration; it is not implied by Nexus. + ## Loading Path Native setup mmaps `model.safetensors` read-only and accepts F32, BF16, or F16 @@ -331,10 +374,11 @@ deterministic generated noise, artifact loading, runtime identity, one variant per captured graph, full/split equivalence, finite logical actions, and teardown. SM110 requires all 72 encoder scales, all 720 decoder scales for ten steps, and all -320 raw action values to be bit-exact. SM120 additionally requires all 109 -vision scales to be bit-exact. Build the Python extension and C++ producer from -the same source revision before producer parity testing; stale compiled kernels -are not a valid reference. +320 raw action values to be bit-exact. It also compares the final F32 logical +action chunk at cosine `>= 0.999999` and `rtol=atol=1e-6`. SM120 additionally +requires all 109 vision scales to be bit-exact. Build the Python extension and +C++ producer from the same source revision before producer parity testing; +stale compiled kernels are not a valid reference. For the complete service loop, profile the CUDA profiler range around 1,000 iterations of prompt/state/image/noise update, replay, and action output: @@ -368,7 +412,7 @@ results were: | RTX 5090 SM120 | 3 | static FP8 E4M3 | 21.22 ms | 21.25 ms | | Thor SM110 | 1 | static FP8 E4M3 | 38.82 ms | 39.01 ms | | Thor SM110 | 2 | static FP8 E4M3 | 46.55 ms | 46.87 ms | -| Thor SM110 | 3 | static FP8 E4M3 | 56.10 ms | 56.24 ms | +| Thor SM110 | 3 | static FP8 E4M3 | 57.04 ms | 57.25 ms | The SM120 label is backed by runtime identity, v2 calibration metadata, producer bit-exact raw actions, and a graph-node trace. One one-view replay @@ -377,3 +421,19 @@ kernels, and the fused E4M3 norm/gating kernels. BF16 FA2 attention kernels in the same trace are intentional boundaries, not evidence of BF16 GEMM fallback. Latency varies with clocks, prompt capacity, build flags, and system load; it is validation evidence rather than a public performance guarantee. + +For the current three-view Thor audit, the same fixed clocks, 64-token prompt +window, inputs, warmup, and 100-sample service scope produced: + +| Producer | Scope | p50 | p99 | +|---|---|---:|---:| +| Python Torch | image/infer wall time | 56.46 ms | 56.70 ms | +| Python Torch | prompt/state plus image/infer wall time | 57.05 ms | 57.20 ms | +| Native C++ | complete service loop | 57.04 ms | 57.25 ms | + +The comparable complete-service medians differ by `0.01 ms` and p99 values by +`0.05 ms` after rounding. The 320 raw F16 action values are bit-exact; the +logical F32 result reaches cosine `1.000000000` with maximum absolute error +`2.98e-8`. Dynamic CPU or memory clocks can widen p99 substantially, so latency +acceptance records clock policy instead of treating an unpinned sample as a +code regression. From 07e2429abc5752dc3dd85d7cfb135346854dc3fa Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 06:52:38 -0400 Subject: [PATCH 76/83] refactor(cpp): extract native config parsing --- cpp/CMakeLists.txt | 5 + cpp/models/pi05/CMakeLists.txt | 3 +- cpp/models/pi05/src/native_open.cpp | 228 +++--------------- .../flashrt/cpp/native/config_object.h | 37 +++ cpp/native/src/config_object.cpp | 204 ++++++++++++++++ cpp/tests/CMakeLists.txt | 4 + cpp/tests/test_native_config.cpp | 40 +++ 7 files changed, 319 insertions(+), 202 deletions(-) create mode 100644 cpp/native/include/flashrt/cpp/native/config_object.h create mode 100644 cpp/native/src/config_object.cpp create mode 100644 cpp/tests/test_native_config.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5a9a3594..06c114f3 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -179,6 +179,11 @@ target_include_directories(flashrt_cpp_vla target_link_libraries(flashrt_cpp_vla INTERFACE flashrt_cpp_modalities flashrt_cpp) +add_library(flashrt_cpp_native STATIC + native/src/config_object.cpp) +target_include_directories(flashrt_cpp_native + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/native/include) + add_library(flashrt_cpp_loader STATIC loader/src/safetensors.cpp loader/src/sha256.cpp) diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index 59e04d9b..40fac3be 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -148,7 +148,8 @@ add_library(flashrt_cpp_pi05_c SHARED src/native_model_runtime.cpp src/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c - PUBLIC flashrt_cpp_pi05 flashrt_runtime) + PUBLIC flashrt_cpp_pi05 flashrt_runtime + PRIVATE flashrt_cpp_native) target_include_directories(flashrt_cpp_pi05_c PUBLIC ${_pi05_public_include} PRIVATE ${_pi05_internal_include}) diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/native_open.cpp index 0f48a0cc..cbab5956 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/native_open.cpp @@ -1,6 +1,7 @@ #include "flashrt/model_runtime.h" #include "flashrt/cpp/loader/safetensors.h" #include "flashrt/cpp/models/pi05/native_weights.h" +#include "flashrt/cpp/native/config_object.h" #include "native_open_internal.h" #include @@ -14,7 +15,6 @@ #include #include #include -#include #include #include #include @@ -23,154 +23,6 @@ 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; @@ -553,38 +405,6 @@ bool validate_pi05_safetensors(const std::string& checkpoint_path) { return true; } -bool string_field(const std::map& obj, - const char* key, - std::string* out, - bool required) { - auto it = obj.find(key); - if (it == obj.end()) { - if (!required) return true; - g_last_error = std::string("missing required field: ") + key; - return false; - } - if (it->second.type != JsonValue::Type::kString || - it->second.text.empty()) { - g_last_error = std::string("field must be a non-empty string: ") + key; - return false; - } - if (out) *out = it->second.text; - return true; -} - -bool integer_field(const std::map& obj, - const char* key, - int64_t* out) { - auto it = obj.find(key); - if (it == obj.end()) return true; - if (it->second.type != JsonValue::Type::kInteger) { - g_last_error = std::string("field must be an integer: ") + key; - return false; - } - if (out) *out = it->second.integer; - return true; -} - int validate_config( const char* config_json, flashrt::models::pi05::NativeOpenConfig* parsed) { @@ -592,10 +412,9 @@ int validate_config( 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(); + flashrt::native::ConfigObject config_object; + if (!config_object.parse(config_json)) { + g_last_error = config_object.error(); return -1; } @@ -606,14 +425,18 @@ int validate_config( std::string precision; std::string calibration_path; std::string stage_plan; - 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) || - !string_field(obj, "precision", &precision, false) || - !string_field(obj, "calibration_path", &calibration_path, false) || - !string_field(obj, "stage_plan", &stage_plan, false)) { + if (!config_object.string_field("io", &io, true) || + !config_object.string_field("checkpoint_path", &checkpoint_path, + true) || + !config_object.string_field("tokenizer_model_path", + &tokenizer_model_path, true) || + !config_object.string_field("state_prompt_mode", &state_prompt_mode, + true) || + !config_object.string_field("precision", &precision, false) || + !config_object.string_field("calibration_path", &calibration_path, + false) || + !config_object.string_field("stage_plan", &stage_plan, false)) { + g_last_error = config_object.error(); return -1; } if (io != "native_v2") { @@ -663,14 +486,17 @@ int validate_config( int64_t vision_pool_factor = 1; int64_t max_frame_width = 1280; int64_t max_frame_height = 720; - if (!integer_field(obj, "max_prompt_tokens", &max_prompt_tokens) || - !integer_field(obj, "state_dim", &state_dim) || - !integer_field(obj, "num_views", &num_views) || - !integer_field(obj, "chunk", &chunk) || - !integer_field(obj, "num_steps", &num_steps) || - !integer_field(obj, "vision_pool_factor", &vision_pool_factor) || - !integer_field(obj, "max_frame_width", &max_frame_width) || - !integer_field(obj, "max_frame_height", &max_frame_height)) { + if (!config_object.integer_field("max_prompt_tokens", + &max_prompt_tokens) || + !config_object.integer_field("state_dim", &state_dim) || + !config_object.integer_field("num_views", &num_views) || + !config_object.integer_field("chunk", &chunk) || + !config_object.integer_field("num_steps", &num_steps) || + !config_object.integer_field("vision_pool_factor", + &vision_pool_factor) || + !config_object.integer_field("max_frame_width", &max_frame_width) || + !config_object.integer_field("max_frame_height", &max_frame_height)) { + g_last_error = config_object.error(); return -1; } if (max_prompt_tokens <= 0 || max_prompt_tokens > INT_MAX) { diff --git a/cpp/native/include/flashrt/cpp/native/config_object.h b/cpp/native/include/flashrt/cpp/native/config_object.h new file mode 100644 index 00000000..d71e8dde --- /dev/null +++ b/cpp/native/include/flashrt/cpp/native/config_object.h @@ -0,0 +1,37 @@ +#pragma once + +#include +#include +#include + +namespace flashrt::native { + +class ConfigObject { +public: + bool parse(const char* json); + + bool string_field(const char* key, + std::string* out, + bool required) const; + bool integer_field(const char* key, int64_t* out) const; + + const std::string& error() const { return error_; } + +private: + friend class ConfigObjectParser; + + enum class ValueType { kString, kInteger, kBool, kNull }; + + struct Value { + ValueType type = ValueType::kNull; + std::string text; + int64_t integer = 0; + }; + + bool fail(std::string message) const; + + std::map values_; + mutable std::string error_; +}; + +} // namespace flashrt::native diff --git a/cpp/native/src/config_object.cpp b/cpp/native/src/config_object.cpp new file mode 100644 index 00000000..5c2ff753 --- /dev/null +++ b/cpp/native/src/config_object.cpp @@ -0,0 +1,204 @@ +#include "flashrt/cpp/native/config_object.h" + +#include +#include +#include +#include +#include + +namespace flashrt::native { + +class ConfigObjectParser { +public: + explicit ConfigObjectParser(const char* source) + : current_(source ? source : "") {} + + bool parse(std::map* values, + std::string* error) { + values_ = values; + error_ = error; + skip_whitespace(); + if (!consume('{')) return fail("config_json must be a JSON object"); + skip_whitespace(); + if (consume('}')) return finish(); + + while (true) { + std::string key; + if (!parse_string(&key)) return false; + skip_whitespace(); + if (!consume(':')) return fail("expected ':' after JSON key"); + skip_whitespace(); + + ConfigObject::Value value; + if (!parse_value(&value)) return false; + (*values_)[key] = std::move(value); + + skip_whitespace(); + if (consume('}')) return finish(); + if (!consume(',')) return fail("expected ',' or '}' in object"); + skip_whitespace(); + } + } + +private: + bool finish() { + skip_whitespace(); + if (*current_) { + return fail("unexpected trailing data after JSON object"); + } + return true; + } + + void skip_whitespace() { + while (*current_ && + std::isspace(static_cast(*current_))) { + ++current_; + } + } + + bool consume(char value) { + if (*current_ != value) return false; + ++current_; + return true; + } + + bool parse_value(ConfigObject::Value* value) { + if (!value) return fail("internal parser error"); + if (*current_ == '"') { + value->type = ConfigObject::ValueType::kString; + return parse_string(&value->text); + } + if (*current_ == '-' || + std::isdigit(static_cast(*current_))) { + value->type = ConfigObject::ValueType::kInteger; + return parse_integer(&value->integer); + } + if (match_literal("true")) { + value->type = ConfigObject::ValueType::kBool; + return true; + } + if (match_literal("false")) { + value->type = ConfigObject::ValueType::kBool; + return true; + } + if (match_literal("null")) { + value->type = ConfigObject::ValueType::kNull; + return true; + } + return fail("unsupported JSON value"); + } + + bool parse_string(std::string* out) { + if (!consume('"')) return fail("expected JSON string"); + std::string result; + while (*current_ && *current_ != '"') { + const unsigned char value = + static_cast(*current_++); + if (value < 0x20) { + return fail("control character in JSON string"); + } + if (value != '\\') { + result.push_back(static_cast(value)); + continue; + } + + const char escape = *current_++; + switch (escape) { + case '"': result.push_back('"'); break; + case '\\': result.push_back('\\'); break; + case '/': result.push_back('/'); break; + case 'b': result.push_back('\b'); break; + case 'f': result.push_back('\f'); break; + case 'n': result.push_back('\n'); break; + case 'r': result.push_back('\r'); break; + case 't': result.push_back('\t'); break; + default: return fail("unsupported JSON string escape"); + } + } + if (!consume('"')) return fail("unterminated JSON string"); + if (out) *out = std::move(result); + return true; + } + + bool parse_integer(int64_t* out) { + const char* begin = current_; + if (*current_ == '-') ++current_; + if (!std::isdigit(static_cast(*current_))) { + return fail("expected JSON integer"); + } + if (*current_ == '0') { + ++current_; + } else { + while (std::isdigit(static_cast(*current_))) { + ++current_; + } + } + if (*current_ == '.' || *current_ == 'e' || *current_ == '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 != current_) { + 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 size = std::strlen(text); + if (std::strncmp(current_, text, size) != 0) return false; + current_ += size; + return true; + } + + bool fail(const char* message) { + if (error_) *error_ = message; + return false; + } + + const char* current_; + std::map* values_ = nullptr; + std::string* error_ = nullptr; +}; + +bool ConfigObject::parse(const char* json) { + values_.clear(); + error_.clear(); + ConfigObjectParser parser(json); + return parser.parse(&values_, &error_); +} + +bool ConfigObject::string_field(const char* key, + std::string* out, + bool required) const { + const auto it = values_.find(key); + if (it == values_.end()) { + if (!required) return true; + return fail(std::string("missing required field: ") + key); + } + if (it->second.type != ValueType::kString || it->second.text.empty()) { + return fail(std::string("field must be a non-empty string: ") + key); + } + if (out) *out = it->second.text; + return true; +} + +bool ConfigObject::integer_field(const char* key, int64_t* out) const { + const auto it = values_.find(key); + if (it == values_.end()) return true; + if (it->second.type != ValueType::kInteger) { + return fail(std::string("field must be an integer: ") + key); + } + if (out) *out = it->second.integer; + return true; +} + +bool ConfigObject::fail(std::string message) const { + error_ = std::move(message); + return false; +} + +} // namespace flashrt::native diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 21fc5f7a..558da751 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -15,6 +15,10 @@ add_executable(test_sha256 test_sha256.cpp) target_link_libraries(test_sha256 PRIVATE flashrt_cpp_loader) add_test(NAME sha256 COMMAND test_sha256) +add_executable(test_native_config test_native_config.cpp) +target_link_libraries(test_native_config PRIVATE flashrt_cpp_native) +add_test(NAME native_config COMMAND test_native_config) + add_executable(test_text_modalities test_text_modalities.cpp) target_link_libraries(test_text_modalities PRIVATE flashrt_cpp_modalities) add_test(NAME text_modalities COMMAND test_text_modalities) diff --git a/cpp/tests/test_native_config.cpp b/cpp/tests/test_native_config.cpp new file mode 100644 index 00000000..343bd9aa --- /dev/null +++ b/cpp/tests/test_native_config.cpp @@ -0,0 +1,40 @@ +#include "flashrt/cpp/native/config_object.h" + +#include +#include +#include + +int main() { + flashrt::native::ConfigObject config; + assert(config.parse( + R"({"io":"native","count":3,"enabled":true,"unused":null})")); + + std::string io; + int64_t count = 0; + assert(config.string_field("io", &io, true)); + assert(io == "native"); + assert(config.integer_field("count", &count)); + assert(count == 3); + assert(config.string_field("optional", nullptr, false)); + assert(config.integer_field("optional", nullptr)); + + assert(!config.string_field("missing", nullptr, true)); + assert(config.error() == "missing required field: missing"); + assert(!config.string_field("count", nullptr, true)); + assert(config.error() == "field must be a non-empty string: count"); + assert(!config.integer_field("io", nullptr)); + assert(config.error() == "field must be an integer: io"); + + assert(config.parse(R"({"value":1,"value":2})")); + int64_t value = 0; + assert(config.integer_field("value", &value)); + assert(value == 2); + + assert(!config.parse(R"({"value":1.5})")); + assert(config.error() == "JSON number must be an integer"); + assert(!config.parse(R"({"value":[]})")); + assert(config.error() == "unsupported JSON value"); + assert(!config.parse(R"({"value":1} trailing)")); + assert(config.error() == "unexpected trailing data after JSON object"); + return 0; +} From 5950ba7c646a1381e8f07d8b361e233f31efbda4 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 06:58:50 -0400 Subject: [PATCH 77/83] refactor(cpp): centralize CUDA graph ownership --- cpp/CMakeLists.txt | 8 ++ cpp/models/pi05/CMakeLists.txt | 3 +- .../cpp/models/pi05/native_graph_owner.h | 6 +- .../cpp/models/pi05/native_graph_runtime.h | 43 ++---- .../cpp/models/pi05/native_thor_graph_owner.h | 6 +- cpp/models/pi05/src/native_graph_owner.cpp | 17 ++- cpp/models/pi05/src/native_graph_runtime.cpp | 112 ++-------------- cpp/models/pi05/src/native_model_runtime.cpp | 6 +- .../pi05/src/native_thor_graph_owner.cpp | 17 ++- .../flashrt/cpp/native/cuda_graph_set.h | 53 ++++++++ cpp/native/src/cuda_graph_set.cpp | 124 ++++++++++++++++++ cpp/tests/CMakeLists.txt | 7 + cpp/tests/test_native_cuda_graph_set.cpp | 61 +++++++++ 13 files changed, 307 insertions(+), 156 deletions(-) create mode 100644 cpp/native/include/flashrt/cpp/native/cuda_graph_set.h create mode 100644 cpp/native/src/cuda_graph_set.cpp create mode 100644 cpp/tests/test_native_cuda_graph_set.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 06c114f3..4a8560aa 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -183,6 +183,14 @@ add_library(flashrt_cpp_native STATIC native/src/config_object.cpp) target_include_directories(flashrt_cpp_native PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/native/include) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + add_library(flashrt_cpp_native_cuda STATIC + native/src/cuda_graph_set.cpp) + target_include_directories(flashrt_cpp_native_cuda + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/native/include) + target_link_libraries(flashrt_cpp_native_cuda + PUBLIC flashrt_cpp_modalities flashrt_exec CUDA::cudart) +endif() add_library(flashrt_cpp_loader STATIC loader/src/safetensors.cpp diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index 40fac3be..b5c2f9a3 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -68,7 +68,8 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels) target_link_libraries(flashrt_cpp_pi05_kernels - PUBLIC flashrt_cpp_pi05 CUDA::cublas CUDA::cublasLt CUDA::cudart) + PUBLIC flashrt_cpp_pi05 flashrt_cpp_native_cuda + CUDA::cublas CUDA::cublasLt CUDA::cudart) set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h index 65141be7..3b78f1e6 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h @@ -32,7 +32,7 @@ class NativeGraphOwner final : public NativeGraphRuntime { frt_ctx context() const override { return graphs_.context(); } frt_graph graph(NativeGraphKind kind) const override { - return graphs_.graph(kind); + return graphs_.graph(static_cast(kind)); } int stream_id() const override { return graphs_.stream_id(); } void* native_stream() const override { return graphs_.native_stream(); } @@ -58,9 +58,9 @@ class NativeGraphOwner final : public NativeGraphRuntime { modalities::Status record_context(void* stream); modalities::Status record_action(void* stream); static modalities::Status record_graph( - void* user, NativeGraphKind kind, void* stream); + void* user, std::size_t slot, void* stream); - NativeGraphCatalog graphs_; + native::CudaGraphSet graphs_; NativeGraphConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h index 2f5c3df9..0bbb64eb 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h @@ -3,6 +3,7 @@ #include "flashrt/cpp/models/pi05/native_device_weights.h" #include "flashrt/cpp/models/pi05/native_workspace.h" +#include "flashrt/cpp/native/cuda_graph_set.h" #include #include @@ -32,41 +33,15 @@ enum class NativeGraphKind : std::size_t { kCount = 3, }; -class NativeGraphCatalog { -public: - using RecordFn = modalities::Status (*)( - void* owner, NativeGraphKind kind, void* stream); - - explicit NativeGraphCatalog(frt_ctx ctx) : ctx_(ctx) {} - ~NativeGraphCatalog(); - - NativeGraphCatalog(const NativeGraphCatalog&) = delete; - NativeGraphCatalog& operator=(const NativeGraphCatalog&) = delete; - - modalities::Status capture( - NativeGraphKind kind, const NativeWorkspace& workspace, - std::initializer_list bindings, - RecordFn record, void* owner); - modalities::Status create_replay_stream(); +const char* native_graph_name(NativeGraphKind kind); - frt_ctx context() const { return ctx_; } - frt_graph graph(NativeGraphKind kind) const; - int stream_id() const { return stream_id_; } - void* native_stream() const { return replay_stream_; } - int replay(NativeGraphKind kind) const; - modalities::Status synchronize() const; - - static const char* name(NativeGraphKind kind); - -private: - struct CaptureCall; - static void record_graph(void* user, void* stream); - - frt_ctx ctx_ = nullptr; - frt_graph graphs_[static_cast(NativeGraphKind::kCount)] = {}; - void* replay_stream_ = nullptr; - int stream_id_ = -1; -}; +modalities::Status capture_native_graph( + native::CudaGraphSet* graphs, + NativeGraphKind kind, + const NativeWorkspace& workspace, + std::initializer_list bindings, + native::CudaGraphSet::RecordFn record, + void* owner); modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, void* stream); diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h index b57c3ee5..bfba72eb 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h @@ -28,7 +28,7 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { frt_ctx context() const override { return graphs_.context(); } frt_graph graph(NativeGraphKind kind) const override { - return graphs_.graph(kind); + return graphs_.graph(static_cast(kind)); } int stream_id() const override { return graphs_.stream_id(); } void* native_stream() const override { return graphs_.native_stream(); } @@ -50,9 +50,9 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { modalities::Status record_context(void* stream); modalities::Status record_action(void* stream); static modalities::Status record_graph( - void* user, NativeGraphKind kind, void* stream); + void* user, std::size_t slot, void* stream); - NativeGraphCatalog graphs_; + native::CudaGraphSet graphs_; NativeGraphConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index a1399348..f4993aca 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -55,7 +55,7 @@ NativeGraphOwner::NativeGraphOwner( frt_ctx ctx, const NativeGraphConfig& config, NativeRtxLinearMode linear_mode) - : graphs_(ctx), + : graphs_(ctx, static_cast(NativeGraphKind::kCount)), config_(config), weights_(ctx), workspace_(ctx), @@ -295,20 +295,23 @@ modalities::Status NativeGraphOwner::initialize( } report("input_init"); - st = graphs_.capture( + st = capture_native_graph( + &graphs_, NativeGraphKind::kInfer, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = graphs_.capture( + st = capture_native_graph( + &graphs_, NativeGraphKind::kDecodeOnly, workspace_, {"encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = graphs_.capture( + st = capture_native_graph( + &graphs_, NativeGraphKind::kContext, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x"}, record_graph, this); @@ -359,9 +362,9 @@ modalities::Status NativeGraphOwner::record(NativeGraphKind kind, } modalities::Status NativeGraphOwner::record_graph( - void* user, NativeGraphKind kind, void* stream) { + void* user, std::size_t slot, void* stream) { auto* owner = static_cast(user); - return owner->record(kind, stream); + return owner->record(static_cast(slot), stream); } modalities::Status NativeGraphOwner::set_prompt_length(int prompt_tokens) { @@ -371,7 +374,7 @@ modalities::Status NativeGraphOwner::set_prompt_length(int prompt_tokens) { } int NativeGraphOwner::replay(NativeGraphKind kind) const { - return graphs_.replay(kind); + return graphs_.replay(static_cast(kind)); } modalities::Status NativeGraphOwner::synchronize() const { diff --git a/cpp/models/pi05/src/native_graph_runtime.cpp b/cpp/models/pi05/src/native_graph_runtime.cpp index 4ee994e3..99581544 100644 --- a/cpp/models/pi05/src/native_graph_runtime.cpp +++ b/cpp/models/pi05/src/native_graph_runtime.cpp @@ -21,32 +21,9 @@ modalities::Status backend(const char* message) { message); } -std::size_t graph_index(NativeGraphKind kind) { - return static_cast(kind); -} - } // namespace -struct NativeGraphCatalog::CaptureCall { - RecordFn record = nullptr; - void* owner = nullptr; - NativeGraphKind kind = NativeGraphKind::kInfer; - modalities::Status status = modalities::Status::ok(); -}; - -NativeGraphCatalog::~NativeGraphCatalog() { - if (replay_stream_) { - cudaStreamSynchronize(static_cast(replay_stream_)); - cudaStreamDestroy(static_cast(replay_stream_)); - replay_stream_ = nullptr; - } - if (ctx_) { - frt_ctx_destroy(ctx_); - ctx_ = nullptr; - } -} - -const char* NativeGraphCatalog::name(NativeGraphKind kind) { +const char* native_graph_name(NativeGraphKind kind) { switch (kind) { case NativeGraphKind::kInfer: return "infer"; case NativeGraphKind::kDecodeOnly: return "decode_only"; @@ -56,88 +33,27 @@ const char* NativeGraphCatalog::name(NativeGraphKind kind) { return nullptr; } -frt_graph NativeGraphCatalog::graph(NativeGraphKind kind) const { - const std::size_t index = graph_index(kind); - return index < graph_index(NativeGraphKind::kCount) - ? graphs_[index] - : nullptr; -} - -void NativeGraphCatalog::record_graph(void* user, void* stream) { - auto* call = static_cast(user); - if (!call || !call->record) return; - call->status = call->record(call->owner, call->kind, stream); -} - -modalities::Status NativeGraphCatalog::capture( - NativeGraphKind kind, const NativeWorkspace& workspace, +modalities::Status capture_native_graph( + native::CudaGraphSet* graphs, + NativeGraphKind kind, + const NativeWorkspace& workspace, std::initializer_list bindings, - RecordFn record, void* owner) { - const std::size_t index = graph_index(kind); - const char* graph_name = name(kind); - if (!ctx_ || !graph_name || index >= graph_index(NativeGraphKind::kCount) || - graphs_[index] || !record || !owner) { + native::CudaGraphSet::RecordFn record, + void* owner) { + if (!graphs || !native_graph_name(kind)) { return invalid("native graph capture request is invalid"); } - frt_graph captured = frt_graph_create(ctx_, graph_name, 1); - if (!captured) return backend("native graph creation failed"); - graphs_[index] = captured; + std::vector resolved; + resolved.reserve(bindings.size()); for (const char* binding : bindings) { const NativeWorkspaceBuffer* buffer = workspace.find(binding); - if (!buffer || - frt_graph_bind(captured, binding, buffer->buffer) != FRT_OK) { - frt_graph_destroy(captured); - graphs_[index] = nullptr; + if (!buffer) { return backend("native graph binding failed"); } + resolved.push_back({binding, buffer->buffer}); } - CaptureCall call; - call.record = record; - call.owner = owner; - call.kind = kind; - const int rc = frt_graph_capture(captured, 0, record_graph, &call); - if (!call.status.ok_status() || rc != FRT_OK || - frt_graph_variant_count(captured) != 1) { - frt_graph_destroy(captured); - graphs_[index] = nullptr; - return call.status.ok_status() - ? backend("native graph capture failed") - : call.status; - } - return modalities::Status::ok(); -} - -modalities::Status NativeGraphCatalog::create_replay_stream() { - if (!ctx_ || replay_stream_) { - return invalid("native replay stream request is invalid"); - } - cudaStream_t stream = nullptr; - if (cudaStreamCreate(&stream) != cudaSuccess) { - return backend("native replay stream creation failed"); - } - const int wrapped = frt_ctx_wrap_stream(ctx_, stream); - if (wrapped < 0) { - cudaStreamDestroy(stream); - return backend("native replay stream wrapping failed"); - } - replay_stream_ = stream; - stream_id_ = wrapped; - return modalities::Status::ok(); -} - -int NativeGraphCatalog::replay(NativeGraphKind kind) const { - frt_graph selected = graph(kind); - if (!selected || stream_id_ < 0) return FRT_ERR_INVALID; - return frt_graph_replay(selected, 0, stream_id_); -} - -modalities::Status NativeGraphCatalog::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)); + return graphs->capture(static_cast(kind), + native_graph_name(kind), resolved, record, owner); } modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 3fa32915..62e08fa5 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -284,15 +284,15 @@ int build_native_model_runtime(const NativeOpenConfig& config, builder, "main", graph->stream_id(), 0, graph->native_stream()) == 0 && frt_runtime_builder_add_graph( - builder, NativeGraphCatalog::name(NativeGraphKind::kInfer), + builder, native_graph_name(NativeGraphKind::kInfer), graph->graph(NativeGraphKind::kInfer), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_graph( - builder, NativeGraphCatalog::name(NativeGraphKind::kDecodeOnly), + builder, native_graph_name(NativeGraphKind::kDecodeOnly), graph->graph(NativeGraphKind::kDecodeOnly), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_graph( - builder, NativeGraphCatalog::name(NativeGraphKind::kContext), + builder, native_graph_name(NativeGraphKind::kContext), graph->graph(NativeGraphKind::kContext), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_buffer( diff --git a/cpp/models/pi05/src/native_thor_graph_owner.cpp b/cpp/models/pi05/src/native_thor_graph_owner.cpp index 66fd6ec1..939d38b0 100644 --- a/cpp/models/pi05/src/native_thor_graph_owner.cpp +++ b/cpp/models/pi05/src/native_thor_graph_owner.cpp @@ -59,7 +59,7 @@ bool calibration_matches(const NativeCalibrationArtifact& artifact, NativeThorGraphOwner::NativeThorGraphOwner( frt_ctx ctx, const NativeGraphConfig& config) - : graphs_(ctx), + : graphs_(ctx, static_cast(NativeGraphKind::kCount)), config_(config), weights_(ctx), workspace_(ctx), @@ -207,20 +207,23 @@ modalities::Status NativeThorGraphOwner::initialize( } report("warmup"); - st = graphs_.capture( + st = capture_native_graph( + &graphs_, NativeGraphKind::kInfer, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = graphs_.capture( + st = capture_native_graph( + &graphs_, NativeGraphKind::kDecodeOnly, workspace_, {"encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = graphs_.capture( + st = capture_native_graph( + &graphs_, NativeGraphKind::kContext, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x"}, record_graph, this); @@ -266,9 +269,9 @@ modalities::Status NativeThorGraphOwner::record(NativeGraphKind kind, } modalities::Status NativeThorGraphOwner::record_graph( - void* user, NativeGraphKind kind, void* stream) { + void* user, std::size_t slot, void* stream) { auto* owner = static_cast(user); - return owner->record(kind, stream); + return owner->record(static_cast(slot), stream); } modalities::Status NativeThorGraphOwner::set_prompt_length(int prompt_tokens) { @@ -276,7 +279,7 @@ modalities::Status NativeThorGraphOwner::set_prompt_length(int prompt_tokens) { } int NativeThorGraphOwner::replay(NativeGraphKind kind) const { - return graphs_.replay(kind); + return graphs_.replay(static_cast(kind)); } modalities::Status NativeThorGraphOwner::synchronize() const { diff --git a/cpp/native/include/flashrt/cpp/native/cuda_graph_set.h b/cpp/native/include/flashrt/cpp/native/cuda_graph_set.h new file mode 100644 index 00000000..f28e46c9 --- /dev/null +++ b/cpp/native/include/flashrt/cpp/native/cuda_graph_set.h @@ -0,0 +1,53 @@ +#pragma once + +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include +#include + +namespace flashrt::native { + +struct CudaGraphBinding { + const char* name = nullptr; + frt_buffer buffer = nullptr; +}; + +class CudaGraphSet { +public: + using RecordFn = modalities::Status (*)( + void* owner, std::size_t slot, void* stream); + + // The graph set takes ownership of context. + CudaGraphSet(frt_ctx context, std::size_t graph_count); + ~CudaGraphSet(); + + CudaGraphSet(const CudaGraphSet&) = delete; + CudaGraphSet& operator=(const CudaGraphSet&) = delete; + + modalities::Status capture( + std::size_t slot, + const char* name, + const std::vector& bindings, + RecordFn record, + void* owner); + modalities::Status create_replay_stream(); + + frt_ctx context() const { return context_; } + frt_graph graph(std::size_t slot) const; + int stream_id() const { return stream_id_; } + void* native_stream() const { return replay_stream_; } + int replay(std::size_t slot) const; + modalities::Status synchronize() const; + +private: + struct CaptureCall; + static void record_graph(void* user, void* stream); + + frt_ctx context_ = nullptr; + std::vector graphs_; + void* replay_stream_ = nullptr; + int stream_id_ = -1; +}; + +} // namespace flashrt::native diff --git a/cpp/native/src/cuda_graph_set.cpp b/cpp/native/src/cuda_graph_set.cpp new file mode 100644 index 00000000..352d4c81 --- /dev/null +++ b/cpp/native/src/cuda_graph_set.cpp @@ -0,0 +1,124 @@ +#include "flashrt/cpp/native/cuda_graph_set.h" + +#include + +namespace flashrt::native { +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 + +struct CudaGraphSet::CaptureCall { + RecordFn record = nullptr; + void* owner = nullptr; + std::size_t slot = 0; + modalities::Status status = modalities::Status::ok(); +}; + +CudaGraphSet::CudaGraphSet(frt_ctx context, std::size_t graph_count) + : context_(context), graphs_(graph_count, nullptr) {} + +CudaGraphSet::~CudaGraphSet() { + if (replay_stream_) { + cudaStreamSynchronize(static_cast(replay_stream_)); + cudaStreamDestroy(static_cast(replay_stream_)); + replay_stream_ = nullptr; + } + if (context_) { + frt_ctx_destroy(context_); + context_ = nullptr; + } +} + +frt_graph CudaGraphSet::graph(std::size_t slot) const { + return slot < graphs_.size() ? graphs_[slot] : nullptr; +} + +void CudaGraphSet::record_graph(void* user, void* stream) { + auto* call = static_cast(user); + if (!call || !call->record) return; + call->status = call->record(call->owner, call->slot, stream); +} + +modalities::Status CudaGraphSet::capture( + std::size_t slot, + const char* name, + const std::vector& bindings, + RecordFn record, + void* owner) { + if (!context_ || slot >= graphs_.size() || !name || !*name || + graphs_[slot] || !record || !owner) { + return invalid("native graph capture request is invalid"); + } + + frt_graph captured = frt_graph_create(context_, name, 1); + if (!captured) return backend("native graph creation failed"); + graphs_[slot] = captured; + for (const CudaGraphBinding& binding : bindings) { + if (!binding.name || !binding.buffer || + frt_graph_bind(captured, binding.name, binding.buffer) != FRT_OK) { + frt_graph_destroy(captured); + graphs_[slot] = nullptr; + return backend("native graph binding failed"); + } + } + + CaptureCall call; + call.record = record; + call.owner = owner; + call.slot = slot; + const int rc = frt_graph_capture(captured, 0, record_graph, &call); + if (!call.status.ok_status() || rc != FRT_OK || + frt_graph_variant_count(captured) != 1) { + frt_graph_destroy(captured); + graphs_[slot] = nullptr; + return call.status.ok_status() + ? backend("native graph capture failed") + : call.status; + } + return modalities::Status::ok(); +} + +modalities::Status CudaGraphSet::create_replay_stream() { + if (!context_ || replay_stream_) { + return invalid("native replay stream request is invalid"); + } + cudaStream_t stream = nullptr; + if (cudaStreamCreate(&stream) != cudaSuccess) { + return backend("native replay stream creation failed"); + } + const int wrapped = frt_ctx_wrap_stream(context_, stream); + if (wrapped < 0) { + cudaStreamDestroy(stream); + return backend("native replay stream wrapping failed"); + } + replay_stream_ = stream; + stream_id_ = wrapped; + return modalities::Status::ok(); +} + +int CudaGraphSet::replay(std::size_t slot) const { + const frt_graph selected = graph(slot); + if (!selected || stream_id_ < 0) return FRT_ERR_INVALID; + return frt_graph_replay(selected, 0, stream_id_); +} + +modalities::Status CudaGraphSet::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 flashrt::native diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 558da751..b1cbc4b6 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -19,6 +19,13 @@ add_executable(test_native_config test_native_config.cpp) target_link_libraries(test_native_config PRIVATE flashrt_cpp_native) add_test(NAME native_config COMMAND test_native_config) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + add_executable(test_native_cuda_graph_set test_native_cuda_graph_set.cpp) + target_link_libraries(test_native_cuda_graph_set + PRIVATE flashrt_cpp_native_cuda) + add_test(NAME native_cuda_graph_set COMMAND test_native_cuda_graph_set) +endif() + add_executable(test_text_modalities test_text_modalities.cpp) target_link_libraries(test_text_modalities PRIVATE flashrt_cpp_modalities) add_test(NAME text_modalities COMMAND test_text_modalities) diff --git a/cpp/tests/test_native_cuda_graph_set.cpp b/cpp/tests/test_native_cuda_graph_set.cpp new file mode 100644 index 00000000..b896f054 --- /dev/null +++ b/cpp/tests/test_native_cuda_graph_set.cpp @@ -0,0 +1,61 @@ +#include "flashrt/cpp/native/cuda_graph_set.h" + +#include + +#include +#include +#include + +namespace { + +struct RecordCall { + void* destination = nullptr; + std::size_t bytes = 0; +}; + +flashrt::modalities::Status record_fill( + void* user, std::size_t slot, void* stream) { + auto* call = static_cast(user); + if (!call || slot != 0 || !stream || !call->destination) { + return flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kInvalidArgument, + "invalid graph test record request"); + } + const cudaError_t result = cudaMemsetAsync( + call->destination, 0x5a, call->bytes, + static_cast(stream)); + return result == cudaSuccess + ? flashrt::modalities::Status::ok() + : flashrt::modalities::Status::error( + flashrt::modalities::StatusCode::kBackend, + cudaGetErrorString(result)); +} + +} // namespace + +int main() { + frt_ctx context = frt_ctx_create(); + assert(context); + flashrt::native::CudaGraphSet graphs(context, 1); + + constexpr std::size_t kBytes = 64; + frt_buffer output = frt_buffer_alloc(context, "output", kBytes); + assert(output); + RecordCall call{frt_buffer_dptr(output), kBytes}; + const std::vector bindings = { + {"output", output}, + }; + assert(graphs.capture(0, "fill", bindings, record_fill, &call) + .ok_status()); + assert(graphs.graph(0)); + assert(frt_graph_variant_count(graphs.graph(0)) == 1); + assert(graphs.create_replay_stream().ok_status()); + assert(graphs.replay(0) == FRT_OK); + assert(graphs.synchronize().ok_status()); + + std::array result{}; + assert(cudaMemcpy(result.data(), frt_buffer_dptr(output), kBytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + for (unsigned char value : result) assert(value == 0x5a); + return 0; +} From 70bbf0b010518d3363e2d38ef385a2b2fbf8a151 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 07:06:55 -0400 Subject: [PATCH 78/83] refactor(pi05): publish typed runtime artifacts --- .../cpp/models/pi05/native_graph_owner.h | 12 ++-- .../cpp/models/pi05/native_graph_runtime.h | 22 ++++++-- .../cpp/models/pi05/native_thor_graph_owner.h | 12 ++-- cpp/models/pi05/src/native_graph_owner.cpp | 4 ++ cpp/models/pi05/src/native_graph_runtime.cpp | 30 ++++++++++ cpp/models/pi05/src/native_model_runtime.cpp | 56 +++++++------------ .../pi05/src/native_thor_graph_owner.cpp | 4 ++ 7 files changed, 92 insertions(+), 48 deletions(-) diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h index 3b78f1e6..ffd75617 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h @@ -37,10 +37,13 @@ class NativeGraphOwner final : public NativeGraphRuntime { int stream_id() const override { return graphs_.stream_id(); } void* native_stream() const override { return graphs_.native_stream(); } const NativeGraphConfig& config() const { return config_; } - NativeDeviceWeightStore& weights() override { return weights_; } - const NativeDeviceWeightStore& weights() const override { return weights_; } - NativeWorkspace& workspace() override { return workspace_; } - const NativeWorkspace& workspace() const override { return workspace_; } + NativeDeviceWeightStore& weights() { return weights_; } + const NativeDeviceWeightStore& weights() const { return weights_; } + NativeWorkspace& workspace() { return workspace_; } + const NativeWorkspace& workspace() const { return workspace_; } + const NativeRuntimeArtifacts& artifacts() const override { + return artifacts_; + } NativeRtxAttentionWorkspace& attention() { return attention_; } const NativeRtxAttentionWorkspace& attention() const { return attention_; } @@ -64,6 +67,7 @@ class NativeGraphOwner final : public NativeGraphRuntime { NativeGraphConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; + NativeRuntimeArtifacts artifacts_; NativeRtxAttentionWorkspace attention_; NativeKernelDriver driver_; NativeRtxLinear linear_; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h index 0bbb64eb..abdd8466 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h @@ -46,6 +46,23 @@ modalities::Status capture_native_graph( modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, void* stream); +struct NativeRuntimeArtifacts { + const NativeWorkspaceBuffer* images = nullptr; + const NativeWorkspaceBuffer* noise = nullptr; + const NativeWorkspaceBuffer* encoder = nullptr; + const NativeWorkspaceBuffer* previous_actions = nullptr; + const NativeWorkspaceBuffer* prefix_weights = nullptr; + const NativeWorkspaceBuffer* guidance_weight = nullptr; + const NativeWorkspaceBuffer* prompt_embedding = nullptr; + const NativeDeviceWeight* embedding_table = nullptr; +}; + +modalities::Status resolve_native_runtime_artifacts( + const NativeWorkspace& workspace, + const NativeDeviceWeightStore& weights, + NativeWeightDType embedding_dtype, + NativeRuntimeArtifacts* artifacts); + class NativeGraphRuntime { public: virtual ~NativeGraphRuntime() = default; @@ -55,10 +72,7 @@ class NativeGraphRuntime { frt_graph infer_graph() const { return graph(NativeGraphKind::kInfer); } virtual int stream_id() const = 0; virtual void* native_stream() const = 0; - virtual NativeDeviceWeightStore& weights() = 0; - virtual const NativeDeviceWeightStore& weights() const = 0; - virtual NativeWorkspace& workspace() = 0; - virtual const NativeWorkspace& workspace() const = 0; + virtual const NativeRuntimeArtifacts& artifacts() const = 0; virtual modalities::Status set_prompt_length(int prompt_tokens) = 0; virtual int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const = 0; virtual modalities::Status synchronize() const = 0; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h index bfba72eb..266fb209 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h @@ -32,10 +32,13 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { } int stream_id() const override { return graphs_.stream_id(); } void* native_stream() const override { return graphs_.native_stream(); } - NativeDeviceWeightStore& weights() override { return weights_; } - const NativeDeviceWeightStore& weights() const override { return weights_; } - NativeWorkspace& workspace() override { return workspace_; } - const NativeWorkspace& workspace() const override { return workspace_; } + NativeDeviceWeightStore& weights() { return weights_; } + const NativeDeviceWeightStore& weights() const { return weights_; } + NativeWorkspace& workspace() { return workspace_; } + const NativeWorkspace& workspace() const { return workspace_; } + const NativeRuntimeArtifacts& artifacts() const override { + return artifacts_; + } modalities::Status set_prompt_length(int prompt_tokens) override; int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const override; @@ -56,6 +59,7 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { NativeGraphConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; + NativeRuntimeArtifacts artifacts_; NativeThorKernelDriver driver_; NativeThorFp8Forward forward_; NativeThorWeightScales weight_scales_; diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/native_graph_owner.cpp index f4993aca..946a2800 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/native_graph_owner.cpp @@ -280,6 +280,10 @@ modalities::Status NativeGraphOwner::initialize( } report("workspace_style"); + st = resolve_native_runtime_artifacts( + workspace_, weights_, NativeWeightDType::kBf16, &artifacts_); + if (!st.ok_status()) return st; + for (const char* name : {"observation_images_normalized", "prompt_embedding", "encoder_x", "diffusion_noise"}) { diff --git a/cpp/models/pi05/src/native_graph_runtime.cpp b/cpp/models/pi05/src/native_graph_runtime.cpp index 99581544..027ada34 100644 --- a/cpp/models/pi05/src/native_graph_runtime.cpp +++ b/cpp/models/pi05/src/native_graph_runtime.cpp @@ -81,6 +81,36 @@ modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, : backend("native prompt graph copy failed"); } +modalities::Status resolve_native_runtime_artifacts( + const NativeWorkspace& workspace, + const NativeDeviceWeightStore& weights, + NativeWeightDType embedding_dtype, + NativeRuntimeArtifacts* artifacts) { + if (!artifacts) { + return invalid("native runtime artifacts destination is null"); + } + NativeRuntimeArtifacts result; + result.images = workspace.find("observation_images_normalized"); + result.noise = workspace.find("diffusion_noise"); + result.encoder = workspace.find("encoder_x"); + result.previous_actions = workspace.find("rtc_prev_action_chunk"); + result.prefix_weights = workspace.find("rtc_prefix_weights"); + result.guidance_weight = workspace.find("rtc_guidance_weight"); + result.prompt_embedding = workspace.find("prompt_embedding"); + result.embedding_table = weights.find("embedding_weight"); + if (!result.images || !result.noise || !result.encoder || + !result.previous_actions || !result.prefix_weights || + !result.guidance_weight || !result.prompt_embedding || + !result.embedding_table || + result.embedding_table->dtype != embedding_dtype || + result.embedding_table->shape.size() != 2 || + result.embedding_table->shape[1] != kEncoderWidth) { + return backend("native graph export buffers are incomplete"); + } + *artifacts = result; + return modalities::Status::ok(); +} + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/native_model_runtime.cpp index 62e08fa5..01fd3e12 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/native_model_runtime.cpp @@ -8,6 +8,7 @@ #include "flashrt/cpp/models/pi05/model_runtime.h" #include "flashrt/cpp/models/pi05/native_calibration.h" #include "flashrt/cpp/models/pi05/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/spec.h" #if defined(FLASHRT_CPP_WITH_FA2) #include "flashrt/cpp/models/pi05/native_graph_owner.h" #endif @@ -28,12 +29,6 @@ namespace models { namespace pi05 { namespace { -const NativeWorkspaceBuffer* workspace_buffer( - const NativeGraphRuntime& owner, - const char* name) { - return owner.workspace().find(name); -} - void release_graph_owner(void* owner) { delete static_cast(owner); } @@ -248,30 +243,15 @@ int build_native_model_runtime(const NativeOpenConfig& config, return -2; } - 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 != (thor_fp8 ? NativeWeightDType::kFloat16 - : NativeWeightDType::kBf16) || - embedding->shape.size() != 2 || embedding->shape[1] != 2048) { - if (error) *error = "native graph export buffers are incomplete"; - return -6; - } + const NativeRuntimeArtifacts& artifacts = graph->artifacts(); + const NativeWorkspaceBuffer* images = artifacts.images; + const NativeWorkspaceBuffer* noise = artifacts.noise; + const NativeWorkspaceBuffer* encoder = artifacts.encoder; + const NativeWorkspaceBuffer* previous = artifacts.previous_actions; + const NativeWorkspaceBuffer* prefix_weights = artifacts.prefix_weights; + const NativeWorkspaceBuffer* guidance = artifacts.guidance_weight; + const NativeWorkspaceBuffer* prompt = artifacts.prompt_embedding; + const NativeDeviceWeight* embedding = artifacts.embedding_table; frt_runtime_builder builder = frt_runtime_builder_create(graph->context()); if (!builder) { @@ -353,7 +333,8 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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, "model_action_dim", + std::to_string(kModelActionDim)) && add_identity(builder, "robot_action_dim", std::to_string(config.action_q01.size())); if (ok && fp8) { @@ -385,8 +366,10 @@ int build_native_model_runtime(const NativeOpenConfig& config, 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 image_shape[] = { + config.num_views, kImageSize, kImageSize, 3}; + const int64_t raw_action_shape[] = { + config.chunk, kModelActionDim}; const int64_t action_shape[] = { config.chunk, static_cast(config.action_q01.size())}; const std::uint64_t action_bytes = @@ -455,7 +438,7 @@ int build_native_model_runtime(const NativeOpenConfig& 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.model_action_dim = kModelActionDim; 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(); @@ -478,12 +461,13 @@ int build_native_model_runtime(const NativeOpenConfig& config, frt_buffer_bytes(embedding->buffer); runtime_config.prompt_embedding_table_dtype = runtime_dtype; runtime_config.prompt_embedding_vocab_size = embedding->shape[0]; - runtime_config.prompt_embedding_hidden_dim = 2048; + runtime_config.prompt_embedding_hidden_dim = kEncoderWidth; 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 = runtime_dtype; runtime_config.max_prompt_tokens = config.max_prompt_tokens; - runtime_config.prompt_embedding_scale = std::sqrt(2048.0f); + runtime_config.prompt_embedding_scale = + std::sqrt(static_cast(kEncoderWidth)); 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(); diff --git a/cpp/models/pi05/src/native_thor_graph_owner.cpp b/cpp/models/pi05/src/native_thor_graph_owner.cpp index 939d38b0..46217b69 100644 --- a/cpp/models/pi05/src/native_thor_graph_owner.cpp +++ b/cpp/models/pi05/src/native_thor_graph_owner.cpp @@ -180,6 +180,10 @@ modalities::Status NativeThorGraphOwner::initialize( if (!st.ok_status()) return st; report("workspace_style"); + st = resolve_native_runtime_artifacts( + workspace_, weights_, NativeWeightDType::kFloat16, &artifacts_); + if (!st.ok_status()) return st; + for (const char* name : {"observation_images_normalized", "prompt_embedding", "diffusion_noise"}) { const NativeWorkspaceBuffer* input = workspace_.find(name); From 87608fbe902b915b28ba2e1d4984f01837b6f016 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 07:28:33 -0400 Subject: [PATCH 79/83] test(pi05): retain public contract coverage --- README.md | 6 +- benchmarks/pi05_production_profile.py | 254 ----------- cpp/CMakeLists.txt | 2 - cpp/models/pi05/tests/CMakeLists.txt | 212 +-------- cpp/tests/gate_pi05_kernel_sequence.py | 119 ----- cpp/tests/pi05_native_diffusion_probe.cpp | 195 -------- cpp/tests/pi05_native_e2e_probe.cpp | 177 -------- cpp/tests/pi05_native_encoder_layer_probe.cpp | 133 ------ cpp/tests/pi05_native_encoder_probe.cpp | 143 ------ cpp/tests/pi05_native_encoder_qkv_probe.cpp | 99 ----- cpp/tests/pi05_native_graph_probe.cpp | 123 ------ cpp/tests/pi05_native_quant_probe.cpp | 277 ------------ cpp/tests/pi05_native_rope_probe.cpp | 90 ---- cpp/tests/pi05_native_style_probe.cpp | 85 ---- cpp/tests/pi05_native_vision_probe.cpp | 415 ------------------ cpp/tests/pi05_native_weight_probe.cpp | 200 --------- cpp/tests/pi05_tokenizer_corpus_probe.cpp | 104 ----- cpp/tests/test_pi05_native_bf16_forward.cpp | 192 -------- cpp/tests/test_pi05_native_device_weights.cpp | 97 ---- .../test_pi05_native_forward_primitives.cpp | 348 --------------- cpp/tests/test_pi05_native_kernel_driver.cpp | 139 ------ cpp/tests/test_pi05_native_quantization.cpp | 56 --- cpp/tests/test_pi05_native_rtx_attention.cpp | 73 --- .../test_pi05_native_rtx_attention_driver.cpp | 157 ------- .../test_pi05_native_style_precompute.cpp | 27 -- .../test_pi05_native_weight_materializer.cpp | 345 --------------- cpp/tests/test_pi05_native_weight_ops.cpp | 316 ------------- cpp/tests/test_pi05_native_weight_packer.cpp | 141 ------ cpp/tests/test_pi05_native_workspace.cpp | 237 ---------- docs/pi05_io_contract.md | 12 +- docs/pi05_thor_native_fp8.md | 12 +- tests/bench_pi05_thor_views.py | 335 -------------- 32 files changed, 39 insertions(+), 5082 deletions(-) delete mode 100644 benchmarks/pi05_production_profile.py delete mode 100644 cpp/tests/gate_pi05_kernel_sequence.py delete mode 100644 cpp/tests/pi05_native_diffusion_probe.cpp delete mode 100644 cpp/tests/pi05_native_e2e_probe.cpp delete mode 100644 cpp/tests/pi05_native_encoder_layer_probe.cpp delete mode 100644 cpp/tests/pi05_native_encoder_probe.cpp delete mode 100644 cpp/tests/pi05_native_encoder_qkv_probe.cpp delete mode 100644 cpp/tests/pi05_native_graph_probe.cpp delete mode 100644 cpp/tests/pi05_native_quant_probe.cpp delete mode 100644 cpp/tests/pi05_native_rope_probe.cpp delete mode 100644 cpp/tests/pi05_native_style_probe.cpp delete mode 100644 cpp/tests/pi05_native_vision_probe.cpp delete mode 100644 cpp/tests/pi05_native_weight_probe.cpp delete mode 100644 cpp/tests/pi05_tokenizer_corpus_probe.cpp delete mode 100644 cpp/tests/test_pi05_native_bf16_forward.cpp delete mode 100644 cpp/tests/test_pi05_native_device_weights.cpp delete mode 100644 cpp/tests/test_pi05_native_forward_primitives.cpp delete mode 100644 cpp/tests/test_pi05_native_kernel_driver.cpp delete mode 100644 cpp/tests/test_pi05_native_quantization.cpp delete mode 100644 cpp/tests/test_pi05_native_rtx_attention.cpp delete mode 100644 cpp/tests/test_pi05_native_rtx_attention_driver.cpp delete mode 100644 cpp/tests/test_pi05_native_style_precompute.cpp delete mode 100644 cpp/tests/test_pi05_native_weight_materializer.cpp delete mode 100644 cpp/tests/test_pi05_native_weight_ops.cpp delete mode 100644 cpp/tests/test_pi05_native_weight_packer.cpp delete mode 100644 cpp/tests/test_pi05_native_workspace.cpp delete mode 100644 tests/bench_pi05_thor_views.py diff --git a/README.md b/README.md index a503c5f4..b8e228d9 100644 --- a/README.md +++ b/README.md @@ -578,9 +578,9 @@ NVFP4 weights directly from the Orbax checkpoint (no torch dependency at runtime) and uses the same two-phase multi-sample calibration flow as the torch FP4 path. Treat the table as Thor correctness / availability evidence, not a broad performance claim across every view count or host. -Reproduce with -[`tests/bench_pi05_thor_views.py`](tests/bench_pi05_thor_views.py) -(defaults now include `jax_fp4`). +The view-count benchmark is maintained in the +[MindOn PI0.5 validation suite](https://github.com/LiangSu8899/MindOn-dev/tree/main/validation/pi05_cpp/perf) +(`bench_pi05_thor_views.py`; defaults include `jax_fp4`). **What's next**: - Decoder FP4 (S2 precision-validated set — 72 weight tensors, ~-6 ms estimated) diff --git a/benchmarks/pi05_production_profile.py b/benchmarks/pi05_production_profile.py deleted file mode 100644 index 99db92c2..00000000 --- a/benchmarks/pi05_production_profile.py +++ /dev/null @@ -1,254 +0,0 @@ -#!/usr/bin/env python3 -"""Production-profile Pi0.5 benchmark: Python predict vs model-runtime tick. - -This is intentionally wall-clock, not graph-only: - - cold: - import/runtime setup, load_model, first predict (prompt/calibrate/capture), - export, and native model-runtime adapter creation. - - steady: - host uint8 images -> preprocessing/staging -> replay -> D2H/postprocess -> - float32 robot actions. - -The current native Pi0.5 model-runtime exposes diffusion noise as a SWAP port. -There is not yet a native C++ RNG/noise-fill verb, so this benchmark reports -the native loop with a fixed/prewritten noise window separately from the Python -public predict loop, which generates fresh noise each call. -""" - -from __future__ import annotations - -import argparse -import importlib.util -import statistics -import sys -import time -from pathlib import Path - - -ROOT = Path(__file__).resolve().parents[1] - - -def _stats(xs: list[float]) -> str: - xs = sorted(float(x) for x in xs) - if not xs: - return "n=0" - - def pct(p: float) -> float: - return xs[int(p * (len(xs) - 1))] - - return ( - f"n={len(xs)} p50={pct(0.50):.3f} p90={pct(0.90):.3f} " - f"p95={pct(0.95):.3f} mean={statistics.fmean(xs):.3f} " - f"min={xs[0]:.3f} max={xs[-1]:.3f}" - ) - - -def _load_gate_helpers(): - gate = ROOT / "cpp/tests/gate_pi05_model_runtime_export.py" - spec = importlib.util.spec_from_file_location("pi05_model_runtime_gate", gate) - if spec is None or spec.loader is None: - raise RuntimeError(f"cannot import helper module: {gate}") - mod = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = mod - spec.loader.exec_module(mod) - return mod - - -def _time_s(fn): - t0 = time.perf_counter() - value = fn() - return value, time.perf_counter() - t0 - - -def _bench_ms(label: str, iters: int, fn) -> list[float]: - times = [] - for i in range(iters): - t0 = time.perf_counter() - fn(i) - times.append((time.perf_counter() - t0) * 1000.0) - print(f"{label:<30}: {_stats(times)}", flush=True) - return times - - -def main() -> None: - ap = argparse.ArgumentParser() - ap.add_argument("--checkpoint", required=True) - ap.add_argument("--num-views", type=int, default=3) - ap.add_argument("--steps", type=int, default=10) - ap.add_argument("--prompt", default="pick up the red block") - ap.add_argument("--fp8", action="store_true", - help="use production FP8/BF16 path") - ap.add_argument("--seed", type=int, default=0) - ap.add_argument("--warmup", type=int, default=5) - ap.add_argument("--iters", type=int, default=50) - ap.add_argument("--lib", default=str( - ROOT / "cpp/build-container/libflashrt_cpp_pi05_c.so")) - args = ap.parse_args() - - helper, import_s = _time_s(_load_gate_helpers) - np = helper.np - torch = helper.torch - flash_rt = helper.flash_rt - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - - print("===== PI0.5 PRODUCTION PROFILE =====", flush=True) - print(f"device : {torch.cuda.get_device_name(0)} " - f"{torch.cuda.get_device_capability(0)}", flush=True) - print(f"precision : {'fp8/bf16' if args.fp8 else 'fp16'}", flush=True) - print(f"checkpoint : {args.checkpoint}", flush=True) - print(f"num_views/steps : {args.num_views}/{args.steps}", flush=True) - print("scope : host uint8 images -> robot float32 actions", - flush=True) - - image_seq = [ - helper._make_images(args.num_views, args.seed + i) - for i in range(max(args.iters, args.warmup) + 1) - ] - obs_seq = [helper._make_obs(images) for images in image_seq] - view_seq = [helper._make_image_views(images) for images in image_seq] - - def load_model(): - return flash_rt.load_model( - args.checkpoint, - framework="torch", - config="pi05", - hardware="auto", - num_views=args.num_views, - num_steps=args.steps, - cache_frames=1, - use_fp8=bool(args.fp8), - use_fp16=not args.fp8, - ) - - model, load_s = _time_s(load_model) - torch.cuda.synchronize() - - _, first_predict_s = _time_s( - lambda: model.predict(image_seq[0], prompt=args.prompt)) - torch.cuda.synchronize() - - pipe = model._pipe - pl = pipe.pipeline - assert getattr(pl, "_graph", None) is not None, "Pi05 graph was not captured" - - def make_model_runtime(): - export = pl.export_runtime(identity={"bench": "pi05_production_profile"}) - lib = helper._load_lib(args.lib) - dtype_id, torch_dtype = helper._dtype_from_frontend(pipe) - action_mean, action_stddev = helper._action_affine(pipe.norm_stats) - - cfg = helper.Pi05RuntimeConfig() - cfg.struct_size = helper.ctypes.sizeof(helper.Pi05RuntimeConfig) - cfg.num_views = args.num_views - cfg.chunk = int(pl.chunk_size) - cfg.model_action_dim = 32 - cfg.robot_action_dim = helper.LIBERO_ACTION_DIM - cfg.action_mean = action_mean.ctypes.data_as( - helper.ctypes.POINTER(helper.ctypes.c_float)) - cfg.n_action_mean = action_mean.size - cfg.action_stddev = action_stddev.ctypes.data_as( - helper.ctypes.POINTER(helper.ctypes.c_float)) - cfg.n_action_stddev = action_stddev.size - cfg.graph_name = b"infer" - cfg.image_buffer_name = b"observation_images_normalized" - cfg.action_buffer_name = b"diffusion_noise" - cfg.image_dtype = dtype_id - cfg.action_dtype = dtype_id - - m_ptr = helper.ctypes.c_void_p() - rc = lib.frt_pi05_model_runtime_create( - helper.ctypes.c_void_p(export.ptr), - helper.ctypes.byref(cfg), - helper.ctypes.byref(m_ptr), - ) - if rc != 0: - export.release() - raise RuntimeError(f"frt_pi05_model_runtime_create failed rc={rc}") - m = helper.ctypes.cast( - m_ptr, helper.ctypes.POINTER(helper.FrtModelRuntimeV1)).contents - return export, m, torch_dtype - - (export, m, torch_dtype), export_adopt_s = _time_s(make_model_runtime) - - try: - start_noise = helper._seed_noise(pipe, pl, args.seed + 1009) - - # Fixed-noise equivalence check between Python manual replay and the - # generic native model-runtime tick. - py_raw = helper._python_replay(pipe, pl, obs_seq[0], start_noise) - helper._upload_bytes(pl.input_noise_buf, start_noise) - helper._model_set_images(m, view_seq[0]) - cxx_actions = helper._model_step_get_actions(m, int(pl.chunk_size)) - cxx_raw = helper._read(pl.input_noise_buf) - py_raw_f = helper._raw_to_float(py_raw, torch_dtype).reshape( - int(pl.chunk_size), 32) - py_actions = helper.unnormalize_actions(py_raw_f, pipe.norm_stats)[ - :, :helper.LIBERO_ACTION_DIM].astype(np.float32) - raw_exact = bool(np.array_equal(py_raw, cxx_raw)) - act_max = float(np.max(np.abs(py_actions - cxx_actions))) - act_ok = bool(np.allclose(py_actions, cxx_actions, - rtol=1e-4, atol=1e-3)) - - print("\n===== COLD START =====", flush=True) - print(f"python imports : {import_s:.3f} s", flush=True) - print(f"load_model : {load_s:.3f} s", flush=True) - print(f"first predict/setup : {first_predict_s:.3f} s", flush=True) - print(f"export + native adopt : {export_adopt_s * 1000.0:.3f} ms", - flush=True) - print(f"python ready total : {import_s + load_s + first_predict_s:.3f} s", - flush=True) - print(f"hybrid ready total : " - f"{import_s + load_s + first_predict_s + export_adopt_s:.3f} s", - flush=True) - - print("\n===== FIXED-NOISE EQUIVALENCE =====", flush=True) - print(f"raw action exact : {raw_exact}", flush=True) - print(f"robot action allclose : {act_ok} max_abs={act_max:.6g}", - flush=True) - - for i in range(args.warmup): - idx = i % len(image_seq) - model.predict(image_seq[idx], prompt=args.prompt) - helper._upload_bytes(pl.input_noise_buf, start_noise) - helper._model_set_images(m, view_seq[idx]) - helper._model_step_get_actions(m, int(pl.chunk_size)) - torch.cuda.synchronize() - pipe.latency_records.clear() - - print("\n===== STEADY WALL-CLOCK =====", flush=True) - _bench_ms( - "python predict fresh-noise", - args.iters, - lambda i: model.predict( - image_seq[i % len(image_seq)], prompt=args.prompt), - ) - torch.cuda.synchronize() - if pipe.latency_records: - print(f"python predict internal : {_stats(list(pipe.latency_records))}", - flush=True) - - def native_fixed(i: int) -> None: - idx = i % len(view_seq) - helper._upload_bytes(pl.input_noise_buf, start_noise) - helper._model_set_images(m, view_seq[idx]) - helper._model_step_get_actions(m, int(pl.chunk_size)) - - _bench_ms("model-runtime fixed-noise", args.iters, native_fixed) - - print("\nNOTE", flush=True) - print("model-runtime fixed-noise includes host uint8 image staging, graph " - "replay, D2H action readback, and C++ postprocess. It does not " - "include native C++ RNG because the current ABI exposes noise as " - "a SWAP window; adding a native RNG/fill stage is the next gap for " - "a fully no-Python production loop.", flush=True) - finally: - m.release(m.owner) - export.release() - - -if __name__ == "__main__": - main() diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4a8560aa..b56783a3 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -26,8 +26,6 @@ option(FLASHRT_CPP_WITH_SENTENCEPIECE "Enable native SentencePiece text tokeniza option(FLASHRT_CPP_WITH_FA2 "Enable the Python-free vendored FA2 driver" OFF) option(FLASHRT_CPP_WITH_THOR_FP8 "Enable the native Thor SM110 FP8 kernel backend" OFF) -option(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS - "Build detailed PI0.5 real-checkpoint diagnostic probes" OFF) set(FLASHRT_CPP_FA2_LIBRARY "" CACHE FILEPATH "Path to the Python-free libflashrt_fa2_raw shared library") set(FLASHRT_CPP_CUTLASS_DIR diff --git a/cpp/models/pi05/tests/CMakeLists.txt b/cpp/models/pi05/tests/CMakeLists.txt index 876ce6a6..23f49cb8 100644 --- a/cpp/models/pi05/tests/CMakeLists.txt +++ b/cpp/models/pi05/tests/CMakeLists.txt @@ -1,4 +1,4 @@ -# PI0.5 contract, backend, lifecycle, and opt-in diagnostic targets. +# PI0.5 public contract, service, and lifecycle tests. set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") set(_pi05_test_source_dir "${FLASHRT_CPP_SOURCE_DIR}/tests") @@ -21,184 +21,11 @@ target_link_libraries(test_pi05_runtime PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) add_test(NAME pi05_runtime COMMAND test_pi05_runtime) -add_executable(test_pi05_native_weight_ops - ${_pi05_test_source_dir}/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(test_pi05_native_calibration ${_pi05_test_source_dir}/test_pi05_native_calibration.cpp) target_link_libraries(test_pi05_native_calibration PRIVATE flashrt_cpp_pi05) add_test(NAME pi05_native_calibration COMMAND test_pi05_native_calibration) -if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_native_weight_probe - ${_pi05_test_source_dir}/pi05_native_weight_probe.cpp) - target_link_libraries(pi05_native_weight_probe PRIVATE flashrt_cpp_pi05) -endif() - -if(FLASHRT_CPP_WITH_CUDA_KERNELS) - add_executable(test_pi05_native_quantization - ${_pi05_test_source_dir}/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) - - if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_native_quant_probe - ${_pi05_test_source_dir}/pi05_native_quant_probe.cpp) - target_link_libraries(pi05_native_quant_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - endif() - - add_executable(test_pi05_native_weight_packer - ${_pi05_test_source_dir}/test_pi05_native_weight_packer.cpp) - target_link_libraries(test_pi05_native_weight_packer - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_weight_packer - COMMAND test_pi05_native_weight_packer) - - add_executable(test_pi05_native_kernel_driver - ${_pi05_test_source_dir}/test_pi05_native_kernel_driver.cpp) - target_link_libraries(test_pi05_native_kernel_driver - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_kernel_driver - COMMAND test_pi05_native_kernel_driver) - - add_executable(test_pi05_native_forward_primitives - ${_pi05_test_source_dir}/test_pi05_native_forward_primitives.cpp) - target_link_libraries(test_pi05_native_forward_primitives - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_forward_primitives - COMMAND test_pi05_native_forward_primitives) - - add_executable(test_pi05_native_bf16_forward - ${_pi05_test_source_dir}/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) - - if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_native_encoder_qkv_probe - ${_pi05_test_source_dir}/pi05_native_encoder_qkv_probe.cpp) - target_link_libraries(pi05_native_encoder_qkv_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - endif() - - add_executable(test_pi05_native_style_precompute - ${_pi05_test_source_dir}/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) - - if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_native_style_probe - ${_pi05_test_source_dir}/pi05_native_style_probe.cpp) - target_link_libraries(pi05_native_style_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - endif() - - add_executable(test_pi05_native_workspace - ${_pi05_test_source_dir}/test_pi05_native_workspace.cpp) - target_link_libraries(test_pi05_native_workspace - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_workspace - COMMAND test_pi05_native_workspace) - - add_executable(test_pi05_native_rtx_attention - ${_pi05_test_source_dir}/test_pi05_native_rtx_attention.cpp) - target_link_libraries(test_pi05_native_rtx_attention - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_rtx_attention - COMMAND test_pi05_native_rtx_attention) - - if(FLASHRT_CPP_WITH_FA2) - if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_native_encoder_layer_probe - ${_pi05_test_source_dir}/pi05_native_encoder_layer_probe.cpp) - target_link_libraries(pi05_native_encoder_layer_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_encoder_probe - ${_pi05_test_source_dir}/pi05_native_encoder_probe.cpp) - target_link_libraries(pi05_native_encoder_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_vision_probe - ${_pi05_test_source_dir}/pi05_native_vision_probe.cpp) - target_link_libraries(pi05_native_vision_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_diffusion_probe - ${_pi05_test_source_dir}/pi05_native_diffusion_probe.cpp) - target_link_libraries(pi05_native_diffusion_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - add_executable(pi05_native_graph_probe - ${_pi05_test_source_dir}/pi05_native_graph_probe.cpp) - target_link_libraries(pi05_native_graph_probe - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - - if(FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_e2e_probe - ${_pi05_test_source_dir}/pi05_native_e2e_probe.cpp) - target_link_libraries(pi05_native_e2e_probe - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - endif() - endif() - - if(FLASHRT_CPP_WITH_SENTENCEPIECE) - add_executable(pi05_native_dlopen_probe - ${_pi05_test_source_dir}/pi05_native_dlopen_probe.cpp) - target_include_directories(pi05_native_dlopen_probe - PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../runtime/include - ${FLASHRT_CPP_SOURCE_DIR}/../exec/include) - target_link_libraries(pi05_native_dlopen_probe - PRIVATE ${CMAKE_DL_LIBS}) - endif() - - add_executable(test_pi05_native_rtx_attention_driver - ${_pi05_test_source_dir}/test_pi05_native_rtx_attention_driver.cpp) - target_link_libraries(test_pi05_native_rtx_attention_driver - PRIVATE flashrt_cpp_pi05_kernels flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_rtx_attention_driver - COMMAND test_pi05_native_rtx_attention_driver) - endif() - - if(FLASHRT_CPP_WITH_SENTENCEPIECE AND - (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) - add_executable(pi05_native_open_probe - ${_pi05_test_source_dir}/pi05_native_open_probe.cpp) - target_link_libraries(pi05_native_open_probe - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - endif() - - if(FLASHRT_CPP_WITH_SENTENCEPIECE AND - FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_tokenizer_corpus_probe - ${_pi05_test_source_dir}/pi05_tokenizer_corpus_probe.cpp) - target_link_libraries(pi05_tokenizer_corpus_probe - PRIVATE flashrt_cpp_pi05) - endif() - - if(FLASHRT_CPP_WITH_SENTENCEPIECE AND - (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) - add_executable(pi05_native_fp8_calibration_probe - ${_pi05_test_source_dir}/pi05_native_fp8_calibration_probe.cpp) - target_link_libraries(pi05_native_fp8_calibration_probe - PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) - endif() - - if(FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS) - add_executable(pi05_native_rope_probe - ${_pi05_test_source_dir}/pi05_native_rope_probe.cpp) - target_link_libraries(pi05_native_rope_probe - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - endif() -endif() - add_executable(test_pi05_prompt_format ${_pi05_test_source_dir}/test_pi05_prompt_format.cpp) target_link_libraries(test_pi05_prompt_format PRIVATE flashrt_cpp_pi05) @@ -209,6 +36,29 @@ add_executable(test_pi05_prompt_embed 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_KERNELS AND FLASHRT_CPP_WITH_SENTENCEPIECE AND + (FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8)) + add_executable(pi05_native_open_probe + ${_pi05_test_source_dir}/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_fp8_calibration_probe + ${_pi05_test_source_dir}/pi05_native_fp8_calibration_probe.cpp) + target_link_libraries(pi05_native_fp8_calibration_probe + PRIVATE flashrt_cpp_pi05_c flashrt_exec CUDA::cudart) +endif() + +if(FLASHRT_CPP_WITH_CUDA_KERNELS AND FLASHRT_CPP_WITH_SENTENCEPIECE AND + FLASHRT_CPP_WITH_FA2) + add_executable(pi05_native_dlopen_probe + ${_pi05_test_source_dir}/pi05_native_dlopen_probe.cpp) + target_include_directories(pi05_native_dlopen_probe + PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../runtime/include + ${FLASHRT_CPP_SOURCE_DIR}/../exec/include) + target_link_libraries(pi05_native_dlopen_probe PRIVATE ${CMAKE_DL_LIBS}) +endif() + if(FLASHRT_CPP_WITH_CUDA_STAGING) add_executable(test_device_staging ${_pi05_test_source_dir}/test_device_staging.cpp) @@ -216,20 +66,6 @@ if(FLASHRT_CPP_WITH_CUDA_STAGING) PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities CUDA::cudart) add_test(NAME device_staging COMMAND test_device_staging) - add_executable(test_pi05_native_device_weights - ${_pi05_test_source_dir}/test_pi05_native_device_weights.cpp) - target_link_libraries(test_pi05_native_device_weights - PRIVATE flashrt_cpp_pi05 flashrt_exec CUDA::cudart) - add_test(NAME pi05_native_device_weights - COMMAND test_pi05_native_device_weights) - - add_executable(test_pi05_native_weight_materializer - ${_pi05_test_source_dir}/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 ${_pi05_test_source_dir}/test_pi05_c_api.cpp) target_link_libraries(test_pi05_c_api diff --git a/cpp/tests/gate_pi05_kernel_sequence.py b/cpp/tests/gate_pi05_kernel_sequence.py deleted file mode 100644 index 1c8ae859..00000000 --- a/cpp/tests/gate_pi05_kernel_sequence.py +++ /dev/null @@ -1,119 +0,0 @@ -"""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/pi05_native_diffusion_probe.cpp b/cpp/tests/pi05_native_diffusion_probe.cpp deleted file mode 100644 index 02556842..00000000 --- a/cpp/tests/pi05_native_diffusion_probe.cpp +++ /dev/null @@ -1,195 +0,0 @@ -#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/cpp/tests/pi05_native_e2e_probe.cpp b/cpp/tests/pi05_native_e2e_probe.cpp deleted file mode 100644 index 0b3e5249..00000000 --- a/cpp/tests/pi05_native_e2e_probe.cpp +++ /dev/null @@ -1,177 +0,0 @@ -#include "flashrt/model_runtime.h" - -#include - -#include -#include -#include -#include -#include -#include -#include - -extern "C" int frt_model_runtime_open_v1(const char* config_json, - frt_model_runtime_v1** out); -extern "C" const char* frt_pi05_native_open_last_error(); - -namespace { - -bool read_file(const std::string& path, std::vector* out) { - std::ifstream input(path, std::ios::binary); - if (!input) return false; - input.seekg(0, std::ios::end); - const std::streamoff size = input.tellg(); - if (size < 0) return false; - input.seekg(0, std::ios::beg); - std::vector data(static_cast(size)); - if (size && !input.read(reinterpret_cast(data.data()), size)) { - return false; - } - *out = std::move(data); - return true; -} - -bool write_file(const std::string& path, const void* data, std::size_t bytes) { - std::ofstream output(path, std::ios::binary | std::ios::trunc); - if (!output) return false; - output.write(static_cast(data), - static_cast(bytes)); - return output.good(); -} - -std::string json_string(const std::string& value) { - std::string output = "\""; - for (char c : value) { - if (c == '\\' || c == '"') output.push_back('\\'); - output.push_back(c); - } - output.push_back('"'); - return output; -} - -int model_error(frt_model_runtime_v1* model, const char* message) { - std::cerr << message; - if (model && model->verbs.last_error) { - const char* detail = model->verbs.last_error(model->self); - if (detail && *detail) std::cerr << ": " << detail; - } - std::cerr << '\n'; - if (model) model->release(model->owner); - return 1; -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 5) { - std::cerr << "usage: pi05_native_e2e_probe CHECKPOINT TOKENIZER " - "FIXTURE_DIR OUTPUT_DIR\n"; - return 2; - } - const std::string checkpoint = argv[1]; - const std::string tokenizer = argv[2]; - const std::string fixture = argv[3]; - const std::string output = argv[4]; - - std::ostringstream json; - json << "{\"io\":\"native_v2\",\"checkpoint_path\":" - << json_string(checkpoint) << ",\"tokenizer_model_path\":" - << json_string(tokenizer) - << ",\"state_prompt_mode\":\"fixed\"," - "\"max_prompt_tokens\":200,\"state_dim\":8," - "\"num_views\":2,\"chunk\":10,\"num_steps\":10," - "\"vision_pool_factor\":1}"; - frt_model_runtime_v1* model = nullptr; - const int open_rc = frt_model_runtime_open_v1(json.str().c_str(), &model); - if (open_rc != 0 || !model) { - std::cerr << "native open failed: rc=" << open_rc << " error=" - << frt_pi05_native_open_last_error() << '\n'; - return 1; - } - const char* names[] = { - "prompt", "state", "images", "noise", "actions", "actions_raw"}; - if (model->n_ports != 6) return model_error(model, "unexpected port count"); - for (std::uint64_t i = 0; i < model->n_ports; ++i) { - if (!model->ports[i].name || - std::strcmp(model->ports[i].name, names[i]) != 0) { - return model_error(model, "unexpected port schema"); - } - } - if (model->ports[4].dtype != FRT_RT_DTYPE_F32 || - model->ports[4].update != FRT_RT_PORT_STAGED || - model->ports[4].buffer != nullptr || - model->ports[4].bytes != 10 * 7 * sizeof(float) || - model->ports[5].dtype != FRT_RT_DTYPE_BF16 || - model->ports[5].update != FRT_RT_PORT_SWAP || - model->ports[5].buffer != model->ports[3].buffer) { - return model_error(model, "native action port contract mismatch"); - } - - std::vector prompt; - std::vector state; - std::vector image0; - std::vector image1; - std::vector noise; - if (!read_file(fixture + "/prompt.txt", &prompt) || - !read_file(fixture + "/state.f32", &state) || - !read_file(fixture + "/image_0.rgb", &image0) || - !read_file(fixture + "/image_1.rgb", &image1) || - !read_file(fixture + "/noise.bf16", &noise) || - state.size() != 8 * sizeof(float) || - image0.size() != 224 * 224 * 3 || image1.size() != image0.size() || - noise.size() != 10 * 32 * sizeof(std::uint16_t)) { - return model_error(model, "invalid E2E fixture"); - } - if (model->verbs.set_input(model->self, 0, prompt.data(), prompt.size(), - -1) != 0 || - model->verbs.set_input(model->self, 1, state.data(), state.size(), - -1) != 0) { - return model_error(model, "prompt/state staging failed"); - } - - frt_image_view views[2]{}; - const std::vector* images[] = {&image0, &image1}; - for (int i = 0; i < 2; ++i) { - views[i].struct_size = sizeof(frt_image_view); - views[i].pixel_format = FRT_RT_PIXEL_RGB8; - views[i].data = images[i]->data(); - views[i].bytes = images[i]->size(); - views[i].width = 224; - views[i].height = 224; - views[i].stride_bytes = 224 * 3; - } - if (model->verbs.set_input(model->self, 2, views, sizeof(views), -1) != 0) { - return model_error(model, "image staging failed"); - } - frt_buffer noise_buffer = model->ports[3].buffer; - if (!noise_buffer || frt_buffer_bytes(noise_buffer) != noise.size() || - cudaMemcpy(frt_buffer_dptr(noise_buffer), noise.data(), noise.size(), - cudaMemcpyHostToDevice) != cudaSuccess) { - return model_error(model, "noise upload failed"); - } - if (model->verbs.step(model->self) != 0) { - return model_error(model, "native infer failed"); - } - - std::vector actions(10 * 7); - std::uint64_t written = 0; - if (model->verbs.get_output(model->self, 4, actions.data(), - actions.size() * sizeof(float), &written, - -1) != 0 || - written != actions.size() * sizeof(float)) { - return model_error(model, "action output failed"); - } - std::vector raw(noise.size()); - if (cudaMemcpy(raw.data(), frt_buffer_dptr(model->ports[5].buffer), - raw.size(), cudaMemcpyDeviceToHost) != cudaSuccess) { - return model_error(model, "raw action download failed"); - } - if (!write_file(output + "/native_raw.bf16", raw.data(), raw.size()) || - !write_file(output + "/native_actions.f32", actions.data(), - actions.size() * sizeof(float))) { - return model_error(model, "native output write failed"); - } - model->release(model->owner); - std::cout << "PASS native real-episode fixture\n"; - return 0; -} diff --git a/cpp/tests/pi05_native_encoder_layer_probe.cpp b/cpp/tests/pi05_native_encoder_layer_probe.cpp deleted file mode 100644 index 0aff2e74..00000000 --- a/cpp/tests/pi05_native_encoder_layer_probe.cpp +++ /dev/null @@ -1,133 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_bf16_forward.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" - -#include - -#include -#include -#include -#include - -namespace { - -struct CaptureArgs { - const flashrt::models::pi05::NativeBf16Forward* forward = nullptr; - const flashrt::models::pi05::NativeDeviceWeightStore* weights = nullptr; - flashrt::models::pi05::NativeWorkspace* workspace = nullptr; - flashrt::models::pi05::NativeRtxAttentionWorkspace* attention = nullptr; - const flashrt::models::pi05::NativeRtxAttentionDriver* attention_driver = - nullptr; - bool recorded = false; -}; - -void record_layer(void* user, void* stream) { - auto* args = static_cast(user); - args->recorded = args->forward - ->encoder_layer(0, *args->weights, args->workspace, args->attention, - args->attention_driver, - reinterpret_cast(stream)) - .ok_status(); -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_encoder_layer_probe CHECKPOINT OUTPUT\n"; - return 2; - } - using namespace flashrt::models::pi05; - flashrt::loader::SafetensorsFile source; - if (!source.open(std::string(argv[1]) + "/model.safetensors")) { - std::cerr << source.error() << '\n'; - return 2; - } - frt_ctx ctx = frt_ctx_create(); - if (!ctx) return 1; - NativeDeviceWeightStore weights(ctx); - NativeWeightMaterializer materializer(source, &weights); - flashrt::modalities::Status st = materializer.materialize_encoder_layer(0); - NativeWorkspace workspace(ctx); - NativeRtxAttentionWorkspace attention(ctx); - if (!st.ok_status() || - !workspace.allocate(NativeWorkspaceConfig{}).ok_status() || - !attention.allocate(NativeRtxAttentionConfig{}).ok_status() || - !attention.set_fixed_prompt_length(200).ok_status()) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - const auto* encoder_x = workspace.find("encoder_x"); - std::vector host_x(712 * 2048, 0); - for (int row = 0; row < 712; ++row) { - for (int column = 0; column < 512; ++column) { - const float value = float((row + column) % 15 - 7) / 8.0f; - host_x[static_cast(row) * 2048 + column] = - flashrt::modalities::float_to_bfloat16(value); - } - } - if (!encoder_x || - cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), - host_x.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) != cudaSuccess) { - frt_ctx_destroy(ctx); - return 1; - } - NativeKernelDriver driver; - NativeRtxAttentionDriver attention_driver(&attention); - NativeBf16Forward forward(&driver); - frt_graph graph = frt_graph_create(ctx, "native_encoder_layer", 712); - cudaStream_t stream = nullptr; - if (!graph || cudaStreamCreate(&stream) != cudaSuccess || - frt_graph_bind(graph, "encoder_x", encoder_x->buffer) != FRT_OK) { - frt_ctx_destroy(ctx); - return 1; - } - CaptureArgs capture{&forward, &weights, &workspace, &attention, - &attention_driver, false}; - if (frt_graph_capture(graph, 712, record_layer, &capture) != FRT_OK || - !capture.recorded) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - const int stream_id = frt_ctx_wrap_stream(ctx, stream); - for (int i = 0; i < 100; ++i) { - if (cudaMemcpyAsync(frt_buffer_dptr(encoder_x->buffer), host_x.data(), - host_x.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice, stream) != cudaSuccess || - frt_graph_replay(graph, 712, stream_id) != FRT_OK) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - } - if (frt_graph_variant_count(graph) != 1 || - cudaStreamSynchronize(stream) != cudaSuccess) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - std::vector output(712 * 2048); - if (cudaMemcpy(output.data(), frt_buffer_dptr(encoder_x->buffer), - output.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) != cudaSuccess) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - std::ofstream file(argv[2], std::ios::binary | std::ios::trunc); - file.write(reinterpret_cast(output.data()), - static_cast(output.size() * sizeof(std::uint16_t))); - const bool ok = file.good(); - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - if (!ok) return 1; - std::cout << "PASS native encoder layer 0\n"; - return 0; -} diff --git a/cpp/tests/pi05_native_encoder_probe.cpp b/cpp/tests/pi05_native_encoder_probe.cpp deleted file mode 100644 index c3930e6b..00000000 --- a/cpp/tests/pi05_native_encoder_probe.cpp +++ /dev/null @@ -1,143 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_bf16_forward.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" - -#include - -#include -#include -#include -#include - -namespace { - -struct CaptureArgs { - const flashrt::models::pi05::NativeBf16Forward* forward = nullptr; - const flashrt::models::pi05::NativeDeviceWeightStore* weights = nullptr; - flashrt::models::pi05::NativeWorkspace* workspace = nullptr; - flashrt::models::pi05::NativeRtxAttentionWorkspace* attention = nullptr; - const flashrt::models::pi05::NativeRtxAttentionDriver* attention_driver = - nullptr; - bool recorded = false; -}; - -void record_encoder(void* user, void* stream) { - auto* args = static_cast(user); - args->recorded = args->forward - ->encoder(*args->weights, args->workspace, args->attention, - args->attention_driver, - reinterpret_cast(stream)) - .ok_status(); -} - -bool write_buffer(std::ofstream* file, const void* device, std::size_t elements) { - std::vector host(elements); - if (cudaMemcpy(host.data(), device, elements * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) != cudaSuccess) { - return false; - } - file->write(reinterpret_cast(host.data()), - static_cast(host.size() * - sizeof(std::uint16_t))); - return file->good(); -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_encoder_probe CHECKPOINT OUTPUT\n"; - return 2; - } - using namespace flashrt::models::pi05; - flashrt::loader::SafetensorsFile source; - if (!source.open(std::string(argv[1]) + "/model.safetensors")) { - std::cerr << source.error() << '\n'; - return 2; - } - frt_ctx ctx = frt_ctx_create(); - if (!ctx) return 1; - NativeDeviceWeightStore weights(ctx); - NativeWeightMaterializer materializer(source, &weights); - flashrt::modalities::Status st = flashrt::modalities::Status::ok(); - for (int layer = 0; layer < 18 && st.ok_status(); ++layer) { - st = materializer.materialize_encoder_layer(layer); - } - NativeWorkspace workspace(ctx); - NativeRtxAttentionWorkspace attention(ctx); - if (!st.ok_status() || - !workspace.allocate(NativeWorkspaceConfig{}).ok_status() || - !attention.allocate(NativeRtxAttentionConfig{}).ok_status() || - !attention.set_fixed_prompt_length(200).ok_status()) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - const auto* encoder_x = workspace.find("encoder_x"); - const auto* encoder_q = attention.find("attn_enc_Q"); - std::vector host_x(712 * 2048, 0); - for (int row = 0; row < 712; ++row) { - for (int column = 0; column < 512; ++column) { - const float value = float((row + column) % 15 - 7) / 8.0f; - host_x[static_cast(row) * 2048 + column] = - flashrt::modalities::float_to_bfloat16(value); - } - } - if (!encoder_x || !encoder_q || - cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), - host_x.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) != cudaSuccess) { - frt_ctx_destroy(ctx); - return 1; - } - NativeKernelDriver driver; - NativeRtxAttentionDriver attention_driver(&attention); - NativeBf16Forward forward(&driver); - frt_graph graph = frt_graph_create(ctx, "native_encoder", 712); - cudaStream_t stream = nullptr; - if (!graph || cudaStreamCreate(&stream) != cudaSuccess || - frt_graph_bind(graph, "encoder_x", encoder_x->buffer) != FRT_OK || - frt_graph_bind(graph, "encoder_q", encoder_q->buffer) != FRT_OK) { - frt_ctx_destroy(ctx); - return 1; - } - CaptureArgs capture{&forward, &weights, &workspace, &attention, - &attention_driver, false}; - if (frt_graph_capture(graph, 712, record_encoder, &capture) != FRT_OK || - !capture.recorded) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - const int stream_id = frt_ctx_wrap_stream(ctx, stream); - for (int i = 0; i < 100; ++i) { - if (cudaMemcpyAsync(frt_buffer_dptr(encoder_x->buffer), host_x.data(), - host_x.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice, stream) != cudaSuccess || - frt_graph_replay(graph, 712, stream_id) != FRT_OK) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - } - if (frt_graph_variant_count(graph) != 1 || - cudaStreamSynchronize(stream) != cudaSuccess) { - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - return 1; - } - std::ofstream file(argv[2], std::ios::binary | std::ios::trunc); - const bool ok = file && - write_buffer(&file, frt_buffer_dptr(encoder_x->buffer), 712 * 2048) && - write_buffer(&file, frt_buffer_dptr(encoder_q->buffer), 712 * 2048) && - write_buffer(&file, attention.encoder_k_layer_dptr(17), 712 * 256) && - write_buffer(&file, attention.encoder_v_layer_dptr(17), 712 * 256); - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - if (!ok) return 1; - std::cout << "PASS native encoder 18 layers\n"; - return 0; -} diff --git a/cpp/tests/pi05_native_encoder_qkv_probe.cpp b/cpp/tests/pi05_native_encoder_qkv_probe.cpp deleted file mode 100644 index ce6b22b6..00000000 --- a/cpp/tests/pi05_native_encoder_qkv_probe.cpp +++ /dev/null @@ -1,99 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_bf16_forward.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" - -#include - -#include -#include -#include -#include - -namespace { - -bool write_device(const std::string& path, const void* device, - std::size_t elements) { - std::vector host(elements); - if (cudaMemcpy(host.data(), device, host.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) != cudaSuccess) { - return false; - } - std::ofstream file(path, std::ios::binary | std::ios::trunc); - file.write(reinterpret_cast(host.data()), - static_cast(host.size() * sizeof(std::uint16_t))); - return file.good(); -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_encoder_qkv_probe CHECKPOINT OUTPUT\n"; - return 2; - } - using namespace flashrt::models::pi05; - flashrt::loader::SafetensorsFile source; - if (!source.open(std::string(argv[1]) + "/model.safetensors")) { - std::cerr << source.error() << '\n'; - return 2; - } - frt_ctx ctx = frt_ctx_create(); - if (!ctx) return 1; - NativeDeviceWeightStore weights(ctx); - NativeWeightMaterializer materializer(source, &weights); - flashrt::modalities::Status st = materializer.materialize_encoder_layer(17); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - NativeWorkspace workspace(ctx); - if (!workspace.allocate(NativeWorkspaceConfig{}).ok_status()) { - frt_ctx_destroy(ctx); - return 1; - } - NativeRtxAttentionWorkspace attention(ctx); - if (!attention.allocate(NativeRtxAttentionConfig{}).ok_status()) { - frt_ctx_destroy(ctx); - return 1; - } - const auto* encoder_x = workspace.find("encoder_x"); - if (!encoder_x) { - frt_ctx_destroy(ctx); - return 1; - } - std::vector host_x(712 * 2048, 0); - for (int row = 0; row < 712; ++row) { - for (int column = 0; column < 512; ++column) { - const float value = float((row + column) % 15 - 7) / 8.0f; - host_x[static_cast(row) * 2048 + column] = - flashrt::modalities::float_to_bfloat16(value); - } - } - if (cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), - host_x.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) != cudaSuccess) { - frt_ctx_destroy(ctx); - return 1; - } - NativeKernelDriver driver; - NativeBf16Forward forward(&driver); - st = forward.encoder_qkv(17, weights, &workspace, &attention, 0); - if (!st.ok_status() || cudaDeviceSynchronize() != cudaSuccess) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - const auto* query = attention.find("attn_enc_Q"); - const std::string prefix = argv[2]; - const bool ok = query && - write_device(prefix + ".q.bin", frt_buffer_dptr(query->buffer), - 712 * 2048) && - write_device(prefix + ".k.bin", attention.encoder_k_layer_dptr(17), - 712 * 256) && - write_device(prefix + ".v.bin", attention.encoder_v_layer_dptr(17), - 712 * 256); - frt_ctx_destroy(ctx); - if (!ok) return 1; - std::cout << "PASS native encoder QKV layer 17\n"; - return 0; -} diff --git a/cpp/tests/pi05_native_graph_probe.cpp b/cpp/tests/pi05_native_graph_probe.cpp deleted file mode 100644 index 6a5d828f..00000000 --- a/cpp/tests/pi05_native_graph_probe.cpp +++ /dev/null @@ -1,123 +0,0 @@ -#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"); - const frt_graph infer = owner->graph(NativeGraphKind::kInfer); - const frt_graph decode = owner->graph(NativeGraphKind::kDecodeOnly); - const frt_graph context = owner->graph(NativeGraphKind::kContext); - if (!images || !prompt || !noise || !infer || !decode || !context || - frt_graph_variant_count(infer) != 1 || - frt_graph_variant_count(decode) != 1 || - frt_graph_variant_count(context) != 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; - if (cudaMemcpy(frt_buffer_dptr(noise->buffer), host_noise.data(), - frt_buffer_bytes(noise->buffer), - cudaMemcpyHostToDevice) != cudaSuccess || - owner->replay(NativeGraphKind::kContext) != FRT_OK || - owner->replay(NativeGraphKind::kDecodeOnly) != FRT_OK || - !owner->synchronize().ok_status() || - download(noise->buffer) != expected) { - 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(replay % 2 == 0 - ? NativeGraphKind::kInfer - : NativeGraphKind::kContext) != FRT_OK || - (replay % 2 != 0 && - owner->replay(NativeGraphKind::kDecodeOnly) != FRT_OK)) { - return 1; - } - } - if (!owner->synchronize().ok_status() || - frt_graph_variant_count(infer) != 1 || - frt_graph_variant_count(decode) != 1 || - frt_graph_variant_count(context) != 1 || - owner->workspace().allocation_count() != allocation_count || - download(noise->buffer) != expected) { - return 1; - } - std::cout << "PASS native full/split graphs 100 replays\n"; - return 0; -} diff --git a/cpp/tests/pi05_native_quant_probe.cpp b/cpp/tests/pi05_native_quant_probe.cpp deleted file mode 100644 index da58b579..00000000 --- a/cpp/tests/pi05_native_quant_probe.cpp +++ /dev/null @@ -1,277 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_quantization.h" -#include "flashrt/cpp/models/pi05/native_device_weights.h" -#include "flashrt/cpp/models/pi05/native_kernel_driver.h" -#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" -#include "flashrt/exec.h" - -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace { - -using flashrt::loader::SafetensorsFile; -using flashrt::models::pi05::NativeF16Tensor; -using flashrt::models::pi05::NativeFloatTensor; -using flashrt::models::pi05::NativeFp8Tensor; -using flashrt::models::pi05::NativeInt8Tensor; -using flashrt::models::pi05::NativeSourceTensorView; -using flashrt::modalities::Status; - -constexpr const char* kDecoder = - "paligemma_with_expert.gemma_expert.model.layers.0"; -constexpr const char* kEncoder = - "paligemma_with_expert.paligemma.model.language_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(); -} - -bool encoder_gate_up(const SafetensorsFile& file, NativeF16Tensor* out) { - NativeSourceTensorView gate; - NativeSourceTensorView up; - NativeFloatTensor norm; - return flashrt::models::pi05::load_native_source_tensor( - file, std::string(kEncoder) + ".mlp.gate_proj.weight", &gate) - .ok_status() && - flashrt::models::pi05::load_native_source_tensor( - file, std::string(kEncoder) + ".mlp.up_proj.weight", &up) - .ok_status() && - load(file, std::string(kEncoder) + - ".post_attention_layernorm.weight", &norm) && - flashrt::models::pi05::native_source_pair_to_f16( - gate, up, &norm, false, 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) { - std::cout << std::dec; - 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::uint32_t first_scale_bits = 0; - if (!scales.empty()) { - static_assert(sizeof(first_scale_bits) == sizeof(scales.front())); - std::memcpy(&first_scale_bits, scales.data(), sizeof(first_scale_bits)); - } - 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)) - << " first_scale_bits=" << std::setw(8) - << first_scale_bits << '\n'; -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3 && argc != 4) { - std::cerr << "usage: pi05_native_quant_probe CHECKPOINT OP [OUTPUT]\n"; - return 2; - } - SafetensorsFile file; - if (!file.open(std::string(argv[1]) + "/model.safetensors")) { - std::cerr << file.error() << '\n'; - return 2; - } - const std::string op = argv[2]; - if (op == "decoder_qkv0_fp8_kn_gpu" || - op == "vision_attn_qkv0_fp8_kn_gpu" || - op == "vision_ffn_down4_fp8_kn_gpu") { - frt_ctx ctx = frt_ctx_create(); - if (!ctx) { - std::cerr << "failed to create FlashRT context\n"; - return 1; - } - std::vector shape; - std::vector values; - std::vector scales; - std::uint64_t input_hash = 0; - std::string error; - { - flashrt::models::pi05::NativeDeviceWeightStore weights(ctx); - flashrt::models::pi05::NativeWeightMaterializer materializer( - file, &weights); - const bool vision_qkv = op.find("vision_attn") == 0; - const bool vision_down = op.find("vision_ffn") == 0; - const std::string name = - vision_qkv ? "vision_attn_qkv_w_0" : - vision_down ? "vision_ffn_down_w_4" : - "decoder_attn_qkv_w_0"; - Status st = vision_qkv - ? materializer.materialize_vision_layer(0) - : vision_down - ? materializer.materialize_vision_layer(4) - : materializer.materialize_decoder_layer(0, true); - flashrt::models::pi05::NativeKernelDriver driver; - if (st.ok_status()) st = driver.status(); - flashrt::models::pi05::NativeRtxWeightPacker packer( - &weights, &driver); - if (st.ok_status()) st = packer.pack_weight(name); - if (!st.ok_status()) { - error = st.message; - } else { - const auto* input = weights.find(name); - const auto* output = weights.find("fp8." + name); - const auto* scale = weights.find("fp8." + name + ".scale"); - if (!input || !output || !scale) { - error = "GPU FP8 weight pack output is missing"; - } else { - std::vector input_values( - frt_buffer_bytes(input->buffer)); - shape = output->shape; - values.resize(frt_buffer_bytes(output->buffer)); - scales.resize(1); - if (cudaMemcpy( - input_values.data(), frt_buffer_dptr(input->buffer), - input_values.size(), cudaMemcpyDeviceToHost) != - cudaSuccess || - cudaMemcpy( - values.data(), frt_buffer_dptr(output->buffer), - values.size(), cudaMemcpyDeviceToHost) != - cudaSuccess || - cudaMemcpy( - scales.data(), frt_buffer_dptr(scale->buffer), - sizeof(float), cudaMemcpyDeviceToHost) != - cudaSuccess) { - error = "GPU FP8 weight download failed"; - } else { - input_hash = fnv1a( - input_values.data(), input_values.size()); - } - } - } - } - frt_ctx_destroy(ctx); - if (!error.empty()) { - std::cerr << error << '\n'; - return 1; - } - std::cout << "input_fnv=" << std::hex << std::setw(16) - << std::setfill('0') << input_hash << ' '; - if (argc == 4) { - std::ofstream output(argv[3], std::ios::binary | std::ios::trunc); - output.write( - reinterpret_cast(values.data()), - static_cast(values.size())); - if (!output) return 1; - } - print_result( - shape, values.data(), values.size(), scales); - return 0; - } - if (op == "encoder_gate_up0_fp8") { - NativeF16Tensor weight; - NativeFp8Tensor output; - if (!encoder_gate_up(file, &weight)) return 1; - std::cout << "input_fnv=" << std::hex << std::setw(16) - << std::setfill('0') - << fnv1a(weight.values.data(), - weight.values.size() * sizeof(std::uint16_t)) - << ' '; - const Status st = flashrt::models::pi05::native_quantize_fp8_e4m3( - weight, false, &output); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - return 1; - } - if (argc == 4) { - std::ofstream file(argv[3], std::ios::binary | std::ios::trunc); - file.write(reinterpret_cast(output.values.data()), - static_cast(output.values.size())); - if (!file) return 1; - } - print_result(output.shape, output.values.data(), output.values.size(), - {output.scale}); - return 0; - } - NativeFloatTensor weight; - if (!decoder_qkv(file, &weight)) return 1; - if (op == "decoder_qkv0_fp8_kn" || op == "decoder_qkv0_fp8_nk") { - NativeFp8Tensor output; - const bool transpose = op.back() == 'k'; - const Status st = flashrt::models::pi05::native_quantize_fp8_e4m3( - weight, transpose, &output); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - return 1; - } - print_result(output.shape, output.values.data(), output.values.size(), - {output.scale}); - return 0; - } - if (op == "decoder_qkv0_int8") { - NativeInt8Tensor output; - const Status st = - flashrt::models::pi05::native_quantize_int8_per_output( - weight, &output); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - return 1; - } - print_result(output.shape, output.values.data(), output.values.size(), - output.scales); - return 0; - } - std::cerr << "unknown quantization probe operation: " << op << '\n'; - return 2; -} diff --git a/cpp/tests/pi05_native_rope_probe.cpp b/cpp/tests/pi05_native_rope_probe.cpp deleted file mode 100644 index 4e94b581..00000000 --- a/cpp/tests/pi05_native_rope_probe.cpp +++ /dev/null @@ -1,90 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_workspace.h" - -#include - -#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 || argc > 8) { - std::cerr << "usage: pi05_native_rope_probe VIEWS MAX_PROMPT CHUNK " - "POOL PROMPT [thor_fp8 [OUTPUT_PREFIX]]\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]); - if (argc >= 7) { - if (std::string(argv[6]) != "thor_fp8") return 2; - config.flavor = - flashrt::models::pi05::NativeWorkspaceFlavor::kThorFp8; - } - 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() || - !(argc >= 7 ? workspace.set_fixed_prompt_length(prompt) - : 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")); - if (argc == 8) { - std::ofstream encoder_file( - std::string(argv[7]) + ".encoder", std::ios::binary); - std::ofstream decoder_file( - std::string(argv[7]) + ".decoder", std::ios::binary); - encoder_file.write(reinterpret_cast(encoder.data()), - encoder.size() * sizeof(encoder[0])); - decoder_file.write(reinterpret_cast(decoder.data()), - decoder.size() * sizeof(decoder[0])); - if (!encoder_file || !decoder_file) { - frt_ctx_destroy(ctx); - return 1; - } - } - std::cout << "encoder_shape=" << workspace.encoder_sequence() << ",256" - << " encoder_fnv=" << std::hex << std::setw(16) - << std::setfill('0') << fnv1a(encoder) - << " decoder_shape=" << std::dec << config.chunk_size << ",256" - << " decoder_fnv=" << std::hex << std::setw(16) - << fnv1a(decoder) << '\n'; - frt_ctx_destroy(ctx); - return 0; -} diff --git a/cpp/tests/pi05_native_style_probe.cpp b/cpp/tests/pi05_native_style_probe.cpp deleted file mode 100644 index 0848ebfc..00000000 --- a/cpp/tests/pi05_native_style_probe.cpp +++ /dev/null @@ -1,85 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_style_precompute.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" - -#include - -#include -#include -#include - -namespace { - -bool write_buffer(const std::string& path, - const flashrt::models::pi05::NativeWorkspaceBuffer& buffer) { - const std::size_t bytes = frt_buffer_bytes(buffer.buffer); - std::vector host(bytes); - if (cudaMemcpy(host.data(), frt_buffer_dptr(buffer.buffer), bytes, - cudaMemcpyDeviceToHost) != cudaSuccess) { - return false; - } - std::ofstream file(path, std::ios::binary | std::ios::trunc); - file.write(reinterpret_cast(host.data()), - static_cast(host.size())); - return file.good(); -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_style_probe CHECKPOINT OUTPUT_PREFIX\n"; - return 2; - } - flashrt::loader::SafetensorsFile source; - if (!source.open(std::string(argv[1]) + "/model.safetensors")) { - std::cerr << source.error() << '\n'; - return 2; - } - frt_ctx ctx = frt_ctx_create(); - if (!ctx) return 1; - flashrt::models::pi05::NativeDeviceWeightStore weights(ctx); - flashrt::models::pi05::NativeWeightMaterializer materializer(source, - &weights); - for (int layer = 0; layer < 18; ++layer) { - const flashrt::modalities::Status st = - materializer.materialize_decoder_layer(layer, false); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - } - flashrt::modalities::Status st = - materializer.materialize_decoder_globals(10); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - flashrt::models::pi05::NativeWorkspace workspace(ctx); - flashrt::models::pi05::NativeWorkspaceConfig config; - if (!workspace.allocate(config).ok_status()) { - frt_ctx_destroy(ctx); - return 1; - } - flashrt::models::pi05::NativeKernelDriver driver; - flashrt::models::pi05::NativeStylePrecomputer precomputer(&driver); - st = precomputer.run(weights, &workspace, 0); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - frt_ctx_destroy(ctx); - return 1; - } - const std::string prefix = argv[2]; - for (const char* name : {"decoder_time_emb", "decoder_style_attn", - "decoder_style_ffn", "decoder_style_final"}) { - const auto* buffer = workspace.find(name); - if (!buffer || !write_buffer(prefix + "." + name + ".bin", *buffer)) { - frt_ctx_destroy(ctx); - return 1; - } - } - std::cout << "PASS native decoder style precompute\n"; - frt_ctx_destroy(ctx); - return 0; -} diff --git a/cpp/tests/pi05_native_vision_probe.cpp b/cpp/tests/pi05_native_vision_probe.cpp deleted file mode 100644 index 83a3b98b..00000000 --- a/cpp/tests/pi05_native_vision_probe.cpp +++ /dev/null @@ -1,415 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_bf16_forward.h" -#include "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" - -#include - -#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(); -} - -bool download_buffer( - const void* device, - std::size_t elements, - std::vector* host) { - if (!device || !host) return false; - host->resize(elements); - return cudaMemcpy( - host->data(), device, elements * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess; -} - -bool upload_scales( - flashrt::models::pi05::NativeWorkspace* workspace, - const std::vector& scales) { - const auto* output = workspace->find("rtx_fp8_vision_scales"); - return output && output->shape == - std::vector({scales.size()}) && - cudaMemcpy(frt_buffer_dptr(output->buffer), scales.data(), - scales.size() * sizeof(float), - cudaMemcpyHostToDevice) == cudaSuccess; -} - -bool read_images( - const char* path, - std::vector* images) { - if (!path || !images) return false; - std::ifstream file(path, std::ios::binary | std::ios::ate); - if (!file || - file.tellg() != static_cast( - images->size() * sizeof(std::uint16_t))) { - return false; - } - file.seekg(0); - file.read( - reinterpret_cast(images->data()), - static_cast( - images->size() * sizeof(std::uint16_t))); - return file.good(); -} - -flashrt::modalities::Status write_vision_layer_diagnostics( - const char* path, - const flashrt::models::pi05::NativeKernelDriver& driver, - const flashrt::models::pi05::NativeBf16Forward& forward, - const flashrt::models::pi05::NativeDeviceWeightStore& weights, - flashrt::models::pi05::NativeWorkspace* workspace, - flashrt::models::pi05::NativeRtxAttentionWorkspace* attention, - const flashrt::models::pi05::NativeRtxAttentionDriver& attention_driver, - cudaStream_t stream) { - using flashrt::models::pi05::NativeWorkspaceBuffer; - 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 auto* patch_weight = weights.find("vision_patch_embedding_w"); - const auto* patch_bias = weights.find("vision_patch_embedding_b"); - const auto* norm_weight = weights.find("vision_pre_attn_norm_w_0"); - const auto* norm_bias = weights.find("vision_pre_attn_norm_b_0"); - if (!images || !patches || !position || !x || !x_norm || - !patch_weight || !patch_bias || !norm_weight || !norm_bias) { - return flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kInvalidArgument, - "native vision diagnostic buffers are incomplete"); - } - const auto ptr = [](const auto* value) { - return frt_buffer_dptr(value->buffer); - }; - const int sequence = workspace->vision_sequence(); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - flashrt::modalities::Status st = driver.patch_im2col_16bit( - ptr(images), ptr(patches), workspace->num_views(), native_stream); - if (st.ok_status()) { - st = driver.bf16_nn( - ptr(patches), ptr(patch_weight), ptr(x), sequence, 1152, 588, - native_stream); - } - if (st.ok_status()) { - st = driver.bias_residual_bf16( - ptr(x), ptr(position), ptr(patch_bias), sequence, 1152, - native_stream); - } - if (st.ok_status()) { - st = driver.layer_norm_bf16( - ptr(x), ptr(norm_weight), ptr(norm_bias), ptr(x_norm), sequence, - 1152, 1.0e-5f, native_stream); - } - if (!st.ok_status()) return st; - - int diagnostic_layer = 0; - if (const char* value = - std::getenv("FLASHRT_VISION_DIAGNOSTIC_LAYER")) { - char* end = nullptr; - const long parsed = std::strtol(value, &end, 10); - if (!end || *end != '\0' || parsed < 0 || parsed >= 27) { - return flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kInvalidArgument, - "native vision diagnostic layer is invalid"); - } - diagnostic_layer = static_cast(parsed); - } - - std::ofstream output(path, std::ios::binary | std::ios::trunc); - const std::size_t elements = - static_cast(sequence) * 1152; - if (!output || cudaStreamSynchronize(stream) != cudaSuccess || - !write_buffer(&output, ptr(x), elements)) { - return flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kBackend, - "native vision diagnostic output failed"); - } - std::vector layer_qkv; - std::vector layer_attention; - std::vector layer_hidden; - for (int layer = 0; layer < 27; ++layer) { - st = forward.vision_layer( - layer, weights, workspace, attention, &attention_driver, - native_stream); - if (!st.ok_status()) return st; - if (cudaStreamSynchronize(stream) != cudaSuccess || - !write_buffer(&output, ptr(x), elements)) { - return flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kBackend, - "native vision layer diagnostic output failed"); - } - if (layer == diagnostic_layer) { - const NativeWorkspaceBuffer* qkv = - workspace->find("vision_QKV"); - const NativeWorkspaceBuffer* hidden = - workspace->find("vision_hidden"); - const auto* attention_output = - attention->find("attn_vis_O"); - if (!qkv || !hidden || !attention_output || - !download_buffer( - ptr(qkv), static_cast(sequence) * 3456, - &layer_qkv) || - !download_buffer( - frt_buffer_dptr(attention_output->buffer), - static_cast(sequence) * 1152, - &layer_attention) || - !download_buffer( - ptr(hidden), static_cast(sequence) * 4304, - &layer_hidden)) { - return flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kBackend, - "native vision sublayer diagnostics failed"); - } - } - } - for (const auto* values : { - &layer_qkv, - &layer_attention, - &layer_hidden}) { - output.write( - reinterpret_cast(values->data()), - static_cast( - values->size() * sizeof(std::uint16_t))); - if (!output) { - return flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kBackend, - "native vision diagnostic write failed"); - } - } - return flashrt::modalities::Status::ok(); -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3 && argc != 5) { - std::cerr << "usage: pi05_native_vision_probe CHECKPOINT OUTPUT " - "[CALIBRATION INPUT_BF16]\n"; - return 2; - } - using namespace flashrt::models::pi05; - const bool fp8 = argc == 5; - NativeCalibrationArtifact calibration; - if (fp8) { - const flashrt::modalities::Status calibration_status = - load_native_calibration_artifact(argv[3], &calibration); - if (!calibration_status.ok_status() || - calibration.activation_dtype != "bfloat16" || - calibration.hardware != "sm120") { - std::cerr << (calibration_status.ok_status() - ? "SM120 FP8 calibration is required" - : calibration_status.message) - << '\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; - 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); - } - NativeKernelDriver driver; - if (st.ok_status()) st = driver.status(); - if (fp8 && st.ok_status()) { - NativeRtxWeightPacker packer(&weights, &driver); - for (int layer = 0; layer < 27 && st.ok_status(); ++layer) { - for (const char* stem : { - "vision_attn_qkv_w_", "vision_attn_o_w_", - "vision_ffn_up_w_", "vision_ffn_down_w_"}) { - st = packer.pack_weight( - std::string(stem) + std::to_string(layer)); - if (!st.ok_status()) break; - } - } - if (st.ok_status()) { - st = packer.pack_weight( - "encoder_multi_modal_projector_w", "vision_projector_w"); - } - } - NativeWorkspace workspace(ctx); - NativeRtxAttentionWorkspace attention(ctx); - NativeWorkspaceConfig workspace_config; - if (fp8) { - workspace_config.num_views = calibration.num_views; - workspace_config.max_prompt_tokens = calibration.max_prompt_tokens; - workspace_config.chunk_size = calibration.chunk_size; - workspace_config.num_steps = calibration.num_steps; - workspace_config.vision_pool_factor = - calibration.vision_pool_factor; - workspace_config.flavor = NativeWorkspaceFlavor::kRtxFp8; - } - NativeRtxAttentionConfig attention_config; - if (st.ok_status()) st = workspace.allocate(workspace_config); - if (st.ok_status()) { - st = workspace.expand_vision_position_embedding(weights); - } - attention_config.num_views = workspace.num_views(); - attention_config.encoder_sequence = workspace.encoder_sequence(); - attention_config.encoder_vision_sequence = - workspace.encoder_vision_sequence(); - attention_config.chunk_size = workspace.chunk_size(); - if (st.ok_status()) st = attention.allocate(attention_config); - if (st.ok_status() && fp8 && - !upload_scales(&workspace, calibration.vision_scales)) { - st = flashrt::modalities::Status::error( - flashrt::modalities::StatusCode::kBackend, - "native vision scale upload failed"); - } - if (!st.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( - static_cast(workspace.num_views()) * 224 * 224 * 3); - if (fp8) { - if (!read_images(argv[4], &host_images)) { - std::cerr << "native vision input is invalid\n"; - frt_ctx_destroy(ctx); - return 1; - } - } else { - 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; - } - NativeRtxAttentionDriver attention_driver(&attention); - NativeRtxLinear linear( - &driver, fp8 ? NativeRtxLinearMode::kFp8Static - : NativeRtxLinearMode::kBf16); - NativeBf16Forward forward(&driver, &linear); - const int vision_sequence = workspace.vision_sequence(); - if (fp8 && std::getenv("FLASHRT_VISION_LAYER_DIAGNOSTICS")) { - cudaStream_t diagnostic_stream = nullptr; - if (cudaStreamCreate(&diagnostic_stream) != cudaSuccess) { - frt_ctx_destroy(ctx); - return 1; - } - st = write_vision_layer_diagnostics( - argv[2], driver, forward, weights, &workspace, &attention, - attention_driver, diagnostic_stream); - cudaStreamDestroy(diagnostic_stream); - frt_ctx_destroy(ctx); - if (!st.ok_status()) { - std::cerr << st.message << '\n'; - return 1; - } - std::cout << "PASS native vision layer diagnostics\n"; - return 0; - } - frt_graph graph = frt_graph_create( - ctx, "native_vision", vision_sequence); - 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, vision_sequence, 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, vision_sequence, 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), - static_cast(vision_sequence) * 1152) && - write_buffer( - &file, frt_buffer_dptr(encoder_x->buffer), - static_cast(workspace.encoder_vision_sequence()) * - 2048); - frt_graph_destroy(graph); - cudaStreamDestroy(stream); - frt_ctx_destroy(ctx); - if (!ok) return 1; - std::cout << "PASS native vision 27 layers\n"; - return 0; -} diff --git a/cpp/tests/pi05_native_weight_probe.cpp b/cpp/tests/pi05_native_weight_probe.cpp deleted file mode 100644 index 872a041d..00000000 --- a/cpp/tests/pi05_native_weight_probe.cpp +++ /dev/null @@ -1,200 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_weight_ops.h" - -#include -#include -#include -#include -#include -#include - -namespace { - -using flashrt::loader::SafetensorsFile; -using flashrt::models::pi05::NativeBf16Tensor; -using flashrt::models::pi05::NativeFloatTensor; -using flashrt::models::pi05::NativeSourceTensorView; -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 source_view(const SafetensorsFile& file, const std::string& key, - NativeSourceTensorView* out) { - const Status st = - flashrt::models::pi05::load_native_source_tensor(file, key, out); - if (!st.ok_status()) std::cerr << st.message << '\n'; - return st.ok_status(); -} - -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 finish(const NativeFloatTensor& input, NativeBf16Tensor* out) { - return flashrt::models::pi05::native_to_bf16(input, out).ok_status(); -} - -bool patch(const SafetensorsFile& file, NativeBf16Tensor* out) { - NativeSourceTensorView source; - return source_view(file, std::string(kVision) + - ".embeddings.patch_embedding.weight", - &source) && - flashrt::models::pi05::native_source_patch_oihw_to_hwio_bf16( - source, out).ok_status(); -} - -bool qkv(const SafetensorsFile& file, const std::string& prefix, - bool fold_rms, NativeBf16Tensor* out) { - NativeSourceTensorView q; - NativeSourceTensorView k; - NativeSourceTensorView v; - if (!source_view(file, prefix + ".self_attn.q_proj.weight", &q) || - !source_view(file, prefix + ".self_attn.k_proj.weight", &k) || - !source_view(file, prefix + ".self_attn.v_proj.weight", &v)) { - return false; - } - NativeFloatTensor norm; - const NativeFloatTensor* norm_ptr = nullptr; - if (fold_rms) { - if (!load(file, prefix + ".input_layernorm.weight", &norm)) return false; - norm_ptr = &norm; - } - return flashrt::models::pi05::native_source_qkv_to_bf16( - q, k, v, 8, 1, norm_ptr, out).ok_status(); -} - -bool gate_up(const SafetensorsFile& file, NativeBf16Tensor* out) { - NativeSourceTensorView gate; - NativeSourceTensorView up; - return source_view(file, std::string(kDecoder) + - ".mlp.gate_proj.weight", &gate) && - source_view(file, std::string(kDecoder) + - ".mlp.up_proj.weight", &up) && - flashrt::models::pi05::native_source_pair_transpose_concat_bf16( - gate, up, out).ok_status(); -} - -bool action_out(const SafetensorsFile& file, int num_steps, - NativeBf16Tensor* out) { - NativeSourceTensorView source; - return source_view(file, "action_out_proj.weight", &source) && - flashrt::models::pi05::native_source_round_scale_to_bf16( - source, -1.0f / static_cast(num_steps), - true, out).ok_status(); -} - -bool rounded_transpose(const SafetensorsFile& file, - const std::string& key, - NativeBf16Tensor* out) { - NativeSourceTensorView source; - return source_view(file, key, &source) && - flashrt::models::pi05::native_source_to_bf16(source, true, out) - .ok_status(); -} - -bool rounded_copy(const SafetensorsFile& file, - const std::string& key, - NativeBf16Tensor* out) { - NativeSourceTensorView source; - return source_view(file, key, &source) && - flashrt::models::pi05::native_source_to_bf16(source, false, out) - .ok_status(); -} - -bool folded_transpose(const SafetensorsFile& file, - const std::string& key, - const std::string& norm_key, - NativeBf16Tensor* out) { - NativeSourceTensorView source; - NativeFloatTensor norm; - return source_view(file, key, &source) && load(file, norm_key, &norm) && - flashrt::models::pi05::native_source_fold_rms_columns_transpose( - source, norm, out) - .ok_status(); -} - -bool time_embeds(int num_steps, NativeBf16Tensor* out) { - NativeFloatTensor generated; - return flashrt::models::pi05::native_pi05_time_embeddings(num_steps, 1024, - &generated) - .ok_status() && - finish(generated, out); -} - -std::uint64_t fnv1a(const std::vector& values) { - std::uint64_t hash = 14695981039346656037ull; - const auto* bytes = reinterpret_cast(values.data()); - for (std::size_t i = 0; i < values.size() * sizeof(std::uint16_t); ++i) { - hash ^= bytes[i]; - hash *= 1099511628211ull; - } - return hash; -} - -} // namespace - -int main(int argc, char** argv) { - if (argc != 3) { - std::cerr << "usage: pi05_native_weight_probe CHECKPOINT OP\n"; - return 2; - } - SafetensorsFile file; - if (!file.open(std::string(argv[1]) + "/model.safetensors")) { - std::cerr << file.error() << '\n'; - return 2; - } - NativeBf16Tensor output; - const std::string op = argv[2]; - bool ok = false; - if (op == "patch") { - ok = patch(file, &output); - } else if (op == "encoder_qkv0") { - ok = qkv(file, kEncoder, true, &output); - } else if (op == "decoder_qkv0") { - ok = qkv(file, kDecoder, false, &output); - } else if (op == "decoder_gate_up0") { - ok = gate_up(file, &output); - } else if (op == "encoder_o0_fast") { - ok = rounded_transpose( - file, std::string(kEncoder) + ".self_attn.o_proj.weight", - &output); - } else if (op == "encoder_gate0_fast") { - ok = folded_transpose( - file, std::string(kEncoder) + ".mlp.gate_proj.weight", - std::string(kEncoder) + ".post_attention_layernorm.weight", - &output); - } else if (op == "decoder_mod_bias0_fast") { - ok = rounded_copy( - file, std::string(kDecoder) + - ".input_layernorm.dense.bias", - &output); - } else if (op == "action_out10") { - ok = action_out(file, 10, &output); - } else if (op == "action_out5") { - ok = action_out(file, 5, &output); - } else if (op == "time_embeds10") { - ok = time_embeds(10, &output); - } else if (op == "time_embeds5") { - ok = time_embeds(5, &output); - } - if (!ok) { - std::cerr << "weight probe operation failed: " << op << '\n'; - return 1; - } - std::cout << "shape="; - for (std::size_t i = 0; i < output.shape.size(); ++i) { - if (i) std::cout << ','; - std::cout << output.shape[i]; - } - std::cout << " fnv=" << std::hex << std::setw(16) << std::setfill('0') - << fnv1a(output.values) << '\n'; - return 0; -} diff --git a/cpp/tests/pi05_tokenizer_corpus_probe.cpp b/cpp/tests/pi05_tokenizer_corpus_probe.cpp deleted file mode 100644 index d650a919..00000000 --- a/cpp/tests/pi05_tokenizer_corpus_probe.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#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/test_pi05_native_bf16_forward.cpp b/cpp/tests/test_pi05_native_bf16_forward.cpp deleted file mode 100644 index a37e22ce..00000000 --- a/cpp/tests/test_pi05_native_bf16_forward.cpp +++ /dev/null @@ -1,192 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_bf16_forward.h" -#include "flashrt/cpp/modalities/types.h" -#include "flashrt/exec.h" - -#include - -#include -#include -#include -#include - -namespace { - -using flashrt::models::pi05::NativeBf16Forward; -using flashrt::models::pi05::NativeDeviceWeightStore; -using flashrt::models::pi05::NativeKernelDriver; -using flashrt::models::pi05::NativeRtxAttentionWorkspace; -using flashrt::models::pi05::NativeWorkspace; - -std::vector download(const void* device, std::size_t elements) { - std::vector result(elements); - assert(cudaMemcpy(result.data(), device, - result.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - return result; -} - -struct CaptureArgs { - const NativeBf16Forward* forward = nullptr; - const NativeDeviceWeightStore* weights = nullptr; - NativeWorkspace* workspace = nullptr; - NativeRtxAttentionWorkspace* attention = nullptr; - bool recorded = false; -}; - -void record_encoder_qkv(void* user, void* stream) { - auto* args = static_cast(user); - args->recorded = args->forward - ->encoder_qkv(17, *args->weights, args->workspace, args->attention, - reinterpret_cast(stream)) - .ok_status(); -} - -} // namespace - -int main() { - int count = 0; - if (cudaGetDeviceCount(&count) != cudaSuccess || !count) { - cudaGetLastError(); - std::printf("SKIP - no CUDA device\n"); - return 0; - } - using namespace flashrt::models::pi05; - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - NativeWorkspace workspace(ctx); - NativeWorkspaceConfig workspace_config; - workspace_config.num_views = 1; - workspace_config.max_prompt_tokens = 1; - workspace_config.chunk_size = 2; - workspace_config.num_steps = 2; - workspace_config.vision_pool_factor = 4; - assert(workspace.allocate(workspace_config).ok_status()); - assert(workspace.encoder_sequence() == 17); - - NativeRtxAttentionWorkspace attention(ctx); - NativeRtxAttentionConfig attention_config; - attention_config.num_views = 1; - attention_config.encoder_sequence = 17; - attention_config.encoder_vision_sequence = 16; - attention_config.chunk_size = 2; - assert(attention.allocate(attention_config).ok_status()); - - NativeKernelDriver driver; - NativeBf16Forward forward(&driver); - NativeDeviceWeightStore weights(ctx); - assert(!forward.encoder_qkv(17, weights, &workspace, &attention, 0) - .ok_status()); - - NativeBf16Tensor qkv_weight; - qkv_weight.shape = {2048, 2560}; - qkv_weight.values.assign(2048 * 2560, 0); - const std::uint16_t one = - flashrt::modalities::float_to_bfloat16(1.0f); - for (int column = 0; column < 2048; ++column) { - qkv_weight.values[static_cast(column) * 2560 + column] = - one; - } - for (int column = 0; column < 256; ++column) { - qkv_weight.values[static_cast(column) * 2560 + - 2048 + column] = one; - qkv_weight.values[static_cast(256 + column) * 2560 + - 2304 + column] = one; - } - assert(weights.upload("encoder_attn_qkv_w_17", qkv_weight).ok_status()); - - const auto* encoder_x = workspace.find("encoder_x"); - assert(encoder_x); - std::vector host_x(17 * 2048, 0); - for (int row = 0; row < 17; ++row) { - for (int column = 0; column < 512; ++column) { - const float value = float((row + column) % 15 - 7) / 8.0f; - host_x[static_cast(row) * 2048 + column] = - flashrt::modalities::float_to_bfloat16(value); - } - } - assert(cudaMemcpy(frt_buffer_dptr(encoder_x->buffer), host_x.data(), - host_x.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); - - cudaStream_t stream = nullptr; - assert(cudaStreamCreate(&stream) == cudaSuccess); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - assert(forward.encoder_qkv(17, weights, &workspace, &attention, - native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - const auto* query_buffer = attention.find("attn_enc_Q"); - assert(query_buffer); - const std::vector expected_q = download( - frt_buffer_dptr(query_buffer->buffer), 17 * 2048); - const std::vector expected_k = - download(attention.encoder_k_layer_dptr(17), 17 * 256); - const std::vector expected_v = - download(attention.encoder_v_layer_dptr(17), 17 * 256); - - assert(cudaMemset(frt_buffer_dptr(query_buffer->buffer), 0, - expected_q.size() * sizeof(std::uint16_t)) == - cudaSuccess); - assert(cudaMemset(attention.encoder_k_layer_dptr(17), 0, - expected_k.size() * sizeof(std::uint16_t)) == - cudaSuccess); - assert(cudaMemset(attention.encoder_v_layer_dptr(17), 0, - expected_v.size() * sizeof(std::uint16_t)) == - cudaSuccess); - const auto* x_norm = workspace.find("encoder_x_norm"); - const auto* qkv = workspace.find("encoder_QKV"); - const auto* rms = workspace.find("encoder_rms_ones"); - const auto* rope = workspace.find("encoder_rope_weights"); - const auto* weight = weights.find("encoder_attn_qkv_w_17"); - assert(x_norm && qkv && rms && rope && weight); - assert(driver.rms_norm_bf16( - frt_buffer_dptr(encoder_x->buffer), - frt_buffer_dptr(rms->buffer), - frt_buffer_dptr(x_norm->buffer), 17, 2048, 1e-6f, - native_stream) - .ok_status()); - assert(driver.bf16_nn( - frt_buffer_dptr(x_norm->buffer), - frt_buffer_dptr(weight->buffer), - frt_buffer_dptr(qkv->buffer), 17, 2560, 2048, - native_stream) - .ok_status()); - assert(driver.qkv_split_rope_bf16( - frt_buffer_dptr(qkv->buffer), frt_buffer_dptr(rope->buffer), - frt_buffer_dptr(query_buffer->buffer), - attention.encoder_k_layer_dptr(17), - attention.encoder_v_layer_dptr(17), 17, 2048, 256, 256, - 256, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - assert(download(frt_buffer_dptr(query_buffer->buffer), 17 * 2048) == - expected_q); - assert(download(attention.encoder_k_layer_dptr(17), 17 * 256) == - expected_k); - assert(download(attention.encoder_v_layer_dptr(17), 17 * 256) == - expected_v); - - frt_graph graph = frt_graph_create(ctx, "native_encoder_qkv", 17); - assert(graph); - assert(frt_graph_bind(graph, "encoder_x", encoder_x->buffer) == FRT_OK); - assert(frt_graph_bind(graph, "encoder_q", query_buffer->buffer) == FRT_OK); - CaptureArgs capture{&forward, &weights, &workspace, &attention, false}; - assert(frt_graph_capture(graph, 17, record_encoder_qkv, &capture) == FRT_OK); - assert(capture.recorded); - const int stream_id = frt_ctx_wrap_stream(ctx, stream); - assert(stream_id >= 0); - for (int i = 0; i < 100; ++i) { - assert(frt_graph_replay(graph, 17, stream_id) == FRT_OK); - } - assert(frt_graph_variant_count(graph) == 1); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - assert(download(attention.encoder_k_layer_dptr(17), 17 * 256) == - expected_k); - - frt_graph_destroy(graph); - assert(cudaStreamDestroy(stream) == cudaSuccess); - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native BF16 encoder QKV\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_device_weights.cpp b/cpp/tests/test_pi05_native_device_weights.cpp deleted file mode 100644 index 0c884818..00000000 --- a/cpp/tests/test_pi05_native_device_weights.cpp +++ /dev/null @@ -1,97 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_device_weights.h" - -#include - -#include -#include -#include - -namespace { - -bool has_cuda_device() { - int count = 0; - const cudaError_t rc = cudaGetDeviceCount(&count); - if (rc != cudaSuccess) { - cudaGetLastError(); - return false; - } - return count > 0; -} - -} // namespace - -int main() { - if (!has_cuda_device()) { - std::printf("SKIP - no CUDA device\n"); - return 0; - } - using flashrt::models::pi05::NativeBf16Tensor; - using flashrt::models::pi05::NativeDeviceWeightStore; - using flashrt::models::pi05::NativeWeightDType; - - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - { - NativeDeviceWeightStore store(ctx); - NativeBf16Tensor tensor; - tensor.shape = {2, 3}; - tensor.values = { - flashrt::modalities::float_to_bfloat16(1.0f), - flashrt::modalities::float_to_bfloat16(2.0f), - flashrt::modalities::float_to_bfloat16(3.0f), - flashrt::modalities::float_to_bfloat16(4.0f), - flashrt::modalities::float_to_bfloat16(5.0f), - flashrt::modalities::float_to_bfloat16(6.0f), - }; - assert(store.upload("encoder.layer0.qkv", tensor).ok_status()); - assert(store.size() == 1); - const auto* weight = store.find("encoder.layer0.qkv"); - assert(weight && weight->buffer); - assert(weight->shape == tensor.shape); - assert(weight->dtype == NativeWeightDType::kBf16); - assert(frt_buffer_bytes(weight->buffer) == - tensor.values.size() * sizeof(std::uint16_t)); - assert(std::string(frt_buffer_name(weight->buffer)) == - "encoder.layer0.qkv"); - - std::vector copied(tensor.values.size()); - assert(cudaMemcpy(copied.data(), frt_buffer_dptr(weight->buffer), - copied.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - assert(copied == tensor.values); - - const std::vector int8_values = {-127, -1, 0, 127}; - assert(store.upload_bytes("encoder.layer0.qkv.int8", {2, 2}, - NativeWeightDType::kInt8, - int8_values.data(), int8_values.size()) - .ok_status()); - const auto* int8_weight = store.find("encoder.layer0.qkv.int8"); - assert(int8_weight && int8_weight->dtype == NativeWeightDType::kInt8); - std::vector int8_copied(int8_values.size()); - assert(cudaMemcpy(int8_copied.data(), - frt_buffer_dptr(int8_weight->buffer), - int8_copied.size(), cudaMemcpyDeviceToHost) == - cudaSuccess); - assert(int8_copied == int8_values); - - const std::vector scales = {0.25f, 0.5f}; - assert(store.upload_bytes("encoder.layer0.qkv.scale", {2}, - NativeWeightDType::kFloat32, - scales.data(), - scales.size() * sizeof(float)) - .ok_status()); - assert(store.find("encoder.layer0.qkv.scale")->dtype == - NativeWeightDType::kFloat32); - assert(!store.upload_bytes("bad.bytes", {3}, - NativeWeightDType::kFp8E4M3, - int8_values.data(), int8_values.size()) - .ok_status()); - assert(!store.upload("encoder.layer0.qkv", tensor).ok_status()); - tensor.shape = {3, 3}; - assert(!store.upload("bad", tensor).ok_status()); - } - - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native device weights\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_forward_primitives.cpp b/cpp/tests/test_pi05_native_forward_primitives.cpp deleted file mode 100644 index 8648592a..00000000 --- a/cpp/tests/test_pi05_native_forward_primitives.cpp +++ /dev/null @@ -1,348 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_kernel_driver.h" -#include "flashrt/cpp/modalities/types.h" -#include "flashrt/exec.h" - -#include - -#include -#include -#include -#include -#include - -namespace { - -using flashrt::models::pi05::NativeKernelDriver; - -frt_buffer allocate(frt_ctx ctx, const char* name, std::size_t elements) { - frt_buffer buffer = - frt_buffer_alloc(ctx, name, elements * sizeof(std::uint16_t)); - assert(buffer); - return buffer; -} - -std::vector bf16(const std::vector& values) { - std::vector result(values.size()); - for (std::size_t i = 0; i < values.size(); ++i) { - result[i] = flashrt::modalities::float_to_bfloat16(values[i]); - } - return result; -} - -void upload(frt_buffer buffer, const std::vector& values) { - assert(cudaMemcpy(frt_buffer_dptr(buffer), values.data(), - values.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); -} - -std::vector download(frt_buffer buffer, std::size_t elements) { - std::vector bits(elements); - assert(cudaMemcpy(bits.data(), frt_buffer_dptr(buffer), - bits.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - std::vector result(elements); - for (std::size_t i = 0; i < elements; ++i) { - result[i] = flashrt::modalities::bfloat16_to_float(bits[i]); - } - return result; -} - -void expect_close(const std::vector& actual, - const std::vector& expected, float tolerance) { - assert(actual.size() == expected.size()); - for (std::size_t i = 0; i < actual.size(); ++i) { - assert(std::fabs(actual[i] - expected[i]) <= tolerance); - } -} - -struct CaptureArgs { - const NativeKernelDriver* driver = nullptr; - void* values = nullptr; - const void* weight = nullptr; - void* norm_output = nullptr; - const void* qkv = nullptr; - const void* rope = nullptr; - void* query = nullptr; - void* key = nullptr; - void* value = nullptr; - const void* devpos = nullptr; - bool recorded = false; -}; - -void record_primitives(void* user, void* stream) { - auto* args = static_cast(user); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - args->recorded = - args->driver - ->rms_norm_bf16(args->values, args->weight, args->norm_output, - 2, 4, 1e-6f, native_stream) - .ok_status() && - args->driver - ->qkv_split_rope_devpos_bf16( - args->qkv, args->rope, args->query, args->key, args->value, - args->devpos, 2, 4, 2, 2, 2, native_stream) - .ok_status(); -} - -} // namespace - -int main() { - int device_count = 0; - if (cudaGetDeviceCount(&device_count) != cudaSuccess || !device_count) { - cudaGetLastError(); - std::printf("SKIP - no CUDA device\n"); - return 0; - } - - NativeKernelDriver driver; - assert(driver.status().ok_status()); - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - cudaStream_t stream = nullptr; - assert(cudaStreamCreate(&stream) == cudaSuccess); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - - const std::vector host_x = {1, 2, 3, 4, -1, 0, 1, 2}; - const std::vector host_weight = {1, 1.5f, 0.5f, 2}; - const std::vector host_bias = {0.1f, -0.2f, 0.3f, -0.4f}; - frt_buffer x = allocate(ctx, "primitive_x", 8); - frt_buffer weight = allocate(ctx, "primitive_weight", 4); - frt_buffer bias = allocate(ctx, "primitive_bias", 4); - frt_buffer output = allocate(ctx, "primitive_output", 8); - frt_buffer gate = allocate(ctx, "primitive_gate", 8); - upload(x, bf16(host_x)); - upload(weight, bf16(host_weight)); - upload(bias, bf16(host_bias)); - - assert(driver.rms_norm_bf16( - frt_buffer_dptr(x), frt_buffer_dptr(weight), - frt_buffer_dptr(output), 2, 4, 1e-6f, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector rms_expected(8); - for (int row = 0; row < 2; ++row) { - float sum = 0; - for (int col = 0; col < 4; ++col) { - const float value = host_x[row * 4 + col]; - sum += value * value; - } - const float scale = 1.0f / std::sqrt(sum / 4 + 1e-6f); - for (int col = 0; col < 4; ++col) { - rms_expected[row * 4 + col] = - host_x[row * 4 + col] * scale * host_weight[col]; - } - } - expect_close(download(output, 8), rms_expected, 0.025f); - - assert(driver.layer_norm_bf16( - frt_buffer_dptr(x), frt_buffer_dptr(weight), - frt_buffer_dptr(bias), frt_buffer_dptr(output), 2, 4, - 1e-5f, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector layer_expected(8); - for (int row = 0; row < 2; ++row) { - float mean = 0; - for (int col = 0; col < 4; ++col) mean += host_x[row * 4 + col]; - mean /= 4; - float variance = 0; - for (int col = 0; col < 4; ++col) { - const float delta = host_x[row * 4 + col] - mean; - variance += delta * delta; - } - const float scale = 1.0f / std::sqrt(variance / 4 + 1e-5f); - for (int col = 0; col < 4; ++col) { - layer_expected[row * 4 + col] = - (host_x[row * 4 + col] - mean) * scale * host_weight[col] + - host_bias[col]; - } - } - expect_close(download(output, 8), layer_expected, 0.025f); - - std::vector style(24, 0.0f); - for (int row = 0; row < 2; ++row) { - for (int col = 0; col < 4; ++col) { - style[row * 12 + col] = 0.25f; - style[row * 12 + 4 + col] = -0.5f; - style[row * 12 + 8 + col] = 0.75f; - } - } - frt_buffer style_buffer = allocate(ctx, "primitive_style", 24); - upload(style_buffer, bf16(style)); - assert(driver.ada_rms_norm_style_bf16( - frt_buffer_dptr(x), frt_buffer_dptr(weight), - frt_buffer_dptr(style_buffer), frt_buffer_dptr(output), - frt_buffer_dptr(gate), 2, 4, 1e-6f, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector ada_expected(8); - for (std::size_t i = 0; i < ada_expected.size(); ++i) { - ada_expected[i] = rms_expected[i] * 1.25f - 0.5f; - } - expect_close(download(output, 8), ada_expected, 0.035f); - expect_close(download(gate, 8), std::vector(8, 0.75f), 0.0f); - - frt_buffer residual = allocate(ctx, "primitive_residual", 8); - upload(residual, bf16(std::vector(8, 1.0f))); - upload(output, bf16(std::vector(8, 2.0f))); - upload(gate, bf16(std::vector(8, 0.5f))); - assert(driver.gate_mul_residual_bf16( - frt_buffer_dptr(residual), frt_buffer_dptr(output), - frt_buffer_dptr(gate), 8, native_stream) - .ok_status()); - assert(driver.bias_residual_bf16( - frt_buffer_dptr(residual), frt_buffer_dptr(output), - frt_buffer_dptr(bias), 2, 4, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector residual_expected(8); - for (int i = 0; i < 8; ++i) { - residual_expected[i] = 4.0f + host_bias[i % 4]; - } - expect_close(download(residual, 8), residual_expected, 0.025f); - - upload(residual, bf16(std::vector(8, 1.0f))); - upload(output, bf16(std::vector(8, 2.0f))); - assert(driver.residual_add_bf16( - frt_buffer_dptr(residual), frt_buffer_dptr(output), 8, - native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - expect_close(download(residual, 8), std::vector(8, 3.0f), 0.0f); - - const std::vector activation_input = - {-3, -2, -1, 0, 0.5f, 1, 2, 3}; - upload(gate, bf16(activation_input)); - upload(output, bf16(std::vector(8, 1.5f))); - assert(driver.gate_gelu_bf16( - frt_buffer_dptr(gate), frt_buffer_dptr(output), - frt_buffer_dptr(residual), 8, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector gated_expected(8); - for (int i = 0; i < 8; ++i) { - const float value = activation_input[i]; - const float gelu = value / - (1.0f + std::exp(-1.5957691216057308f * value * - (1.0f + 0.044715f * value * value))); - gated_expected[i] = gelu * 1.5f; - } - expect_close(download(residual, 8), gated_expected, 0.025f); - upload(output, bf16(activation_input)); - assert(driver.gelu_bf16(frt_buffer_dptr(output), 8, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector gelu_expected(8); - for (int i = 0; i < 8; ++i) { - const float value = activation_input[i]; - gelu_expected[i] = value * 0.5f * - (1.0f + std::tanh(0.7978845608f * - (value + 0.044715f * value * value * value))); - } - expect_close(download(output, 8), gelu_expected, 0.025f); - - std::vector pool_input(4 * 4 * 2); - for (int row = 0; row < 4; ++row) { - for (int col = 0; col < 4; ++col) { - pool_input[(row * 4 + col) * 2] = float(row * 4 + col); - pool_input[(row * 4 + col) * 2 + 1] = float(row * 4 + col + 1); - } - } - frt_buffer pool_values = allocate(ctx, "primitive_pool_values", 32); - frt_buffer pool_output = allocate(ctx, "primitive_pool_output", 8); - upload(pool_values, bf16(pool_input)); - assert(driver.avg_pool_vision_bf16( - frt_buffer_dptr(pool_values), frt_buffer_dptr(pool_output), - 1, 4, 4, 2, 2, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - expect_close(download(pool_output, 8), - {2.5f, 3.5f, 4.5f, 5.5f, 10.5f, 11.5f, 12.5f, 13.5f}, - 0.0f); - - const std::vector host_qkv = { - 1, 2, 3, 4, 5, 6, 7, 8, - 9, 10, 11, 12, 13, 14, 15, 16}; - frt_buffer qkv = allocate(ctx, "primitive_qkv", 16); - frt_buffer rope = allocate(ctx, "primitive_rope", 4); - frt_buffer query = allocate(ctx, "primitive_query", 8); - frt_buffer key = allocate(ctx, "primitive_key", 8); - frt_buffer value = allocate(ctx, "primitive_value", 8); - frt_buffer devpos = frt_buffer_alloc(ctx, "primitive_devpos", sizeof(int)); - assert(devpos); - upload(qkv, bf16(host_qkv)); - upload(rope, bf16({1, 0, 0, 1})); - const int position = 1; - assert(cudaMemcpy(frt_buffer_dptr(devpos), &position, sizeof(position), - cudaMemcpyHostToDevice) == cudaSuccess); - assert(cudaMemset(frt_buffer_dptr(key), 0, 8 * sizeof(std::uint16_t)) == - cudaSuccess); - assert(cudaMemset(frt_buffer_dptr(value), 0, 8 * sizeof(std::uint16_t)) == - cudaSuccess); - assert(driver.qkv_split_rope_devpos_bf16( - frt_buffer_dptr(qkv), frt_buffer_dptr(rope), - frt_buffer_dptr(query), frt_buffer_dptr(key), - frt_buffer_dptr(value), frt_buffer_dptr(devpos), 2, 4, 2, - 2, 2, native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - expect_close(download(query, 8), {1, 2, 3, 4, -10, 9, -12, 11}, 0.0f); - expect_close(download(key, 8), {0, 0, 5, 6, -14, 13, 0, 0}, 0.0f); - expect_close(download(value, 8), {0, 0, 7, 8, 15, 16, 0, 0}, 0.0f); - - const std::size_t image_elements = 224 * 224 * 3; - frt_buffer image = allocate(ctx, "primitive_image", image_elements); - frt_buffer patches = allocate(ctx, "primitive_patches", image_elements); - std::vector image_bits(image_elements); - for (std::size_t i = 0; i < image_bits.size(); ++i) { - image_bits[i] = static_cast(i); - } - upload(image, image_bits); - assert(driver.patch_im2col_16bit( - frt_buffer_dptr(image), frt_buffer_dptr(patches), 1, - native_stream) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - std::vector patch_bits(image_elements); - assert(cudaMemcpy(patch_bits.data(), frt_buffer_dptr(patches), - patch_bits.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - for (int patch = 0; patch < 256; ++patch) { - const int patch_row = patch / 16; - const int patch_col = patch % 16; - for (int feature = 0; feature < 588; ++feature) { - const int pixel_row = feature / 42; - const int pixel_col = (feature % 42) / 3; - const int channel = feature % 3; - const std::size_t source = - static_cast(patch_row * 14 + pixel_row) * 224 * 3 + - (patch_col * 14 + pixel_col) * 3 + channel; - assert(patch_bits[patch * 588 + feature] == image_bits[source]); - } - } - - frt_graph graph = frt_graph_create(ctx, "native_forward_primitives", 7); - assert(graph); - CaptureArgs capture{ - &driver, frt_buffer_dptr(x), frt_buffer_dptr(weight), - frt_buffer_dptr(output), frt_buffer_dptr(qkv), frt_buffer_dptr(rope), - frt_buffer_dptr(query), frt_buffer_dptr(key), frt_buffer_dptr(value), - frt_buffer_dptr(devpos), false}; - assert(frt_graph_capture(graph, 7, record_primitives, &capture) == FRT_OK); - assert(capture.recorded); - const int stream_id = frt_ctx_wrap_stream(ctx, stream); - assert(stream_id >= 0); - for (int i = 0; i < 100; ++i) { - assert(frt_graph_replay(graph, 7, stream_id) == FRT_OK); - } - assert(frt_graph_variant_count(graph) == 1); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - - frt_graph_destroy(graph); - assert(cudaStreamDestroy(stream) == cudaSuccess); - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native forward primitives\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_kernel_driver.cpp b/cpp/tests/test_pi05_native_kernel_driver.cpp deleted file mode 100644 index 1bc11324..00000000 --- a/cpp/tests/test_pi05_native_kernel_driver.cpp +++ /dev/null @@ -1,139 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_kernel_driver.h" -#include "flashrt/cpp/modalities/types.h" -#include "flashrt/exec.h" - -#include - -#include -#include -#include -#include - -namespace { - -struct CaptureArgs { - flashrt::models::pi05::NativeKernelDriver* driver = nullptr; - void* a = nullptr; - void* b = nullptr; - void* output = nullptr; - const void* bias = nullptr; - bool recorded = false; -}; - -bool has_cuda_device() { - int count = 0; - const cudaError_t rc = cudaGetDeviceCount(&count); - if (rc != cudaSuccess) { - cudaGetLastError(); - return false; - } - return count > 0; -} - -void record_gemm(void* user, void* stream) { - auto* args = static_cast(user); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - args->recorded = - args->driver - ->bf16_nn(args->a, args->b, args->output, 2, 2, 3, - native_stream) - .ok_status() && - args->driver - ->add_bias_bf16(args->output, args->bias, 2, 2, native_stream) - .ok_status() && - args->driver->silu_bf16(args->output, 4, native_stream).ok_status(); -} - -} // namespace - -int main() { - if (!has_cuda_device()) { - std::printf("SKIP - no CUDA device\n"); - return 0; - } - using flashrt::modalities::bfloat16_to_float; - using flashrt::modalities::float_to_bfloat16; - - flashrt::models::pi05::NativeKernelDriver driver; - assert(driver.status().ok_status()); - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - frt_buffer a = frt_buffer_alloc(ctx, "a", 6 * sizeof(std::uint16_t)); - frt_buffer b = frt_buffer_alloc(ctx, "b", 6 * sizeof(std::uint16_t)); - frt_buffer output = - frt_buffer_alloc(ctx, "output", 4 * sizeof(std::uint16_t)); - frt_buffer bias = frt_buffer_alloc(ctx, "bias", 2 * sizeof(std::uint16_t)); - assert(a && b && output && bias); - const std::vector host_a = { - float_to_bfloat16(1), float_to_bfloat16(2), float_to_bfloat16(3), - float_to_bfloat16(4), float_to_bfloat16(5), float_to_bfloat16(6)}; - const std::vector host_b = { - float_to_bfloat16(1), float_to_bfloat16(2), - float_to_bfloat16(3), float_to_bfloat16(4), - float_to_bfloat16(5), float_to_bfloat16(6)}; - const std::vector host_bias = { - float_to_bfloat16(-24), float_to_bfloat16(-30)}; - assert(cudaMemcpy(frt_buffer_dptr(a), host_a.data(), - host_a.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); - assert(cudaMemcpy(frt_buffer_dptr(b), host_b.data(), - host_b.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); - assert(cudaMemcpy(frt_buffer_dptr(bias), host_bias.data(), - host_bias.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); - - cudaStream_t stream = nullptr; - assert(cudaStreamCreate(&stream) == cudaSuccess); - assert(driver.bf16_nn(frt_buffer_dptr(a), frt_buffer_dptr(b), - frt_buffer_dptr(output), 2, 2, 3, - reinterpret_cast(stream)) - .ok_status()); - assert(driver.add_bias_bf16(frt_buffer_dptr(output), - frt_buffer_dptr(bias), 2, 2, - reinterpret_cast(stream)) - .ok_status()); - assert(driver.silu_bf16(frt_buffer_dptr(output), 4, - reinterpret_cast(stream)) - .ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - - frt_graph graph = frt_graph_create(ctx, "native_bf16_gemm", 1); - assert(graph); - assert(frt_graph_bind(graph, "a", a) == FRT_OK); - assert(frt_graph_bind(graph, "b", b) == FRT_OK); - assert(frt_graph_bind(graph, "output", output) == FRT_OK); - assert(frt_graph_bind(graph, "bias", bias) == FRT_OK); - CaptureArgs capture{&driver, frt_buffer_dptr(a), frt_buffer_dptr(b), - frt_buffer_dptr(output), frt_buffer_dptr(bias), false}; - assert(frt_graph_capture(graph, 1, record_gemm, &capture) == FRT_OK); - assert(capture.recorded); - assert(frt_graph_variant_count(graph) == 1); - const int stream_id = frt_ctx_wrap_stream(ctx, stream); - assert(stream_id >= 0); - for (int i = 0; i < 100; ++i) { - assert(frt_graph_replay(graph, 1, stream_id) == FRT_OK); - } - assert(frt_graph_variant_count(graph) == 1); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - - std::vector host_output(4); - assert(cudaMemcpy(host_output.data(), frt_buffer_dptr(output), - host_output.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - const float expected[] = { - -2.0f / (1.0f + std::exp(2.0f)), - -2.0f / (1.0f + std::exp(2.0f)), - 25.0f / (1.0f + std::exp(-25.0f)), - 34.0f / (1.0f + std::exp(-34.0f))}; - for (std::size_t i = 0; i < host_output.size(); ++i) { - assert(std::fabs(bfloat16_to_float(host_output[i]) - expected[i]) < - 0.02f); - } - frt_graph_destroy(graph); - assert(cudaStreamDestroy(stream) == cudaSuccess); - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native kernel driver capture\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_quantization.cpp b/cpp/tests/test_pi05_native_quantization.cpp deleted file mode 100644 index 9340c230..00000000 --- a/cpp/tests/test_pi05_native_quantization.cpp +++ /dev/null @@ -1,56 +0,0 @@ -#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})); - - NativeF16Tensor fp16_input{{2, 3}, {}}; - for (float value : {-448.0f, -1.0f, 0.0f, 1.0f, 224.0f, 448.0f}) { - fp16_input.values.push_back( - flashrt::modalities::float_to_float16(value)); - } - assert(native_quantize_fp8_e4m3(fp16_input, true, &fp8).ok_status()); - assert(fp8.shape == std::vector({3, 2})); - assert(fp8.scale == 1.0f); - assert(fp8.values == std::vector( - {0xfe, 0x38, 0xb8, 0x76, 0x00, 0x7e})); - - NativeF16Tensor reciprocal_input{{1, 2}, {}}; - for (float value : {-0.10406494140625f, 3.0078125f}) { - reciprocal_input.values.push_back( - flashrt::modalities::float_to_float16(value)); - } - assert(native_quantize_fp8_e4m3( - reciprocal_input, false, &fp8).ok_status()); - assert(fp8.scale == 0.0067138671875f); - assert(fp8.values == std::vector({0xd7, 0x7e})); - - NativeFloatTensor int8_input{{2, 3}, {1, 2, 3, 4, 5, 6}}; - NativeInt8Tensor int8; - assert(native_quantize_int8_per_output(int8_input, &int8).ok_status()); - assert(int8.shape == std::vector({3, 2})); - assert(int8.values == - std::vector({32, 127, 51, 127, 64, 127})); - assert(int8.scales.size() == 3); - assert(std::fabs(int8.scales[0] - 4.0f / 127.0f) < 1e-9f); - assert(std::fabs(int8.scales[1] - 5.0f / 127.0f) < 1e-9f); - assert(std::fabs(int8.scales[2] - 6.0f / 127.0f) < 1e-9f); - - NativeFloatTensor invalid{{2}, {1, 2}}; - assert(!native_quantize_fp8_e4m3(invalid, false, &fp8).ok_status()); - assert(!native_quantize_int8_per_output(invalid, &int8).ok_status()); - std::printf("PASS - Pi0.5 native weight quantization\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_rtx_attention.cpp b/cpp/tests/test_pi05_native_rtx_attention.cpp deleted file mode 100644 index e5e8e179..00000000 --- a/cpp/tests/test_pi05_native_rtx_attention.cpp +++ /dev/null @@ -1,73 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_rtx_attention.h" - -#include - -#include -#include - -int main() { - int count = 0; - if (cudaGetDeviceCount(&count) != cudaSuccess || count == 0) { - cudaGetLastError(); - std::printf("SKIP - no CUDA device\n"); - return 0; - } - using namespace flashrt::models::pi05; - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - { - NativeRtxAttentionWorkspace attention(ctx); - NativeRtxAttentionConfig bad; - bad.encoder_layers = 17; - assert(!attention.allocate(bad).ok_status()); - NativeRtxAttentionConfig config; - assert(attention.allocate(config).ok_status()); - assert(attention.size() == 22); - assert(attention.allocated_bytes() > 0); - assert(attention.encoder_splits() == 12); - assert(attention.decoder_splits() == 12); - assert(attention.kv_layer_stride_bytes() == 722 * 256 * 2); - assert(attention.find("attn_enc_K")->shape == - std::vector({18, 722, 1, 256})); - assert(attention.find("attn_enc_lse")->shape == - std::vector({1, 8, 768})); - assert(attention.find("attn_dec_lse")->shape == - std::vector({1, 8, 128})); - void* base = frt_buffer_dptr(attention.find("attn_enc_K")->buffer); - assert(attention.encoder_k_layer_dptr(0) == base); - assert(static_cast(attention.encoder_k_layer_dptr(17)) == - static_cast(base) + - 17 * attention.kv_layer_stride_bytes()); - assert(!attention.encoder_k_layer_dptr(18)); - - void* seqused_ptr = - frt_buffer_dptr(attention.find("attn_enc_seqused")->buffer); - const std::size_t bytes = attention.allocated_bytes(); - for (int i = 0; i < 1000; ++i) { - assert(attention.set_fixed_prompt_length(i % 201).ok_status()); - assert(frt_buffer_dptr( - attention.find("attn_enc_seqused")->buffer) == - seqused_ptr); - assert(attention.allocated_bytes() == bytes); - } - std::int32_t enc = 0; - std::int32_t dec = 0; - std::int32_t pos = 0; - assert(cudaMemcpy(&enc, frt_buffer_dptr( - attention.find("attn_enc_seqused")->buffer), - sizeof(enc), cudaMemcpyDeviceToHost) == cudaSuccess); - assert(cudaMemcpy(&dec, frt_buffer_dptr( - attention.find("attn_dec_seqused")->buffer), - sizeof(dec), cudaMemcpyDeviceToHost) == cudaSuccess); - assert(cudaMemcpy(&pos, frt_buffer_dptr( - attention.find("attn_dec_devpos")->buffer), - sizeof(pos), cudaMemcpyDeviceToHost) == cudaSuccess); - assert(enc == 512 + (999 % 201)); - assert(dec == enc + 10); - assert(pos == enc); - assert(!attention.set_fixed_prompt_length(201).ok_status()); - } - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native RTX attention workspace\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_rtx_attention_driver.cpp b/cpp/tests/test_pi05_native_rtx_attention_driver.cpp deleted file mode 100644 index 80f01126..00000000 --- a/cpp/tests/test_pi05_native_rtx_attention_driver.cpp +++ /dev/null @@ -1,157 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_rtx_attention_driver.h" -#include "flashrt/cpp/modalities/types.h" -#include "flashrt/exec.h" - -#include - -#include -#include -#include -#include -#include - -namespace { - -using flashrt::models::pi05::NativeAttentionBuffer; -using flashrt::models::pi05::NativeRtxAttentionDriver; - -struct CaptureArgs { - const NativeRtxAttentionDriver* driver = nullptr; - bool recorded = false; -}; - -void record_attention(void* user, void* stream) { - auto* args = static_cast(user); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - args->recorded = - args->driver->vision(native_stream).ok_status() && - args->driver->encoder(0, native_stream).ok_status() && - args->driver->decoder(0, native_stream).ok_status(); -} - -std::size_t elements(const NativeAttentionBuffer* buffer) { - assert(buffer); - std::size_t count = 1; - for (std::uint64_t dim : buffer->shape) count *= dim; - return count; -} - -void upload_constant(const NativeAttentionBuffer* buffer, float value) { - std::vector host( - elements(buffer), flashrt::modalities::float_to_bfloat16(value)); - assert(cudaMemcpy(frt_buffer_dptr(buffer->buffer), host.data(), - host.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); -} - -void upload_kv_rows(const NativeAttentionBuffer* buffer, int total_kv) { - std::vector host(elements(buffer), 0); - for (int row = 0; row < total_kv; ++row) { - const std::uint16_t value = - flashrt::modalities::float_to_bfloat16(float(row + 1)); - for (int column = 0; column < 256; ++column) { - host[static_cast(row) * 256 + column] = value; - } - } - assert(cudaMemcpy(frt_buffer_dptr(buffer->buffer), host.data(), - host.size() * sizeof(std::uint16_t), - cudaMemcpyHostToDevice) == cudaSuccess); -} - -void expect_constant(void* device, std::size_t count, float expected) { - std::vector host(count); - assert(cudaMemcpy(host.data(), device, count * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - for (std::uint16_t value : host) { - assert(std::fabs(flashrt::modalities::bfloat16_to_float(value) - - expected) < 0.02f); - } -} - -} // namespace - -int main() { - int count = 0; - if (cudaGetDeviceCount(&count) != cudaSuccess || count == 0) { - cudaGetLastError(); - std::printf("SKIP - no CUDA device\n"); - return 0; - } - cudaDeviceProp properties{}; - assert(cudaGetDeviceProperties(&properties, 0) == cudaSuccess); - if (properties.major < 8) { - std::printf("SKIP - BF16 FA2 needs compute capability 8.0+\n"); - return 0; - } - - using namespace flashrt::models::pi05; - NativeRtxAttentionDriver invalid_driver(nullptr); - assert(!invalid_driver.status().ok_status()); - - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - NativeRtxAttentionWorkspace workspace(ctx); - NativeRtxAttentionConfig config; - config.num_views = 1; - config.encoder_sequence = 128; - config.encoder_vision_sequence = 2; - config.chunk_size = 2; - assert(workspace.allocate(config).ok_status()); - assert(workspace.decoder_splits() == 3); - assert(workspace.set_fixed_prompt_length(1).ok_status()); - - NativeRtxAttentionDriver driver(&workspace); - assert(driver.status().ok_status()); - assert(driver.num_sms() == properties.multiProcessorCount); - - upload_constant(workspace.find("attn_vis_Q"), 0.0f); - upload_constant(workspace.find("attn_vis_K"), 0.0f); - upload_constant(workspace.find("attn_vis_V"), 2.0f); - upload_constant(workspace.find("attn_enc_Q"), 0.0f); - upload_constant(workspace.find("attn_dec_Q"), 0.0f); - upload_kv_rows(workspace.find("attn_enc_K"), 130); - upload_kv_rows(workspace.find("attn_enc_V"), 130); - - cudaStream_t stream = nullptr; - assert(cudaStreamCreate(&stream) == cudaSuccess); - const std::uintptr_t native_stream = - reinterpret_cast(stream); - assert(driver.vision(native_stream).ok_status()); - assert(driver.encoder(0, native_stream).ok_status()); - assert(driver.decoder(0, native_stream).ok_status()); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - expect_constant(driver.vision_output(), 256 * 16 * 72, 2.0f); - expect_constant(driver.encoder_output(), 128 * 8 * 256, 2.0f); - expect_constant(driver.decoder_output(), 2 * 8 * 256, 3.0f); - - frt_graph graph = frt_graph_create(ctx, "native_rtx_attention", 1); - assert(graph); - assert(frt_graph_bind(graph, "vis_q", - workspace.find("attn_vis_Q")->buffer) == FRT_OK); - assert(frt_graph_bind(graph, "enc_q", - workspace.find("attn_enc_Q")->buffer) == FRT_OK); - assert(frt_graph_bind(graph, "dec_q", - workspace.find("attn_dec_Q")->buffer) == FRT_OK); - CaptureArgs capture{&driver, false}; - assert(frt_graph_capture(graph, 1, record_attention, &capture) == FRT_OK); - assert(capture.recorded); - assert(frt_graph_variant_count(graph) == 1); - const int stream_id = frt_ctx_wrap_stream(ctx, stream); - assert(stream_id >= 0); - assert(workspace.set_fixed_prompt_length(0).ok_status()); - for (int i = 0; i < 100; ++i) { - assert(frt_graph_replay(graph, 1, stream_id) == FRT_OK); - } - assert(frt_graph_variant_count(graph) == 1); - assert(cudaStreamSynchronize(stream) == cudaSuccess); - expect_constant(driver.vision_output(), 256 * 16 * 72, 2.0f); - expect_constant(driver.encoder_output(), 128 * 8 * 256, 1.5f); - expect_constant(driver.decoder_output(), 2 * 8 * 256, 2.5f); - - frt_graph_destroy(graph); - assert(cudaStreamDestroy(stream) == cudaSuccess); - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native RTX FA2 driver\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_style_precompute.cpp b/cpp/tests/test_pi05_native_style_precompute.cpp deleted file mode 100644 index 0ec5b5dd..00000000 --- a/cpp/tests/test_pi05_native_style_precompute.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_style_precompute.h" - -#include - -#include -#include - -int main() { - int count = 0; - if (cudaGetDeviceCount(&count) != cudaSuccess || count == 0) { - cudaGetLastError(); - std::printf("SKIP - no CUDA device\n"); - return 0; - } - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - flashrt::models::pi05::NativeWorkspace workspace(ctx); - flashrt::models::pi05::NativeWorkspaceConfig config; - assert(workspace.allocate(config).ok_status()); - flashrt::models::pi05::NativeDeviceWeightStore weights(ctx); - flashrt::models::pi05::NativeKernelDriver driver; - flashrt::models::pi05::NativeStylePrecomputer precomputer(&driver); - assert(!precomputer.run(weights, &workspace, 0).ok_status()); - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native style precompute validation\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_weight_materializer.cpp b/cpp/tests/test_pi05_native_weight_materializer.cpp deleted file mode 100644 index 9703229b..00000000 --- a/cpp/tests/test_pi05_native_weight_materializer.cpp +++ /dev/null @@ -1,345 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" -#include "flashrt/cpp/models/pi05/native_weight_packer.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -struct Entry { - std::string key; - std::vector shape; - std::vector values; -}; - -bool has_cuda_device() { - int count = 0; - const cudaError_t rc = cudaGetDeviceCount(&count); - if (rc != cudaSuccess) { - cudaGetLastError(); - return false; - } - return count > 0; -} - -std::string temp_path() { - char path[] = "/tmp/frt_pi05_materializer_XXXXXX"; - const int fd = ::mkstemp(path); - assert(fd >= 0); - ::close(fd); - return path; -} - -std::vector sequence(std::size_t count, float start) { - std::vector values(count); - for (std::size_t i = 0; i < count; ++i) { - values[i] = start + static_cast(i) * 0.01f; - } - return values; -} - -void write_checkpoint(const std::string& path, - const std::vector& entries) { - std::string header = "{"; - std::uint64_t offset = 0; - for (std::size_t i = 0; i < entries.size(); ++i) { - const Entry& entry = entries[i]; - if (i) header += ','; - header += '"' + entry.key + "\":{\"dtype\":\"F32\",\"shape\":["; - for (std::size_t d = 0; d < entry.shape.size(); ++d) { - if (d) header += ','; - header += std::to_string(entry.shape[d]); - } - header += "],\"data_offsets\":[" + std::to_string(offset) + ','; - offset += entry.values.size() * sizeof(float); - header += std::to_string(offset) + "]}"; - } - header += '}'; - std::ofstream file(path, std::ios::binary | std::ios::trunc); - const std::uint64_t n = header.size(); - for (int i = 0; i < 8; ++i) { - const char byte = static_cast((n >> (8 * i)) & 0xffu); - file.write(&byte, 1); - } - file.write(header.data(), static_cast(header.size())); - for (const Entry& entry : entries) { - file.write(reinterpret_cast(entry.values.data()), - static_cast(entry.values.size() * - sizeof(float))); - } - assert(file.good()); -} - -} // namespace - -int main() { - if (!has_cuda_device()) { - std::printf("SKIP - no CUDA device\n"); - return 0; - } - const std::string prefix = - "paligemma_with_expert.paligemma.model.language_model.layers.0"; - const std::string decoder = - "paligemma_with_expert.gemma_expert.model.layers.0"; - const std::string vision = - "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; - const std::string vision_layer = vision + ".encoder.layers.0"; - const std::vector entries = { - {prefix + ".input_layernorm.weight", {4}, {-0.5f, 0.0f, 0.5f, 1.0f}}, - {prefix + ".self_attn.q_proj.weight", {16, 4}, sequence(64, 0.1f)}, - {prefix + ".self_attn.k_proj.weight", {4, 4}, sequence(16, 1.0f)}, - {prefix + ".self_attn.v_proj.weight", {4, 4}, sequence(16, 2.0f)}, - {prefix + ".self_attn.o_proj.weight", {4, 8}, sequence(32, 3.0f)}, - {prefix + ".post_attention_layernorm.weight", {4}, - {-0.25f, 0.0f, 0.25f, 0.5f}}, - {prefix + ".mlp.gate_proj.weight", {6, 4}, sequence(24, 4.0f)}, - {prefix + ".mlp.up_proj.weight", {6, 4}, sequence(24, 5.0f)}, - {prefix + ".mlp.down_proj.weight", {4, 6}, sequence(24, 6.0f)}, - {decoder + ".self_attn.q_proj.weight", {16, 4}, sequence(64, 7.0f)}, - {decoder + ".self_attn.k_proj.weight", {4, 4}, sequence(16, 8.0f)}, - {decoder + ".self_attn.v_proj.weight", {4, 4}, sequence(16, 9.0f)}, - {decoder + ".self_attn.o_proj.weight", {4, 16}, sequence(64, 10.0f)}, - {decoder + ".mlp.gate_proj.weight", {6, 4}, sequence(24, 11.0f)}, - {decoder + ".mlp.up_proj.weight", {6, 4}, sequence(24, 12.0f)}, - {decoder + ".mlp.down_proj.weight", {4, 6}, sequence(24, 13.0f)}, - {decoder + ".input_layernorm.dense.weight", {12, 4}, - sequence(48, 14.0f)}, - {decoder + ".input_layernorm.dense.bias", {12}, sequence(12, 15.0f)}, - {decoder + ".post_attention_layernorm.dense.weight", {12, 4}, - sequence(48, 16.0f)}, - {decoder + ".post_attention_layernorm.dense.bias", {12}, - sequence(12, 17.0f)}, - {vision + ".embeddings.patch_embedding.weight", {2, 2, 2, 1}, - sequence(8, 18.0f)}, - {vision + ".embeddings.patch_embedding.bias", {2}, - sequence(2, 19.0f)}, - {vision + ".embeddings.position_embedding.weight", {3, 2}, - sequence(6, 20.0f)}, - {vision + ".post_layernorm.weight", {2}, sequence(2, 21.0f)}, - {vision + ".post_layernorm.bias", {2}, sequence(2, 22.0f)}, - {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear." - "weight", {4, 2}, sequence(8, 23.0f)}, - {"paligemma_with_expert.paligemma.model.multi_modal_projector.linear." - "bias", {4}, sequence(4, 24.0f)}, - {vision_layer + ".self_attn.q_proj.weight", {2, 2}, - sequence(4, 25.0f)}, - {vision_layer + ".self_attn.q_proj.bias", {2}, - sequence(2, 26.0f)}, - {vision_layer + ".self_attn.k_proj.weight", {2, 2}, - sequence(4, 27.0f)}, - {vision_layer + ".self_attn.k_proj.bias", {2}, - sequence(2, 28.0f)}, - {vision_layer + ".self_attn.v_proj.weight", {2, 2}, - sequence(4, 29.0f)}, - {vision_layer + ".self_attn.v_proj.bias", {2}, - sequence(2, 30.0f)}, - {vision_layer + ".self_attn.out_proj.weight", {2, 2}, - sequence(4, 31.0f)}, - {vision_layer + ".self_attn.out_proj.bias", {2}, - sequence(2, 32.0f)}, - {vision_layer + ".mlp.fc1.weight", {3, 2}, - sequence(6, 33.0f)}, - {vision_layer + ".mlp.fc1.bias", {3}, sequence(3, 34.0f)}, - {vision_layer + ".mlp.fc2.weight", {2, 3}, - sequence(6, 35.0f)}, - {vision_layer + ".mlp.fc2.bias", {2}, sequence(2, 36.0f)}, - {vision_layer + ".layer_norm1.weight", {2}, sequence(2, 37.0f)}, - {vision_layer + ".layer_norm1.bias", {2}, sequence(2, 38.0f)}, - {vision_layer + ".layer_norm2.weight", {2}, sequence(2, 39.0f)}, - {vision_layer + ".layer_norm2.bias", {2}, sequence(2, 40.0f)}, - {"paligemma_with_expert.gemma_expert.model.norm.dense.weight", - {3, 2}, sequence(6, 41.0f)}, - {"paligemma_with_expert.gemma_expert.model.norm.dense.bias", - {3}, sequence(3, 42.0f)}, - {"time_mlp_in.weight", {2, 2}, sequence(4, 43.0f)}, - {"time_mlp_in.bias", {2}, sequence(2, 44.0f)}, - {"time_mlp_out.weight", {2, 2}, sequence(4, 45.0f)}, - {"time_mlp_out.bias", {2}, sequence(2, 46.0f)}, - {"action_in_proj.weight", {2, 1}, sequence(2, 47.0f)}, - {"action_in_proj.bias", {2}, sequence(2, 48.0f)}, - {"action_out_proj.weight", {1, 2}, sequence(2, 49.0f)}, - {"action_out_proj.bias", {1}, sequence(1, 50.0f)}, - {"paligemma_with_expert.paligemma.lm_head.weight", - {4, 2}, sequence(8, 51.0f)}, - }; - const std::string path = temp_path(); - write_checkpoint(path, entries); - - flashrt::loader::SafetensorsFile source; - assert(source.open(path)); - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - { - flashrt::models::pi05::NativeDeviceWeightStore destination(ctx); - flashrt::models::pi05::NativeWeightMaterializer materializer( - source, &destination); - const auto encoder_status = materializer.materialize_encoder_layer(0); - if (!encoder_status.ok_status()) { - std::fprintf(stderr, "encoder materialization failed: %s\n", - encoder_status.message.c_str()); - } - assert(encoder_status.ok_status()); - assert(destination.size() == 5); - const auto* qkv = destination.find("encoder_attn_qkv_w_0"); - assert(qkv && qkv->shape == std::vector({4, 24})); - const auto* gate = destination.find("encoder_ffn_gate_w_0"); - assert(gate && gate->shape == std::vector({4, 6})); - const auto* down = destination.find("encoder_ffn_down_w_0"); - assert(down && down->shape == std::vector({6, 4})); - assert(!materializer.materialize_encoder_layer(0).ok_status()); - assert(!materializer.materialize_encoder_layer(18).ok_status()); - assert(materializer.materialize_decoder_layer(0, true).ok_status()); - assert(destination.size() == 15); - const auto* decoder_qkv = destination.find("decoder_attn_qkv_w_0"); - assert(decoder_qkv && - decoder_qkv->shape == std::vector({4, 24})); - const auto* gate_up = destination.find("decoder_ffn_gate_up_w_0"); - assert(gate_up && - gate_up->shape == std::vector({4, 12})); - const auto* attn_mod = - destination.find("decoder_pre_attn_norm_mod_w_0"); - assert(attn_mod && - attn_mod->shape == std::vector({4, 12})); - assert(!materializer.materialize_decoder_layer(0, true).ok_status()); - assert(!materializer.materialize_decoder_layer(18, true).ok_status()); - assert(materializer.materialize_vision_globals().ok_status()); - assert(materializer.materialize_vision_layer(0).ok_status()); - assert(destination.size() == 34); - const auto* patch = destination.find("vision_patch_embedding_w"); - assert(patch && patch->shape == - std::vector({2, 1, 2, 2})); - const auto* vision_qkv = destination.find("vision_attn_qkv_w_0"); - assert(vision_qkv && - vision_qkv->shape == std::vector({2, 6})); - assert(!materializer.materialize_vision_layer(27).ok_status()); - assert(!materializer.materialize_decoder_globals(0).ok_status()); - assert(materializer.materialize_decoder_globals(10).ok_status()); - assert(destination.size() == 45); - assert(destination.find("decoder_final_norm_mod_w")->shape == - std::vector({2, 3})); - assert(destination.find("decoder_time_embeds")->shape == - std::vector({10, 1024})); - assert(destination.find("decoder_action_out_proj_w")->shape == - std::vector({2, 1})); - assert(materializer.materialize_embedding().ok_status()); - assert(destination.size() == 46); - assert(destination.find("embedding_weight")->shape == - std::vector({4, 2})); - flashrt::models::pi05::NativeWeightPacker packer(&destination); - assert(packer.pack_fp8("decoder_attn_qkv_w_0", false).ok_status()); - assert(packer.pack_int8("decoder_attn_qkv_w_0").ok_status()); - assert(destination.size() == 50); - assert(destination.find("fp8.decoder_attn_qkv_w_0")->shape == - std::vector({4, 24})); - assert(destination.find("int8.decoder_attn_qkv_w_0")->shape == - std::vector({24, 4})); - } - frt_ctx_destroy(ctx); - assert(::unlink(path.c_str()) == 0); - - const char* real_checkpoint = std::getenv("FLASH_RT_PI05_CHECKPOINT"); - if (real_checkpoint && real_checkpoint[0]) { - flashrt::loader::SafetensorsFile real_source; - assert(real_source.open(std::string(real_checkpoint) + - "/model.safetensors")); - frt_ctx real_ctx = frt_ctx_create(); - assert(real_ctx); - { - flashrt::models::pi05::NativeDeviceWeightStore destination( - real_ctx); - flashrt::models::pi05::NativeWeightMaterializer materializer( - real_source, &destination); - assert(materializer.materialize_encoder_layer(0).ok_status()); - assert(destination.size() == 5); - assert(destination.find("encoder_attn_qkv_w_0")->shape == - std::vector({2048, 2560})); - assert(destination.find("encoder_ffn_gate_w_0")->shape == - std::vector({2048, 16384})); - assert(materializer.materialize_decoder_layer(0, true).ok_status()); - assert(destination.size() == 15); - assert(destination.find("decoder_attn_qkv_w_0")->shape == - std::vector({1024, 2560})); - assert(destination.find("decoder_ffn_gate_up_w_0")->shape == - std::vector({1024, 8192})); - assert(materializer.materialize_vision_globals().ok_status()); - assert(materializer.materialize_vision_layer(0).ok_status()); - assert(destination.size() == 34); - assert(destination.find("vision_patch_embedding_w")->shape == - std::vector({14, 14, 3, 1152})); - assert(destination.find("vision_attn_qkv_w_0")->shape == - std::vector({1152, 3456})); - assert(materializer.materialize_decoder_globals(10).ok_status()); - assert(destination.size() == 45); - assert(destination.find("decoder_final_norm_mod_w")->shape == - std::vector({1024, 3072})); - assert(destination.find("decoder_time_embeds")->shape == - std::vector({10, 1024})); - assert(destination.find("decoder_action_in_proj_w")->shape == - std::vector({32, 1024})); - assert(destination.find("decoder_action_out_proj_w")->shape == - std::vector({1024, 32})); - flashrt::models::pi05::NativeWeightPacker packer(&destination); - assert(packer.pack_fp8("decoder_attn_qkv_w_0", false) - .ok_status()); - assert(packer.pack_int8("decoder_attn_qkv_w_0").ok_status()); - assert(destination.size() == 49); - assert(destination.find("fp8.decoder_attn_qkv_w_0")->shape == - std::vector({1024, 2560})); - assert(destination.find("int8.decoder_attn_qkv_w_0")->shape == - std::vector({2560, 1024})); - } - frt_ctx_destroy(real_ctx); - - const char* full = std::getenv("FLASH_RT_PI05_FULL_MATERIALIZE"); - if (full && std::strcmp(full, "1") == 0) { - frt_ctx full_ctx = frt_ctx_create(); - assert(full_ctx); - { - flashrt::models::pi05::NativeDeviceWeightStore destination( - full_ctx); - flashrt::models::pi05::NativeWeightMaterializer materializer( - real_source, &destination); - flashrt::models::pi05::NativeMaterializationOptions options; - assert(materializer.materialize_all(options).ok_status()); - assert(destination.size() == 613); - assert(destination.find("vision_attn_qkv_w_26")->shape == - std::vector({1152, 3456})); - assert(destination.find("encoder_ffn_down_w_17")->shape == - std::vector({16384, 2048})); - assert(destination.find("decoder_ffn_gate_up_w_17")->shape == - std::vector({1024, 8192})); - assert(destination.find("embedding_weight")->shape == - std::vector({257152, 2048})); - const char* pack = - std::getenv("FLASH_RT_PI05_FULL_PACK_FP8"); - if (pack && std::strcmp(pack, "1") == 0) { - flashrt::models::pi05::NativeWeightPacker packer( - &destination); - assert(packer.pack_all_fp8(false).ok_status()); - assert(destination.size() == 1137); - assert(destination.find( - "fp8.vision_attn_qkv_w_26")->dtype == - flashrt::models::pi05::NativeWeightDType:: - kFp8E4M3); - assert(destination.find( - "fp8.encoder_ffn_gate_up_w_17")->shape == - std::vector({2048, 32768})); - assert(destination.find( - "fp8.decoder_ffn_down_w_17.scale")->dtype == - flashrt::models::pi05::NativeWeightDType::kFloat32); - } - } - frt_ctx_destroy(full_ctx); - } - } - std::printf("PASS - Pi0.5 native layer materializer\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_weight_ops.cpp b/cpp/tests/test_pi05_native_weight_ops.cpp deleted file mode 100644 index 0072916b..00000000 --- a/cpp/tests/test_pi05_native_weight_ops.cpp +++ /dev/null @@ -1,316 +0,0 @@ -#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::NativeF16Tensor; -using flashrt::models::pi05::NativeFloatTensor; -using flashrt::models::pi05::NativeSourceDType; -using flashrt::models::pi05::NativeSourceTensorView; - -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); - } -} - -void expect_bf16(const NativeBf16Tensor& 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(tensor.values[i++] == - flashrt::modalities::float_to_bfloat16(value)); - } -} - -void expect_f16(const NativeF16Tensor& 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(tensor.values[i++] == flashrt::modalities::float_to_float16(value)); - } -} - -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}); - - NativeSourceTensorView mapped; - NativeBf16Tensor mapped_bf16; - NativeF16Tensor mapped_f16; - assert(load_native_source_tensor(file, "f32", &mapped).ok_status()); - assert(mapped.dtype == NativeSourceDType::kF32); - assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); - expect_bf16(mapped_bf16, {2}, {1.25f, -2.5f}); - assert(native_source_to_f16(mapped, false, &mapped_f16).ok_status()); - expect_f16(mapped_f16, {2}, {1.25f, -2.5f}); - assert(load_native_source_tensor(file, "bf16", &mapped).ok_status()); - assert(mapped.dtype == NativeSourceDType::kBf16); - assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); - expect_bf16(mapped_bf16, {2}, {3.0f, -4.0f}); - assert(native_source_to_f16(mapped, false, &mapped_f16).ok_status()); - expect_f16(mapped_f16, {2}, {3.0f, -4.0f}); - assert(load_native_source_tensor(file, "f16", &mapped).ok_status()); - assert(mapped.dtype == NativeSourceDType::kF16); - assert(native_source_to_bf16(mapped, false, &mapped_bf16).ok_status()); - expect_bf16(mapped_bf16, {2}, {5.0f, -6.0f}); - assert(native_source_to_f16(mapped, false, &mapped_f16).ok_status()); - expect_f16(mapped_f16, {2}, {5.0f, -6.0f}); - assert(::unlink(path.c_str()) == 0); - - NativeFloatTensor matrix{{2, 3}, {1, 2, 3, 4, 5, 6}}; - NativeFloatTensor result; - assert(native_transpose_2d(matrix, &result).ok_status()); - expect(result, {3, 2}, {1, 4, 2, 5, 3, 6}); - - NativeFloatTensor patch{{2, 2, 2, 1}, {0, 1, 2, 3, 4, 5, 6, 7}}; - assert(native_patch_oihw_to_hwio(patch, &result).ok_status()); - expect(result, {2, 1, 2, 2}, {0, 4, 2, 6, 1, 5, 3, 7}); - - NativeFloatTensor qk{{8, 1}, {0, 1, 2, 3, 4, 5, 6, 7}}; - assert(native_interleave_qk_rows(qk, 2, &result).ok_status()); - expect(result, {8, 1}, {0, 2, 1, 3, 4, 6, 5, 7}); - - NativeFloatTensor norm{{3}, {-0.5f, 0.0f, 1.0f}}; - assert(native_fold_rms_columns(matrix, norm, &result).ok_status()); - expect(result, {2, 3}, {0.5f, 2.0f, 6.0f, 2.0f, 5.0f, 12.0f}); - - NativeFloatTensor q{{2, 2}, {1, 2, 3, 4}}; - NativeFloatTensor k{{1, 2}, {5, 6}}; - NativeFloatTensor v{{1, 2}, {7, 8}}; - assert(native_concat_rows_transpose({&q, &k, &v}, &result).ok_status()); - expect(result, {2, 4}, {1, 3, 5, 7, 2, 4, 6, 8}); - - NativeFloatTensor left{{2, 2}, {1, 2, 3, 4}}; - NativeFloatTensor right{{2, 2}, {5, 6, 7, 8}}; - assert(native_concat_columns(left, right, &result).ok_status()); - expect(result, {2, 4}, {1, 2, 5, 6, 3, 4, 7, 8}); - - NativeFloatTensor a{{2}, {1, 2}}; - NativeFloatTensor b{{3}, {3, 4, 5}}; - assert(native_concat_vectors({&a, &b}, &result).ok_status()); - expect(result, {5}, {1, 2, 3, 4, 5}); - - assert(native_scale(matrix, -0.1f, &result).ok_status()); - expect(result, {2, 3}, {-0.1f, -0.2f, -0.3f, - -0.4f, -0.5f, -0.6f}); - - assert(native_pi05_time_embeddings(2, 4, &result).ok_status()); - assert(result.shape == std::vector({2, 4})); - assert(result.values.size() == 8); - assert(!native_pi05_time_embeddings(0, 4, &result).ok_status()); - assert(!native_pi05_time_embeddings(2, 3, &result).ok_status()); - - NativeF16Tensor time_f16; - assert(native_pi05_time_embeddings_f16(2, 4, &time_f16).ok_status()); - assert(time_f16.shape == std::vector({2, 4})); - assert(time_f16.values.size() == 8); - const double two_pi = 6.2831853071795864769; - const double first_angles[] = {two_pi / 4.0e-3, two_pi / 4.0}; - for (std::size_t i = 0; i < 2; ++i) { - assert(time_f16.values[i] == flashrt::modalities::float_to_float16( - static_cast(std::sin(first_angles[i])))); - assert(time_f16.values[2 + i] == - flashrt::modalities::float_to_float16( - static_cast(std::cos(first_angles[i])))); - } - assert(native_pi05_time_embeddings_f16(10, 1024, &time_f16).ok_status()); - assert(time_f16.shape == std::vector({10, 1024})); - assert(time_f16.values[6 * 1024 + 7] == 0xb1c0u); - assert(time_f16.values[8 * 1024 + 7] == 0x2dc6u); - assert(time_f16.values[9 * 1024 + 3] == 0x292au); - - 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])); - } - - NativeF16Tensor converted_f16; - assert(native_to_f16(matrix, &converted_f16).ok_status()); - expect_f16(converted_f16, {2, 3}, {1, 2, 3, 4, 5, 6}); - - NativeBf16Tensor direct; - const float source_f32[] = {1, 2, 3, 4, 5, 6, 7, 8}; - std::uint16_t source_bf16[8]; - std::uint16_t source_f16[8]; - for (std::size_t i = 0; i < 8; ++i) { - source_bf16[i] = flashrt::modalities::float_to_bfloat16(source_f32[i]); - source_f16[i] = flashrt::modalities::float_to_float16(source_f32[i]); - } - const NativeSourceTensorView source_views[] = { - {source_f32, {2, 4}, NativeSourceDType::kF32}, - {source_bf16, {2, 4}, NativeSourceDType::kBf16}, - {source_f16, {2, 4}, NativeSourceDType::kF16}, - }; - NativeFloatTensor source_norm{{4}, {0, 0, 0, 0}}; - for (const NativeSourceTensorView& source_view : source_views) { - assert(native_source_to_bf16(source_view, true, &direct).ok_status()); - expect_bf16(direct, {4, 2}, {1, 5, 2, 6, 3, 7, 4, 8}); - assert(native_source_fold_rms_columns_transpose( - source_view, source_norm, &direct).ok_status()); - expect_bf16(direct, {4, 2}, {1, 5, 2, 6, 3, 7, 4, 8}); - assert(native_source_round_scale_to_bf16( - source_view, 2.0f, true, &direct).ok_status()); - expect_bf16(direct, {4, 2}, {2, 10, 4, 12, 6, 14, 8, 16}); - assert(native_source_qkv_to_bf16( - source_view, source_view, source_view, 1, 1, nullptr, - &direct).ok_status()); - expect_bf16(direct, {4, 6}, - {1, 5, 1, 5, 1, 5, - 2, 6, 2, 6, 2, 6, - 3, 7, 3, 7, 3, 7, - 4, 8, 4, 8, 4, 8}); - - NativeF16Tensor direct_f16; - assert(native_source_to_f16(source_view, true, &direct_f16).ok_status()); - expect_f16(direct_f16, {4, 2}, {1, 5, 2, 6, 3, 7, 4, 8}); - assert(native_source_qkv_to_f16( - source_view, source_view, source_view, 1, 1, nullptr, - true, &direct_f16).ok_status()); - expect_f16(direct_f16, {4, 6}, - {1, 5, 1, 5, 1, 5, - 2, 6, 2, 6, 2, 6, - 3, 7, 3, 7, 3, 7, - 4, 8, 4, 8, 4, 8}); - assert(native_source_qkv_to_f16( - source_view, source_view, source_view, 1, 1, nullptr, - false, &direct_f16).ok_status()); - expect_f16(direct_f16, {6, 4}, - {1, 2, 3, 4, 5, 6, 7, 8, - 1, 2, 3, 4, 5, 6, 7, 8, - 1, 2, 3, 4, 5, 6, 7, 8}); - } - - const float fold_source[] = {1.003f, -2.007f, 3.011f, -4.015f}; - const NativeSourceTensorView fold_view{ - fold_source, {2, 2}, NativeSourceDType::kF32}; - NativeFloatTensor fold_norm{{2}, {0.125f, -0.375f}}; - NativeF16Tensor direct_f16; - assert(native_source_qkv_to_f16( - fold_view, fold_view, fold_view, 1, 1, &fold_norm, - false, &direct_f16).ok_status()); - expect_f16(direct_f16, {6, 2}, - {1.003f * 1.125f, -2.007f * 0.625f, - 3.011f * 1.125f, -4.015f * 0.625f, - 1.003f * 1.125f, -2.007f * 0.625f, - 3.011f * 1.125f, -4.015f * 0.625f, - 1.003f * 1.125f, -2.007f * 0.625f, - 3.011f * 1.125f, -4.015f * 0.625f}); - - assert(native_source_pair_to_f16( - source_views[0], source_views[0], nullptr, false, - &direct_f16).ok_status()); - expect_f16(direct_f16, {4, 4}, - {1, 2, 3, 4, 5, 6, 7, 8, - 1, 2, 3, 4, 5, 6, 7, 8}); - assert(native_source_pair_to_f16( - source_views[0], source_views[0], nullptr, true, - &direct_f16).ok_status()); - expect_f16(direct_f16, {4, 4}, - {1, 5, 1, 5, 2, 6, 2, 6, - 3, 7, 3, 7, 4, 8, 4, 8}); - - const float vector_a[] = {1, 2}; - const std::uint16_t vector_b[] = { - flashrt::modalities::float_to_bfloat16(3), - flashrt::modalities::float_to_bfloat16(4)}; - const NativeSourceTensorView vector_a_view{ - vector_a, {2}, NativeSourceDType::kF32}; - const NativeSourceTensorView vector_b_view{ - vector_b, {2}, NativeSourceDType::kBf16}; - assert(native_source_concat_vectors_to_f16( - {&vector_a_view, &vector_b_view}, &direct_f16).ok_status()); - expect_f16(direct_f16, {4}, {1, 2, 3, 4}); - - const float patch_source[] = {0, 1, 2, 3, 4, 5, 6, 7}; - const NativeSourceTensorView patch_view{ - patch_source, {2, 2, 2, 1}, NativeSourceDType::kF32}; - assert(native_source_patch_oihw_to_hwio_f16( - patch_view, &direct_f16).ok_status()); - expect_f16(direct_f16, {2, 1, 2, 2}, {0, 4, 2, 6, 1, 5, 3, 7}); - - assert(!native_interleave_qk_rows(matrix, 2, &result).ok_status()); - assert(!native_concat_columns(matrix, k, &result).ok_status()); - std::printf("PASS - Pi0.5 native weight transforms\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_weight_packer.cpp b/cpp/tests/test_pi05_native_weight_packer.cpp deleted file mode 100644 index 83e52180..00000000 --- a/cpp/tests/test_pi05_native_weight_packer.cpp +++ /dev/null @@ -1,141 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_weight_packer.h" - -#include - -#include -#include -#include -#include - -namespace { - -bool has_cuda_device() { - int count = 0; - const cudaError_t rc = cudaGetDeviceCount(&count); - if (rc != cudaSuccess) { - cudaGetLastError(); - return false; - } - return count > 0; -} - -template -std::vector download(const flashrt::models::pi05::NativeDeviceWeight& weight) { - std::vector result(frt_buffer_bytes(weight.buffer) / sizeof(T)); - assert(cudaMemcpy(result.data(), frt_buffer_dptr(weight.buffer), - result.size() * sizeof(T), cudaMemcpyDeviceToHost) == - cudaSuccess); - return result; -} - -void upload(flashrt::models::pi05::NativeDeviceWeightStore* store, - const std::string& name, - const flashrt::models::pi05::NativeBf16Tensor& tensor) { - assert(store->upload(name, tensor).ok_status()); -} - -} // namespace - -int main() { - if (!has_cuda_device()) { - std::printf("SKIP - no CUDA device\n"); - return 0; - } - using namespace flashrt::models::pi05; - NativeFloatTensor source{{2, 3}, {1, 2, 3, 4, 5, 6}}; - NativeBf16Tensor bf16; - assert(native_to_bf16(source, &bf16).ok_status()); - NativeFloatTensor rounded; - assert(native_round_to_bf16_float(source, &rounded).ok_status()); - - frt_ctx ctx = frt_ctx_create(); - assert(ctx); - { - NativeDeviceWeightStore store(ctx); - assert(store.upload("weight", bf16).ok_status()); - NativeBf16Tensor copied; - assert(store.download_bf16("weight", &copied).ok_status()); - assert(copied.values == bf16.values); - - NativeWeightPacker packer(&store); - assert(packer.pack_fp8("weight", false).ok_status()); - assert(packer.pack_fp8("weight", true).ok_status() == false); - assert(packer.pack_int8("weight").ok_status()); - assert(store.size() == 5); - - NativeFp8Tensor expected_fp8; - assert(native_quantize_fp8_e4m3(rounded, false, &expected_fp8) - .ok_status()); - const auto* fp8 = store.find("fp8.weight"); - assert(fp8 && fp8->dtype == NativeWeightDType::kFp8E4M3); - assert(download(*fp8) == expected_fp8.values); - assert(download(*store.find("fp8.weight.scale")) == - std::vector({expected_fp8.scale})); - - NativeInt8Tensor expected_int8; - assert(native_quantize_int8_per_output(rounded, &expected_int8) - .ok_status()); - const auto* int8 = store.find("int8.weight"); - assert(int8 && int8->dtype == NativeWeightDType::kInt8); - assert(download(*int8) == expected_int8.values); - assert(download(*store.find("int8.weight.scale")) == - expected_int8.scales); - assert(!packer.pack_int8("missing").ok_status()); - } - frt_ctx_destroy(ctx); - - ctx = frt_ctx_create(); - assert(ctx); - { - NativeDeviceWeightStore store(ctx); - NativeBf16Tensor tiny; - tiny.shape = {1, 1}; - tiny.values = {flashrt::modalities::float_to_bfloat16(1.0f)}; - for (int layer = 0; layer < 27; ++layer) { - for (const char* stem : { - "vision_attn_qkv_w_", "vision_attn_o_w_", - "vision_ffn_up_w_", "vision_ffn_down_w_"}) { - upload(&store, std::string(stem) + std::to_string(layer), - tiny); - } - } - upload(&store, "encoder_multi_modal_projector_w", tiny); - for (int layer = 0; layer < 18; ++layer) { - const std::string suffix = std::to_string(layer); - for (const std::string& name : { - "encoder_attn_qkv_w_" + suffix, - "encoder_attn_o_w_" + suffix, - "encoder_ffn_gate_w_" + suffix, - "encoder_ffn_up_w_" + suffix, - "encoder_ffn_down_w_" + suffix, - "decoder_attn_qkv_w_" + suffix, - "decoder_attn_o_w_" + suffix, - "decoder_ffn_gate_w_" + suffix, - "decoder_ffn_up_w_" + suffix, - "decoder_ffn_gate_up_w_" + suffix, - "decoder_ffn_down_w_" + suffix}) { - upload(&store, name, tiny); - } - } - assert(store.size() == 307); - NativeWeightPacker packer(&store); - assert(packer.pack_all_fp8(false).ok_status()); - assert(packer.pack_vision_int8().ok_status()); - assert(packer.pack_encoder_int8().ok_status()); - assert(packer.pack_decoder_int8().ok_status()); - assert(store.size() == 1407); - assert(store.find("fp8.vision_projector_w")->dtype == - NativeWeightDType::kFp8E4M3); - assert(store.find("fp8.encoder_ffn_gate_up_w_17")->shape == - std::vector({1, 2})); - assert(store.find("int8.vision_ffn_down_w_26.scale")->dtype == - NativeWeightDType::kFloat32); - assert(store.find("int8.encoder_ffn_up_w_17")->dtype == - NativeWeightDType::kInt8); - assert(store.find("int8.decoder_ffn_down_w_17")->dtype == - NativeWeightDType::kInt8); - } - frt_ctx_destroy(ctx); - std::printf("PASS - Pi0.5 native weight packer\n"); - return 0; -} diff --git a/cpp/tests/test_pi05_native_workspace.cpp b/cpp/tests/test_pi05_native_workspace.cpp deleted file mode 100644 index 459d73b4..00000000 --- a/cpp/tests/test_pi05_native_workspace.cpp +++ /dev/null @@ -1,237 +0,0 @@ -#include "flashrt/cpp/models/pi05/native_workspace.h" - -#include - -#include -#include -#include -#include - -namespace { - -bool has_cuda_device() { - int count = 0; - const cudaError_t rc = cudaGetDeviceCount(&count); - if (rc != cudaSuccess) { - cudaGetLastError(); - return false; - } - return count > 0; -} - -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); - const std::uint16_t expected = - buffer.dtype == flashrt::modalities::DType::kFloat16 - ? flashrt::modalities::float_to_float16(1.0f) - : flashrt::modalities::float_to_bfloat16(1.0f); - for (std::uint16_t value : values) assert(value == expected); -} - -} // 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()); - invalid.vision_pool_factor = 1; - invalid.num_views = 3; - invalid.max_prompt_tokens = std::numeric_limits::max() - 769; - invalid.chunk_size = 2; - assert(!workspace.allocate(invalid).ok_status()); - NativeWorkspaceConfig config; - assert(workspace.allocate(config).ok_status()); - assert(workspace.logical_size() == 35); - assert(workspace.allocation_count() == 34); - assert(workspace.allocated_bytes() > 0); - assert(workspace.vision_sequence() == 512); - assert(workspace.encoder_vision_sequence() == 512); - assert(workspace.encoder_sequence() == 712); - assert(workspace.find("prompt_embedding")->shape == - std::vector({200, 2048})); - const auto* vision_x = workspace.find("vision_x"); - const auto* pooled = workspace.find("vision_x_pooled"); - assert(vision_x && pooled && pooled->alias); - assert(vision_x->buffer == pooled->buffer); - assert(workspace.find("decoder_style_attn")->shape == - std::vector({10, 18, 10, 3072})); - assert(workspace.find("rtc_prefix_weights")->dtype == - flashrt::modalities::DType::kFloat32); - check_ones(*workspace.find("encoder_rms_ones")); - check_ones(*workspace.find("decoder_rms_ones")); - assert(workspace.update_decoder_rope(37).ok_status()); - assert(!workspace.update_decoder_rope(201).ok_status()); - void* decoder_rope_ptr = - frt_buffer_dptr(workspace.find("decoder_rope_weights")->buffer); - const std::size_t allocation_count = workspace.allocation_count(); - const std::size_t allocated_bytes = workspace.allocated_bytes(); - for (int i = 0; i < 1000; ++i) { - assert(workspace.update_decoder_rope(i % 201).ok_status()); - assert(frt_buffer_dptr( - workspace.find("decoder_rope_weights")->buffer) == - decoder_rope_ptr); - assert(workspace.allocation_count() == allocation_count); - assert(workspace.allocated_bytes() == allocated_bytes); - } - - NativeDeviceWeightStore weights(ctx); - NativeBf16Tensor position; - position.shape = {256, 1152}; - position.values.resize(256 * 1152); - for (std::size_t i = 0; i < position.values.size(); ++i) { - position.values[i] = flashrt::modalities::float_to_bfloat16( - static_cast(i % 97) / 97.0f); - } - assert(weights.upload("vision_position_embedding", position) - .ok_status()); - assert(workspace.expand_vision_position_embedding(weights) - .ok_status()); - const auto* expanded = workspace.find("vision_pos_embed_expanded"); - std::vector expanded_values(position.values.size() * 2); - assert(cudaMemcpy(expanded_values.data(), - frt_buffer_dptr(expanded->buffer), - expanded_values.size() * sizeof(std::uint16_t), - cudaMemcpyDeviceToHost) == cudaSuccess); - assert(std::vector(expanded_values.begin(), - expanded_values.begin() + - position.values.size()) == - position.values); - assert(std::vector( - expanded_values.begin() + position.values.size(), - expanded_values.end()) == position.values); - assert(!workspace.allocate(config).ok_status()); - } - frt_ctx_destroy(ctx); - - ctx = frt_ctx_create(); - assert(ctx); - { - NativeWorkspace workspace(ctx); - NativeWorkspaceConfig config; - config.num_views = 3; - config.max_prompt_tokens = 256; - config.chunk_size = 50; - config.num_steps = 5; - config.vision_pool_factor = 2; - assert(workspace.allocate(config).ok_status()); - assert(workspace.logical_size() == 35); - assert(workspace.allocation_count() == 35); - assert(workspace.vision_sequence() == 768); - assert(workspace.encoder_vision_sequence() == 192); - assert(workspace.encoder_sequence() == 448); - const auto* pooled = workspace.find("vision_x_pooled"); - assert(pooled && !pooled->alias); - assert(pooled->shape == std::vector({192, 1152})); - assert(pooled->buffer != workspace.find("vision_x")->buffer); - assert(workspace.find("decoder_time_emb")->shape == - std::vector({5, 50, 1024})); - } - frt_ctx_destroy(ctx); - - ctx = frt_ctx_create(); - assert(ctx); - { - NativeWorkspace workspace(ctx); - NativeWorkspaceConfig invalid; - invalid.flavor = NativeWorkspaceFlavor::kThorFp8; - invalid.vision_pool_factor = 2; - assert(!workspace.allocate(invalid).ok_status()); - invalid.vision_pool_factor = 1; - invalid.max_prompt_tokens = 199; - assert(!workspace.allocate(invalid).ok_status()); - - NativeWorkspaceConfig config; - config.flavor = NativeWorkspaceFlavor::kThorFp8; - config.enable_calibration = true; - assert(workspace.allocate(config).ok_status()); - assert(workspace.vision_sequence() == 512); - assert(workspace.encoder_sequence() == 712); - assert(workspace.total_keys() == 722); - assert(workspace.find("observation_images_normalized")->dtype == - flashrt::modalities::DType::kFloat16); - assert(workspace.find("encoder_x_fp8")->dtype == - flashrt::modalities::DType::kUInt8); - assert(workspace.find("encoder_logits")->shape == - std::vector({712 * 8, 722})); - assert(workspace.find("encoder_k_cache")->shape == - std::vector({18, 722, 256})); - assert(workspace.find("decoder_activation_scales")->shape == - std::vector({10, 18, 4})); - assert(workspace.find("encoder_sample_scales")); - assert(workspace.find("decoder_sample_scales")); - check_ones(*workspace.find("encoder_rms_ones")); - check_ones(*workspace.find("decoder_rms_ones")); - std::vector encoder_rope(712 * 256); - assert(cudaMemcpy( - encoder_rope.data(), - frt_buffer_dptr( - workspace.find("encoder_rope_weights")->buffer), - encoder_rope.size() * sizeof(encoder_rope[0]), - cudaMemcpyDeviceToHost) == cudaSuccess); - assert(encoder_rope[468 * 256 + 2 * 30 + 1] == 0xb7fdu); - assert(encoder_rope[624 * 256 + 2 * 34 + 1] == 0xb7fdu); - - const auto* prompt = workspace.find("prompt_embedding"); - std::vector row(2048); - for (std::size_t i = 0; i < row.size(); ++i) { - row[i] = flashrt::modalities::float_to_float16( - static_cast(i % 31)); - } - auto* prompt_base = static_cast( - frt_buffer_dptr(prompt->buffer)); - assert(cudaMemcpy(prompt_base + 4 * row.size() * sizeof(row[0]), - row.data(), row.size() * sizeof(row[0]), - cudaMemcpyHostToDevice) == cudaSuccess); - assert(workspace.set_fixed_prompt_length(5).ok_status()); - std::vector padded(row.size()); - assert(cudaMemcpy(padded.data(), - prompt_base + 5 * row.size() * sizeof(row[0]), - padded.size() * sizeof(padded[0]), - cudaMemcpyDeviceToHost) == cudaSuccess); - assert(padded == row); - const char* control_names[] = { - "attn_enc_seqused", "attn_dec_seqused", "attn_dec_devpos"}; - const std::int32_t expected_controls[] = {518, 528, 518}; - for (int i = 0; i < 3; ++i) { - std::int32_t value = 0; - assert(cudaMemcpy( - &value, - frt_buffer_dptr(workspace.find(control_names[i])->buffer), - sizeof(value), cudaMemcpyDeviceToHost) == cudaSuccess); - assert(value == expected_controls[i]); - } - const std::size_t allocations = workspace.allocation_count(); - for (int i = 0; i < 1000; ++i) { - assert(workspace.set_fixed_prompt_length(i % 200).ok_status()); - assert(workspace.allocation_count() == allocations); - } - - NativeDeviceWeightStore weights(ctx); - NativeF16Tensor 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_float16( - static_cast(i % 97) / 97.0f); - } - assert(weights.upload("vision_position_embedding", position) - .ok_status()); - assert(workspace.expand_vision_position_embedding(weights) - .ok_status()); - } - 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 a43e360c..77f7516c 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -490,8 +490,8 @@ ctest --test-dir cpp/build --output-on-failure Detailed real-checkpoint oracles are maintained in the [MindOn validation suite](https://github.com/LiangSu8899/MindOn-dev/tree/main/validation/pi05_cpp). -Build their stage probes with -`FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS=ON`; default FlashRT builds retain contract, +Build private stage probes from that suite's CMake overlay against the exact +FlashRT source revision under test. Default FlashRT builds retain contract, lifecycle, calibration, final-result, and hot-path tests. Real-checkpoint oracle example: @@ -499,7 +499,7 @@ Real-checkpoint oracle example: ``` python /oracles/gate_pi05_native_weight_ops.py \ --checkpoint \ - --probe cpp/build/pi05_native_weight_probe + --probe /pi05_native_weight_probe ``` ``` @@ -582,7 +582,7 @@ python /oracles/gate_pi05_tokenizer_corpus.py \ --dataset \ --checkpoint \ --tokenizer \ - --probe /pi05_tokenizer_corpus_probe \ + --probe /pi05_tokenizer_corpus_probe \ --count 10000 ``` @@ -602,7 +602,7 @@ python /oracles/gate_pi05_native_e2e.py \ --checkpoint \ --tokenizer \ --dataset \ - --probe /pi05_native_e2e_probe \ + --probe /pi05_native_e2e_probe \ --episode 0 --frame 0 ``` @@ -636,7 +636,7 @@ 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 \ +python /perf/gate_pi05_kernel_sequence.py \ --native .csv --python .csv ``` diff --git a/docs/pi05_thor_native_fp8.md b/docs/pi05_thor_native_fp8.md index c9c9c214..c6af8aeb 100644 --- a/docs/pi05_thor_native_fp8.md +++ b/docs/pi05_thor_native_fp8.md @@ -59,10 +59,10 @@ useful for the same test matrix but is not a latency reference. For SM120, build the shared FA2 raw library first, then configure the C++ tree with `CMAKE_CUDA_ARCHITECTURES=120`, `FLASHRT_CPP_WITH_FA2=ON`, and -`FLASHRT_CPP_FA2_LIBRARY=`. Detailed real-checkpoint -diagnostic probes are opt-in through -`FLASHRT_CPP_BUILD_PI05_DIAGNOSTICS=ON`; default builds retain the contract, -lifecycle, calibration, final-result, and hot-path tests. +`FLASHRT_CPP_FA2_LIBRARY=`. Default FlashRT builds retain +the contract, lifecycle, calibration, final-result, and hot-path tests. +Implementation-level probes and unit tests are built from the separate MindOn +validation overlay and are not part of the product target graph. ## Native Configuration @@ -349,8 +349,8 @@ ctest --test-dir "$BUILD_DIR" --output-on-failure Real-checkpoint producer oracles and profiling tools are maintained in the [MindOn validation suite](https://github.com/LiangSu8899/MindOn-dev/tree/main/validation/pi05_cpp). -Configure diagnostic probes when an oracle needs a stage-level binary, then -run the matching script. The public paths below are placeholders: +Configure the validation overlay when an oracle needs a private stage-level +binary, then run the matching script. The public paths below are placeholders: ```bash python /oracles/gate_pi05_native_thor_fp8.py \ diff --git a/tests/bench_pi05_thor_views.py b/tests/bench_pi05_thor_views.py deleted file mode 100644 index f09e3f19..00000000 --- a/tests/bench_pi05_thor_views.py +++ /dev/null @@ -1,335 +0,0 @@ -#!/usr/bin/env python3 -"""Pi0.5 Thor 1-view / 2-view / 3-view latency + cos benchmark. - -Populates the ``Pi0.5 Thor (SM110)`` 1v/2v/3v table in the README, -matching the existing ``Pi0.5 RTX 5090`` table structure. - -Matrix (per view count): - Frontend Calibration Precision - ---------------- -------------- --------- - Pi0.5 Torch FP4 N=8 LIBERO NVFP4 encoder (18 layers + AWQ + P1) - Pi0.5 Torch FP8 N=8 LIBERO FP8 baseline - Pi0.5 JAX FP4 N=8 LIBERO NVFP4 encoder (18 layers + AWQ + P1) - Pi0.5 JAX FP8 N=1 (lazy) FP8 baseline [JAX FP8 does not - support N>=2 dataset - calibrate today] - -Reports per cell: - calibrate_ms - wall-clock of the calibrate() call (0 for lazy N=1) - p50_ms - median of 50 CUDA-graph replays (matches the README - RTX table's methodology) - p95_ms - 95th percentile - cos_vs_fp32_ref - 3v only (we only have a 3-view PyTorch FP32 reference) - -Each cell runs in its own Python subprocess (Thor resource model will -not accept multiple 7 GB VLA weight loads in one process). - -Usage:: - - python3 tests/bench_pi05_thor_views.py - python3 tests/bench_pi05_thor_views.py --views 2 3 # subset - python3 tests/bench_pi05_thor_views.py --frontends torch_fp4 # one row -""" -from __future__ import annotations - -import argparse -import json -import logging -import os -import subprocess -import sys -from pathlib import Path - -logging.basicConfig(level=logging.INFO, format="%(message)s") -logger = logging.getLogger("bench_pi05_views") - -ROOT = Path(__file__).resolve().parents[1] - -PI05_TORCH_CKPT = os.environ.get("PI05_CKPT", "/workspace/pytorch_checkpoints/pi05_libero_converted") -PI05_JAX_CKPT = os.environ.get("PI05_JAX_CKPT", "/workspace/checkpoints/pi05_libero_jax") - - -TORCH_SCRIPT = r""" -import sys, json, time, pathlib, numpy as np, torch -sys.path.insert(0, "ROOTDIR") - -# Clear calibration cache so the reported calibrate_ms is a real -# forward-pass cost, not a disk read. -cdir = pathlib.Path.home() / ".flash_rt" / "calibration" -if cdir.exists(): - for f in cdir.glob("*.json"): f.unlink() - -# Load per-view-count LIBERO obs bundle. -d = np.load(f"/tmp/libero_obs_{NV}v_n8.npz") -obs_list = [] -for i in range(int(d["n"])): - o = {"image": d[f"img_{i}"], "state": d[f"state_{i}"]} - if NV >= 2: o["wrist_image"] = d[f"wrist_{i}"] - if NV >= 3: o["wrist_image_right"] = d[f"wrist_right_{i}"] - obs_list.append(o) - -# Build frontend -if USE_FP4: - from flash_rt.frontends.torch.pi05_thor_fp4 import Pi05TorchFrontendThorFP4 - pipe = Pi05TorchFrontendThorFP4( - CKPT, num_views=NV, autotune=3, - use_fp4_encoder_ffn=True, fp4_layers=tuple(range(18)), - use_awq=True, use_p1_split_gu=True) -else: - from flash_rt.frontends.torch.pi05_thor import Pi05TorchFrontendThor - pipe = Pi05TorchFrontendThor(CKPT, num_views=NV, autotune=3) - -# 3v uses pytorch_reference.npz tokens + noise to allow cos vs FP32 ref. -# 1v/2v have no FP32 ref so we use a plain string prompt; cos is not -# reported for those. -if NV == 3: - ref = np.load("/tmp/pytorch_reference.npz", allow_pickle=True) - tok_mask = ref["arg8_tokenized_prompt_mask"][0] - tokens = ref["arg7_tokenized_prompt"][0][:int(tok_mask.sum())].astype(np.int64) - pipe.set_prompt(tokens.tolist()) - deploy_obs = { - "image": ref["arg0_base_rgb"][0], - "wrist_image": ref["arg1_left_wrist_rgb"][0], - "wrist_image_right": ref["arg2_right_wrist_rgb"][0], - } - ref_raw = ref["pytorch_raw_output"][0].astype(np.float32) - ref_noise = ref["arg9_noise"][0].astype(np.float16) -else: - pipe.set_prompt("pick up the red block and place it in the tray") - deploy_obs = obs_list[0] # use a LIBERO obs as deployment - ref_raw = None - ref_noise = None - -# Calibrate with 8 LIBERO samples. -t0 = time.perf_counter() -pipe.calibrate(obs_list, percentile=99.9) -cal_ms = (time.perf_counter() - t0) * 1000 - -# Warmup (Graph capture happens here implicitly in first infer). -for _ in range(5): pipe.infer(deploy_obs) - -# 3v: measure cos vs FP32 reference using matched noise. -# Pi0.5 torch frontend (the current Pi0.5 torch frontend) draws _g_noise via -# np.random.randn(Sa, 32) on CPU then H2D-copies; the legacy -# torch.Tensor.normal_ monkey-patch no longer fires. Mirror the JAX -# path's _PatchedRNG pattern (lines below). -cos_ref = None; maxdiff = None -if NV == 3 and ref_raw is not None: - Sa = ref_noise.shape[0] - _orig_randn = np.random.randn - class _PatchedRNG: - on = False - def __call__(self, *a, **kw): - if self.on and a == (Sa, 32): - return ref_noise.astype(np.float64) - return _orig_randn(*a, **kw) - p = _PatchedRNG() - np.random.randn = p - p.on = True - pipe.infer(deploy_obs) - p.on = False - np.random.randn = _orig_randn - out = pipe._g_noise.float().cpu().numpy() - a, b = out.flatten().astype(np.float64), ref_raw.flatten().astype(np.float64) - cos_ref = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)) - maxdiff = float(np.max(np.abs(out - ref_raw))) - -# Latency: 50 CUDA-graph replays. Match the README's RTX table -# (--warmup 50 --iters 500) approach of reporting replay P50 + P95. -lat = [] -for _ in range(50): - t0 = time.perf_counter() - pipe.infer(deploy_obs) - lat.append((time.perf_counter() - t0) * 1000) -lat.sort() -p50 = lat[25] -p95 = lat[47] - -out = { - "frontend": ("torch_fp4" if USE_FP4 else "torch_fp8"), - "num_views": NV, - "calibrate_ms": round(cal_ms, 1), - "p50_ms": round(p50, 2), - "p95_ms": round(p95, 2), -} -if cos_ref is not None: out["cos_vs_fp32_ref"] = round(cos_ref, 6) -if maxdiff is not None: out["maxdiff_vs_ref"] = round(maxdiff, 4) -print("__RESULT__ " + json.dumps(out)) -""" - - -JAX_SCRIPT = r""" -import sys, json, time, pathlib, numpy as np -sys.path.insert(0, "ROOTDIR") - -# JAX Pi0.5: FP8 (N=1 implicit) or FP4 (N=8 LIBERO multi-frame). -# Clear cache so calibrate_ms reflects a real forward-pass cost. -cdir = pathlib.Path.home() / ".flash_rt" / "calibration" -if cdir.exists(): - for f in cdir.glob("*.json"): f.unlink() - -d = np.load(f"/tmp/libero_obs_{NV}v_n8.npz") -obs_list = [] -for i in range(int(d["n"])): - o = {"image": d[f"img_{i}"], "state": d[f"state_{i}"]} - # The JAX Pi0.5 frontend's infer() reads wrist_image / wrist_image_right - # unconditionally — duplicate the base image when the view count is - # smaller so the 1v / 2v bench cells work through the same interface. - o["wrist_image"] = d[f"wrist_{i}"] if NV >= 2 else o["image"] - o["wrist_image_right"] = (d[f"wrist_right_{i}"] if NV >= 3 - else o["wrist_image"]) - obs_list.append(o) -deploy_obs = dict(obs_list[0]) - -if USE_FP4: - from flash_rt.frontends.jax.pi05_thor_fp4 import Pi05JaxFrontendThorFP4 - pipe = Pi05JaxFrontendThorFP4( - JAX_CKPT, num_views=NV, autotune=3, weight_cache=True, - use_fp4_encoder_ffn=True, fp4_layers=tuple(range(18)), - use_awq=True, use_p1_split_gu=True) -else: - from flash_rt.frontends.jax.pi05_thor import Pi05JaxFrontendThor - pipe = Pi05JaxFrontendThor(JAX_CKPT, num_views=NV, autotune=3) - -# 3v matches pytorch_reference.npz so we can report cos. -if NV == 3: - ref = np.load("/tmp/pytorch_reference.npz", allow_pickle=True) - tok_mask = ref["arg8_tokenized_prompt_mask"][0] - tokens = ref["arg7_tokenized_prompt"][0][:int(tok_mask.sum())].astype(np.int64) - pipe.set_prompt(tokens.tolist()) - deploy_obs = { - "image": ref["arg0_base_rgb"][0], - "wrist_image": ref["arg1_left_wrist_rgb"][0], - "wrist_image_right": ref["arg2_right_wrist_rgb"][0], - } - ref_raw = ref["pytorch_raw_output"][0].astype(np.float32) - ref_noise = ref["arg9_noise"][0].astype(np.float16) -else: - pipe.set_prompt("pick up the red block and place it in the tray") - ref_raw = None - ref_noise = None - -# FP4: N=8 LIBERO multi-frame; FP8: N=1 implicit (JAX FP8 path lacks N>=2). -t0 = time.perf_counter() -pipe.calibrate(obs_list if USE_FP4 else [deploy_obs], percentile=99.9) -cal_ms = (time.perf_counter() - t0) * 1000 - -for _ in range(5): pipe.infer(deploy_obs) - -cos_ref = None; maxdiff = None -if NV == 3 and ref_raw is not None: - _orig_randn = np.random.randn - class _PatchedRNG: - on = False - def __call__(self, *a, **kw): - if self.on and a == (10, 32): - return ref_noise.astype(np.float64) - return _orig_randn(*a, **kw) - p = _PatchedRNG() - np.random.randn = p - p.on = True - pipe.infer(deploy_obs) - p.on = False - np.random.randn = _orig_randn - out = pipe.g_noise.download_new((10, 32), np.float16).astype(np.float32) - a, b = out.flatten().astype(np.float64), ref_raw.flatten().astype(np.float64) - cos_ref = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)) - maxdiff = float(np.max(np.abs(out - ref_raw))) - -lat = [] -for _ in range(50): - t0 = time.perf_counter() - pipe.infer(deploy_obs) - lat.append((time.perf_counter() - t0) * 1000) -lat.sort() -p50 = lat[25]; p95 = lat[47] - -out = { - "frontend": ("jax_fp4" if USE_FP4 else "jax_fp8"), - "num_views": NV, - "calibrate_ms": round(cal_ms, 1), - "p50_ms": round(p50, 2), - "p95_ms": round(p95, 2), -} -if cos_ref is not None: out["cos_vs_fp32_ref"] = round(cos_ref, 6) -if maxdiff is not None: out["maxdiff_vs_ref"] = round(maxdiff, 4) -print("__RESULT__ " + json.dumps(out)) -""" - - -def _run(*, frontend: str, nv: int, ckpt: str) -> dict: - if frontend.startswith("torch"): - body = TORCH_SCRIPT - use_fp4 = (frontend == "torch_fp4") - header = ( - f"CKPT = {ckpt!r}\n" - f"NV = {int(nv)}\n" - f"USE_FP4 = {bool(use_fp4)}\n" - ) - elif frontend.startswith("jax"): - body = JAX_SCRIPT - use_fp4 = (frontend == "jax_fp4") - header = ( - f"JAX_CKPT = {ckpt!r}\n" - f"NV = {int(nv)}\n" - f"USE_FP4 = {bool(use_fp4)}\n" - ) - else: - raise ValueError(f"unknown frontend {frontend!r}") - - script = header + body.replace("ROOTDIR", str(ROOT)) - r = subprocess.run( - [sys.executable, "-c", script], capture_output=True, text=True, - timeout=1800) - for line in r.stdout.splitlines(): - if line.startswith("__RESULT__ "): - return json.loads(line[len("__RESULT__ "):]) - return { - "frontend": frontend, "num_views": nv, "error": ( - f"no __RESULT__; stderr tail: " - + "\n".join(r.stderr.splitlines()[-8:])) - } - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--views", nargs="+", type=int, default=[1, 2, 3]) - ap.add_argument("--frontends", nargs="+", - default=["torch_fp4", "torch_fp8", "jax_fp4"]) - args = ap.parse_args() - - rows = [] - for fr in args.frontends: - for nv in args.views: - logger.info("── %s num_views=%d ──", fr, nv) - ckpt = (PI05_JAX_CKPT if fr.startswith("jax") else PI05_TORCH_CKPT) - r = _run(frontend=fr, nv=nv, ckpt=ckpt) - logger.info(" %s", json.dumps(r)) - rows.append(r) - - # Summary table. - print("\n" + "=" * 95) - print(f"{'frontend':>10} {'nv':>3} {'calibrate':>10} " - f"{'p50':>9} {'p95':>9} {'cos_ref':>10} {'maxdiff':>8}") - print("-" * 95) - for r in rows: - if "error" in r: - print(f"{r.get('frontend','?'):>10} {r.get('num_views','?'):>3} " - f"ERROR {r['error'][:60]}") - continue - cos = r.get("cos_vs_fp32_ref", "-") - mdf = r.get("maxdiff_vs_ref", "-") - cos_s = f"{cos:.6f}" if isinstance(cos, float) else str(cos) - mdf_s = f"{mdf:.4f}" if isinstance(mdf, float) else str(mdf) - print(f"{r['frontend']:>10} {r['num_views']:>3} " - f"{r['calibrate_ms']:>8.1f} ms " - f"{r['p50_ms']:>6.2f} ms {r['p95_ms']:>6.2f} ms " - f"{cos_s:>10} {mdf_s:>8}") - print("=" * 95) - print(json.dumps(rows, indent=2)) - - return 0 if all("error" not in r for r in rows) else 1 - - -if __name__ == "__main__": - sys.exit(main()) From b5b87e31bedbf8bac5562f144d1a58df8407a3d4 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 07:37:26 -0400 Subject: [PATCH 80/83] refactor(pi05): organize model and backend sources --- cpp/models/pi05/CMakeLists.txt | 82 +++++++++---------- .../native_calibration_session.h | 0 .../pi05/{ => backend}/native_graph_runtime.h | 4 +- .../sm110}/native_thor_calibration_session.h | 2 +- .../sm110}/native_thor_fp8_forward.h | 6 +- .../sm110}/native_thor_graph_owner.h | 6 +- .../sm110}/native_thor_kernel_driver.h | 0 .../sm110}/native_thor_style_precompute.h | 6 +- .../sm110}/native_thor_weight_materializer.h | 4 +- .../sm120}/native_bf16_forward.h | 10 +-- .../{ => backends/sm120}/native_graph_owner.h | 6 +- .../sm120}/native_kernel_driver.h | 0 .../sm120}/native_rtx_attention.h | 0 .../sm120}/native_rtx_attention_driver.h | 2 +- .../sm120}/native_rtx_autotune.h | 2 +- .../sm120}/native_rtx_calibration_session.h | 2 +- .../{ => backends/sm120}/native_rtx_linear.h | 6 +- .../sm120}/native_rtx_weight_packer.h | 4 +- .../sm120}/native_style_precompute.h | 4 +- .../flashrt/cpp/models/pi05/{ => model}/io.h | 2 +- .../models/pi05/{ => model}/prompt_embed.h | 0 .../models/pi05/{ => model}/prompt_format.h | 0 .../cpp/models/pi05/{ => model}/runtime.h | 4 +- .../cpp/models/pi05/{ => model}/spec.h | 0 .../pi05/{ => support}/native_calibration.h | 0 .../{ => support}/native_device_weights.h | 2 +- .../pi05/{ => support}/native_quantization.h | 2 +- .../models/pi05/{ => support}/native_rope.h | 0 .../native_weight_materializer.h | 2 +- .../pi05/{ => support}/native_weight_ops.h | 0 .../pi05/{ => support}/native_weight_packer.h | 4 +- .../pi05/{ => support}/native_weights.h | 0 .../pi05/{ => support}/native_workspace.h | 2 +- .../native_calibration_session.cpp | 2 +- .../{ => backend}/native_graph_runtime.cpp | 4 +- .../native_thor_calibration_session.cpp | 14 ++-- .../sm110}/native_thor_fp8_forward.cpp | 2 +- .../sm110}/native_thor_graph_owner.cpp | 6 +- .../sm110}/native_thor_kernel_driver.cu | 2 +- .../sm110}/native_thor_setup_ops.cu | 0 .../sm110}/native_thor_style_precompute.cpp | 2 +- .../native_thor_weight_materializer.cpp | 2 +- .../sm120}/native_bf16_forward.cpp | 2 +- .../sm120}/native_graph_owner.cpp | 10 +-- .../sm120}/native_kernel_driver.cu | 2 +- .../sm120}/native_rtx_attention.cpp | 2 +- .../sm120}/native_rtx_attention_driver.cu | 2 +- .../sm120}/native_rtx_autotune.cpp | 2 +- .../sm120}/native_rtx_calibration_session.cpp | 10 +-- .../sm120}/native_rtx_linear.cpp | 2 +- .../sm120}/native_rtx_weight_packer.cpp | 2 +- .../sm120}/native_style_precompute.cu | 2 +- cpp/models/pi05/src/{ => model}/io.cpp | 2 +- .../pi05/src/{ => model}/prompt_embed.cpp | 4 +- .../pi05/src/{ => model}/prompt_format.cpp | 2 +- cpp/models/pi05/src/{ => model}/runtime.cpp | 2 +- cpp/models/pi05/src/{ => model}/spec.cpp | 2 +- cpp/models/pi05/src/{ => service}/c_api.cpp | 0 .../pi05/src/{ => service}/config_map.h | 2 +- .../pi05/src/{ => service}/model_runtime.cpp | 0 .../native_calibration_c_api.cpp | 8 +- .../{ => service}/native_model_runtime.cpp | 10 +-- .../pi05/src/{ => service}/native_open.cpp | 2 +- .../src/{ => service}/native_open_internal.h | 0 .../src/{ => support}/native_calibration.cpp | 2 +- .../{ => support}/native_device_weights.cpp | 2 +- .../src/{ => support}/native_quantization.cu | 2 +- .../native_quantization_unavailable.cpp | 2 +- .../pi05/src/{ => support}/native_rope.cu | 2 +- .../{ => support}/native_rope_unavailable.cpp | 2 +- .../native_weight_materializer.cpp | 2 +- .../src/{ => support}/native_weight_ops.cpp | 2 +- .../{ => support}/native_weight_packer.cpp | 2 +- .../pi05/src/{ => support}/native_weights.cpp | 2 +- .../src/{ => support}/native_workspace.cpp | 4 +- .../pi05_native_fp8_calibration_probe.cpp | 2 +- cpp/tests/test_device_staging.cpp | 2 +- cpp/tests/test_modalities.cpp | 4 +- cpp/tests/test_pi05_native_calibration.cpp | 4 +- cpp/tests/test_pi05_native_open.cpp | 2 +- cpp/tests/test_pi05_prompt_embed.cpp | 2 +- cpp/tests/test_pi05_prompt_format.cpp | 2 +- cpp/tests/test_pi05_runtime.cpp | 2 +- docs/pi05_io_contract.md | 6 +- 84 files changed, 159 insertions(+), 159 deletions(-) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backend}/native_calibration_session.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backend}/native_graph_runtime.h (95%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm110}/native_thor_calibration_session.h (95%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm110}/native_thor_fp8_forward.h (88%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm110}/native_thor_graph_owner.h (92%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm110}/native_thor_kernel_driver.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm110}/native_thor_style_precompute.h (78%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm110}/native_thor_weight_materializer.h (94%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_bf16_forward.h (89%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_graph_owner.h (93%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_kernel_driver.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_rtx_attention.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_rtx_attention_driver.h (94%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_rtx_autotune.h (87%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_rtx_calibration_session.h (95%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_rtx_linear.h (92%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_rtx_weight_packer.h (87%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => backends/sm120}/native_style_precompute.h (83%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => model}/io.h (97%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => model}/prompt_embed.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => model}/prompt_format.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => model}/runtime.h (98%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => model}/spec.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_calibration.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_device_weights.h (96%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_quantization.h (94%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_rope.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_weight_materializer.h (96%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_weight_ops.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_weight_packer.h (91%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_weights.h (100%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/{ => support}/native_workspace.h (98%) rename cpp/models/pi05/src/{ => backend}/native_calibration_session.cpp (98%) rename cpp/models/pi05/src/{ => backend}/native_graph_runtime.cpp (97%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_calibration_session.cpp (97%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_fp8_forward.cpp (99%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_graph_owner.cpp (97%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_kernel_driver.cu (99%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_setup_ops.cu (100%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_style_precompute.cpp (99%) rename cpp/models/pi05/src/{ => backends/sm110}/native_thor_weight_materializer.cpp (99%) rename cpp/models/pi05/src/{ => backends/sm120}/native_bf16_forward.cpp (99%) rename cpp/models/pi05/src/{ => backends/sm120}/native_graph_owner.cpp (97%) rename cpp/models/pi05/src/{ => backends/sm120}/native_kernel_driver.cu (99%) rename cpp/models/pi05/src/{ => backends/sm120}/native_rtx_attention.cpp (99%) rename cpp/models/pi05/src/{ => backends/sm120}/native_rtx_attention_driver.cu (99%) rename cpp/models/pi05/src/{ => backends/sm120}/native_rtx_autotune.cpp (97%) rename cpp/models/pi05/src/{ => backends/sm120}/native_rtx_calibration_session.cpp (98%) rename cpp/models/pi05/src/{ => backends/sm120}/native_rtx_linear.cpp (99%) rename cpp/models/pi05/src/{ => backends/sm120}/native_rtx_weight_packer.cpp (98%) rename cpp/models/pi05/src/{ => backends/sm120}/native_style_precompute.cu (99%) rename cpp/models/pi05/src/{ => model}/io.cpp (98%) rename cpp/models/pi05/src/{ => model}/prompt_embed.cpp (98%) rename cpp/models/pi05/src/{ => model}/prompt_format.cpp (98%) rename cpp/models/pi05/src/{ => model}/runtime.cpp (99%) rename cpp/models/pi05/src/{ => model}/spec.cpp (96%) rename cpp/models/pi05/src/{ => service}/c_api.cpp (100%) rename cpp/models/pi05/src/{ => service}/config_map.h (99%) rename cpp/models/pi05/src/{ => service}/model_runtime.cpp (100%) rename cpp/models/pi05/src/{ => service}/native_calibration_c_api.cpp (97%) rename cpp/models/pi05/src/{ => service}/native_model_runtime.cpp (98%) rename cpp/models/pi05/src/{ => service}/native_open.cpp (99%) rename cpp/models/pi05/src/{ => service}/native_open_internal.h (100%) rename cpp/models/pi05/src/{ => support}/native_calibration.cpp (99%) rename cpp/models/pi05/src/{ => support}/native_device_weights.cpp (99%) rename cpp/models/pi05/src/{ => support}/native_quantization.cu (98%) rename cpp/models/pi05/src/{ => support}/native_quantization_unavailable.cpp (93%) rename cpp/models/pi05/src/{ => support}/native_rope.cu (97%) rename cpp/models/pi05/src/{ => support}/native_rope_unavailable.cpp (86%) rename cpp/models/pi05/src/{ => support}/native_weight_materializer.cpp (99%) rename cpp/models/pi05/src/{ => support}/native_weight_ops.cpp (99%) rename cpp/models/pi05/src/{ => support}/native_weight_packer.cpp (99%) rename cpp/models/pi05/src/{ => support}/native_weights.cpp (98%) rename cpp/models/pi05/src/{ => support}/native_workspace.cpp (99%) diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index b5c2f9a3..08d286f5 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -9,28 +9,28 @@ set(_pi05_public_include "${CMAKE_CURRENT_SOURCE_DIR}/include") set(_pi05_internal_include "${CMAKE_CURRENT_SOURCE_DIR}/internal/include") set(_pi05_sources - src/spec.cpp - src/native_weights.cpp - src/native_weight_ops.cpp - src/native_device_weights.cpp - src/native_weight_packer.cpp - src/native_weight_materializer.cpp - src/native_calibration.cpp - src/native_calibration_session.cpp - src/native_workspace.cpp - src/native_rtx_attention.cpp - src/prompt_format.cpp - src/prompt_embed.cpp - src/io.cpp - src/runtime.cpp) + src/model/spec.cpp + src/model/prompt_format.cpp + src/model/prompt_embed.cpp + src/model/io.cpp + src/model/runtime.cpp + src/support/native_weights.cpp + src/support/native_weight_ops.cpp + src/support/native_device_weights.cpp + src/support/native_weight_packer.cpp + src/support/native_weight_materializer.cpp + src/support/native_calibration.cpp + src/support/native_workspace.cpp + src/backend/native_calibration_session.cpp + src/backends/sm120/native_rtx_attention.cpp) if(FLASHRT_CPP_WITH_CUDA_KERNELS) list(APPEND _pi05_sources - src/native_quantization.cu - src/native_rope.cu) + src/support/native_quantization.cu + src/support/native_rope.cu) else() list(APPEND _pi05_sources - src/native_quantization_unavailable.cpp - src/native_rope_unavailable.cpp) + src/support/native_quantization_unavailable.cpp + src/support/native_rope_unavailable.cpp) endif() add_library(flashrt_cpp_pi05 STATIC ${_pi05_sources}) @@ -41,18 +41,18 @@ target_link_libraries(flashrt_cpp_pi05 PUBLIC flashrt_cpp_modalities flashrt_cpp_vla flashrt_cpp_loader) if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") # Preserve producer-side add/cast/multiply operation boundaries. - set_property(SOURCE src/native_weight_ops.cpp APPEND PROPERTY + set_property(SOURCE src/support/native_weight_ops.cpp APPEND PROPERTY COMPILE_OPTIONS -ffp-contract=off) endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) add_library(flashrt_cpp_pi05_kernels STATIC - src/native_bf16_forward.cpp - src/native_graph_runtime.cpp - src/native_kernel_driver.cu - src/native_rtx_linear.cpp - src/native_rtx_weight_packer.cpp - src/native_style_precompute.cu + src/backend/native_graph_runtime.cpp + src/backends/sm120/native_bf16_forward.cpp + src/backends/sm120/native_kernel_driver.cu + src/backends/sm120/native_rtx_linear.cpp + src/backends/sm120/native_rtx_weight_packer.cpp + src/backends/sm120/native_style_precompute.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/activation.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu @@ -80,7 +80,7 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) if(FLASHRT_CPP_WITH_THOR_FP8) add_library(flashrt_cpp_pi05_thor_precise OBJECT - src/native_thor_setup_ops.cu + src/backends/sm110/native_thor_setup_ops.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention/fmha_fp16_strided.cu) target_include_directories(flashrt_cpp_pi05_thor_precise PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention @@ -98,12 +98,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) target_sources(flashrt_cpp_pi05_kernels PRIVATE $ - src/native_thor_kernel_driver.cu - src/native_thor_weight_materializer.cpp - src/native_thor_fp8_forward.cpp - src/native_thor_graph_owner.cpp - src/native_thor_calibration_session.cpp - src/native_thor_style_precompute.cpp + src/backends/sm110/native_thor_kernel_driver.cu + src/backends/sm110/native_thor_weight_materializer.cpp + src/backends/sm110/native_thor_fp8_forward.cpp + src/backends/sm110/native_thor_graph_owner.cpp + src/backends/sm110/native_thor_calibration_session.cpp + src/backends/sm110/native_thor_style_precompute.cpp ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/cutlass_sm100.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/attention_cublas.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/decoder_fused.cu @@ -126,12 +126,12 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) if(FLASHRT_CPP_WITH_FA2) target_sources(flashrt_cpp_pi05_kernels PRIVATE - src/native_graph_owner.cpp - src/native_rtx_autotune.cpp - src/native_rtx_attention_driver.cu) + src/backends/sm120/native_graph_owner.cpp + src/backends/sm120/native_rtx_autotune.cpp + src/backends/sm120/native_rtx_attention_driver.cu) if(FLASHRT_CPP_WITH_SENTENCEPIECE) target_sources(flashrt_cpp_pi05_kernels PRIVATE - src/native_rtx_calibration_session.cpp) + src/backends/sm120/native_rtx_calibration_session.cpp) endif() target_include_directories(flashrt_cpp_pi05_kernels PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../csrc) @@ -143,11 +143,11 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) endif() add_library(flashrt_cpp_pi05_c SHARED - src/c_api.cpp - src/model_runtime.cpp - src/native_calibration_c_api.cpp - src/native_model_runtime.cpp - src/native_open.cpp) + src/service/c_api.cpp + src/service/model_runtime.cpp + src/service/native_calibration_c_api.cpp + src/service/native_model_runtime.cpp + src/service/native_open.cpp) target_link_libraries(flashrt_cpp_pi05_c PUBLIC flashrt_cpp_pi05 flashrt_runtime PRIVATE flashrt_cpp_native) diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_calibration_session.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration_session.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_calibration_session.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_graph_runtime.h similarity index 95% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_graph_runtime.h index abdd8466..97b7d9ce 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_graph_runtime.h @@ -1,8 +1,8 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H -#include "flashrt/cpp/models/pi05/native_device_weights.h" -#include "flashrt/cpp/models/pi05/native_workspace.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" #include "flashrt/cpp/native/cuda_graph_set.h" #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h similarity index 95% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h index 6c7d6fdf..3a0b9c80 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_calibration_session.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h @@ -1,7 +1,7 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_CALIBRATION_SESSION_H -#include "flashrt/cpp/models/pi05/native_calibration_session.h" +#include "flashrt/cpp/models/pi05/backend/native_calibration_session.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h similarity index 88% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h index 7efd7126..2898c765 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_fp8_forward.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h @@ -1,9 +1,9 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_FP8_FORWARD_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_FP8_FORWARD_H -#include "flashrt/cpp/models/pi05/native_thor_kernel_driver.h" -#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" -#include "flashrt/cpp/models/pi05/native_workspace.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h similarity index 92% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h index 266fb209..ae355b2c 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h @@ -1,9 +1,9 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H -#include "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_graph_runtime.h" -#include "flashrt/cpp/models/pi05/native_thor_fp8_forward.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_kernel_driver.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h similarity index 78% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h index b152321d..872b16bc 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_style_precompute.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h @@ -1,9 +1,9 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_STYLE_PRECOMPUTE_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_STYLE_PRECOMPUTE_H -#include "flashrt/cpp/models/pi05/native_device_weights.h" -#include "flashrt/cpp/models/pi05/native_thor_kernel_driver.h" -#include "flashrt/cpp/models/pi05/native_workspace.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h similarity index 94% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h index 8bedf5e8..80ac2fb8 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_thor_weight_materializer.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h @@ -2,8 +2,8 @@ #define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_WEIGHT_MATERIALIZER_H #include "flashrt/cpp/loader/safetensors.h" -#include "flashrt/cpp/models/pi05/native_device_weights.h" -#include "flashrt/cpp/models/pi05/native_quantization.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/support/native_quantization.h" #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h similarity index 89% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h index 6fb9bb0f..1f0431a9 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_bf16_forward.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h @@ -1,11 +1,11 @@ #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_linear.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" +#include "flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h similarity index 93% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h index ffd75617..56ffe90c 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h @@ -1,9 +1,9 @@ #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 "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_kernel_driver.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h similarity index 94% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h index 62259246..95855f17 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_attention_driver.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h @@ -1,7 +1,7 @@ #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 "flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h similarity index 87% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h index eba22efb..e12686e5 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_autotune.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h @@ -1,7 +1,7 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_AUTOTUNE_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_AUTOTUNE_H -#include "flashrt/cpp/models/pi05/native_rtx_linear.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h similarity index 95% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h index fb2c24ae..c04ddcad 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_calibration_session.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h @@ -1,7 +1,7 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_CALIBRATION_SESSION_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_CALIBRATION_SESSION_H -#include "flashrt/cpp/models/pi05/native_calibration_session.h" +#include "flashrt/cpp/models/pi05/backend/native_calibration_session.h" #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h similarity index 92% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h index ec1bc58a..7b0db7c3 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_linear.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h @@ -1,9 +1,9 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_LINEAR_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_LINEAR_H -#include "flashrt/cpp/models/pi05/native_device_weights.h" -#include "flashrt/cpp/models/pi05/native_kernel_driver.h" -#include "flashrt/cpp/models/pi05/native_workspace.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h similarity index 87% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h index 3d2be1c8..4ec4ceb0 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rtx_weight_packer.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h @@ -1,8 +1,8 @@ #ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_WEIGHT_PACKER_H #define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_WEIGHT_PACKER_H -#include "flashrt/cpp/models/pi05/native_device_weights.h" -#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h" #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_style_precompute.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h similarity index 83% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_style_precompute.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h index df8c99c0..0c47f37c 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_style_precompute.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h @@ -1,8 +1,8 @@ #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" +#include "flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/io.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/io.h similarity index 97% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/io.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/io.h index 66ef4673..807fdb01 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/io.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/io.h @@ -1,7 +1,7 @@ #ifndef FLASHRT_MODELS_PI05_CPP_RUNTIME_IO_H #define FLASHRT_MODELS_PI05_CPP_RUNTIME_IO_H -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/model/spec.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_embed.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_embed.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_embed.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_embed.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_format.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_format.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/prompt_format.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_format.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/runtime.h similarity index 98% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/runtime.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/runtime.h index 2c65967d..b9f1da9f 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/runtime.h @@ -2,8 +2,8 @@ #define FLASHRT_CPP_MODELS_PI05_RUNTIME_H #include "flashrt/cpp/families/vla/runtime.h" -#include "flashrt/cpp/models/pi05/io.h" -#include "flashrt/cpp/models/pi05/prompt_embed.h" +#include "flashrt/cpp/models/pi05/model/io.h" +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/spec.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/spec.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/spec.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_calibration.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_calibration.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_calibration.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h similarity index 96% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h index 401051cf..b130dacb 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_device_weights.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h @@ -1,7 +1,7 @@ #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/cpp/models/pi05/support/native_weight_ops.h" #include "flashrt/exec.h" #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_quantization.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_quantization.h similarity index 94% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_quantization.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_quantization.h index bc3c1e9d..67009f2f 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_quantization.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_quantization.h @@ -1,7 +1,7 @@ #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 "flashrt/cpp/models/pi05/support/native_weight_ops.h" #include #include diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rope.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_rope.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_rope.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_rope.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_materializer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_materializer.h similarity index 96% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_materializer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_materializer.h index ece1efb3..856acd0d 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_materializer.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_materializer.h @@ -2,7 +2,7 @@ #define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H #include "flashrt/cpp/loader/safetensors.h" -#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_ops.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_ops.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_ops.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_ops.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_packer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h similarity index 91% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_packer.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h index 281885d7..d200d035 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weight_packer.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h @@ -1,8 +1,8 @@ #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" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/support/native_quantization.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weights.h similarity index 100% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_weights.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weights.h diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_workspace.h similarity index 98% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_workspace.h index bc4bc3a0..efc58ddf 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/native_workspace.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_workspace.h @@ -2,7 +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/cpp/models/pi05/support/native_device_weights.h" #include "flashrt/exec.h" #include diff --git a/cpp/models/pi05/src/native_calibration_session.cpp b/cpp/models/pi05/src/backend/native_calibration_session.cpp similarity index 98% rename from cpp/models/pi05/src/native_calibration_session.cpp rename to cpp/models/pi05/src/backend/native_calibration_session.cpp index 07df2c5d..e3a6b36f 100644 --- a/cpp/models/pi05/src/native_calibration_session.cpp +++ b/cpp/models/pi05/src/backend/native_calibration_session.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_calibration_session.h" +#include "flashrt/cpp/models/pi05/backend/native_calibration_session.h" #include #include diff --git a/cpp/models/pi05/src/native_graph_runtime.cpp b/cpp/models/pi05/src/backend/native_graph_runtime.cpp similarity index 97% rename from cpp/models/pi05/src/native_graph_runtime.cpp rename to cpp/models/pi05/src/backend/native_graph_runtime.cpp index 027ada34..01fe07cb 100644 --- a/cpp/models/pi05/src/native_graph_runtime.cpp +++ b/cpp/models/pi05/src/backend/native_graph_runtime.cpp @@ -1,6 +1,6 @@ -#include "flashrt/cpp/models/pi05/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/model/spec.h" #include diff --git a/cpp/models/pi05/src/native_thor_calibration_session.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_calibration_session.cpp similarity index 97% rename from cpp/models/pi05/src/native_thor_calibration_session.cpp rename to cpp/models/pi05/src/backends/sm110/native_thor_calibration_session.cpp index 36696a28..730679ad 100644 --- a/cpp/models/pi05/src/native_thor_calibration_session.cpp +++ b/cpp/models/pi05/src/backends/sm110/native_thor_calibration_session.cpp @@ -1,13 +1,13 @@ -#include "flashrt/cpp/models/pi05/native_thor_calibration_session.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h" #include "flashrt/cpp/loader/safetensors.h" #include "flashrt/cpp/loader/sha256.h" -#include "flashrt/cpp/models/pi05/io.h" -#include "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_thor_fp8_forward.h" -#include "flashrt/cpp/models/pi05/native_thor_style_precompute.h" -#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" -#include "flashrt/cpp/models/pi05/prompt_embed.h" +#include "flashrt/cpp/models/pi05/model/io.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h" +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" #include diff --git a/cpp/models/pi05/src/native_thor_fp8_forward.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_fp8_forward.cpp similarity index 99% rename from cpp/models/pi05/src/native_thor_fp8_forward.cpp rename to cpp/models/pi05/src/backends/sm110/native_thor_fp8_forward.cpp index 1db58508..92be7fa8 100644 --- a/cpp/models/pi05/src/native_thor_fp8_forward.cpp +++ b/cpp/models/pi05/src/backends/sm110/native_thor_fp8_forward.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_thor_fp8_forward.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h" #include diff --git a/cpp/models/pi05/src/native_thor_graph_owner.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_graph_owner.cpp similarity index 97% rename from cpp/models/pi05/src/native_thor_graph_owner.cpp rename to cpp/models/pi05/src/backends/sm110/native_thor_graph_owner.cpp index 46217b69..159940d9 100644 --- a/cpp/models/pi05/src/native_thor_graph_owner.cpp +++ b/cpp/models/pi05/src/backends/sm110/native_thor_graph_owner.cpp @@ -1,7 +1,7 @@ -#include "flashrt/cpp/models/pi05/native_thor_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h" -#include "flashrt/cpp/models/pi05/native_thor_style_precompute.h" -#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h" #include diff --git a/cpp/models/pi05/src/native_thor_kernel_driver.cu b/cpp/models/pi05/src/backends/sm110/native_thor_kernel_driver.cu similarity index 99% rename from cpp/models/pi05/src/native_thor_kernel_driver.cu rename to cpp/models/pi05/src/backends/sm110/native_thor_kernel_driver.cu index 0ab2b298..1a9067c1 100644 --- a/cpp/models/pi05/src/native_thor_kernel_driver.cu +++ b/cpp/models/pi05/src/backends/sm110/native_thor_kernel_driver.cu @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_thor_kernel_driver.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h" #include "activation.cuh" #include "attention_cublas.cuh" diff --git a/cpp/models/pi05/src/native_thor_setup_ops.cu b/cpp/models/pi05/src/backends/sm110/native_thor_setup_ops.cu similarity index 100% rename from cpp/models/pi05/src/native_thor_setup_ops.cu rename to cpp/models/pi05/src/backends/sm110/native_thor_setup_ops.cu diff --git a/cpp/models/pi05/src/native_thor_style_precompute.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_style_precompute.cpp similarity index 99% rename from cpp/models/pi05/src/native_thor_style_precompute.cpp rename to cpp/models/pi05/src/backends/sm110/native_thor_style_precompute.cpp index 25f996bd..80601069 100644 --- a/cpp/models/pi05/src/native_thor_style_precompute.cpp +++ b/cpp/models/pi05/src/backends/sm110/native_thor_style_precompute.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_thor_style_precompute.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h" #include diff --git a/cpp/models/pi05/src/native_thor_weight_materializer.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_weight_materializer.cpp similarity index 99% rename from cpp/models/pi05/src/native_thor_weight_materializer.cpp rename to cpp/models/pi05/src/backends/sm110/native_thor_weight_materializer.cpp index 163059fa..f0028027 100644 --- a/cpp/models/pi05/src/native_thor_weight_materializer.cpp +++ b/cpp/models/pi05/src/backends/sm110/native_thor_weight_materializer.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_thor_weight_materializer.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h" #include #include diff --git a/cpp/models/pi05/src/native_bf16_forward.cpp b/cpp/models/pi05/src/backends/sm120/native_bf16_forward.cpp similarity index 99% rename from cpp/models/pi05/src/native_bf16_forward.cpp rename to cpp/models/pi05/src/backends/sm120/native_bf16_forward.cpp index b84502aa..5c5f6094 100644 --- a/cpp/models/pi05/src/native_bf16_forward.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_bf16_forward.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_bf16_forward.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h" #include #include diff --git a/cpp/models/pi05/src/native_graph_owner.cpp b/cpp/models/pi05/src/backends/sm120/native_graph_owner.cpp similarity index 97% rename from cpp/models/pi05/src/native_graph_owner.cpp rename to cpp/models/pi05/src/backends/sm120/native_graph_owner.cpp index 946a2800..397acc47 100644 --- a/cpp/models/pi05/src/native_graph_owner.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_graph_owner.cpp @@ -1,9 +1,9 @@ -#include "flashrt/cpp/models/pi05/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h" -#include "flashrt/cpp/models/pi05/native_style_precompute.h" -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" -#include "flashrt/cpp/models/pi05/native_rtx_autotune.h" -#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h" +#include "flashrt/cpp/models/pi05/support/native_weight_materializer.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h" #include diff --git a/cpp/models/pi05/src/native_kernel_driver.cu b/cpp/models/pi05/src/backends/sm120/native_kernel_driver.cu similarity index 99% rename from cpp/models/pi05/src/native_kernel_driver.cu rename to cpp/models/pi05/src/backends/sm120/native_kernel_driver.cu index 2a4a7964..c02c38f8 100644 --- a/cpp/models/pi05/src/native_kernel_driver.cu +++ b/cpp/models/pi05/src/backends/sm120/native_kernel_driver.cu @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h" #include "activation.cuh" #include "elementwise.cuh" diff --git a/cpp/models/pi05/src/native_rtx_attention.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_attention.cpp similarity index 99% rename from cpp/models/pi05/src/native_rtx_attention.cpp rename to cpp/models/pi05/src/backends/sm120/native_rtx_attention.cpp index f77a6b91..40e2d812 100644 --- a/cpp/models/pi05/src/native_rtx_attention.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_attention.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rtx_attention.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h" #ifdef FLASHRT_CPP_WITH_CUDA_STAGING #include diff --git a/cpp/models/pi05/src/native_rtx_attention_driver.cu b/cpp/models/pi05/src/backends/sm120/native_rtx_attention_driver.cu similarity index 99% rename from cpp/models/pi05/src/native_rtx_attention_driver.cu rename to cpp/models/pi05/src/backends/sm120/native_rtx_attention_driver.cu index 2850b769..7aaca702 100644 --- a/cpp/models/pi05/src/native_rtx_attention_driver.cu +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_attention_driver.cu @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rtx_attention_driver.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h" #include "attention/fa2_wrapper.h" diff --git a/cpp/models/pi05/src/native_rtx_autotune.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_autotune.cpp similarity index 97% rename from cpp/models/pi05/src/native_rtx_autotune.cpp rename to cpp/models/pi05/src/backends/sm120/native_rtx_autotune.cpp index 31cbd319..eaf19fc2 100644 --- a/cpp/models/pi05/src/native_rtx_autotune.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_autotune.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rtx_autotune.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/src/native_rtx_calibration_session.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp similarity index 98% rename from cpp/models/pi05/src/native_rtx_calibration_session.cpp rename to cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp index 53f4b223..4c1ffb28 100644 --- a/cpp/models/pi05/src/native_rtx_calibration_session.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp @@ -1,10 +1,10 @@ -#include "flashrt/cpp/models/pi05/native_rtx_calibration_session.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h" #include "flashrt/cpp/loader/sha256.h" -#include "flashrt/cpp/models/pi05/io.h" -#include "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_graph_owner.h" -#include "flashrt/cpp/models/pi05/prompt_embed.h" +#include "flashrt/cpp/models/pi05/model/io.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" #include diff --git a/cpp/models/pi05/src/native_rtx_linear.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_linear.cpp similarity index 99% rename from cpp/models/pi05/src/native_rtx_linear.cpp rename to cpp/models/pi05/src/backends/sm120/native_rtx_linear.cpp index 9e018c7a..c09ebd27 100644 --- a/cpp/models/pi05/src/native_rtx_linear.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_linear.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rtx_linear.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h" #include diff --git a/cpp/models/pi05/src/native_rtx_weight_packer.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_weight_packer.cpp similarity index 98% rename from cpp/models/pi05/src/native_rtx_weight_packer.cpp rename to cpp/models/pi05/src/backends/sm120/native_rtx_weight_packer.cpp index 5511337f..a08e8c48 100644 --- a/cpp/models/pi05/src/native_rtx_weight_packer.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_weight_packer.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rtx_weight_packer.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h" #include diff --git a/cpp/models/pi05/src/native_style_precompute.cu b/cpp/models/pi05/src/backends/sm120/native_style_precompute.cu similarity index 99% rename from cpp/models/pi05/src/native_style_precompute.cu rename to cpp/models/pi05/src/backends/sm120/native_style_precompute.cu index 07c496ab..003a07da 100644 --- a/cpp/models/pi05/src/native_style_precompute.cu +++ b/cpp/models/pi05/src/backends/sm120/native_style_precompute.cu @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_style_precompute.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h" #include diff --git a/cpp/models/pi05/src/io.cpp b/cpp/models/pi05/src/model/io.cpp similarity index 98% rename from cpp/models/pi05/src/io.cpp rename to cpp/models/pi05/src/model/io.cpp index 27eb5c6e..5ad26d3a 100644 --- a/cpp/models/pi05/src/io.cpp +++ b/cpp/models/pi05/src/model/io.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/io.h" +#include "flashrt/cpp/models/pi05/model/io.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/src/prompt_embed.cpp b/cpp/models/pi05/src/model/prompt_embed.cpp similarity index 98% rename from cpp/models/pi05/src/prompt_embed.cpp rename to cpp/models/pi05/src/model/prompt_embed.cpp index 50a9c046..1f8fdaae 100644 --- a/cpp/models/pi05/src/prompt_embed.cpp +++ b/cpp/models/pi05/src/model/prompt_embed.cpp @@ -1,6 +1,6 @@ -#include "flashrt/cpp/models/pi05/prompt_embed.h" +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" -#include "flashrt/cpp/models/pi05/prompt_format.h" +#include "flashrt/cpp/models/pi05/model/prompt_format.h" #ifdef FLASHRT_CPP_WITH_CUDA_STAGING #include diff --git a/cpp/models/pi05/src/prompt_format.cpp b/cpp/models/pi05/src/model/prompt_format.cpp similarity index 98% rename from cpp/models/pi05/src/prompt_format.cpp rename to cpp/models/pi05/src/model/prompt_format.cpp index 5e1c36a1..7f0cf64f 100644 --- a/cpp/models/pi05/src/prompt_format.cpp +++ b/cpp/models/pi05/src/model/prompt_format.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/prompt_format.h" +#include "flashrt/cpp/models/pi05/model/prompt_format.h" #include #include diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/model/runtime.cpp similarity index 99% rename from cpp/models/pi05/src/runtime.cpp rename to cpp/models/pi05/src/model/runtime.cpp index ecdbae61..7faba9d5 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/model/runtime.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/runtime.h" +#include "flashrt/cpp/models/pi05/model/runtime.h" #include #include diff --git a/cpp/models/pi05/src/spec.cpp b/cpp/models/pi05/src/model/spec.cpp similarity index 96% rename from cpp/models/pi05/src/spec.cpp rename to cpp/models/pi05/src/model/spec.cpp index f6eb0067..617bf36e 100644 --- a/cpp/models/pi05/src/spec.cpp +++ b/cpp/models/pi05/src/model/spec.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/model/spec.h" #include diff --git a/cpp/models/pi05/src/c_api.cpp b/cpp/models/pi05/src/service/c_api.cpp similarity index 100% rename from cpp/models/pi05/src/c_api.cpp rename to cpp/models/pi05/src/service/c_api.cpp diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/service/config_map.h similarity index 99% rename from cpp/models/pi05/src/config_map.h rename to cpp/models/pi05/src/service/config_map.h index a0bc1d21..fb78d803 100644 --- a/cpp/models/pi05/src/config_map.h +++ b/cpp/models/pi05/src/service/config_map.h @@ -5,7 +5,7 @@ #define FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H #include "flashrt/cpp/models/pi05/c_api.h" -#include "flashrt/cpp/models/pi05/runtime.h" +#include "flashrt/cpp/models/pi05/model/runtime.h" #include diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/service/model_runtime.cpp similarity index 100% rename from cpp/models/pi05/src/model_runtime.cpp rename to cpp/models/pi05/src/service/model_runtime.cpp diff --git a/cpp/models/pi05/src/native_calibration_c_api.cpp b/cpp/models/pi05/src/service/native_calibration_c_api.cpp similarity index 97% rename from cpp/models/pi05/src/native_calibration_c_api.cpp rename to cpp/models/pi05/src/service/native_calibration_c_api.cpp index 28373b67..6757b7ac 100644 --- a/cpp/models/pi05/src/native_calibration_c_api.cpp +++ b/cpp/models/pi05/src/service/native_calibration_c_api.cpp @@ -1,19 +1,19 @@ #include "flashrt/cpp/models/pi05/c_api.h" #include "config_map.h" -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/model/spec.h" #include "native_open_internal.h" #if defined(FLASHRT_CPP_HAS_SENTENCEPIECE) && \ (defined(FLASHRT_CPP_WITH_FA2) || \ defined(FLASHRT_CPP_WITH_THOR_FP8)) #define FLASHRT_CPP_HAS_NATIVE_CALIBRATION 1 -#include "flashrt/cpp/models/pi05/native_calibration_session.h" +#include "flashrt/cpp/models/pi05/backend/native_calibration_session.h" #if defined(FLASHRT_CPP_WITH_FA2) -#include "flashrt/cpp/models/pi05/native_rtx_calibration_session.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h" #endif #if defined(FLASHRT_CPP_WITH_THOR_FP8) -#include "flashrt/cpp/models/pi05/native_thor_calibration_session.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h" #endif #include #endif diff --git a/cpp/models/pi05/src/native_model_runtime.cpp b/cpp/models/pi05/src/service/native_model_runtime.cpp similarity index 98% rename from cpp/models/pi05/src/native_model_runtime.cpp rename to cpp/models/pi05/src/service/native_model_runtime.cpp index 01fd3e12..5a550771 100644 --- a/cpp/models/pi05/src/native_model_runtime.cpp +++ b/cpp/models/pi05/src/service/native_model_runtime.cpp @@ -6,14 +6,14 @@ #include "config_map.h" #include "flashrt/cpp/loader/sha256.h" #include "flashrt/cpp/models/pi05/model_runtime.h" -#include "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_graph_runtime.h" -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/model/spec.h" #if defined(FLASHRT_CPP_WITH_FA2) -#include "flashrt/cpp/models/pi05/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h" #endif #if defined(FLASHRT_CPP_WITH_THOR_FP8) -#include "flashrt/cpp/models/pi05/native_thor_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h" #endif #include diff --git a/cpp/models/pi05/src/native_open.cpp b/cpp/models/pi05/src/service/native_open.cpp similarity index 99% rename from cpp/models/pi05/src/native_open.cpp rename to cpp/models/pi05/src/service/native_open.cpp index cbab5956..d73050a6 100644 --- a/cpp/models/pi05/src/native_open.cpp +++ b/cpp/models/pi05/src/service/native_open.cpp @@ -1,6 +1,6 @@ #include "flashrt/model_runtime.h" #include "flashrt/cpp/loader/safetensors.h" -#include "flashrt/cpp/models/pi05/native_weights.h" +#include "flashrt/cpp/models/pi05/support/native_weights.h" #include "flashrt/cpp/native/config_object.h" #include "native_open_internal.h" diff --git a/cpp/models/pi05/src/native_open_internal.h b/cpp/models/pi05/src/service/native_open_internal.h similarity index 100% rename from cpp/models/pi05/src/native_open_internal.h rename to cpp/models/pi05/src/service/native_open_internal.h diff --git a/cpp/models/pi05/src/native_calibration.cpp b/cpp/models/pi05/src/support/native_calibration.cpp similarity index 99% rename from cpp/models/pi05/src/native_calibration.cpp rename to cpp/models/pi05/src/support/native_calibration.cpp index b460e360..e32b96a2 100644 --- a/cpp/models/pi05/src/native_calibration.cpp +++ b/cpp/models/pi05/src/support/native_calibration.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" #include "flashrt/cpp/loader/safetensors.h" diff --git a/cpp/models/pi05/src/native_device_weights.cpp b/cpp/models/pi05/src/support/native_device_weights.cpp similarity index 99% rename from cpp/models/pi05/src/native_device_weights.cpp rename to cpp/models/pi05/src/support/native_device_weights.cpp index a8bddb53..ef0aac41 100644 --- a/cpp/models/pi05/src/native_device_weights.cpp +++ b/cpp/models/pi05/src/support/native_device_weights.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_device_weights.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" #ifdef FLASHRT_CPP_WITH_CUDA_STAGING #include diff --git a/cpp/models/pi05/src/native_quantization.cu b/cpp/models/pi05/src/support/native_quantization.cu similarity index 98% rename from cpp/models/pi05/src/native_quantization.cu rename to cpp/models/pi05/src/support/native_quantization.cu index d2a6366c..cb4eba3b 100644 --- a/cpp/models/pi05/src/native_quantization.cu +++ b/cpp/models/pi05/src/support/native_quantization.cu @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_quantization.h" +#include "flashrt/cpp/models/pi05/support/native_quantization.h" #include diff --git a/cpp/models/pi05/src/native_quantization_unavailable.cpp b/cpp/models/pi05/src/support/native_quantization_unavailable.cpp similarity index 93% rename from cpp/models/pi05/src/native_quantization_unavailable.cpp rename to cpp/models/pi05/src/support/native_quantization_unavailable.cpp index d043b01d..cec2fb4f 100644 --- a/cpp/models/pi05/src/native_quantization_unavailable.cpp +++ b/cpp/models/pi05/src/support/native_quantization_unavailable.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_quantization.h" +#include "flashrt/cpp/models/pi05/support/native_quantization.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/src/native_rope.cu b/cpp/models/pi05/src/support/native_rope.cu similarity index 97% rename from cpp/models/pi05/src/native_rope.cu rename to cpp/models/pi05/src/support/native_rope.cu index ca57bfe0..e9a926c0 100644 --- a/cpp/models/pi05/src/native_rope.cu +++ b/cpp/models/pi05/src/support/native_rope.cu @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rope.h" +#include "flashrt/cpp/models/pi05/support/native_rope.h" #include #include diff --git a/cpp/models/pi05/src/native_rope_unavailable.cpp b/cpp/models/pi05/src/support/native_rope_unavailable.cpp similarity index 86% rename from cpp/models/pi05/src/native_rope_unavailable.cpp rename to cpp/models/pi05/src/support/native_rope_unavailable.cpp index 84655921..ec43576e 100644 --- a/cpp/models/pi05/src/native_rope_unavailable.cpp +++ b/cpp/models/pi05/src/support/native_rope_unavailable.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_rope.h" +#include "flashrt/cpp/models/pi05/support/native_rope.h" namespace flashrt { namespace models { diff --git a/cpp/models/pi05/src/native_weight_materializer.cpp b/cpp/models/pi05/src/support/native_weight_materializer.cpp similarity index 99% rename from cpp/models/pi05/src/native_weight_materializer.cpp rename to cpp/models/pi05/src/support/native_weight_materializer.cpp index 7531dbbb..21dd765f 100644 --- a/cpp/models/pi05/src/native_weight_materializer.cpp +++ b/cpp/models/pi05/src/support/native_weight_materializer.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_weight_materializer.h" +#include "flashrt/cpp/models/pi05/support/native_weight_materializer.h" #include #include diff --git a/cpp/models/pi05/src/native_weight_ops.cpp b/cpp/models/pi05/src/support/native_weight_ops.cpp similarity index 99% rename from cpp/models/pi05/src/native_weight_ops.cpp rename to cpp/models/pi05/src/support/native_weight_ops.cpp index 85388399..aff9888e 100644 --- a/cpp/models/pi05/src/native_weight_ops.cpp +++ b/cpp/models/pi05/src/support/native_weight_ops.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_weight_ops.h" +#include "flashrt/cpp/models/pi05/support/native_weight_ops.h" #include #include diff --git a/cpp/models/pi05/src/native_weight_packer.cpp b/cpp/models/pi05/src/support/native_weight_packer.cpp similarity index 99% rename from cpp/models/pi05/src/native_weight_packer.cpp rename to cpp/models/pi05/src/support/native_weight_packer.cpp index 54e3d9b3..362c1ad6 100644 --- a/cpp/models/pi05/src/native_weight_packer.cpp +++ b/cpp/models/pi05/src/support/native_weight_packer.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_weight_packer.h" +#include "flashrt/cpp/models/pi05/support/native_weight_packer.h" #include #include diff --git a/cpp/models/pi05/src/native_weights.cpp b/cpp/models/pi05/src/support/native_weights.cpp similarity index 98% rename from cpp/models/pi05/src/native_weights.cpp rename to cpp/models/pi05/src/support/native_weights.cpp index 5abedc13..bf684832 100644 --- a/cpp/models/pi05/src/native_weights.cpp +++ b/cpp/models/pi05/src/support/native_weights.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/native_weights.h" +#include "flashrt/cpp/models/pi05/support/native_weights.h" #include diff --git a/cpp/models/pi05/src/native_workspace.cpp b/cpp/models/pi05/src/support/native_workspace.cpp similarity index 99% rename from cpp/models/pi05/src/native_workspace.cpp rename to cpp/models/pi05/src/support/native_workspace.cpp index 01c7ef2e..2d80f06d 100644 --- a/cpp/models/pi05/src/native_workspace.cpp +++ b/cpp/models/pi05/src/support/native_workspace.cpp @@ -1,5 +1,5 @@ -#include "flashrt/cpp/models/pi05/native_workspace.h" -#include "flashrt/cpp/models/pi05/native_rope.h" +#include "flashrt/cpp/models/pi05/support/native_workspace.h" +#include "flashrt/cpp/models/pi05/support/native_rope.h" #ifdef FLASHRT_CPP_WITH_CUDA_STAGING #include diff --git a/cpp/tests/pi05_native_fp8_calibration_probe.cpp b/cpp/tests/pi05_native_fp8_calibration_probe.cpp index f0089adb..74ed6c98 100644 --- a/cpp/tests/pi05_native_fp8_calibration_probe.cpp +++ b/cpp/tests/pi05_native_fp8_calibration_probe.cpp @@ -1,7 +1,7 @@ #include "flashrt/model_runtime.h" #include "flashrt/cpp/modalities/types.h" #include "flashrt/cpp/models/pi05/c_api.h" -#include "flashrt/cpp/models/pi05/native_calibration.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" #include diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index 57855ef4..3c89b24f 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -1,7 +1,7 @@ #include "flashrt/cpp/modalities/action.h" #include "flashrt/cpp/modalities/text.h" #include "flashrt/cpp/modalities/vision.h" -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/model/spec.h" #include diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index f4d012e1..5cc44045 100644 --- a/cpp/tests/test_modalities.cpp +++ b/cpp/tests/test_modalities.cpp @@ -1,7 +1,7 @@ #include "flashrt/cpp/modalities/action.h" #include "flashrt/cpp/modalities/vision.h" -#include "flashrt/cpp/models/pi05/io.h" -#include "flashrt/cpp/models/pi05/spec.h" +#include "flashrt/cpp/models/pi05/model/io.h" +#include "flashrt/cpp/models/pi05/model/spec.h" #include #include diff --git a/cpp/tests/test_pi05_native_calibration.cpp b/cpp/tests/test_pi05_native_calibration.cpp index 65448417..3e9f6875 100644 --- a/cpp/tests/test_pi05_native_calibration.cpp +++ b/cpp/tests/test_pi05_native_calibration.cpp @@ -1,5 +1,5 @@ -#include "flashrt/cpp/models/pi05/native_calibration.h" -#include "flashrt/cpp/models/pi05/native_calibration_session.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backend/native_calibration_session.h" #include #include diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp index bc6005d7..758d3720 100644 --- a/cpp/tests/test_pi05_native_open.cpp +++ b/cpp/tests/test_pi05_native_open.cpp @@ -1,5 +1,5 @@ #include "flashrt/model_runtime.h" -#include "flashrt/cpp/models/pi05/native_weights.h" +#include "flashrt/cpp/models/pi05/support/native_weights.h" #include #include diff --git a/cpp/tests/test_pi05_prompt_embed.cpp b/cpp/tests/test_pi05_prompt_embed.cpp index 89c05560..c625a615 100644 --- a/cpp/tests/test_pi05_prompt_embed.cpp +++ b/cpp/tests/test_pi05_prompt_embed.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/prompt_embed.h" +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" #ifdef FLASHRT_CPP_WITH_CUDA_STAGING #include diff --git a/cpp/tests/test_pi05_prompt_format.cpp b/cpp/tests/test_pi05_prompt_format.cpp index 7184b0ea..f7431909 100644 --- a/cpp/tests/test_pi05_prompt_format.cpp +++ b/cpp/tests/test_pi05_prompt_format.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/prompt_format.h" +#include "flashrt/cpp/models/pi05/model/prompt_format.h" #include #include diff --git a/cpp/tests/test_pi05_runtime.cpp b/cpp/tests/test_pi05_runtime.cpp index fe594dd7..5e2a1e18 100644 --- a/cpp/tests/test_pi05_runtime.cpp +++ b/cpp/tests/test_pi05_runtime.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/runtime.h" +#include "flashrt/cpp/models/pi05/model/runtime.h" #include #include diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 77f7516c..54ec9f45 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -29,9 +29,9 @@ 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` +- Native verb implementation: `cpp/models/pi05/src/service/model_runtime.cpp` +- C++ modality binding: `cpp/models/pi05/src/model/runtime.cpp`, + `cpp/models/pi05/src/model/io.cpp`, `cpp/models/pi05/src/model/spec.cpp` `io="native"` and `io="native_v2"` are declaration-only handoffs. Their discovery manifest carries `declaration_only: true`; callers must pass them to From 01b33c372513763fe5683c4038f24b5c98bcc77a Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 07:55:44 -0400 Subject: [PATCH 81/83] refactor(pi05): isolate hardware backend targets --- CONTRIBUTING.md | 15 ++++ cpp/models/pi05/CMakeLists.txt | 135 ++++++++++++++++++++------------- docs/pi05_io_contract.md | 16 ++++ docs/pr_review_checklist.md | 4 + 4 files changed, 116 insertions(+), 54 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8dae9c99..72ada51f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -237,6 +237,21 @@ reviewers hold every PR to: - Schema shared by multiple producers needs checked-in canonical records; every producer compares independently against that golden face. +### Native C++ Model Ownership + +Keep a native model producer split by ownership: model semantics, model-private +support, a coarse backend session seam, hardware backend implementations, and +the public service/plugin face. Do not expose those private classes as installed +headers or move model topology into generic runtime/exec code. + +Each hardware backend target must compile only its own implementation and link +the explicitly shared mechanisms it needs. An aggregate compatibility target +may forward link dependencies, but it must not become a source bucket that +injects every hardware implementation and architecture flag into every build. +Validation targets must be registered from actual backend capabilities; do not +compile one backend's private classes for another device just to preserve a test +count. + ### Calibration And Precision FP8/NVFP4 changes must preserve the calibration cache contract described in diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index 08d286f5..de910620 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -21,8 +21,7 @@ set(_pi05_sources src/support/native_weight_materializer.cpp src/support/native_calibration.cpp src/support/native_workspace.cpp - src/backend/native_calibration_session.cpp - src/backends/sm120/native_rtx_attention.cpp) + src/backend/native_calibration_session.cpp) if(FLASHRT_CPP_WITH_CUDA_KERNELS) list(APPEND _pi05_sources src/support/native_quantization.cu @@ -46,37 +45,75 @@ if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) - add_library(flashrt_cpp_pi05_kernels STATIC - src/backend/native_graph_runtime.cpp - src/backends/sm120/native_bf16_forward.cpp - src/backends/sm120/native_kernel_driver.cu - src/backends/sm120/native_rtx_linear.cpp - src/backends/sm120/native_rtx_weight_packer.cpp - src/backends/sm120/native_style_precompute.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/activation.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/elementwise.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/fusion.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/norm.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/patch_embed.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/quantize.cu - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/rope.cu) - target_include_directories(flashrt_cpp_pi05_kernels - PRIVATE ${_pi05_public_include} - ${_pi05_internal_include} - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm - ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels) - target_link_libraries(flashrt_cpp_pi05_kernels - PUBLIC flashrt_cpp_pi05 flashrt_cpp_native_cuda - 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>) + function(_flashrt_configure_pi05_cuda_target target) + target_include_directories(${target} PRIVATE + ${_pi05_public_include} + ${_pi05_internal_include} + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels) + target_link_libraries(${target} PUBLIC + flashrt_cpp_pi05 flashrt_cpp_native_cuda + CUDA::cublas CUDA::cublasLt CUDA::cudart) + set_target_properties(${target} PROPERTIES + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON) + target_compile_options(${target} PRIVATE + $<$: + --expt-relaxed-constexpr -O3 + --ftz=true --prec-div=false --prec-sqrt=false>) + endfunction() + + add_library(flashrt_cpp_pi05_kernels INTERFACE) + set(_pi05_backend_targets) + + if(FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8) + add_library(flashrt_cpp_pi05_backend_cuda STATIC + src/backend/native_graph_runtime.cpp + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/activation.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/elementwise.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/norm.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/patch_embed.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/quantize.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/rope.cu) + _flashrt_configure_pi05_cuda_target(flashrt_cpp_pi05_backend_cuda) + + if(FLASHRT_CPP_WITH_THOR_FP8) + target_compile_options(flashrt_cpp_pi05_backend_cuda PRIVATE + $<$: + --expt-extended-lambda --use_fast_math + -gencode=arch=compute_110a,code=sm_110a>) + set_target_properties(flashrt_cpp_pi05_backend_cuda PROPERTIES + CUDA_ARCHITECTURES 110a) + endif() + endif() + + if(FLASHRT_CPP_WITH_FA2) + add_library(flashrt_cpp_pi05_backend_sm120 STATIC + src/backends/sm120/native_bf16_forward.cpp + src/backends/sm120/native_graph_owner.cpp + src/backends/sm120/native_kernel_driver.cu + src/backends/sm120/native_rtx_attention.cpp + src/backends/sm120/native_rtx_attention_driver.cu + src/backends/sm120/native_rtx_autotune.cpp + src/backends/sm120/native_rtx_linear.cpp + src/backends/sm120/native_rtx_weight_packer.cpp + src/backends/sm120/native_style_precompute.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/dit_bf16.cu + ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/fusion.cu) + if(FLASHRT_CPP_WITH_SENTENCEPIECE) + target_sources(flashrt_cpp_pi05_backend_sm120 PRIVATE + src/backends/sm120/native_rtx_calibration_session.cpp) + endif() + _flashrt_configure_pi05_cuda_target(flashrt_cpp_pi05_backend_sm120) + target_include_directories(flashrt_cpp_pi05_backend_sm120 PRIVATE + ${FLASHRT_CPP_SOURCE_DIR}/../csrc) + target_link_libraries(flashrt_cpp_pi05_backend_sm120 PUBLIC + flashrt_cpp_pi05_backend_cuda flashrt_cpp_fa2_external) + target_compile_definitions(flashrt_cpp_pi05_backend_sm120 + PUBLIC FLASHRT_CPP_WITH_FA2=1) + list(APPEND _pi05_backend_targets flashrt_cpp_pi05_backend_sm120) + endif() if(FLASHRT_CPP_WITH_THOR_FP8) add_library(flashrt_cpp_pi05_thor_precise OBJECT @@ -96,7 +133,7 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) CUDA_STANDARD_REQUIRED ON POSITION_INDEPENDENT_CODE ON) - target_sources(flashrt_cpp_pi05_kernels PRIVATE + add_library(flashrt_cpp_pi05_backend_sm110 STATIC $ src/backends/sm110/native_thor_kernel_driver.cu src/backends/sm110/native_thor_weight_materializer.cpp @@ -108,38 +145,28 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/attention_cublas.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/decoder_fused.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/softmax.cu) - target_include_directories(flashrt_cpp_pi05_kernels PRIVATE + _flashrt_configure_pi05_cuda_target(flashrt_cpp_pi05_backend_sm110) + target_include_directories(flashrt_cpp_pi05_backend_sm110 PRIVATE ${FLASHRT_CPP_SOURCE_DIR}/../csrc/attention ${FLASHRT_CPP_CUTLASS_DIR}/examples/77_blackwell_fmha ${FLASHRT_CPP_CUTLASS_DIR}/include ${FLASHRT_CPP_CUTLASS_DIR}/tools/util/include) - target_compile_definitions(flashrt_cpp_pi05_kernels + target_link_libraries(flashrt_cpp_pi05_backend_sm110 PUBLIC + flashrt_cpp_pi05_backend_cuda) + target_compile_definitions(flashrt_cpp_pi05_backend_sm110 PUBLIC FLASHRT_CPP_WITH_THOR_FP8=1) - target_compile_options(flashrt_cpp_pi05_kernels PRIVATE + target_compile_options(flashrt_cpp_pi05_backend_sm110 PRIVATE $<$: --expt-extended-lambda --use_fast_math -gencode=arch=compute_110a,code=sm_110a>) - set_target_properties(flashrt_cpp_pi05_kernels PROPERTIES + set_target_properties(flashrt_cpp_pi05_backend_sm110 PROPERTIES CUDA_ARCHITECTURES 110a CUDA_RESOLVE_DEVICE_SYMBOLS ON) + list(APPEND _pi05_backend_targets flashrt_cpp_pi05_backend_sm110) endif() - if(FLASHRT_CPP_WITH_FA2) - target_sources(flashrt_cpp_pi05_kernels PRIVATE - src/backends/sm120/native_graph_owner.cpp - src/backends/sm120/native_rtx_autotune.cpp - src/backends/sm120/native_rtx_attention_driver.cu) - if(FLASHRT_CPP_WITH_SENTENCEPIECE) - target_sources(flashrt_cpp_pi05_kernels PRIVATE - src/backends/sm120/native_rtx_calibration_session.cpp) - endif() - target_include_directories(flashrt_cpp_pi05_kernels PRIVATE - ${FLASHRT_CPP_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() + target_link_libraries(flashrt_cpp_pi05_kernels INTERFACE + ${_pi05_backend_targets}) endif() add_library(flashrt_cpp_pi05_c SHARED diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md index 54ec9f45..cac778c3 100644 --- a/docs/pi05_io_contract.md +++ b/docs/pi05_io_contract.md @@ -13,6 +13,22 @@ 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. +## Private Source Ownership + +The native producer is organized by ownership rather than execution order: + +- `src/model/` owns PI0.5 topology and prompt/state/IO semantics; +- `src/support/` owns model-private weights, calibration data, and workspace; +- `src/backend/` defines the coarse model-private backend session seam; +- `src/backends/sm120/` and `src/backends/sm110/` own hardware implementations; +- `src/service/` owns C ABI construction, native open, and runtime publication. + +The matching private headers use the same layout. None of these headers is an +installed API. `flashrt_cpp_pi05_backend_cuda` contains only mechanisms shared +by both CUDA backends; the SM120 and SM110 targets compile only their own +sources and dependencies. `flashrt_cpp_pi05_kernels` is an internal interface +target kept as the aggregate link entry, not a source bucket. + ## Adopted-export Legacy Face The Python-produced Pi0.5 `io="native"` export declares three host-visible diff --git a/docs/pr_review_checklist.md b/docs/pr_review_checklist.md index cd5e81ad..25e39b24 100644 --- a/docs/pr_review_checklist.md +++ b/docs/pr_review_checklist.md @@ -758,6 +758,10 @@ gates: 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. +- [ ] Each native hardware target compiles only its own implementation; shared + CUDA mechanisms contain no model policy or backend-specific flags. +- [ ] Aggregate model targets only forward link dependencies, and validation + cases are registered against the backend capability they exercise. - [ ] 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. From 65f3fbcc55b44be3bd39f786cb1d754ed57be539 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 08:00:39 -0400 Subject: [PATCH 82/83] refactor(pi05): remove unused weight packer --- cpp/models/pi05/CMakeLists.txt | 1 - .../pi05/support/native_device_weights.h | 3 - .../pi05/support/native_weight_packer.h | 40 ---- .../src/support/native_device_weights.cpp | 35 ---- .../pi05/src/support/native_weight_packer.cpp | 194 ------------------ 5 files changed, 273 deletions(-) delete mode 100644 cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h delete mode 100644 cpp/models/pi05/src/support/native_weight_packer.cpp diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index de910620..a6f5c809 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -17,7 +17,6 @@ set(_pi05_sources src/support/native_weights.cpp src/support/native_weight_ops.cpp src/support/native_device_weights.cpp - src/support/native_weight_packer.cpp src/support/native_weight_materializer.cpp src/support/native_calibration.cpp src/support/native_workspace.cpp diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h index b130dacb..e1671de4 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h @@ -49,9 +49,6 @@ class NativeDeviceWeightStore { const std::string& name, const std::vector& shape, NativeWeightDType dtype); - 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/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h deleted file mode 100644 index d200d035..00000000 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_packer.h +++ /dev/null @@ -1,40 +0,0 @@ -#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_PACKER_H -#define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_PACKER_H - -#include "flashrt/cpp/models/pi05/support/native_device_weights.h" -#include "flashrt/cpp/models/pi05/support/native_quantization.h" - -namespace flashrt { -namespace models { -namespace pi05 { - -class NativeWeightPacker { -public: - explicit NativeWeightPacker(NativeDeviceWeightStore* weights) - : weights_(weights) {} - - modalities::Status pack_fp8(const std::string& name, bool transpose); - modalities::Status pack_fp8_as(const std::string& source_name, - const std::string& packed_name, - bool transpose); - modalities::Status pack_int8(const std::string& name); - modalities::Status merge_bf16_columns(const std::string& left_name, - const std::string& right_name, - const std::string& output_name); - modalities::Status pack_all_fp8(bool transpose); - modalities::Status pack_vision_int8(); - modalities::Status pack_encoder_int8(); - modalities::Status pack_decoder_int8(); - -private: - modalities::Status load_bf16(const std::string& name, - NativeFloatTensor* out) const; - - NativeDeviceWeightStore* weights_ = nullptr; -}; - -} // namespace pi05 -} // namespace models -} // namespace flashrt - -#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_PACKER_H diff --git a/cpp/models/pi05/src/support/native_device_weights.cpp b/cpp/models/pi05/src/support/native_device_weights.cpp index ef0aac41..6c0d4e28 100644 --- a/cpp/models/pi05/src/support/native_device_weights.cpp +++ b/cpp/models/pi05/src/support/native_device_weights.cpp @@ -157,41 +157,6 @@ 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/support/native_weight_packer.cpp b/cpp/models/pi05/src/support/native_weight_packer.cpp deleted file mode 100644 index 362c1ad6..00000000 --- a/cpp/models/pi05/src/support/native_weight_packer.cpp +++ /dev/null @@ -1,194 +0,0 @@ -#include "flashrt/cpp/models/pi05/support/native_weight_packer.h" - -#include -#include - -namespace flashrt { -namespace models { -namespace pi05 { -namespace { - -modalities::Status invalid(const char* message) { - return modalities::Status::error(modalities::StatusCode::kInvalidArgument, - message); -} - -} // namespace - -modalities::Status NativeWeightPacker::load_bf16( - const std::string& name, - NativeFloatTensor* out) const { - if (!weights_ || !out) return invalid("native weight packer is invalid"); - NativeBf16Tensor source; - modalities::Status st = weights_->download_bf16(name, &source); - if (!st.ok_status()) return st; - NativeFloatTensor result; - result.shape = source.shape; - result.values.resize(source.values.size()); - for (std::size_t i = 0; i < source.values.size(); ++i) { - result.values[i] = - modalities::bfloat16_to_float(source.values[i]); - } - *out = std::move(result); - return modalities::Status::ok(); -} - -modalities::Status NativeWeightPacker::pack_fp8( - const std::string& name, - bool transpose) { - return pack_fp8_as(name, name, transpose); -} - -modalities::Status NativeWeightPacker::pack_fp8_as( - const std::string& source_name, - const std::string& packed_name, - bool transpose) { - NativeFloatTensor source; - modalities::Status st = load_bf16(source_name, &source); - if (!st.ok_status()) return st; - NativeFp8Tensor packed; - st = native_quantize_fp8_e4m3(source, transpose, &packed); - if (!st.ok_status()) return st; - const std::string prefix = "fp8." + packed_name; - st = weights_->upload_bytes(prefix, packed.shape, - NativeWeightDType::kFp8E4M3, - packed.values.data(), packed.values.size()); - if (!st.ok_status()) return st; - return weights_->upload_bytes(prefix + ".scale", {1}, - NativeWeightDType::kFloat32, - &packed.scale, sizeof(packed.scale)); -} - -modalities::Status NativeWeightPacker::pack_int8(const std::string& name) { - NativeFloatTensor source; - modalities::Status st = load_bf16(name, &source); - if (!st.ok_status()) return st; - NativeInt8Tensor packed; - st = native_quantize_int8_per_output(source, &packed); - if (!st.ok_status()) return st; - const std::string prefix = "int8." + name; - st = weights_->upload_bytes(prefix, packed.shape, - NativeWeightDType::kInt8, - packed.values.data(), packed.values.size()); - if (!st.ok_status()) return st; - return weights_->upload_bytes( - prefix + ".scale", {static_cast(packed.scales.size())}, - NativeWeightDType::kFloat32, packed.scales.data(), - packed.scales.size() * sizeof(float)); -} - -modalities::Status NativeWeightPacker::merge_bf16_columns( - const std::string& left_name, - const std::string& right_name, - const std::string& output_name) { - NativeFloatTensor left; - NativeFloatTensor right; - NativeFloatTensor merged; - modalities::Status st = load_bf16(left_name, &left); - if (!st.ok_status()) return st; - st = load_bf16(right_name, &right); - if (!st.ok_status()) return st; - st = native_concat_columns(left, right, &merged); - if (!st.ok_status()) return st; - NativeBf16Tensor bf16; - st = native_to_bf16(merged, &bf16); - if (!st.ok_status()) return st; - return weights_->upload(output_name, bf16); -} - -modalities::Status NativeWeightPacker::pack_all_fp8(bool transpose) { - if (!weights_) return invalid("native weight packer is invalid"); - modalities::Status st; - for (int layer = 0; layer < 27; ++layer) { - for (const char* stem : {"vision_attn_qkv_w_", "vision_attn_o_w_", - "vision_ffn_up_w_", - "vision_ffn_down_w_"}) { - st = pack_fp8(std::string(stem) + std::to_string(layer), - transpose); - if (!st.ok_status()) return st; - } - } - st = pack_fp8_as("encoder_multi_modal_projector_w", - "vision_projector_w", transpose); - if (!st.ok_status()) return st; - - for (int layer = 0; layer < 18; ++layer) { - const std::string suffix = std::to_string(layer); - const std::string gate_up = "encoder_ffn_gate_up_w_" + suffix; - st = merge_bf16_columns("encoder_ffn_gate_w_" + suffix, - "encoder_ffn_up_w_" + suffix, gate_up); - if (!st.ok_status()) return st; - for (const std::string& name : { - "encoder_attn_qkv_w_" + suffix, - "encoder_attn_o_w_" + suffix, - gate_up, - "encoder_ffn_down_w_" + suffix}) { - st = pack_fp8(name, transpose); - if (!st.ok_status()) return st; - } - } - for (int layer = 0; layer < 18; ++layer) { - const std::string suffix = std::to_string(layer); - for (const std::string& name : { - "decoder_attn_qkv_w_" + suffix, - "decoder_attn_o_w_" + suffix, - "decoder_ffn_gate_up_w_" + suffix, - "decoder_ffn_down_w_" + suffix}) { - st = pack_fp8(name, transpose); - if (!st.ok_status()) return st; - } - } - return modalities::Status::ok(); -} - -modalities::Status NativeWeightPacker::pack_vision_int8() { - if (!weights_) return invalid("native weight packer is invalid"); - for (int layer = 0; layer < 27; ++layer) { - for (const char* stem : {"vision_attn_qkv_w_", "vision_attn_o_w_", - "vision_ffn_up_w_", - "vision_ffn_down_w_"}) { - const modalities::Status st = - pack_int8(std::string(stem) + std::to_string(layer)); - if (!st.ok_status()) return st; - } - } - return modalities::Status::ok(); -} - -modalities::Status NativeWeightPacker::pack_encoder_int8() { - if (!weights_) return invalid("native weight packer is invalid"); - for (int layer = 0; layer < 18; ++layer) { - const std::string suffix = std::to_string(layer); - for (const std::string& name : { - "encoder_attn_qkv_w_" + suffix, - "encoder_attn_o_w_" + suffix, - "encoder_ffn_gate_w_" + suffix, - "encoder_ffn_up_w_" + suffix, - "encoder_ffn_down_w_" + suffix}) { - const modalities::Status st = pack_int8(name); - if (!st.ok_status()) return st; - } - } - return modalities::Status::ok(); -} - -modalities::Status NativeWeightPacker::pack_decoder_int8() { - if (!weights_) return invalid("native weight packer is invalid"); - for (int layer = 0; layer < 18; ++layer) { - const std::string suffix = std::to_string(layer); - for (const std::string& name : { - "decoder_attn_qkv_w_" + suffix, - "decoder_attn_o_w_" + suffix, - "decoder_ffn_gate_w_" + suffix, - "decoder_ffn_up_w_" + suffix, - "decoder_ffn_down_w_" + suffix}) { - const modalities::Status st = pack_int8(name); - if (!st.ok_status()) return st; - } - } - return modalities::Status::ok(); -} - -} // namespace pi05 -} // namespace models -} // namespace flashrt From d31da4b9f7b2aeddbace04bdbd0964065544f428 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 17 Jul 2026 08:15:52 -0400 Subject: [PATCH 83/83] refactor(pi05): name model backend sessions --- cpp/models/pi05/CMakeLists.txt | 6 +- .../{native_graph_runtime.h => session.h} | 38 +++---- .../{native_thor_graph_owner.h => session.h} | 34 +++--- .../sm120/{native_graph_owner.h => session.h} | 46 ++++---- .../{native_graph_runtime.cpp => session.cpp} | 26 ++--- ...ative_thor_graph_owner.cpp => session.cpp} | 70 ++++++------ .../sm120/native_rtx_calibration_session.cpp | 10 +- .../{native_graph_owner.cpp => session.cpp} | 104 +++++++++--------- .../pi05/src/service/native_model_runtime.cpp | 44 ++++---- 9 files changed, 189 insertions(+), 189 deletions(-) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/{native_graph_runtime.h => session.h} (65%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/{native_thor_graph_owner.h => session.h} (64%) rename cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/{native_graph_owner.h => session.h} (59%) rename cpp/models/pi05/src/backend/{native_graph_runtime.cpp => session.cpp} (85%) rename cpp/models/pi05/src/backends/sm110/{native_thor_graph_owner.cpp => session.cpp} (83%) rename cpp/models/pi05/src/backends/sm120/{native_graph_owner.cpp => session.cpp} (81%) diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt index a6f5c809..5102e35f 100644 --- a/cpp/models/pi05/CMakeLists.txt +++ b/cpp/models/pi05/CMakeLists.txt @@ -67,7 +67,7 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) if(FLASHRT_CPP_WITH_FA2 OR FLASHRT_CPP_WITH_THOR_FP8) add_library(flashrt_cpp_pi05_backend_cuda STATIC - src/backend/native_graph_runtime.cpp + src/backend/session.cpp ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/gemm_runner.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/activation.cu ${FLASHRT_CPP_SOURCE_DIR}/../csrc/kernels/elementwise.cu @@ -90,7 +90,7 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) if(FLASHRT_CPP_WITH_FA2) add_library(flashrt_cpp_pi05_backend_sm120 STATIC src/backends/sm120/native_bf16_forward.cpp - src/backends/sm120/native_graph_owner.cpp + src/backends/sm120/session.cpp src/backends/sm120/native_kernel_driver.cu src/backends/sm120/native_rtx_attention.cpp src/backends/sm120/native_rtx_attention_driver.cu @@ -137,7 +137,7 @@ if(FLASHRT_CPP_WITH_CUDA_KERNELS) src/backends/sm110/native_thor_kernel_driver.cu src/backends/sm110/native_thor_weight_materializer.cpp src/backends/sm110/native_thor_fp8_forward.cpp - src/backends/sm110/native_thor_graph_owner.cpp + src/backends/sm110/session.cpp src/backends/sm110/native_thor_calibration_session.cpp src/backends/sm110/native_thor_style_precompute.cpp ${FLASHRT_CPP_SOURCE_DIR}/../csrc/gemm/cutlass_sm100.cu diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_graph_runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/session.h similarity index 65% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_graph_runtime.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/session.h index 97b7d9ce..ebf0ba70 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_graph_runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/session.h @@ -1,5 +1,5 @@ -#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H -#define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H +#ifndef FLASHRT_CPP_MODELS_PI05_BACKEND_SESSION_H +#define FLASHRT_CPP_MODELS_PI05_BACKEND_SESSION_H #include "flashrt/cpp/models/pi05/support/native_device_weights.h" #include "flashrt/cpp/models/pi05/support/native_workspace.h" @@ -12,32 +12,32 @@ namespace flashrt { namespace models { namespace pi05 { -enum class NativeGraphPrecision { +enum class BackendPrecision { kBf16, kFp8E4M3, }; -struct NativeGraphConfig { +struct BackendConfig { int num_views = 2; int max_prompt_tokens = 200; int chunk_size = 10; int num_steps = 10; int vision_pool_factor = 1; - NativeGraphPrecision precision = NativeGraphPrecision::kBf16; + BackendPrecision precision = BackendPrecision::kBf16; }; -enum class NativeGraphKind : std::size_t { +enum class GraphKind : std::size_t { kInfer = 0, kDecodeOnly = 1, kContext = 2, kCount = 3, }; -const char* native_graph_name(NativeGraphKind kind); +const char* backend_graph_name(GraphKind kind); -modalities::Status capture_native_graph( +modalities::Status capture_backend_graph( native::CudaGraphSet* graphs, - NativeGraphKind kind, + GraphKind kind, const NativeWorkspace& workspace, std::initializer_list bindings, native::CudaGraphSet::RecordFn record, @@ -46,7 +46,7 @@ modalities::Status capture_native_graph( modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, void* stream); -struct NativeRuntimeArtifacts { +struct BackendArtifacts { const NativeWorkspaceBuffer* images = nullptr; const NativeWorkspaceBuffer* noise = nullptr; const NativeWorkspaceBuffer* encoder = nullptr; @@ -57,24 +57,24 @@ struct NativeRuntimeArtifacts { const NativeDeviceWeight* embedding_table = nullptr; }; -modalities::Status resolve_native_runtime_artifacts( +modalities::Status resolve_backend_artifacts( const NativeWorkspace& workspace, const NativeDeviceWeightStore& weights, NativeWeightDType embedding_dtype, - NativeRuntimeArtifacts* artifacts); + BackendArtifacts* artifacts); -class NativeGraphRuntime { +class BackendSession { public: - virtual ~NativeGraphRuntime() = default; + virtual ~BackendSession() = default; virtual frt_ctx context() const = 0; - virtual frt_graph graph(NativeGraphKind kind) const = 0; - frt_graph infer_graph() const { return graph(NativeGraphKind::kInfer); } + virtual frt_graph graph(GraphKind kind) const = 0; + frt_graph infer_graph() const { return graph(GraphKind::kInfer); } virtual int stream_id() const = 0; virtual void* native_stream() const = 0; - virtual const NativeRuntimeArtifacts& artifacts() const = 0; + virtual const BackendArtifacts& artifacts() const = 0; virtual modalities::Status set_prompt_length(int prompt_tokens) = 0; - virtual int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const = 0; + virtual int replay(GraphKind kind = GraphKind::kInfer) const = 0; virtual modalities::Status synchronize() const = 0; }; @@ -82,4 +82,4 @@ class NativeGraphRuntime { } // namespace models } // namespace flashrt -#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_RUNTIME_H +#endif // FLASHRT_CPP_MODELS_PI05_BACKEND_SESSION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/session.h similarity index 64% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/session.h index ae355b2c..9e4a5a74 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/session.h @@ -1,8 +1,8 @@ -#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H -#define FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H +#ifndef FLASHRT_CPP_MODELS_PI05_BACKENDS_SM110_SESSION_H +#define FLASHRT_CPP_MODELS_PI05_BACKENDS_SM110_SESSION_H #include "flashrt/cpp/models/pi05/support/native_calibration.h" -#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backend/session.h" #include "flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h" #include @@ -13,21 +13,21 @@ namespace flashrt { namespace models { namespace pi05 { -class NativeThorGraphOwner final : public NativeGraphRuntime { +class Sm110BackendSession final : public BackendSession { public: - static std::unique_ptr create( + static std::unique_ptr create( const std::string& checkpoint_path, - const NativeGraphConfig& config, + const BackendConfig& config, const NativeCalibrationArtifact& calibration, modalities::Status* status); - ~NativeThorGraphOwner() override; + ~Sm110BackendSession() override; - NativeThorGraphOwner(const NativeThorGraphOwner&) = delete; - NativeThorGraphOwner& operator=(const NativeThorGraphOwner&) = delete; + Sm110BackendSession(const Sm110BackendSession&) = delete; + Sm110BackendSession& operator=(const Sm110BackendSession&) = delete; frt_ctx context() const override { return graphs_.context(); } - frt_graph graph(NativeGraphKind kind) const override { + frt_graph graph(GraphKind kind) const override { return graphs_.graph(static_cast(kind)); } int stream_id() const override { return graphs_.stream_id(); } @@ -36,30 +36,30 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { const NativeDeviceWeightStore& weights() const { return weights_; } NativeWorkspace& workspace() { return workspace_; } const NativeWorkspace& workspace() const { return workspace_; } - const NativeRuntimeArtifacts& artifacts() const override { + const BackendArtifacts& artifacts() const override { return artifacts_; } modalities::Status set_prompt_length(int prompt_tokens) override; - int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const override; + int replay(GraphKind kind = GraphKind::kInfer) const override; modalities::Status synchronize() const override; private: - NativeThorGraphOwner(frt_ctx ctx, const NativeGraphConfig& config); + Sm110BackendSession(frt_ctx ctx, const BackendConfig& config); modalities::Status initialize( const std::string& checkpoint_path, const NativeCalibrationArtifact& calibration); - modalities::Status record(NativeGraphKind kind, void* stream); + modalities::Status record(GraphKind kind, void* stream); modalities::Status record_context(void* stream); modalities::Status record_action(void* stream); static modalities::Status record_graph( void* user, std::size_t slot, void* stream); native::CudaGraphSet graphs_; - NativeGraphConfig config_; + BackendConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; - NativeRuntimeArtifacts artifacts_; + BackendArtifacts artifacts_; NativeThorKernelDriver driver_; NativeThorFp8Forward forward_; NativeThorWeightScales weight_scales_; @@ -70,4 +70,4 @@ class NativeThorGraphOwner final : public NativeGraphRuntime { } // namespace models } // namespace flashrt -#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_THOR_GRAPH_OWNER_H +#endif // FLASHRT_CPP_MODELS_PI05_BACKENDS_SM110_SESSION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/session.h similarity index 59% rename from cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/session.h index 56ffe90c..7136ecb9 100644 --- a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/session.h @@ -1,9 +1,9 @@ -#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H -#define FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H +#ifndef FLASHRT_CPP_MODELS_PI05_BACKENDS_SM120_SESSION_H +#define FLASHRT_CPP_MODELS_PI05_BACKENDS_SM120_SESSION_H #include "flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h" #include "flashrt/cpp/models/pi05/support/native_calibration.h" -#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backend/session.h" #include #include @@ -12,62 +12,62 @@ namespace flashrt { namespace models { namespace pi05 { -class NativeGraphOwner final : public NativeGraphRuntime { +class Sm120BackendSession final : public BackendSession { public: - static std::unique_ptr create( - const std::string& checkpoint_path, const NativeGraphConfig& config, + static std::unique_ptr create( + const std::string& checkpoint_path, const BackendConfig& config, modalities::Status* status); - static std::unique_ptr create( - const std::string& checkpoint_path, const NativeGraphConfig& config, + static std::unique_ptr create( + const std::string& checkpoint_path, const BackendConfig& config, const NativeCalibrationArtifact& calibration, modalities::Status* status); - static std::unique_ptr create_calibration( - const std::string& checkpoint_path, const NativeGraphConfig& config, + static std::unique_ptr create_calibration( + const std::string& checkpoint_path, const BackendConfig& config, modalities::Status* status); - ~NativeGraphOwner() override; + ~Sm120BackendSession() override; - NativeGraphOwner(const NativeGraphOwner&) = delete; - NativeGraphOwner& operator=(const NativeGraphOwner&) = delete; + Sm120BackendSession(const Sm120BackendSession&) = delete; + Sm120BackendSession& operator=(const Sm120BackendSession&) = delete; frt_ctx context() const override { return graphs_.context(); } - frt_graph graph(NativeGraphKind kind) const override { + frt_graph graph(GraphKind kind) const override { return graphs_.graph(static_cast(kind)); } int stream_id() const override { return graphs_.stream_id(); } void* native_stream() const override { return graphs_.native_stream(); } - const NativeGraphConfig& config() const { return config_; } + const BackendConfig& config() const { return config_; } NativeDeviceWeightStore& weights() { return weights_; } const NativeDeviceWeightStore& weights() const { return weights_; } NativeWorkspace& workspace() { return workspace_; } const NativeWorkspace& workspace() const { return workspace_; } - const NativeRuntimeArtifacts& artifacts() const override { + const BackendArtifacts& artifacts() const override { return artifacts_; } NativeRtxAttentionWorkspace& attention() { return attention_; } const NativeRtxAttentionWorkspace& attention() const { return attention_; } modalities::Status set_prompt_length(int prompt_tokens) override; - int replay(NativeGraphKind kind = NativeGraphKind::kInfer) const override; + int replay(GraphKind kind = GraphKind::kInfer) const override; modalities::Status synchronize() const override; private: - NativeGraphOwner(frt_ctx ctx, const NativeGraphConfig& config, - NativeRtxLinearMode linear_mode); + Sm120BackendSession(frt_ctx ctx, const BackendConfig& config, + NativeRtxLinearMode linear_mode); modalities::Status initialize( const std::string& checkpoint_path, const NativeCalibrationArtifact* calibration); - modalities::Status record(NativeGraphKind kind, void* stream); + modalities::Status record(GraphKind kind, void* stream); modalities::Status record_context(void* stream); modalities::Status record_action(void* stream); static modalities::Status record_graph( void* user, std::size_t slot, void* stream); native::CudaGraphSet graphs_; - NativeGraphConfig config_; + BackendConfig config_; NativeDeviceWeightStore weights_; NativeWorkspace workspace_; - NativeRuntimeArtifacts artifacts_; + BackendArtifacts artifacts_; NativeRtxAttentionWorkspace attention_; NativeKernelDriver driver_; NativeRtxLinear linear_; @@ -79,4 +79,4 @@ class NativeGraphOwner final : public NativeGraphRuntime { } // namespace models } // namespace flashrt -#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_GRAPH_OWNER_H +#endif // FLASHRT_CPP_MODELS_PI05_BACKENDS_SM120_SESSION_H diff --git a/cpp/models/pi05/src/backend/native_graph_runtime.cpp b/cpp/models/pi05/src/backend/session.cpp similarity index 85% rename from cpp/models/pi05/src/backend/native_graph_runtime.cpp rename to cpp/models/pi05/src/backend/session.cpp index 01fe07cb..b479ef02 100644 --- a/cpp/models/pi05/src/backend/native_graph_runtime.cpp +++ b/cpp/models/pi05/src/backend/session.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backend/session.h" #include "flashrt/cpp/models/pi05/model/spec.h" @@ -23,24 +23,24 @@ modalities::Status backend(const char* message) { } // namespace -const char* native_graph_name(NativeGraphKind kind) { +const char* backend_graph_name(GraphKind kind) { switch (kind) { - case NativeGraphKind::kInfer: return "infer"; - case NativeGraphKind::kDecodeOnly: return "decode_only"; - case NativeGraphKind::kContext: return "context"; - case NativeGraphKind::kCount: break; + case GraphKind::kInfer: return "infer"; + case GraphKind::kDecodeOnly: return "decode_only"; + case GraphKind::kContext: return "context"; + case GraphKind::kCount: break; } return nullptr; } -modalities::Status capture_native_graph( +modalities::Status capture_backend_graph( native::CudaGraphSet* graphs, - NativeGraphKind kind, + GraphKind kind, const NativeWorkspace& workspace, std::initializer_list bindings, native::CudaGraphSet::RecordFn record, void* owner) { - if (!graphs || !native_graph_name(kind)) { + if (!graphs || !backend_graph_name(kind)) { return invalid("native graph capture request is invalid"); } std::vector resolved; @@ -53,7 +53,7 @@ modalities::Status capture_native_graph( resolved.push_back({binding, buffer->buffer}); } return graphs->capture(static_cast(kind), - native_graph_name(kind), resolved, record, owner); + backend_graph_name(kind), resolved, record, owner); } modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, @@ -81,15 +81,15 @@ modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, : backend("native prompt graph copy failed"); } -modalities::Status resolve_native_runtime_artifacts( +modalities::Status resolve_backend_artifacts( const NativeWorkspace& workspace, const NativeDeviceWeightStore& weights, NativeWeightDType embedding_dtype, - NativeRuntimeArtifacts* artifacts) { + BackendArtifacts* artifacts) { if (!artifacts) { return invalid("native runtime artifacts destination is null"); } - NativeRuntimeArtifacts result; + BackendArtifacts result; result.images = workspace.find("observation_images_normalized"); result.noise = workspace.find("diffusion_noise"); result.encoder = workspace.find("encoder_x"); diff --git a/cpp/models/pi05/src/backends/sm110/native_thor_graph_owner.cpp b/cpp/models/pi05/src/backends/sm110/session.cpp similarity index 83% rename from cpp/models/pi05/src/backends/sm110/native_thor_graph_owner.cpp rename to cpp/models/pi05/src/backends/sm110/session.cpp index 159940d9..8c7eafad 100644 --- a/cpp/models/pi05/src/backends/sm110/native_thor_graph_owner.cpp +++ b/cpp/models/pi05/src/backends/sm110/session.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm110/session.h" #include "flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h" #include "flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h" @@ -46,7 +46,7 @@ modalities::Status copy_scales(const NativeWorkspace& workspace, } bool calibration_matches(const NativeCalibrationArtifact& artifact, - const NativeGraphConfig& config) { + const BackendConfig& config) { return artifact.num_views == config.num_views && artifact.max_prompt_tokens == config.max_prompt_tokens && artifact.chunk_size == config.chunk_size && @@ -56,20 +56,20 @@ bool calibration_matches(const NativeCalibrationArtifact& artifact, } // namespace -NativeThorGraphOwner::NativeThorGraphOwner( +Sm110BackendSession::Sm110BackendSession( frt_ctx ctx, - const NativeGraphConfig& config) - : graphs_(ctx, static_cast(NativeGraphKind::kCount)), + const BackendConfig& config) + : graphs_(ctx, static_cast(GraphKind::kCount)), config_(config), weights_(ctx), workspace_(ctx), forward_(&driver_) {} -NativeThorGraphOwner::~NativeThorGraphOwner() = default; +Sm110BackendSession::~Sm110BackendSession() = default; -std::unique_ptr NativeThorGraphOwner::create( +std::unique_ptr Sm110BackendSession::create( const std::string& checkpoint_path, - const NativeGraphConfig& config, + const BackendConfig& config, const NativeCalibrationArtifact& calibration, modalities::Status* status) { if (config.num_views < 1 || config.num_views > 3 || @@ -97,23 +97,23 @@ std::unique_ptr NativeThorGraphOwner::create( if (status) *status = backend("Thor graph context creation failed"); return nullptr; } - std::unique_ptr owner( - new (std::nothrow) NativeThorGraphOwner(ctx, config)); - if (!owner) { + std::unique_ptr session( + new (std::nothrow) Sm110BackendSession(ctx, config)); + if (!session) { frt_ctx_destroy(ctx); - if (status) *status = backend("Thor graph owner allocation failed"); + if (status) *status = backend("SM110 backend session allocation failed"); return nullptr; } - st = owner->initialize(checkpoint_path, calibration); + st = session->initialize(checkpoint_path, calibration); if (!st.ok_status()) { if (status) *status = st; return nullptr; } if (status) *status = modalities::Status::ok(); - return owner; + return session; } -modalities::Status NativeThorGraphOwner::initialize( +modalities::Status Sm110BackendSession::initialize( const std::string& checkpoint_path, const NativeCalibrationArtifact& calibration) { const bool profile_setup = std::getenv("FLASHRT_PROFILE_NATIVE_SETUP"); @@ -180,7 +180,7 @@ modalities::Status NativeThorGraphOwner::initialize( if (!st.ok_status()) return st; report("workspace_style"); - st = resolve_native_runtime_artifacts( + st = resolve_backend_artifacts( workspace_, weights_, NativeWeightDType::kFloat16, &artifacts_); if (!st.ok_status()) return st; @@ -198,7 +198,7 @@ modalities::Status NativeThorGraphOwner::initialize( } // CUTLASS, cuBLAS, and FMHA initialize tactics outside CUDA capture. - st = record(NativeGraphKind::kInfer, nullptr); + st = record(GraphKind::kInfer, nullptr); if (!st.ok_status()) return st; if (cudaDeviceSynchronize() != cudaSuccess) { return backend("Thor graph warmup synchronization failed"); @@ -211,24 +211,24 @@ modalities::Status NativeThorGraphOwner::initialize( } report("warmup"); - st = capture_native_graph( + st = capture_backend_graph( &graphs_, - NativeGraphKind::kInfer, workspace_, + GraphKind::kInfer, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = capture_native_graph( + st = capture_backend_graph( &graphs_, - NativeGraphKind::kDecodeOnly, workspace_, + GraphKind::kDecodeOnly, workspace_, {"encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = capture_native_graph( + st = capture_backend_graph( &graphs_, - NativeGraphKind::kContext, workspace_, + GraphKind::kContext, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x"}, record_graph, this); if (!st.ok_status()) return st; @@ -246,7 +246,7 @@ modalities::Status NativeThorGraphOwner::initialize( return modalities::Status::ok(); } -modalities::Status NativeThorGraphOwner::record_context(void* stream) { +modalities::Status Sm110BackendSession::record_context(void* stream) { modalities::Status st = copy_prompt_to_encoder(&workspace_, stream); if (!st.ok_status()) return st; const std::uintptr_t stream_id = reinterpret_cast(stream); @@ -256,37 +256,37 @@ modalities::Status NativeThorGraphOwner::record_context(void* stream) { weights_, &workspace_, encoder_alphas_, stream_id); } -modalities::Status NativeThorGraphOwner::record_action(void* stream) { +modalities::Status Sm110BackendSession::record_action(void* stream) { return forward_.diffusion( weights_, &workspace_, reinterpret_cast(stream)); } -modalities::Status NativeThorGraphOwner::record(NativeGraphKind kind, +modalities::Status Sm110BackendSession::record(GraphKind kind, void* stream) { - if (kind == NativeGraphKind::kContext) return record_context(stream); - if (kind == NativeGraphKind::kDecodeOnly) return record_action(stream); - if (kind != NativeGraphKind::kInfer) { + if (kind == GraphKind::kContext) return record_context(stream); + if (kind == GraphKind::kDecodeOnly) return record_action(stream); + if (kind != GraphKind::kInfer) { return invalid("Thor graph kind is invalid"); } modalities::Status st = record_context(stream); return st.ok_status() ? record_action(stream) : st; } -modalities::Status NativeThorGraphOwner::record_graph( +modalities::Status Sm110BackendSession::record_graph( void* user, std::size_t slot, void* stream) { - auto* owner = static_cast(user); - return owner->record(static_cast(slot), stream); + auto* session = static_cast(user); + return session->record(static_cast(slot), stream); } -modalities::Status NativeThorGraphOwner::set_prompt_length(int prompt_tokens) { +modalities::Status Sm110BackendSession::set_prompt_length(int prompt_tokens) { return workspace_.set_fixed_prompt_length(prompt_tokens); } -int NativeThorGraphOwner::replay(NativeGraphKind kind) const { +int Sm110BackendSession::replay(GraphKind kind) const { return graphs_.replay(static_cast(kind)); } -modalities::Status NativeThorGraphOwner::synchronize() const { +modalities::Status Sm110BackendSession::synchronize() const { return graphs_.synchronize(); } diff --git a/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp index 4c1ffb28..a41e493d 100644 --- a/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp @@ -3,7 +3,7 @@ #include "flashrt/cpp/loader/sha256.h" #include "flashrt/cpp/models/pi05/model/io.h" #include "flashrt/cpp/models/pi05/support/native_calibration.h" -#include "flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm120/session.h" #include "flashrt/cpp/models/pi05/model/prompt_embed.h" #include @@ -97,15 +97,15 @@ struct NativeRtxCalibrationSession::Impl { return result; }); - NativeGraphConfig graph_config; + BackendConfig graph_config; graph_config.num_views = config.num_views; graph_config.max_prompt_tokens = config.max_prompt_tokens; graph_config.chunk_size = config.chunk_size; graph_config.num_steps = config.num_steps; graph_config.vision_pool_factor = config.vision_pool_factor; - graph_config.precision = NativeGraphPrecision::kFp8E4M3; + graph_config.precision = BackendPrecision::kFp8E4M3; modalities::Status st; - graph = NativeGraphOwner::create_calibration( + graph = Sm120BackendSession::create_calibration( config.checkpoint_path, graph_config, &st); if (!graph) return st; @@ -358,7 +358,7 @@ struct NativeRtxCalibrationSession::Impl { std::string hardware; std::string weights_sha256; std::string tokenizer_sha256; - std::unique_ptr graph; + std::unique_ptr graph; modalities::VisionStaging vision_staging; modalities::TextEmbeddingStaging text_staging; modalities::SentencePieceTokenizer tokenizer; diff --git a/cpp/models/pi05/src/backends/sm120/native_graph_owner.cpp b/cpp/models/pi05/src/backends/sm120/session.cpp similarity index 81% rename from cpp/models/pi05/src/backends/sm120/native_graph_owner.cpp rename to cpp/models/pi05/src/backends/sm120/session.cpp index 397acc47..d2db1f05 100644 --- a/cpp/models/pi05/src/backends/sm120/native_graph_owner.cpp +++ b/cpp/models/pi05/src/backends/sm120/session.cpp @@ -1,4 +1,4 @@ -#include "flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm120/session.h" #include "flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h" #include "flashrt/cpp/models/pi05/support/native_weight_materializer.h" @@ -51,11 +51,11 @@ modalities::Status upload_scales( } // namespace -NativeGraphOwner::NativeGraphOwner( +Sm120BackendSession::Sm120BackendSession( frt_ctx ctx, - const NativeGraphConfig& config, + const BackendConfig& config, NativeRtxLinearMode linear_mode) - : graphs_(ctx, static_cast(NativeGraphKind::kCount)), + : graphs_(ctx, static_cast(GraphKind::kCount)), config_(config), weights_(ctx), workspace_(ctx), @@ -63,13 +63,13 @@ NativeGraphOwner::NativeGraphOwner( linear_(&driver_, linear_mode), forward_(&driver_, &linear_) {} -NativeGraphOwner::~NativeGraphOwner() = default; +Sm120BackendSession::~Sm120BackendSession() = default; -std::unique_ptr NativeGraphOwner::create( +std::unique_ptr Sm120BackendSession::create( const std::string& checkpoint_path, - const NativeGraphConfig& config, + const BackendConfig& config, modalities::Status* status) { - if (config.precision != NativeGraphPrecision::kBf16 || + if (config.precision != BackendPrecision::kBf16 || config.num_views < 1 || config.num_views > 3 || config.max_prompt_tokens < 1 || config.chunk_size < 1 || config.num_steps < 1 || @@ -87,29 +87,29 @@ std::unique_ptr NativeGraphOwner::create( if (status) *status = backend("native graph context creation failed"); return nullptr; } - std::unique_ptr owner( - new (std::nothrow) NativeGraphOwner( + std::unique_ptr session( + new (std::nothrow) Sm120BackendSession( ctx, config, NativeRtxLinearMode::kBf16)); - if (!owner) { + if (!session) { frt_ctx_destroy(ctx); - if (status) *status = backend("native graph owner allocation failed"); + if (status) *status = backend("SM120 backend session allocation failed"); return nullptr; } - modalities::Status st = owner->initialize(checkpoint_path, nullptr); + modalities::Status st = session->initialize(checkpoint_path, nullptr); if (!st.ok_status()) { if (status) *status = st; return nullptr; } if (status) *status = modalities::Status::ok(); - return owner; + return session; } -std::unique_ptr NativeGraphOwner::create( +std::unique_ptr Sm120BackendSession::create( const std::string& checkpoint_path, - const NativeGraphConfig& config, + const BackendConfig& config, const NativeCalibrationArtifact& calibration, modalities::Status* status) { - if (config.precision != NativeGraphPrecision::kFp8E4M3) { + if (config.precision != BackendPrecision::kFp8E4M3) { if (status) *status = invalid("native RTX FP8 graph precision is invalid"); return nullptr; } @@ -118,28 +118,28 @@ std::unique_ptr NativeGraphOwner::create( if (status) *status = backend("native graph context creation failed"); return nullptr; } - std::unique_ptr owner( - new (std::nothrow) NativeGraphOwner( + std::unique_ptr session( + new (std::nothrow) Sm120BackendSession( ctx, config, NativeRtxLinearMode::kFp8Static)); - if (!owner) { + if (!session) { frt_ctx_destroy(ctx); - if (status) *status = backend("native graph owner allocation failed"); + if (status) *status = backend("SM120 backend session allocation failed"); return nullptr; } - modalities::Status st = owner->initialize(checkpoint_path, &calibration); + modalities::Status st = session->initialize(checkpoint_path, &calibration); if (!st.ok_status()) { if (status) *status = st; return nullptr; } if (status) *status = modalities::Status::ok(); - return owner; + return session; } -std::unique_ptr NativeGraphOwner::create_calibration( +std::unique_ptr Sm120BackendSession::create_calibration( const std::string& checkpoint_path, - const NativeGraphConfig& config, + const BackendConfig& config, modalities::Status* status) { - if (config.precision != NativeGraphPrecision::kFp8E4M3) { + if (config.precision != BackendPrecision::kFp8E4M3) { if (status) { *status = invalid("native RTX calibration precision is invalid"); } @@ -150,24 +150,24 @@ std::unique_ptr NativeGraphOwner::create_calibration( if (status) *status = backend("native graph context creation failed"); return nullptr; } - std::unique_ptr owner( - new (std::nothrow) NativeGraphOwner( + std::unique_ptr session( + new (std::nothrow) Sm120BackendSession( ctx, config, NativeRtxLinearMode::kFp8Dynamic)); - if (!owner) { + if (!session) { frt_ctx_destroy(ctx); - if (status) *status = backend("native graph owner allocation failed"); + if (status) *status = backend("SM120 backend session allocation failed"); return nullptr; } - modalities::Status st = owner->initialize(checkpoint_path, nullptr); + modalities::Status st = session->initialize(checkpoint_path, nullptr); if (!st.ok_status()) { if (status) *status = st; return nullptr; } if (status) *status = modalities::Status::ok(); - return owner; + return session; } -modalities::Status NativeGraphOwner::initialize( +modalities::Status Sm120BackendSession::initialize( const std::string& checkpoint_path, const NativeCalibrationArtifact* calibration) { const bool fp8 = linear_.fp8(); @@ -280,7 +280,7 @@ modalities::Status NativeGraphOwner::initialize( } report("workspace_style"); - st = resolve_native_runtime_artifacts( + st = resolve_backend_artifacts( workspace_, weights_, NativeWeightDType::kBf16, &artifacts_); if (!st.ok_status()) return st; @@ -299,24 +299,24 @@ modalities::Status NativeGraphOwner::initialize( } report("input_init"); - st = capture_native_graph( + st = capture_backend_graph( &graphs_, - NativeGraphKind::kInfer, workspace_, + GraphKind::kInfer, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = capture_native_graph( + st = capture_backend_graph( &graphs_, - NativeGraphKind::kDecodeOnly, workspace_, + GraphKind::kDecodeOnly, workspace_, {"encoder_x", "diffusion_noise", "rtc_prev_action_chunk", "rtc_prefix_weights", "rtc_guidance_weight"}, record_graph, this); if (!st.ok_status()) return st; - st = capture_native_graph( + st = capture_backend_graph( &graphs_, - NativeGraphKind::kContext, workspace_, + GraphKind::kContext, workspace_, {"observation_images_normalized", "prompt_embedding", "encoder_x"}, record_graph, this); if (!st.ok_status()) return st; @@ -334,7 +334,7 @@ modalities::Status NativeGraphOwner::initialize( return modalities::Status::ok(); } -modalities::Status NativeGraphOwner::record_context(void* stream) { +modalities::Status Sm120BackendSession::record_context(void* stream) { modalities::Status st = copy_prompt_to_encoder(&workspace_, stream); if (!st.ok_status()) return st; st = forward_.vision( @@ -348,40 +348,40 @@ modalities::Status NativeGraphOwner::record_context(void* stream) { return st; } -modalities::Status NativeGraphOwner::record_action(void* stream) { +modalities::Status Sm120BackendSession::record_action(void* stream) { return forward_.diffusion(weights_, &workspace_, &attention_, attention_driver_.get(), reinterpret_cast(stream)); } -modalities::Status NativeGraphOwner::record(NativeGraphKind kind, +modalities::Status Sm120BackendSession::record(GraphKind kind, void* stream) { - if (kind == NativeGraphKind::kContext) return record_context(stream); - if (kind == NativeGraphKind::kDecodeOnly) return record_action(stream); - if (kind != NativeGraphKind::kInfer) { + if (kind == GraphKind::kContext) return record_context(stream); + if (kind == GraphKind::kDecodeOnly) return record_action(stream); + if (kind != GraphKind::kInfer) { return invalid("native graph kind is invalid"); } modalities::Status st = record_context(stream); return st.ok_status() ? record_action(stream) : st; } -modalities::Status NativeGraphOwner::record_graph( +modalities::Status Sm120BackendSession::record_graph( void* user, std::size_t slot, void* stream) { - auto* owner = static_cast(user); - return owner->record(static_cast(slot), stream); + auto* session = static_cast(user); + return session->record(static_cast(slot), stream); } -modalities::Status NativeGraphOwner::set_prompt_length(int prompt_tokens) { +modalities::Status Sm120BackendSession::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(NativeGraphKind kind) const { +int Sm120BackendSession::replay(GraphKind kind) const { return graphs_.replay(static_cast(kind)); } -modalities::Status NativeGraphOwner::synchronize() const { +modalities::Status Sm120BackendSession::synchronize() const { return graphs_.synchronize(); } diff --git a/cpp/models/pi05/src/service/native_model_runtime.cpp b/cpp/models/pi05/src/service/native_model_runtime.cpp index 5a550771..6f64ea26 100644 --- a/cpp/models/pi05/src/service/native_model_runtime.cpp +++ b/cpp/models/pi05/src/service/native_model_runtime.cpp @@ -7,13 +7,13 @@ #include "flashrt/cpp/loader/sha256.h" #include "flashrt/cpp/models/pi05/model_runtime.h" #include "flashrt/cpp/models/pi05/support/native_calibration.h" -#include "flashrt/cpp/models/pi05/backend/native_graph_runtime.h" +#include "flashrt/cpp/models/pi05/backend/session.h" #include "flashrt/cpp/models/pi05/model/spec.h" #if defined(FLASHRT_CPP_WITH_FA2) -#include "flashrt/cpp/models/pi05/backends/sm120/native_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm120/session.h" #endif #if defined(FLASHRT_CPP_WITH_THOR_FP8) -#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_graph_owner.h" +#include "flashrt/cpp/models/pi05/backends/sm110/session.h" #endif #include @@ -29,12 +29,12 @@ namespace models { namespace pi05 { namespace { -void release_graph_owner(void* owner) { - delete static_cast(owner); +void release_backend_session(void* owner) { + delete static_cast(owner); } int update_prompt_length(void* owner, std::uint64_t prompt_len) { - auto* graph = static_cast(owner); + 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))); @@ -198,17 +198,17 @@ int build_native_model_runtime(const NativeOpenConfig& config, } } - NativeGraphConfig graph_config; + BackendConfig 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; graph_config.precision = precision == Precision::kFp8E4M3Fn - ? NativeGraphPrecision::kFp8E4M3 - : NativeGraphPrecision::kBf16; + ? BackendPrecision::kFp8E4M3 + : BackendPrecision::kBf16; modalities::Status st; - std::unique_ptr graph; + std::unique_ptr graph; const bool thor_fp8 = precision == Precision::kFp8E4M3Fn && properties.major == 11; const bool rtx_fp8 = precision == Precision::kFp8E4M3Fn && @@ -216,15 +216,15 @@ int build_native_model_runtime(const NativeOpenConfig& config, if (precision == Precision::kBf16 || rtx_fp8) { #if defined(FLASHRT_CPP_WITH_FA2) graph = rtx_fp8 - ? NativeGraphOwner::create( + ? Sm120BackendSession::create( config.checkpoint_path, graph_config, calibration, &st) - : NativeGraphOwner::create( + : Sm120BackendSession::create( config.checkpoint_path, graph_config, &st); #endif } else if (thor_fp8) { #if defined(FLASHRT_CPP_WITH_THOR_FP8) - graph = NativeThorGraphOwner::create( + graph = Sm110BackendSession::create( config.checkpoint_path, graph_config, calibration, &st); #endif } @@ -243,7 +243,7 @@ int build_native_model_runtime(const NativeOpenConfig& config, return -2; } - const NativeRuntimeArtifacts& artifacts = graph->artifacts(); + const BackendArtifacts& artifacts = graph->artifacts(); const NativeWorkspaceBuffer* images = artifacts.images; const NativeWorkspaceBuffer* noise = artifacts.noise; const NativeWorkspaceBuffer* encoder = artifacts.encoder; @@ -264,16 +264,16 @@ int build_native_model_runtime(const NativeOpenConfig& config, builder, "main", graph->stream_id(), 0, graph->native_stream()) == 0 && frt_runtime_builder_add_graph( - builder, native_graph_name(NativeGraphKind::kInfer), - graph->graph(NativeGraphKind::kInfer), 0, keys, 1, + builder, backend_graph_name(GraphKind::kInfer), + graph->graph(GraphKind::kInfer), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_graph( - builder, native_graph_name(NativeGraphKind::kDecodeOnly), - graph->graph(NativeGraphKind::kDecodeOnly), 0, keys, 1, + builder, backend_graph_name(GraphKind::kDecodeOnly), + graph->graph(GraphKind::kDecodeOnly), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_graph( - builder, native_graph_name(NativeGraphKind::kContext), - graph->graph(NativeGraphKind::kContext), 0, keys, 1, + builder, backend_graph_name(GraphKind::kContext), + graph->graph(GraphKind::kContext), 0, keys, 1, graph->stream_id()) == 0 && frt_runtime_builder_add_buffer( builder, "observation_images_normalized", images->buffer, @@ -414,13 +414,13 @@ int build_native_model_runtime(const NativeOpenConfig& config, } if (!ok) return fail_builder(builder, error, "native port/stage build failed"); - NativeGraphRuntime* raw_graph = graph.release(); + BackendSession* raw_graph = graph.release(); /* This base is retained only by the verb override below and is never * returned to a consumer. The published object always has real verbs. */ frt_model_runtime_verbs base_verbs = unpublished_verbs(); frt_model_runtime_v1* base = frt_runtime_builder_finish_model( builder, &base_verbs, nullptr, raw_graph, nullptr, - release_graph_owner); + release_backend_session); if (!base) { delete raw_graph; if (error) *error = "native integrated runtime finish failed";