diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 00000000..13d0f3be --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,31 @@ +## 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 +- [ ] 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 +- [ ] 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 +- [ ] External test migration preserves core contract/result coverage diff --git a/CMakeLists.txt b/CMakeLists.txt index 8152df92..183e6f61 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -767,8 +767,8 @@ if(ENABLE_NVFP4) endif() # ── Flash-Attention 2 vendored kernels (fp16 + bf16 fwd SM80) ── -# Built as an object library and linked into a SEPARATE pybind module -# flash_rt_fa2.so — same isolation pattern as flash_rt_fp4.so. The +# Built as an object library and linked into a Python-free raw library; +# flash_rt_fa2.so is a thin pybind adapter over that library. The # main flash_rt_kernels.so stays small (~3.6 MB) and its rebuild no # longer waits on FA2's heavy CUTLASS 3.x template codegen (~8 min). # @@ -831,9 +831,9 @@ if(ENABLE_FA2) csrc/attention/fa2_causal_inst/flash_fwd_split_hdim256_bf16_sm80_causal.cu ) endif() - if("bf16" IN_LIST FA2_DTYPES AND ("128" IN_LIST FA2_HDIMS OR "256" IN_LIST FA2_HDIMS)) - list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu) - endif() + # Always emit the causal C entry. Builds without a causal hdim retain a + # stable raw-library symbol that fails clearly instead of an unresolved ABI. + list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu) add_library(fa2_vendor_obj OBJECT ${FA2_SRCS}) set_target_properties(fa2_vendor_obj PROPERTIES @@ -954,6 +954,21 @@ if(ENABLE_FA2) math(EXPR _FA2_NKERN "${_FA2_NSRC} - 1") message(STATUS "FA2 vendor instantiations: hdim={${_FA2_H}} x dtype={${_FA2_D}} x {split,no-split} = ${_FA2_NKERN} kernel .cu files") + # Keep the raw C entries in a Python-free library. The native C++ runtime + # links this target directly; the pybind module below is only an adapter. + add_library(flashrt_fa2_raw SHARED) + target_sources(flashrt_fa2_raw PRIVATE $) + target_link_libraries(flashrt_fa2_raw PRIVATE CUDA::cudart) + if(NOT MSVC) + target_link_options(flashrt_fa2_raw PRIVATE + "LINKER:--no-undefined" -static-libstdc++ -static-libgcc) + endif() + set_target_properties(flashrt_fa2_raw PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN") + # Independent pybind module. Caller side: # from flash_rt import flash_rt_fa2 as fa2 # fa2.fwd_fp16(...) / fa2.fwd_bf16(...) @@ -961,12 +976,16 @@ if(ENABLE_FA2) set_target_properties(flash_rt_fa2 PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) - target_sources(flash_rt_fa2 PRIVATE $) target_include_directories(flash_rt_fa2 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/csrc ) - target_link_libraries(flash_rt_fa2 PRIVATE CUDA::cudart) - install(TARGETS flash_rt_fa2 DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) + target_link_libraries(flash_rt_fa2 PRIVATE flashrt_fa2_raw CUDA::cudart) + set_target_properties(flash_rt_fa2 PROPERTIES + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH ON) + install(TARGETS flash_rt_fa2 flashrt_fa2_raw + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) message(STATUS "FA2 pybind module: flash_rt_fa2 (separate .so)") endif() @@ -1440,8 +1459,8 @@ if(FLASHRT_ENABLE_MOTUS AND GPU_ARCH STREQUAL "120") ENABLE_TINYFP8_KERNELS=1) endif() -# (ENABLE_FA2 object-lib is linked into flash_rt_fa2.so only — see -# the dedicated pybind11_add_module(flash_rt_fa2 ...) above. The main +# (ENABLE_FA2 object-lib is linked into libflashrt_fa2_raw.so only; the +# dedicated pybind module and native C++ runtime both consume it. The main # flash_rt_kernels.so deliberately does NOT pull FA2 in, so rebuilds # of our hand-written kernels don't re-trigger the FA2 codegen tax.) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8acbe70b..72ada51f 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,28 @@ 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. + +### 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 @@ -235,6 +262,46 @@ 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. +- 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 Use the right metric for the claim: @@ -354,6 +421,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/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/USAGE.md b/USAGE.md index b493e8f7..6bb28ef1 100644 --- a/USAGE.md +++ b/USAGE.md @@ -169,6 +169,38 @@ 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 | +| SM120 | static FP8 E4M3 | native FA2 + SentencePiece | yes (schema v2) | +| Thor SM110 | FP8 E4M3 | Thor FP8/CUTLASS + SentencePiece | yes | + +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](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 @@ -212,7 +244,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 +1187,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 +1195,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/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 1b097a85..b56783a3 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -22,6 +22,16 @@ include(CTest) option(FLASHRT_CPP_WITH_EXEC "Build/link the exec layer for native replay" ON) option(FLASHRT_CPP_WITH_CUDA_STAGING "Enable conservative CUDA H2D/D2H modality staging" ON) option(FLASHRT_CPP_WITH_CUDA_KERNELS "Enable CUDA modality kernels" ON) +option(FLASHRT_CPP_WITH_SENTENCEPIECE "Enable native SentencePiece text tokenization" OFF) +option(FLASHRT_CPP_WITH_FA2 "Enable the Python-free vendored FA2 driver" OFF) +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() @@ -29,6 +39,71 @@ 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_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) + if(FLASHRT_SENTENCEPIECE_INCLUDE_DIR AND FLASHRT_SENTENCEPIECE_LIBRARY) + add_library(flashrt_sentencepiece_external INTERFACE) + target_include_directories(flashrt_sentencepiece_external + INTERFACE ${FLASHRT_SENTENCEPIECE_INCLUDE_DIR}) + target_link_libraries(flashrt_sentencepiece_external + INTERFACE ${FLASHRT_SENTENCEPIECE_LIBRARY}) + set(FLASHRT_CPP_SENTENCEPIECE_TARGET flashrt_sentencepiece_external) + else() + include(FetchContent) + set(SPM_ENABLE_SHARED OFF CACHE BOOL "" FORCE) + set(SPM_BUILD_TEST OFF CACHE BOOL "" FORCE) + set(SPM_BUILD_TESTS OFF CACHE BOOL "" FORCE) + FetchContent_Declare( + sentencepiece + GIT_REPOSITORY https://github.com/google/sentencepiece.git + GIT_TAG v0.2.1) + FetchContent_MakeAvailable(sentencepiece) + set(FLASHRT_SENTENCEPIECE_INCLUDE_DIR ${sentencepiece_SOURCE_DIR}/src) + set(FLASHRT_CPP_SENTENCEPIECE_TARGET sentencepiece-static) + endif() +endif() + if(FLASHRT_CPP_WITH_EXEC AND NOT TARGET flashrt_exec) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../exec ${CMAKE_CURRENT_BINARY_DIR}/exec) @@ -50,12 +125,28 @@ endif() set(FLASHRT_CPP_MODALITY_SRCS modalities/src/types.cpp modalities/src/vision_cpu.cpp + modalities/src/text_cpu.cpp modalities/src/action_cpu.cpp) +if(FLASHRT_CPP_WITH_SENTENCEPIECE) + list(APPEND FLASHRT_CPP_MODALITY_SRCS + modalities/src/sentencepiece_tokenizer.cpp) +else() + list(APPEND FLASHRT_CPP_MODALITY_SRCS + modalities/src/tokenizer_unavailable.cpp) +endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) - list(APPEND FLASHRT_CPP_MODALITY_SRCS modalities/src/vision_cuda.cu) + list(APPEND FLASHRT_CPP_MODALITY_SRCS + modalities/src/vision_cuda.cu + modalities/src/text_cuda.cu) endif() add_library(flashrt_cpp_modalities STATIC ${FLASHRT_CPP_MODALITY_SRCS}) +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) +endif() target_include_directories(flashrt_cpp_modalities PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/modalities/include) if(FLASHRT_CPP_WITH_CUDA_STAGING) @@ -63,6 +154,14 @@ if(FLASHRT_CPP_WITH_CUDA_STAGING) PUBLIC FLASHRT_CPP_WITH_CUDA_STAGING=1) target_link_libraries(flashrt_cpp_modalities PUBLIC CUDA::cudart) endif() +if(FLASHRT_CPP_WITH_SENTENCEPIECE) + target_compile_definitions(flashrt_cpp_modalities + PUBLIC FLASHRT_CPP_HAS_SENTENCEPIECE=1) + target_include_directories(flashrt_cpp_modalities + PUBLIC ${FLASHRT_SENTENCEPIECE_INCLUDE_DIR}) + target_link_libraries(flashrt_cpp_modalities + PUBLIC ${FLASHRT_CPP_SENTENCEPIECE_TARGET}) +endif() if(FLASHRT_CPP_WITH_CUDA_KERNELS) target_compile_definitions(flashrt_cpp_modalities PRIVATE FLASHRT_CPP_WITH_CUDA_KERNELS=1) @@ -78,48 +177,34 @@ target_include_directories(flashrt_cpp_vla target_link_libraries(flashrt_cpp_vla INTERFACE flashrt_cpp_modalities flashrt_cpp) -add_library(flashrt_cpp_pi05 STATIC - models/pi05/src/spec.cpp - models/pi05/src/io.cpp - models/pi05/src/runtime.cpp) -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) +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_pi05_c SHARED - models/pi05/src/c_api.cpp - models/pi05/src/model_runtime.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) +add_library(flashrt_cpp_loader STATIC + loader/src/safetensors.cpp + loader/src/sha256.cpp) +set_property(SOURCE loader/src/sha256.cpp APPEND PROPERTY COMPILE_OPTIONS -O2) +target_include_directories(flashrt_cpp_loader + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/loader/include) +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_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") if(BUILD_TESTING) - add_executable(test_cpp_modalities tests/test_modalities.cpp) - target_link_libraries(test_cpp_modalities - PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) - add_test(NAME cpp_modalities COMMAND test_cpp_modalities) - - add_executable(test_pi05_runtime tests/test_pi05_runtime.cpp) - target_link_libraries(test_pi05_runtime - PRIVATE flashrt_cpp_pi05 flashrt_cpp_modalities) - add_test(NAME pi05_runtime COMMAND test_pi05_runtime) - - 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_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) - endif() + add_subdirectory(tests) endif() +add_subdirectory(models/pi05) 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..e63d7803 --- /dev/null +++ b/cpp/loader/include/flashrt/cpp/loader/safetensors.h @@ -0,0 +1,60 @@ +#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 std::map& metadata() const { + return metadata_; + } + 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_; + std::map metadata_; +}; + +} // namespace loader +} // namespace flashrt diff --git a/cpp/loader/include/flashrt/cpp/loader/sha256.h b/cpp/loader/include/flashrt/cpp/loader/sha256.h new file mode 100644 index 00000000..635b94f3 --- /dev/null +++ b/cpp/loader/include/flashrt/cpp/loader/sha256.h @@ -0,0 +1,15 @@ +#ifndef FLASHRT_CPP_LOADER_SHA256_H +#define FLASHRT_CPP_LOADER_SHA256_H + +#include + +namespace flashrt { +namespace loader { + +bool sha256_file(const std::string& path, std::string* hex, + std::string* error = nullptr); + +} // namespace loader +} // namespace flashrt + +#endif // FLASHRT_CPP_LOADER_SHA256_H diff --git a/cpp/loader/src/safetensors.cpp b/cpp/loader/src/safetensors.cpp new file mode 100644 index 00000000..c8c383aa --- /dev/null +++ b/cpp/loader/src/safetensors.cpp @@ -0,0 +1,495 @@ +#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, + std::map* metadata) { + skip_ws(); + if (!consume('{')) return fail("safetensors header must be an object"); + skip_ws(); + if (consume('}')) return finish(tensors, metadata); + 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 (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; + 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, metadata); + 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, + 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; + 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_; + std::map metadata_; + bool metadata_seen_ = false; +}; + +} // 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_); + metadata_ = std::move(other.metadata_); + 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_, &metadata_)) { + 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(); + metadata_.clear(); +} + +const SafetensorInfo* SafetensorsFile::find(const std::string& name) const { + const auto it = tensors_.find(name); + return it == tensors_.end() ? nullptr : &it->second; +} + +const void* SafetensorsFile::data(const SafetensorInfo& tensor) const { + if (!mapping_ || tensor.data_offset > mapping_bytes_ - data_offset_ || + tensor.bytes > mapping_bytes_ - data_offset_ - tensor.data_offset) { + return nullptr; + } + return static_cast(mapping_) + data_offset_ + + tensor.data_offset; +} + +} // namespace loader +} // namespace flashrt diff --git a/cpp/loader/src/sha256.cpp b/cpp/loader/src/sha256.cpp new file mode 100644 index 00000000..63acb488 --- /dev/null +++ b/cpp/loader/src/sha256.cpp @@ -0,0 +1,212 @@ +#include "flashrt/cpp/loader/sha256.h" + +#include +#include +#include +#include +#include +#include +#include + +#ifdef FLASHRT_CPP_HAS_OPENSSL +#include +#endif + +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; + } +#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; +#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); + *hex = output.str(); + if (error) error->clear(); + return true; +} + +} // namespace loader +} // namespace flashrt diff --git a/cpp/modalities/include/flashrt/cpp/modalities/action.h b/cpp/modalities/include/flashrt/cpp/modalities/action.h index 1596f30e..02d967c5 100644 --- a/cpp/modalities/include/flashrt/cpp/modalities/action.h +++ b/cpp/modalities/include/flashrt/cpp/modalities/action.h @@ -24,6 +24,14 @@ struct ActionPostprocessSpec { bool clamp = false; }; +struct ActionStaging { + void* host_pinned = nullptr; + std::uint64_t bytes = 0; +}; + +Status action_staging_create(ActionStaging* out, std::uint64_t bytes); +void action_staging_destroy(ActionStaging*); + Status postprocess_action_cpu(const ActionPostprocessSpec& spec, TensorView model_output, std::vector* robot_actions); @@ -34,7 +42,8 @@ Status postprocess_action_cpu(const ActionPostprocessSpec& spec, Status postprocess_action(const ActionPostprocessSpec& spec, TensorView model_output, std::vector* robot_actions, - void* stream = nullptr); + void* stream = nullptr, + ActionStaging* staging = nullptr); std::uint64_t required_action_output_bytes(const ActionPostprocessSpec& spec, DType dtype); diff --git a/cpp/modalities/include/flashrt/cpp/modalities/text.h b/cpp/modalities/include/flashrt/cpp/modalities/text.h new file mode 100644 index 00000000..b4a3c39c --- /dev/null +++ b/cpp/modalities/include/flashrt/cpp/modalities/text.h @@ -0,0 +1,44 @@ +#ifndef FLASHRT_MODALITIES_TEXT_H +#define FLASHRT_MODALITIES_TEXT_H + +#include "flashrt/cpp/modalities/types.h" + +#include + +namespace flashrt { +namespace modalities { + +struct EmbeddingGatherSpec { + std::uint64_t vocab_size = 0; + std::uint64_t hidden_dim = 0; + float scale = 1.0f; +}; + +struct TextEmbeddingStaging { + void* device_token_ids = nullptr; + void* device_status = nullptr; + std::uint64_t max_tokens = 0; +}; + +Status gather_token_embeddings_cpu(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output); + +Status text_embedding_staging_create(TextEmbeddingStaging* out, + std::uint64_t max_tokens); +void text_embedding_staging_destroy(TextEmbeddingStaging*); + +Status gather_token_embeddings(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream = nullptr, + TextEmbeddingStaging* staging = nullptr); + +} // namespace modalities +} // namespace flashrt + +#endif // FLASHRT_MODALITIES_TEXT_H diff --git a/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h new file mode 100644 index 00000000..2a4667c0 --- /dev/null +++ b/cpp/modalities/include/flashrt/cpp/modalities/tokenizer.h @@ -0,0 +1,55 @@ +#ifndef FLASHRT_MODALITIES_TOKENIZER_H +#define FLASHRT_MODALITIES_TOKENIZER_H + +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace modalities { + +struct SentencePieceEncodeOptions { + bool add_bos = false; + bool add_eos = false; + bool pad_to_max_tokens = false; + std::uint64_t max_tokens = 0; + std::int32_t pad_id = 0; +}; + +class SentencePieceTokenizer final { +public: + SentencePieceTokenizer(); + ~SentencePieceTokenizer(); + + SentencePieceTokenizer(SentencePieceTokenizer&&) noexcept; + SentencePieceTokenizer& operator=(SentencePieceTokenizer&&) noexcept; + + SentencePieceTokenizer(const SentencePieceTokenizer&) = delete; + SentencePieceTokenizer& operator=(const SentencePieceTokenizer&) = delete; + + Status load_model(const std::string& model_path); + Status encode(const std::string& text, + const SentencePieceEncodeOptions& options, + std::vector* token_ids); + void reserve(std::uint64_t max_tokens); + std::uint64_t workspace_capacity() const; + + std::int32_t bos_id() const; + std::int32_t eos_id() const; + std::int32_t unk_id() const; + std::int32_t pad_id() const; + std::uint64_t vocab_size() const; + bool loaded() const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace modalities +} // namespace flashrt + +#endif // FLASHRT_MODALITIES_TOKENIZER_H diff --git a/cpp/modalities/include/flashrt/cpp/modalities/vision.h b/cpp/modalities/include/flashrt/cpp/modalities/vision.h index 5bea0881..a5cdc330 100644 --- a/cpp/modalities/include/flashrt/cpp/modalities/vision.h +++ b/cpp/modalities/include/flashrt/cpp/modalities/vision.h @@ -11,12 +11,14 @@ namespace modalities { enum class NormalizeMode { kScaleShift, + kDivideShift, kMeanStd, }; struct NormalizeSpec { NormalizeMode mode = NormalizeMode::kScaleShift; float scale = 1.0f / 127.5f; + float divisor = 127.5f; float shift = -1.0f; float mean[3] = {0.0f, 0.0f, 0.0f}; float inv_std[3] = {1.0f, 1.0f, 1.0f}; diff --git a/cpp/modalities/src/action_cpu.cpp b/cpp/modalities/src/action_cpu.cpp index c020814f..06998fbf 100644 --- a/cpp/modalities/src/action_cpu.cpp +++ b/cpp/modalities/src/action_cpu.cpp @@ -31,6 +31,38 @@ bool has_dim(const std::vector& v, int dim) { } // namespace +Status action_staging_create(ActionStaging* out, std::uint64_t bytes) { + if (!out || !bytes) { + return Status::error(StatusCode::kInvalidArgument, + "invalid action staging capacity"); + } + action_staging_destroy(out); +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return Status::error(StatusCode::kUnsupported, + "action staging requires the CUDA build"); +#else + cudaError_t rc = cudaMallocHost(&out->host_pinned, + static_cast(bytes)); + if (rc != cudaSuccess) { + *out = ActionStaging{}; + return Status::error( + StatusCode::kBackend, + std::string("cuda action host staging allocation failed: ") + + cudaGetErrorString(rc)); + } + out->bytes = bytes; + return Status::ok(); +#endif +} + +void action_staging_destroy(ActionStaging* staging) { + if (!staging) return; +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING + if (staging->host_pinned) cudaFreeHost(staging->host_pinned); +#endif + *staging = ActionStaging{}; +} + std::uint64_t required_action_output_bytes(const ActionPostprocessSpec& spec, DType dtype) { if (spec.chunk <= 0 || spec.model_dim <= 0) return 0; @@ -97,7 +129,8 @@ Status postprocess_action_cpu(const ActionPostprocessSpec& spec, Status postprocess_action(const ActionPostprocessSpec& spec, TensorView model_output, std::vector* robot_actions, - void* stream) { + void* stream, + ActionStaging* persistent_staging) { if (model_output.place == MemoryPlace::kHost || model_output.place == MemoryPlace::kHostPinned) { return postprocess_action_cpu(spec, model_output, robot_actions); @@ -117,15 +150,27 @@ Status postprocess_action(const ActionPostprocessSpec& spec, return Status::error(StatusCode::kInsufficientStorage, "device action output storage is too small"); } - std::vector staging(static_cast(bytes)); + std::vector fallback; + void* staging = nullptr; + if (persistent_staging) { + if (!persistent_staging->host_pinned || + persistent_staging->bytes < bytes) { + return Status::error(StatusCode::kInsufficientStorage, + "action staging capacity is too small"); + } + staging = persistent_staging->host_pinned; + } else { + fallback.resize(static_cast(bytes)); + staging = fallback.data(); + } cudaError_t rc = cudaSuccess; if (stream) { auto* cuda_stream = reinterpret_cast(stream); - rc = cudaMemcpyAsync(staging.data(), model_output.data, bytes, + rc = cudaMemcpyAsync(staging, model_output.data, bytes, cudaMemcpyDeviceToHost, cuda_stream); if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); } else { - rc = cudaMemcpy(staging.data(), model_output.data, bytes, + rc = cudaMemcpy(staging, model_output.data, bytes, cudaMemcpyDeviceToHost); } if (rc != cudaSuccess) { @@ -134,7 +179,7 @@ Status postprocess_action(const ActionPostprocessSpec& spec, cudaGetErrorString(rc)); } TensorView host; - host.data = staging.data(); + host.data = staging; host.bytes = bytes; host.dtype = model_output.dtype; host.place = MemoryPlace::kHost; diff --git a/cpp/modalities/src/sentencepiece_tokenizer.cpp b/cpp/modalities/src/sentencepiece_tokenizer.cpp new file mode 100644 index 00000000..363e61e7 --- /dev/null +++ b/cpp/modalities/src/sentencepiece_tokenizer.cpp @@ -0,0 +1,137 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +#include + +#include +#include + +namespace flashrt { +namespace modalities { + +struct SentencePieceTokenizer::Impl { + sentencepiece::SentencePieceProcessor processor; + std::vector encoded; + bool loaded = false; +}; + +SentencePieceTokenizer::SentencePieceTokenizer() + : impl_(new Impl()) {} + +SentencePieceTokenizer::~SentencePieceTokenizer() = default; + +SentencePieceTokenizer::SentencePieceTokenizer( + SentencePieceTokenizer&&) noexcept = default; + +SentencePieceTokenizer& SentencePieceTokenizer::operator=( + SentencePieceTokenizer&&) noexcept = default; + +Status SentencePieceTokenizer::load_model(const std::string& model_path) { + auto status = impl_->processor.Load(model_path); + if (!status.ok()) { + impl_->loaded = false; + return Status::error(StatusCode::kNotFound, status.ToString()); + } + impl_->loaded = true; + return Status::ok(); +} + +Status SentencePieceTokenizer::encode( + const std::string& text, + const SentencePieceEncodeOptions& options, + std::vector* token_ids) { + if (!token_ids) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids output is null"); + } + token_ids->clear(); + if (!impl_->loaded) { + return Status::error(StatusCode::kInvalidArgument, + "SentencePiece model is not loaded"); + } + + impl_->encoded.clear(); + auto status = impl_->processor.Encode(text, &impl_->encoded); + if (!status.ok()) { + return Status::error(StatusCode::kBackend, status.ToString()); + } + const std::uint64_t extra = + (options.add_bos ? 1u : 0u) + (options.add_eos ? 1u : 0u); + if (options.max_tokens && impl_->encoded.size() + extra > + options.max_tokens) { + return Status::error(StatusCode::kShapeMismatch, + "encoded token sequence exceeds max_tokens"); + } + if (impl_->encoded.size() + extra > + static_cast(std::numeric_limits::max())) { + return Status::error(StatusCode::kInsufficientStorage, + "encoded token sequence is too large"); + } + + if (options.add_bos) { + const int bos = impl_->processor.bos_id(); + if (bos < 0) { + return Status::error(StatusCode::kInvalidArgument, + "tokenizer has no BOS id"); + } + token_ids->push_back(static_cast(bos)); + } + token_ids->reserve(impl_->encoded.size() + extra); + for (int id : impl_->encoded) { + token_ids->push_back(static_cast(id)); + } + if (options.add_eos) { + const int eos = impl_->processor.eos_id(); + if (eos < 0) { + return Status::error(StatusCode::kInvalidArgument, + "tokenizer has no EOS id"); + } + token_ids->push_back(static_cast(eos)); + } + + if (options.max_tokens) { + if (options.pad_to_max_tokens) { + token_ids->resize(options.max_tokens, options.pad_id); + } + } else if (options.pad_to_max_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "pad_to_max_tokens requires max_tokens"); + } + return Status::ok(); +} + +void SentencePieceTokenizer::reserve(std::uint64_t max_tokens) { + impl_->encoded.reserve(static_cast(max_tokens)); +} + +std::uint64_t SentencePieceTokenizer::workspace_capacity() const { + return static_cast(impl_->encoded.capacity()); +} + +std::int32_t SentencePieceTokenizer::bos_id() const { + return impl_->loaded ? impl_->processor.bos_id() : -1; +} + +std::int32_t SentencePieceTokenizer::eos_id() const { + return impl_->loaded ? impl_->processor.eos_id() : -1; +} + +std::int32_t SentencePieceTokenizer::unk_id() const { + return impl_->loaded ? impl_->processor.unk_id() : -1; +} + +std::int32_t SentencePieceTokenizer::pad_id() const { + return impl_->loaded ? impl_->processor.pad_id() : -1; +} + +std::uint64_t SentencePieceTokenizer::vocab_size() const { + return impl_->loaded + ? static_cast(impl_->processor.GetPieceSize()) + : 0; +} + +bool SentencePieceTokenizer::loaded() const { + return impl_->loaded; +} + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/text_cpu.cpp b/cpp/modalities/src/text_cpu.cpp new file mode 100644 index 00000000..454b9acb --- /dev/null +++ b/cpp/modalities/src/text_cpu.cpp @@ -0,0 +1,232 @@ +#include "flashrt/cpp/modalities/text.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace modalities { + +#ifdef FLASHRT_CPP_WITH_CUDA_KERNELS +Status gather_token_embeddings_cuda(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream, + TextEmbeddingStaging* staging); +#endif + +namespace { + +float load_scalar(const void* base, std::uint64_t index, DType dtype) { + switch (dtype) { + case DType::kFloat32: + return static_cast(base)[index]; + case DType::kBFloat16: + return bfloat16_to_float( + static_cast(base)[index]); + case DType::kFloat16: + return float16_to_float( + static_cast(base)[index]); + case DType::kUInt8: + return static_cast( + static_cast(base)[index]); + } + return 0.0f; +} + +void store_scalar(void* base, std::uint64_t index, DType dtype, float value) { + switch (dtype) { + case DType::kFloat32: + static_cast(base)[index] = value; + break; + case DType::kBFloat16: + static_cast(base)[index] = float_to_bfloat16(value); + break; + case DType::kFloat16: + static_cast(base)[index] = float_to_float16(value); + break; + case DType::kUInt8: + static_cast(base)[index] = + static_cast(value); + break; + } +} + +Status validate_matrix(const TensorView& tensor, const char* name, + std::uint64_t rows, std::uint64_t cols) { + Status st = validate_host_tensor(tensor, name); + if (!st.ok_status()) return st; + if (tensor.layout != Layout::kFlat || tensor.shape.rank != 2 || + tensor.shape.dims[0] != rows || tensor.shape.dims[1] != cols) { + return Status::error(StatusCode::kShapeMismatch, + std::string(name) + " shape mismatch"); + } + return Status::ok(); +} + +} // namespace + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +Status text_embedding_staging_create(TextEmbeddingStaging* out, + std::uint64_t max_tokens) { + if (!out || !max_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "invalid text embedding staging capacity"); + } + *out = TextEmbeddingStaging{}; + cudaError_t rc = cudaMalloc(&out->device_token_ids, + max_tokens * sizeof(std::int32_t)); + if (rc != cudaSuccess) { + return Status::error( + StatusCode::kBackend, + std::string("text token staging cudaMalloc failed: ") + + cudaGetErrorString(rc)); + } + rc = cudaMalloc(&out->device_status, sizeof(int)); + if (rc != cudaSuccess) { + cudaFree(out->device_token_ids); + *out = TextEmbeddingStaging{}; + return Status::error( + StatusCode::kBackend, + std::string("text status staging cudaMalloc failed: ") + + cudaGetErrorString(rc)); + } + out->max_tokens = max_tokens; + return Status::ok(); +} + +void text_embedding_staging_destroy(TextEmbeddingStaging* s) { + if (!s) return; + if (s->device_token_ids) cudaFree(s->device_token_ids); + if (s->device_status) cudaFree(s->device_status); + *s = TextEmbeddingStaging{}; +} +#else +Status text_embedding_staging_create(TextEmbeddingStaging* out, + std::uint64_t) { + if (out) *out = TextEmbeddingStaging{}; + return Status::error(StatusCode::kUnsupported, + "text embedding staging requires the CUDA build"); +} + +void text_embedding_staging_destroy(TextEmbeddingStaging* s) { + if (s) *s = TextEmbeddingStaging{}; +} +#endif + +Status gather_token_embeddings_cpu(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output) { + if (!token_ids && n_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids is null"); + } + if (!spec.vocab_size || !spec.hidden_dim) { + return Status::error(StatusCode::kInvalidArgument, + "invalid embedding gather dimensions"); + } + Status st = validate_matrix(embedding_table, "embedding_table", + spec.vocab_size, spec.hidden_dim); + if (!st.ok_status()) return st; + st = validate_matrix(output, "embedding_output", n_tokens, + spec.hidden_dim); + if (!st.ok_status()) return st; + + for (std::uint64_t t = 0; t < n_tokens; ++t) { + const std::int32_t token = token_ids[t]; + if (token < 0 || + static_cast(token) >= spec.vocab_size) { + return Status::error(StatusCode::kInvalidArgument, + "token id is out of vocabulary range"); + } + const std::uint64_t src_base = + static_cast(token) * spec.hidden_dim; + const std::uint64_t dst_base = t * spec.hidden_dim; + for (std::uint64_t d = 0; d < spec.hidden_dim; ++d) { + const float value = load_scalar( + embedding_table.data, src_base + d, embedding_table.dtype); + store_scalar(output.data, dst_base + d, output.dtype, + value * spec.scale); + } + } + return Status::ok(); +} + +Status gather_token_embeddings(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream, + TextEmbeddingStaging* staging) { + if (output.place == MemoryPlace::kHost || + output.place == MemoryPlace::kHostPinned) { + (void)stream; + (void)staging; + return gather_token_embeddings_cpu(spec, token_ids, n_tokens, + embedding_table, output); + } + if (output.place != MemoryPlace::kDevice || + embedding_table.place != MemoryPlace::kDevice) { + return Status::error(StatusCode::kUnsupported, + "device text embedding requires device tensors"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)stream; + (void)staging; + return Status::error(StatusCode::kUnsupported, + "device text embedding was not enabled at build time"); +#else + if (!token_ids && n_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids is null"); + } + if (staging && staging->max_tokens < n_tokens) { + return Status::error(StatusCode::kInsufficientStorage, + "text token staging capacity is too small"); + } +#ifdef FLASHRT_CPP_WITH_CUDA_KERNELS + return gather_token_embeddings_cuda(spec, token_ids, n_tokens, + embedding_table, output, stream, + staging); +#else + std::vector host_bytes( + static_cast(n_tokens * spec.hidden_dim * + dtype_size(output.dtype))); + TensorView host_output{host_bytes.data(), + static_cast(host_bytes.size()), + output.dtype, MemoryPlace::kHost, output.layout, + Shape{n_tokens, spec.hidden_dim}}; + TensorView host_table = embedding_table; + if (embedding_table.place != MemoryPlace::kHost && + embedding_table.place != MemoryPlace::kHostPinned) { + return Status::error(StatusCode::kUnsupported, + "CUDA kernel build is required for device embedding tables"); + } + Status st = gather_token_embeddings_cpu(spec, token_ids, n_tokens, + host_table, host_output); + if (!st.ok_status()) return st; + cudaStream_t cuda_stream = reinterpret_cast(stream); + cudaError_t rc = cudaMemcpyAsync(output.data, host_bytes.data(), + host_bytes.size(), cudaMemcpyHostToDevice, + cuda_stream); + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + if (rc != cudaSuccess) { + return Status::error(StatusCode::kBackend, + std::string("cuda H2D text embedding failed: ") + + cudaGetErrorString(rc)); + } + return Status::ok(); +#endif +#endif +} + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/text_cuda.cu b/cpp/modalities/src/text_cuda.cu new file mode 100644 index 00000000..da8bf63d --- /dev/null +++ b/cpp/modalities/src/text_cuda.cu @@ -0,0 +1,208 @@ +#include "flashrt/cpp/modalities/text.h" + +#include +#include + +#include +#include +#include + +namespace flashrt { +namespace modalities { +namespace { + +__device__ __forceinline__ float bf16_to_f32(std::uint16_t value) { + return __uint_as_float(static_cast(value) << 16); +} + +__device__ __forceinline__ std::uint16_t f32_to_bf16(float value) { + std::uint32_t bits = __float_as_uint(value); + const std::uint32_t lsb = (bits >> 16) & 1u; + bits += 0x7fffu + lsb; + return static_cast(bits >> 16); +} + +__device__ __forceinline__ float load_value(const void* base, + std::uint64_t index, + int dtype) { + if (dtype == 1) return static_cast(base)[index]; + if (dtype == 2) return __half2float(static_cast(base)[index]); + if (dtype == 3) { + return bf16_to_f32(static_cast(base)[index]); + } + return static_cast(static_cast(base)[index]); +} + +__device__ __forceinline__ void store_value(void* base, + std::uint64_t index, + int dtype, + float value) { + if (dtype == 1) { + static_cast(base)[index] = value; + } else if (dtype == 2) { + static_cast<__half*>(base)[index] = __float2half_rn(value); + } else if (dtype == 3) { + static_cast(base)[index] = f32_to_bf16(value); + } else { + static_cast(base)[index] = + static_cast(value); + } +} + +int dtype_code(DType dtype) { + switch (dtype) { + case DType::kFloat32: return 1; + case DType::kFloat16: return 2; + case DType::kBFloat16: return 3; + case DType::kUInt8: return 0; + } + return 0; +} + +Status validate_device_matrix(const TensorView& tensor, const char* name, + std::uint64_t rows, std::uint64_t cols) { + if (!tensor.data) { + return Status::error(StatusCode::kInvalidArgument, + std::string(name) + " has null data"); + } + if (tensor.place != MemoryPlace::kDevice) { + return Status::error(StatusCode::kUnsupported, + std::string(name) + " is not device memory"); + } + if (tensor.layout != Layout::kFlat || tensor.shape.rank != 2 || + tensor.shape.dims[0] != rows || tensor.shape.dims[1] != cols) { + return Status::error(StatusCode::kShapeMismatch, + std::string(name) + " shape mismatch"); + } + const std::uint64_t need = rows * cols * dtype_size(tensor.dtype); + if (tensor.bytes < need) { + return Status::error(StatusCode::kInsufficientStorage, + std::string(name) + " storage is too small"); + } + return Status::ok(); +} + +__global__ void gather_kernel(const std::int32_t* ids, + std::uint64_t n_tokens, + std::uint64_t vocab_size, + std::uint64_t hidden_dim, + const void* table, + int table_dtype, + void* output, + int output_dtype, + float scale, + int* bad_token) { + const std::uint64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const std::uint64_t total = n_tokens * hidden_dim; + if (idx >= total) return; + const std::uint64_t token_index = idx / hidden_dim; + const std::uint64_t dim = idx - token_index * hidden_dim; + const std::int32_t token = ids[token_index]; + if (token < 0 || static_cast(token) >= vocab_size) { + atomicCAS(bad_token, 0, 1); + return; + } + const std::uint64_t src = + static_cast(token) * hidden_dim + dim; + const float value = load_value(table, src, table_dtype) * scale; + store_value(output, idx, output_dtype, value); +} + +const char* cuda_error(cudaError_t rc) { + return cudaGetErrorString(rc); +} + +} // namespace + +Status gather_token_embeddings_cuda(const EmbeddingGatherSpec& spec, + const std::int32_t* token_ids, + std::uint64_t n_tokens, + TensorView embedding_table, + TensorView output, + void* stream, + TextEmbeddingStaging* staging) { + if (!token_ids && n_tokens) { + return Status::error(StatusCode::kInvalidArgument, + "token_ids is null"); + } + if (!spec.vocab_size || !spec.hidden_dim) { + return Status::error(StatusCode::kInvalidArgument, + "invalid embedding gather dimensions"); + } + Status st = validate_device_matrix(embedding_table, "embedding_table", + spec.vocab_size, spec.hidden_dim); + if (!st.ok_status()) return st; + st = validate_device_matrix(output, "embedding_output", n_tokens, + spec.hidden_dim); + if (!st.ok_status()) return st; + if (staging && staging->max_tokens < n_tokens) { + return Status::error(StatusCode::kInsufficientStorage, + "text token staging capacity is too small"); + } + + cudaStream_t cuda_stream = reinterpret_cast(stream); + std::int32_t* d_ids = nullptr; + int* d_bad = nullptr; + cudaError_t rc = cudaSuccess; + if (staging) { + d_ids = static_cast(staging->device_token_ids); + d_bad = static_cast(staging->device_status); + } else { + rc = cudaMalloc(&d_ids, n_tokens * sizeof(std::int32_t)); + if (rc != cudaSuccess) { + return Status::error( + StatusCode::kBackend, + std::string("cudaMalloc text token ids failed: ") + + cuda_error(rc)); + } + } + if (!d_bad) rc = cudaMalloc(&d_bad, sizeof(int)); + if (rc == cudaSuccess) { + rc = cudaMemsetAsync(d_bad, 0, sizeof(int), cuda_stream); + } + if (rc == cudaSuccess && n_tokens) { + rc = cudaMemcpyAsync(d_ids, token_ids, n_tokens * sizeof(std::int32_t), + cudaMemcpyHostToDevice, cuda_stream); + } + if (rc != cudaSuccess) { + if (!staging) cudaFree(d_ids); + if (!staging && d_bad) cudaFree(d_bad); + return Status::error(StatusCode::kBackend, + std::string("cuda text token upload failed: ") + + cuda_error(rc)); + } + + const std::uint64_t total = n_tokens * spec.hidden_dim; + if (total) { + const int block = 256; + const int grid = static_cast((total + block - 1) / block); + gather_kernel<<>>( + d_ids, n_tokens, spec.vocab_size, spec.hidden_dim, + embedding_table.data, dtype_code(embedding_table.dtype), + output.data, dtype_code(output.dtype), spec.scale, d_bad); + rc = cudaGetLastError(); + } + + int bad = 0; + if (rc == cudaSuccess) { + rc = cudaMemcpyAsync(&bad, d_bad, sizeof(int), cudaMemcpyDeviceToHost, + cuda_stream); + } + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + if (!staging) cudaFree(d_ids); + if (!staging) cudaFree(d_bad); + if (rc != cudaSuccess) { + return Status::error(StatusCode::kBackend, + std::string("text embedding CUDA failed: ") + + cuda_error(rc)); + } + if (bad) { + return Status::error(StatusCode::kInvalidArgument, + "token id is out of vocabulary range"); + } + return Status::ok(); +} + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/tokenizer_unavailable.cpp b/cpp/modalities/src/tokenizer_unavailable.cpp new file mode 100644 index 00000000..7cb70158 --- /dev/null +++ b/cpp/modalities/src/tokenizer_unavailable.cpp @@ -0,0 +1,52 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +namespace flashrt { +namespace modalities { + +struct SentencePieceTokenizer::Impl {}; + +SentencePieceTokenizer::SentencePieceTokenizer() + : impl_(new Impl()) {} + +SentencePieceTokenizer::~SentencePieceTokenizer() = default; + +SentencePieceTokenizer::SentencePieceTokenizer( + SentencePieceTokenizer&&) noexcept = default; + +SentencePieceTokenizer& SentencePieceTokenizer::operator=( + SentencePieceTokenizer&&) noexcept = default; + +Status SentencePieceTokenizer::load_model(const std::string& model_path) { + (void)model_path; + return Status::error( + StatusCode::kUnsupported, + "native SentencePiece support is not enabled in this build"); +} + +Status SentencePieceTokenizer::encode( + const std::string& text, + const SentencePieceEncodeOptions& options, + std::vector* token_ids) { + (void)text; + (void)options; + if (token_ids) token_ids->clear(); + return Status::error( + StatusCode::kUnsupported, + "native SentencePiece support is not enabled in this build"); +} + +void SentencePieceTokenizer::reserve(std::uint64_t max_tokens) { + (void)max_tokens; +} + +std::uint64_t SentencePieceTokenizer::workspace_capacity() const { return 0; } + +std::int32_t SentencePieceTokenizer::bos_id() const { return -1; } +std::int32_t SentencePieceTokenizer::eos_id() const { return -1; } +std::int32_t SentencePieceTokenizer::unk_id() const { return -1; } +std::int32_t SentencePieceTokenizer::pad_id() const { return -1; } +std::uint64_t SentencePieceTokenizer::vocab_size() const { return 0; } +bool SentencePieceTokenizer::loaded() const { return false; } + +} // namespace modalities +} // namespace flashrt diff --git a/cpp/modalities/src/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 e6c58e49..3a93098d 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 @@ -27,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); @@ -111,6 +119,9 @@ float normalize_value(float raw, int c, const NormalizeSpec& spec) { if (spec.mode == NormalizeMode::kScaleShift) { return raw * spec.scale + spec.shift; } + if (spec.mode == NormalizeMode::kDivideShift) { + return raw / spec.divisor + spec.shift; + } return (raw / 255.0f - spec.mean[c]) * spec.inv_std[c]; } @@ -142,6 +153,13 @@ Status validate_frame(const VisionFrame& frame) { return Status::error(StatusCode::kUnsupported, "unsupported pixel format"); } + if (frame.width > std::numeric_limits::max() / ch || + frame.stride_bytes < 0 || + (frame.stride_bytes > 0 && + frame.stride_bytes < frame.width * ch)) { + return Status::error(StatusCode::kShapeMismatch, + "vision frame stride is smaller than one row"); + } const int stride = frame.stride_bytes > 0 ? frame.stride_bytes : frame.width * ch; const std::uint64_t need = static_cast(stride) * static_cast(frame.height); diff --git a/cpp/modalities/src/vision_cuda.cu b/cpp/modalities/src/vision_cuda.cu index b22d441f..50ee7e73 100644 --- a/cpp/modalities/src/vision_cuda.cu +++ b/cpp/modalities/src/vision_cuda.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -61,6 +62,13 @@ Status validate_frame_for_cuda(const VisionFrame& frame) { return Status::error(StatusCode::kUnsupported, "unsupported pixel format"); } + if (frame.width > std::numeric_limits::max() / ch || + frame.stride_bytes < 0 || + (frame.stride_bytes > 0 && + frame.stride_bytes < frame.width * ch)) { + return Status::error(StatusCode::kShapeMismatch, + "vision frame stride is smaller than one row"); + } const int stride = frame.stride_bytes > 0 ? frame.stride_bytes : frame.width * ch; const std::uint64_t need = static_cast(stride) * @@ -108,11 +116,18 @@ __device__ __forceinline__ float normalize_value(float raw, int c, int norm_mode, float scale, + float divisor, float shift, const float* mean, const float* inv_std) { - if (norm_mode == 0) return raw * scale + shift; - return (raw / 255.0f - mean[c]) * inv_std[c]; + if (norm_mode == 0) { + return __fadd_rn(__fmul_rn(raw, scale), shift); + } + if (norm_mode == 1) { + return __fadd_rn(__fdiv_rn(raw, divisor), shift); + } + return __fmul_rn( + __fsub_rn(__fdiv_rn(raw, 255.0f), mean[c]), inv_std[c]); } __device__ __forceinline__ void store_value(void* out, @@ -141,6 +156,7 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, int th, int norm_mode, float scale, + float divisor, float shift, float mean0, float mean1, @@ -152,10 +168,18 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, const int y = blockIdx.y * blockDim.y + threadIdx.y; if (x >= tw || y >= th) return; - const float fx = (static_cast(x) + 0.5f) * - static_cast(sw) / static_cast(tw) - 0.5f; - const float fy = (static_cast(y) + 0.5f) * - static_cast(sh) / static_cast(th) - 0.5f; + const float fx = __fsub_rn( + __fdiv_rn( + __fmul_rn(static_cast(x) + 0.5f, + static_cast(sw)), + static_cast(tw)), + 0.5f); + const float fy = __fsub_rn( + __fdiv_rn( + __fmul_rn(static_cast(y) + 0.5f, + static_cast(sh)), + static_cast(th)), + 0.5f); const int x0 = max(0, min(sw - 1, static_cast(floorf(fx)))); const int y0 = max(0, min(sh - 1, static_cast(floorf(fy)))); const int x1 = max(0, min(sw - 1, x0 + 1)); @@ -172,11 +196,17 @@ __global__ void resize_normalize_kernel(const std::uint8_t* src, float mean[3] = {mean0, mean1, mean2}; float inv_std[3] = {inv_std0, inv_std1, inv_std2}; for (int c = 0; c < 3; ++c) { - const float top = p00[c] * (1.0f - wx) + p01[c] * wx; - const float bot = p10[c] * (1.0f - wx) + p11[c] * wx; - const float raw = top * (1.0f - wy) + bot * wy; - const float norm = normalize_value(raw, c, norm_mode, scale, shift, - mean, inv_std); + const float top = __fadd_rn( + __fmul_rn(p00[c], __fsub_rn(1.0f, wx)), + __fmul_rn(p01[c], wx)); + const float bot = __fadd_rn( + __fmul_rn(p10[c], __fsub_rn(1.0f, wx)), + __fmul_rn(p11[c], wx)); + const float raw = __fadd_rn( + __fmul_rn(top, __fsub_rn(1.0f, wy)), + __fmul_rn(bot, wy)); + const float norm = normalize_value( + raw, c, norm_mode, scale, divisor, shift, mean, inv_std); const std::uint64_t out_idx = (((static_cast(view) * th + y) * tw + x) * 3ull) + static_cast(c); @@ -305,8 +335,11 @@ Status preprocess_vision_cuda(const VisionPreprocessSpec& spec, spec.output_dtype == DType::kFloat32 ? 1 : (spec.output_dtype == DType::kFloat16 ? 2 : 0), static_cast(v), spec.target_width, spec.target_height, - spec.normalize.mode == NormalizeMode::kScaleShift ? 0 : 1, - spec.normalize.scale, spec.normalize.shift, + spec.normalize.mode == NormalizeMode::kScaleShift + ? 0 + : (spec.normalize.mode == NormalizeMode::kDivideShift ? 1 : 2), + spec.normalize.scale, spec.normalize.divisor, + spec.normalize.shift, spec.normalize.mean[0], spec.normalize.mean[1], spec.normalize.mean[2], spec.normalize.inv_std[0], spec.normalize.inv_std[1], spec.normalize.inv_std[2]); diff --git a/cpp/models/pi05/CMakeLists.txt b/cpp/models/pi05/CMakeLists.txt new file mode 100644 index 00000000..5102e35f --- /dev/null +++ b/cpp/models/pi05/CMakeLists.txt @@ -0,0 +1,190 @@ +# 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/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_materializer.cpp + src/support/native_calibration.cpp + src/support/native_workspace.cpp + src/backend/native_calibration_session.cpp) +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + list(APPEND _pi05_sources + src/support/native_quantization.cu + src/support/native_rope.cu) +else() + list(APPEND _pi05_sources + src/support/native_quantization_unavailable.cpp + src/support/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/support/native_weight_ops.cpp APPEND PROPERTY + COMPILE_OPTIONS -ffp-contract=off) +endif() + +if(FLASHRT_CPP_WITH_CUDA_KERNELS) + 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/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 + ${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/session.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 + 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 + ${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) + + add_library(flashrt_cpp_pi05_backend_sm110 STATIC + $ + 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/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 + ${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) + _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_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_backend_sm110 PRIVATE + $<$: + --expt-extended-lambda --use_fast_math + -gencode=arch=compute_110a,code=sm_110a>) + 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() + + target_link_libraries(flashrt_cpp_pi05_kernels INTERFACE + ${_pi05_backend_targets}) +endif() + +add_library(flashrt_cpp_pi05_c SHARED + 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) +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/include/flashrt/cpp/models/pi05/c_api.h b/cpp/models/pi05/include/flashrt/cpp/models/pi05/c_api.h index 9c727036..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 @@ -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, @@ -43,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; @@ -55,6 +57,37 @@ typedef struct frt_pi05_runtime_config { * capacity is a per-call error, never a fallback allocation. */ int max_frame_width; int max_frame_height; + + /* Optional ABI extension: native prompt staging source. When all fields + * below are present, frt_pi05_runtime_set_prompt* tokenizes text/state and + * writes embeddings into prompt_embedding_data. The captured graph must + * already copy from this stable source buffer into encoder_x; the model + * runtime does not rebind graph pointers. */ + const char* prompt_tokenizer_model_path; + const void* prompt_embedding_table_data; + uint64_t prompt_embedding_table_bytes; + int prompt_embedding_table_dtype; + uint64_t prompt_embedding_vocab_size; + uint64_t prompt_embedding_hidden_dim; + void* prompt_embedding_data; + uint64_t prompt_embedding_bytes; + int prompt_embedding_dtype; + uint64_t max_prompt_tokens; + float prompt_embedding_scale; + + /* Optional ABI extension: raw proprioception normalization for STATE + * STAGED ports. If present, state is mapped with q01/q99 into [-1, 1] + * before Pi0.5 prompt discretization. */ + const float* state_q01; + uint64_t n_state_q01; + const float* state_q99; + uint64_t n_state_q99; + + /* Optional ABI tail: notify an integrated native graph owner after a + * prompt/state update has produced a new valid token count. */ + int (*prompt_length_update)(void* user, uint64_t prompt_len); + void* prompt_length_update_user; + int prompt_embedding_on_device; } frt_pi05_runtime_config; typedef struct frt_pi05_vision_frame { @@ -69,12 +102,33 @@ 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); 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); @@ -87,6 +141,28 @@ 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 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 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, + 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/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/internal/include/flashrt/cpp/models/pi05/backend/native_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/native_calibration_session.h new file mode 100644 index 00000000..6d5cfdbe --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/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/backend/session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/session.h new file mode 100644 index 00000000..ebf0ba70 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backend/session.h @@ -0,0 +1,85 @@ +#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" +#include "flashrt/cpp/native/cuda_graph_set.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class BackendPrecision { + kBf16, + kFp8E4M3, +}; + +struct BackendConfig { + int num_views = 2; + int max_prompt_tokens = 200; + int chunk_size = 10; + int num_steps = 10; + int vision_pool_factor = 1; + BackendPrecision precision = BackendPrecision::kBf16; +}; + +enum class GraphKind : std::size_t { + kInfer = 0, + kDecodeOnly = 1, + kContext = 2, + kCount = 3, +}; + +const char* backend_graph_name(GraphKind kind); + +modalities::Status capture_backend_graph( + native::CudaGraphSet* graphs, + GraphKind kind, + const NativeWorkspace& workspace, + std::initializer_list bindings, + native::CudaGraphSet::RecordFn record, + void* owner); + +modalities::Status copy_prompt_to_encoder(NativeWorkspace* workspace, + void* stream); + +struct BackendArtifacts { + 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_backend_artifacts( + const NativeWorkspace& workspace, + const NativeDeviceWeightStore& weights, + NativeWeightDType embedding_dtype, + BackendArtifacts* artifacts); + +class BackendSession { +public: + virtual ~BackendSession() = default; + + virtual frt_ctx context() const = 0; + 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 BackendArtifacts& artifacts() const = 0; + virtual modalities::Status set_prompt_length(int prompt_tokens) = 0; + virtual int replay(GraphKind kind = GraphKind::kInfer) const = 0; + virtual modalities::Status synchronize() const = 0; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_BACKEND_SESSION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h new file mode 100644 index 00000000..3a0b9c80 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_calibration_session.h @@ -0,0 +1,52 @@ +#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/backend/native_calibration_session.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +using NativeThorCalibrationConfig = NativeCalibrationConfig; + +class NativeThorCalibrationSession final : public NativeCalibrationSession { +public: + static std::unique_ptr create( + const NativeThorCalibrationConfig& config, + double percentile, + modalities::Status* status); + + ~NativeThorCalibrationSession() override; + + 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) override; + modalities::Status finalize( + const std::string& artifact_path) const override; + std::uint64_t sample_count() const override; + +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/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h new file mode 100644 index 00000000..2898c765 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/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/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 + +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/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_kernel_driver.h new file mode 100644 index 00000000..46c1626c --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/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/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_style_precompute.h new file mode 100644 index 00000000..872b16bc --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/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/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 + +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/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/native_thor_weight_materializer.h new file mode 100644 index 00000000..80ac2fb8 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/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/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/support/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/internal/include/flashrt/cpp/models/pi05/backends/sm110/session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/session.h new file mode 100644 index 00000000..9e4a5a74 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm110/session.h @@ -0,0 +1,73 @@ +#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/session.h" +#include "flashrt/cpp/models/pi05/backends/sm110/native_thor_fp8_forward.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class Sm110BackendSession final : public BackendSession { +public: + static std::unique_ptr create( + const std::string& checkpoint_path, + const BackendConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status); + + ~Sm110BackendSession() override; + + Sm110BackendSession(const Sm110BackendSession&) = delete; + Sm110BackendSession& operator=(const Sm110BackendSession&) = delete; + + frt_ctx context() const override { return graphs_.context(); } + 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(); } + NativeDeviceWeightStore& weights() { return weights_; } + const NativeDeviceWeightStore& weights() const { return weights_; } + NativeWorkspace& workspace() { return workspace_; } + const NativeWorkspace& workspace() const { return workspace_; } + const BackendArtifacts& artifacts() const override { + return artifacts_; + } + + modalities::Status set_prompt_length(int prompt_tokens) override; + int replay(GraphKind kind = GraphKind::kInfer) const override; + modalities::Status synchronize() const override; + +private: + Sm110BackendSession(frt_ctx ctx, const BackendConfig& config); + modalities::Status initialize( + const std::string& checkpoint_path, + const NativeCalibrationArtifact& calibration); + 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_; + BackendConfig config_; + NativeDeviceWeightStore weights_; + NativeWorkspace workspace_; + BackendArtifacts artifacts_; + NativeThorKernelDriver driver_; + NativeThorFp8Forward forward_; + NativeThorWeightScales weight_scales_; + std::vector encoder_alphas_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_BACKENDS_SM110_SESSION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h new file mode 100644 index 00000000..1f0431a9 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h @@ -0,0 +1,76 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_BF16_FORWARD_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_BF16_FORWARD_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 { +namespace pi05 { + +class NativeBf16Forward { +public: + explicit NativeBf16Forward(const NativeKernelDriver* 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, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + std::uintptr_t stream) const; +#ifdef FLASHRT_CPP_WITH_FA2 + modalities::Status vision_layer( + int layer, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status vision( + const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status encoder_layer( + int layer, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status encoder( + const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status decoder_layer( + int layer, int step, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status diffusion_step( + int step, const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; + modalities::Status diffusion( + const NativeDeviceWeightStore& weights, NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const; +#endif + +private: + const NativeKernelDriver* driver_ = nullptr; + NativeRtxLinear fallback_linear_; + const NativeRtxLinear* linear_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_BF16_FORWARD_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h new file mode 100644 index 00000000..ebde43e1 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h @@ -0,0 +1,125 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_KERNEL_DRIVER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_KERNEL_DRIVER_H + +#include "flashrt/cpp/modalities/types.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class NativeKernelDriver { +public: + NativeKernelDriver() noexcept; + ~NativeKernelDriver(); + + NativeKernelDriver(const NativeKernelDriver&) = delete; + NativeKernelDriver& operator=(const NativeKernelDriver&) = delete; + + modalities::Status status() const; + modalities::Status bf16_nn(void* a, void* b, void* output, + int m, int n, int k, + std::uintptr_t stream) const; + modalities::Status 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; + 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 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; + 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 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; + 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 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, + std::uintptr_t stream) const; + modalities::Status qkv_split_rope_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + int rows, int query_columns, int key_columns, int value_columns, + int head_dimension, std::uintptr_t stream) const; + modalities::Status qkv_split_rope_devpos_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + const void* device_position, int rows, int query_columns, + int key_columns, int value_columns, int head_dimension, + std::uintptr_t stream) const; + modalities::Status patch_im2col_16bit( + const void* images, void* patches, int num_views, + std::uintptr_t stream) const; + modalities::Status avg_pool_vision_bf16( + const void* values, void* output, int num_views, int height, int width, + int columns, int pool_factor, std::uintptr_t stream) const; + +private: + struct Impl; + std::unique_ptr impl_; + std::string error_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_KERNEL_DRIVER_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h new file mode 100644 index 00000000..cec18243 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h @@ -0,0 +1,76 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H + +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/exec.h" + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeAttentionDType { + kBf16, + kFloat32, + kInt32, +}; + +struct NativeRtxAttentionConfig { + int num_views = 2; + int encoder_sequence = 712; + int encoder_vision_sequence = 512; + int chunk_size = 10; + int encoder_layers = 18; +}; + +struct NativeAttentionBuffer { + frt_buffer buffer = nullptr; + std::vector shape; + NativeAttentionDType dtype = NativeAttentionDType::kBf16; +}; + +class NativeRtxAttentionWorkspace { +public: + explicit NativeRtxAttentionWorkspace(frt_ctx ctx) : ctx_(ctx) {} + + modalities::Status allocate(const NativeRtxAttentionConfig& config); + modalities::Status set_fixed_prompt_length(int prompt_tokens); + const NativeAttentionBuffer* find(const std::string& name) const; + void* encoder_k_layer_dptr(int layer) const; + void* encoder_v_layer_dptr(int layer) const; + + std::size_t size() const { return buffers_.size(); } + std::size_t allocated_bytes() const { return allocated_bytes_; } + std::size_t kv_layer_stride_bytes() const { return kv_layer_stride_bytes_; } + int encoder_splits() const { return encoder_splits_; } + int decoder_splits() const { return decoder_splits_; } + +private: + modalities::Status add(const std::string& name, + std::initializer_list shape, + NativeAttentionDType dtype); + + frt_ctx ctx_ = nullptr; + std::map buffers_; + std::size_t allocated_bytes_ = 0; + std::size_t kv_layer_stride_bytes_ = 0; + int num_views_ = 0; + int encoder_sequence_ = 0; + int encoder_vision_sequence_ = 0; + int chunk_size_ = 0; + int encoder_layers_ = 0; + int encoder_splits_ = 0; + int decoder_splits_ = 0; + frt_buffer prompt_length_buffers_[3] = {nullptr, nullptr, nullptr}; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_RTX_ATTENTION_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention_driver.h new file mode 100644 index 00000000..95855f17 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/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/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h new file mode 100644 index 00000000..e12686e5 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/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/backends/sm120/native_rtx_calibration_session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h new file mode 100644 index 00000000..c04ddcad --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/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/backend/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/backends/sm120/native_rtx_linear.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_linear.h new file mode 100644 index 00000000..7b0db7c3 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/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/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 +#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/backends/sm120/native_rtx_weight_packer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h new file mode 100644 index 00000000..4ec4ceb0 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/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/support/native_device_weights.h" +#include "flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/native_style_precompute.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h new file mode 100644 index 00000000..0c47f37c --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/native_kernel_driver.h" +#include "flashrt/cpp/models/pi05/support/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/internal/include/flashrt/cpp/models/pi05/backends/sm120/session.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/session.h new file mode 100644 index 00000000..7136ecb9 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/backends/sm120/session.h @@ -0,0 +1,82 @@ +#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/session.h" + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +class Sm120BackendSession final : public BackendSession { +public: + 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 BackendConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status); + static std::unique_ptr create_calibration( + const std::string& checkpoint_path, const BackendConfig& config, + modalities::Status* status); + + ~Sm120BackendSession() override; + + Sm120BackendSession(const Sm120BackendSession&) = delete; + Sm120BackendSession& operator=(const Sm120BackendSession&) = delete; + + frt_ctx context() const override { return graphs_.context(); } + 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 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 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(GraphKind kind = GraphKind::kInfer) const override; + modalities::Status synchronize() const override; + +private: + Sm120BackendSession(frt_ctx ctx, const BackendConfig& config, + NativeRtxLinearMode linear_mode); + modalities::Status initialize( + const std::string& checkpoint_path, + const NativeCalibrationArtifact* calibration); + 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_; + BackendConfig config_; + NativeDeviceWeightStore weights_; + NativeWorkspace workspace_; + BackendArtifacts artifacts_; + NativeRtxAttentionWorkspace attention_; + NativeKernelDriver driver_; + NativeRtxLinear linear_; + NativeBf16Forward forward_; + std::unique_ptr attention_driver_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_BACKENDS_SM120_SESSION_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/io.h similarity index 84% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/io.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/io.h index 0ad3f274..807fdb01 100644 --- a/cpp/models/pi05/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 { @@ -23,7 +23,9 @@ class RuntimeIo { int model_action_dim = kModelActionDim, int robot_action_dim = kLiberoActionDim, modalities::DType image_dtype = modalities::DType::kBFloat16, - modalities::VisionStaging* staging = nullptr); + modalities::VisionStaging* staging = nullptr, + modalities::ActionStaging* action_staging = nullptr, + bool strict_rgb8 = true); modalities::Status prepare_vision( const std::vector& frames) const; @@ -42,6 +44,8 @@ class RuntimeIo { modalities::TensorView action_output_; void* stream_ = nullptr; modalities::VisionStaging* staging_ = nullptr; /* borrowed */ + modalities::ActionStaging* action_staging_ = nullptr; /* borrowed */ + bool strict_rgb8_ = true; modalities::VisionPreprocessSpec vision_spec_; modalities::ActionPostprocessSpec action_spec_; }; diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_embed.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_embed.h new file mode 100644 index 00000000..f5bd0d36 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_embed.h @@ -0,0 +1,53 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_PROMPT_EMBED_H +#define FLASHRT_CPP_MODELS_PI05_PROMPT_EMBED_H + +#include "flashrt/cpp/modalities/text.h" +#include "flashrt/cpp/modalities/tokenizer.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +struct PromptEmbeddingSpec { + std::uint64_t vocab_size = 0; + std::uint64_t hidden_dim = 0; + std::uint64_t max_tokens = 0; + float scale = 1.0f; + std::int32_t no_state_suffix_token_id = 108; + bool zero_pad_output = true; +}; + +modalities::Status embed_prompt( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len, + void* stream = nullptr, + modalities::TextEmbeddingStaging* staging = nullptr, + std::string* formatted_workspace = nullptr); + +modalities::Status embed_prompt_cpu( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_PROMPT_EMBED_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_format.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_format.h new file mode 100644 index 00000000..a5ecc9d7 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/prompt_format.h @@ -0,0 +1,32 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H +#define FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +static constexpr int kStatePromptMaxLen = 200; + +std::vector discretize_state_prompt_bins( + const float* state, std::uint64_t n); + +std::string clean_task_prompt(const std::string& prompt); + +std::string format_state_prompt(const std::string& prompt, + const float* state, + std::uint64_t n_state); + +void format_state_prompt_into(const std::string& prompt, + const float* state, + std::uint64_t n_state, + std::string* out); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_PROMPT_FORMAT_H diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/runtime.h similarity index 57% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/runtime.h index 1f817feb..b9f1da9f 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/runtime.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/runtime.h @@ -2,9 +2,11 @@ #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/model/io.h" +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" #include +#include namespace flashrt { namespace models { @@ -12,6 +14,7 @@ namespace pi05 { using ReplayFn = int (*)(frt_graph graph, frt_shape_key key, int stream_id, void* user); +using PromptLengthUpdateFn = int (*)(void* user, std::uint64_t prompt_len); struct RuntimeConfig { int num_views = 3; @@ -33,6 +36,7 @@ struct RuntimeConfig { * construction so prepare_vision never allocates on the hot path. */ int max_frame_width = 1280; int max_frame_height = 720; + bool strict_rgb8 = true; /* Optional host/device overrides. If left null, Runtime derives tensor * views from the export's named buffers. The current CPU reference @@ -41,8 +45,23 @@ struct RuntimeConfig { modalities::TensorView image_input_override; modalities::TensorView action_output_override; + /* Optional native prompt staging. When configured, set_prompt* writes + * token embeddings into prompt_embedding_output. The graph must consume + * this stable source buffer through its own captured copy/update path. */ + std::string prompt_tokenizer_model_path; + modalities::TensorView prompt_embedding_table; + modalities::TensorView prompt_embedding_output; + std::uint64_t prompt_vocab_size = 0; + std::uint64_t prompt_hidden_dim = 0; + std::uint64_t prompt_max_tokens = 0; + float prompt_embedding_scale = 0.0f; + std::vector state_q01; + std::vector state_q99; + ReplayFn replay_fn = nullptr; void* replay_user = nullptr; + PromptLengthUpdateFn prompt_length_update_fn = nullptr; + void* prompt_length_update_user = nullptr; }; class Runtime final : public families::vla::Runtime { @@ -64,6 +83,18 @@ class Runtime final : public families::vla::Runtime { } int set_prompt(const char* text) override; + int set_prompt_state(const char* text, const float* state, + std::uint64_t n_state); + const modalities::Status& prompt_status() const { + return prompt_status_; + } + bool prompt_staging_enabled() const { return prompt_staging_enabled_; } + bool state_normalization_enabled() const { + return !config_.state_q01.empty() && + config_.state_q01.size() == config_.state_q99.size(); + } + std::uint64_t current_prompt_len() const { return current_prompt_len_; } + modalities::Status prepare_vision( const std::vector& frames) override; int replay_tick() override; @@ -73,6 +104,7 @@ class Runtime final : public families::vla::Runtime { void retain_export(); void release_export(); modalities::Status bind(); + modalities::Status bind_prompt_staging(); static int default_replay(frt_graph graph, frt_shape_key key, int stream_id, void* user); @@ -82,7 +114,21 @@ class Runtime final : public families::vla::Runtime { families::vla::Manifest manifest_; modalities::Status status_; modalities::VisionStaging staging_; + modalities::ActionStaging action_staging_; RuntimeIo io_; + modalities::SentencePieceTokenizer prompt_tokenizer_; + modalities::TextEmbeddingStaging prompt_embedding_staging_; + PromptEmbeddingSpec prompt_spec_; + modalities::TensorView prompt_embedding_table_; + modalities::TensorView prompt_embedding_output_; + modalities::Status prompt_status_; + std::vector prompt_token_ids_; + std::vector normalized_state_; + std::string task_prompt_workspace_; + std::string formatted_prompt_workspace_; + std::size_t max_task_prompt_bytes_ = 0; + std::uint64_t current_prompt_len_ = 0; + bool prompt_staging_enabled_ = false; frt_graph graph_ = nullptr; frt_shape_key graph_key_ = 0; int stream_id_ = -1; diff --git a/cpp/models/pi05/include/flashrt/cpp/models/pi05/spec.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/spec.h similarity index 95% rename from cpp/models/pi05/include/flashrt/cpp/models/pi05/spec.h rename to cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/spec.h index 2866cadd..23913960 100644 --- a/cpp/models/pi05/include/flashrt/cpp/models/pi05/spec.h +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/model/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/internal/include/flashrt/cpp/models/pi05/support/native_calibration.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_calibration.h new file mode 100644 index 00000000..06352fcf --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_calibration.h @@ -0,0 +1,52 @@ +#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 activation_dtype = "float16"; + 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 vision_scales; + 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/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 new file mode 100644 index 00000000..e1671de4 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_device_weights.h @@ -0,0 +1,65 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_DEVICE_WEIGHTS_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_DEVICE_WEIGHTS_H + +#include "flashrt/cpp/models/pi05/support/native_weight_ops.h" +#include "flashrt/exec.h" + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeWeightDType { + kBf16, + kFloat16, + kFp8E4M3, + kInt8, + kFloat32, +}; + +struct NativeDeviceWeight { + frt_buffer buffer = nullptr; + std::vector shape; + NativeWeightDType dtype = NativeWeightDType::kBf16; +}; + +class NativeDeviceWeightStore { +public: + explicit NativeDeviceWeightStore(frt_ctx ctx) : ctx_(ctx) {} + + NativeDeviceWeightStore(const NativeDeviceWeightStore&) = delete; + NativeDeviceWeightStore& operator=(const NativeDeviceWeightStore&) = delete; + + modalities::Status upload(const std::string& name, + const NativeBf16Tensor& tensor); + modalities::Status upload(const std::string& name, + const NativeF16Tensor& tensor); + modalities::Status upload_bytes( + const std::string& name, + const std::vector& shape, + NativeWeightDType dtype, + const void* data, + std::size_t bytes); + modalities::Status allocate( + const std::string& name, + const std::vector& shape, + NativeWeightDType dtype); + 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_; + std::mutex upload_mutex_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_DEVICE_WEIGHTS_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_quantization.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_quantization.h new file mode 100644 index 00000000..67009f2f --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_quantization.h @@ -0,0 +1,43 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_QUANTIZATION_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_QUANTIZATION_H + +#include "flashrt/cpp/models/pi05/support/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_fp8_e4m3( + const NativeF16Tensor& fp16_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/internal/include/flashrt/cpp/models/pi05/support/native_rope.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_rope.h new file mode 100644 index 00000000..a7bd0a95 --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/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/internal/include/flashrt/cpp/models/pi05/support/native_weight_materializer.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_materializer.h new file mode 100644 index 00000000..856acd0d --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_materializer.h @@ -0,0 +1,61 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H + +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" + +namespace flashrt { +namespace models { +namespace pi05 { + +struct NativeMaterializationOptions { + int num_steps = 10; + bool merge_decoder_gate_up = true; + bool include_embedding = true; +}; + +class NativeWeightMaterializer { +public: + NativeWeightMaterializer(const loader::SafetensorsFile& source, + NativeDeviceWeightStore* destination) + : source_(source), destination_(destination) {} + + modalities::Status materialize_encoder_layer(int layer); + modalities::Status materialize_decoder_layer(int layer, + bool merge_gate_up); + modalities::Status materialize_vision_layer(int layer); + modalities::Status materialize_vision_globals(); + modalities::Status materialize_decoder_globals(int num_steps); + modalities::Status materialize_embedding(); + modalities::Status materialize_all( + const NativeMaterializationOptions& options); + +private: + modalities::Status load(const std::string& key, NativeFloatTensor* out); + modalities::Status upload(const std::string& name, + const NativeFloatTensor& tensor); + modalities::Status upload_rounded_transpose( + const std::string& source_key, + const std::string& destination_name); + modalities::Status upload_rounded_copy( + const std::string& source_key, + const std::string& destination_name); + modalities::Status upload_folded_transpose( + const std::string& source_key, + const NativeFloatTensor& norm, + const std::string& destination_name); + modalities::Status upload_rounded_scaled( + const std::string& source_key, + const std::string& destination_name, + float scale, + bool transpose); + + const loader::SafetensorsFile& source_; + NativeDeviceWeightStore* destination_ = nullptr; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_MATERIALIZER_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_ops.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_ops.h new file mode 100644 index 00000000..128138bb --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weight_ops.h @@ -0,0 +1,178 @@ +#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; +}; + +struct NativeF16Tensor { + std::vector shape; + std::vector values; +}; + +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, + NativeSourceTensorView* out); + +modalities::Status native_source_to_bf16( + const NativeSourceTensorView& input, + 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, + NativeBf16Tensor* out); + +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_to_f16(const NativeFloatTensor& input, + NativeF16Tensor* out); + +modalities::Status native_round_to_bf16_float( + const NativeFloatTensor& input, + NativeFloatTensor* out); + +modalities::Status native_transpose_2d(const NativeFloatTensor& input, + NativeFloatTensor* out); + +modalities::Status native_patch_oihw_to_hwio( + const NativeFloatTensor& input, + NativeFloatTensor* out); + +modalities::Status native_interleave_qk_rows( + const NativeFloatTensor& input, + std::uint64_t num_heads, + NativeFloatTensor* out); + +modalities::Status native_fold_rms_columns( + const NativeFloatTensor& weight, + const NativeFloatTensor& norm, + NativeFloatTensor* out); + +modalities::Status native_concat_rows_transpose( + const std::vector& inputs, + NativeFloatTensor* out); + +modalities::Status native_concat_columns( + const NativeFloatTensor& left, + const NativeFloatTensor& right, + NativeFloatTensor* out); + +modalities::Status native_concat_vectors( + const std::vector& inputs, + NativeFloatTensor* out); + +modalities::Status native_scale(const NativeFloatTensor& input, + float scale, + NativeFloatTensor* out); + +modalities::Status native_pi05_time_embeddings( + int num_steps, + std::uint64_t embedding_dim, + NativeFloatTensor* out); + +modalities::Status native_pi05_time_embeddings_f16( + int num_steps, + std::uint64_t embedding_dim, + NativeF16Tensor* out); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WEIGHT_OPS_H diff --git a/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weights.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_weights.h new file mode 100644 index 00000000..1c93a8ee --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/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/internal/include/flashrt/cpp/models/pi05/support/native_workspace.h b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_workspace.h new file mode 100644 index 00000000..efc58ddf --- /dev/null +++ b/cpp/models/pi05/internal/include/flashrt/cpp/models/pi05/support/native_workspace.h @@ -0,0 +1,99 @@ +#ifndef FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H +#define FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H + +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" +#include "flashrt/exec.h" + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { + +enum class NativeWorkspaceFlavor { + kBf16, + kRtxFp8, + 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 { + frt_buffer buffer = nullptr; + std::vector shape; + modalities::DType dtype = modalities::DType::kBFloat16; + bool alias = false; +}; + +class NativeWorkspace { +public: + explicit NativeWorkspace(frt_ctx ctx) : ctx_(ctx) {} + + NativeWorkspace(const NativeWorkspace&) = delete; + NativeWorkspace& operator=(const NativeWorkspace&) = delete; + + modalities::Status allocate(const NativeWorkspaceConfig& config); + modalities::Status update_decoder_rope(int prompt_tokens); + modalities::Status set_fixed_prompt_length(int prompt_tokens); + modalities::Status expand_vision_position_embedding( + const NativeDeviceWeightStore& weights); + const NativeWorkspaceBuffer* find(const std::string& name) const; + + std::size_t logical_size() const { return buffers_.size(); } + std::size_t allocation_count() const { return allocation_count_; } + std::size_t allocated_bytes() const { return allocated_bytes_; } + int vision_sequence() const { return vision_sequence_; } + int encoder_vision_sequence() const { return encoder_vision_sequence_; } + int encoder_sequence() const { return encoder_sequence_; } + 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, + std::initializer_list shape, + modalities::DType dtype); + modalities::Status add_alias(const std::string& name, + const std::string& source_name, + std::initializer_list shape); + modalities::Status initialize_rms_ones(); + modalities::Status initialize_rope(); + + frt_ctx ctx_ = nullptr; + std::map buffers_; + std::size_t allocation_count_ = 0; + std::size_t allocated_bytes_ = 0; + int vision_sequence_ = 0; + int encoder_vision_sequence_ = 0; + int encoder_sequence_ = 0; + int num_views_ = 0; + int max_prompt_tokens_ = 0; + int chunk_size_ = 0; + 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_; +}; + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_WORKSPACE_H diff --git a/cpp/models/pi05/src/backend/native_calibration_session.cpp b/cpp/models/pi05/src/backend/native_calibration_session.cpp new file mode 100644 index 00000000..e3a6b36f --- /dev/null +++ b/cpp/models/pi05/src/backend/native_calibration_session.cpp @@ -0,0 +1,133 @@ +#include "flashrt/cpp/models/pi05/backend/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/backend/session.cpp b/cpp/models/pi05/src/backend/session.cpp new file mode 100644 index 00000000..b479ef02 --- /dev/null +++ b/cpp/models/pi05/src/backend/session.cpp @@ -0,0 +1,116 @@ +#include "flashrt/cpp/models/pi05/backend/session.h" + +#include "flashrt/cpp/models/pi05/model/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); +} + +} // namespace + +const char* backend_graph_name(GraphKind kind) { + switch (kind) { + 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_backend_graph( + native::CudaGraphSet* graphs, + GraphKind kind, + const NativeWorkspace& workspace, + std::initializer_list bindings, + native::CudaGraphSet::RecordFn record, + void* owner) { + if (!graphs || !backend_graph_name(kind)) { + return invalid("native graph capture request is invalid"); + } + std::vector resolved; + resolved.reserve(bindings.size()); + for (const char* binding : bindings) { + const NativeWorkspaceBuffer* buffer = workspace.find(binding); + if (!buffer) { + return backend("native graph binding failed"); + } + resolved.push_back({binding, buffer->buffer}); + } + return graphs->capture(static_cast(kind), + backend_graph_name(kind), resolved, record, owner); +} + +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"); +} + +modalities::Status resolve_backend_artifacts( + const NativeWorkspace& workspace, + const NativeDeviceWeightStore& weights, + NativeWeightDType embedding_dtype, + BackendArtifacts* artifacts) { + if (!artifacts) { + return invalid("native runtime artifacts destination is null"); + } + BackendArtifacts 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/backends/sm110/native_thor_calibration_session.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_calibration_session.cpp new file mode 100644 index 00000000..730679ad --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/native_thor_calibration_session.cpp @@ -0,0 +1,420 @@ +#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/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 + +#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; +}; + +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) { + 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; + 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)); + + 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, + 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_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; + } + 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/backends/sm110/native_thor_fp8_forward.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_fp8_forward.cpp new file mode 100644 index 00000000..92be7fa8 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/native_thor_fp8_forward.cpp @@ -0,0 +1,1297 @@ +#include "flashrt/cpp/models/pi05/backends/sm110/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/backends/sm110/native_thor_kernel_driver.cu b/cpp/models/pi05/src/backends/sm110/native_thor_kernel_driver.cu new file mode 100644 index 00000000..1a9067c1 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/native_thor_kernel_driver.cu @@ -0,0 +1,622 @@ +#include "flashrt/cpp/models/pi05/backends/sm110/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/backends/sm110/native_thor_setup_ops.cu b/cpp/models/pi05/src/backends/sm110/native_thor_setup_ops.cu new file mode 100644 index 00000000..e1c6452a --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/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/backends/sm110/native_thor_style_precompute.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_style_precompute.cpp new file mode 100644 index 00000000..80601069 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/native_thor_style_precompute.cpp @@ -0,0 +1,210 @@ +#include "flashrt/cpp/models/pi05/backends/sm110/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/backends/sm110/native_thor_weight_materializer.cpp b/cpp/models/pi05/src/backends/sm110/native_thor_weight_materializer.cpp new file mode 100644 index 00000000..f0028027 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/native_thor_weight_materializer.cpp @@ -0,0 +1,545 @@ +#include "flashrt/cpp/models/pi05/backends/sm110/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/backends/sm110/session.cpp b/cpp/models/pi05/src/backends/sm110/session.cpp new file mode 100644 index 00000000..8c7eafad --- /dev/null +++ b/cpp/models/pi05/src/backends/sm110/session.cpp @@ -0,0 +1,295 @@ +#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" + +#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 BackendConfig& 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 + +Sm110BackendSession::Sm110BackendSession( + frt_ctx ctx, + const BackendConfig& config) + : graphs_(ctx, static_cast(GraphKind::kCount)), + config_(config), + weights_(ctx), + workspace_(ctx), + forward_(&driver_) {} + +Sm110BackendSession::~Sm110BackendSession() = default; + +std::unique_ptr Sm110BackendSession::create( + const std::string& checkpoint_path, + const BackendConfig& 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 session( + new (std::nothrow) Sm110BackendSession(ctx, config)); + if (!session) { + frt_ctx_destroy(ctx); + if (status) *status = backend("SM110 backend session allocation failed"); + return nullptr; + } + st = session->initialize(checkpoint_path, calibration); + if (!st.ok_status()) { + if (status) *status = st; + return nullptr; + } + if (status) *status = modalities::Status::ok(); + return session; +} + +modalities::Status Sm110BackendSession::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"); + + st = resolve_backend_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); + 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(GraphKind::kInfer, 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"); + + st = capture_backend_graph( + &graphs_, + 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_backend_graph( + &graphs_, + 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_backend_graph( + &graphs_, + GraphKind::kContext, workspace_, + {"observation_images_normalized", "prompt_embedding", "encoder_x"}, + record_graph, this); + if (!st.ok_status()) return st; + report("capture"); + + st = graphs_.create_replay_stream(); + if (!st.ok_status()) return st; + 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 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); + st = forward_.vision(weights_, &workspace_, weight_scales_, stream_id); + if (!st.ok_status()) return st; + return forward_.encoder( + weights_, &workspace_, encoder_alphas_, stream_id); +} + +modalities::Status Sm110BackendSession::record_action(void* stream) { + return forward_.diffusion( + weights_, &workspace_, reinterpret_cast(stream)); +} + +modalities::Status Sm110BackendSession::record(GraphKind kind, + void* stream) { + 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 Sm110BackendSession::record_graph( + void* user, std::size_t slot, void* stream) { + auto* session = static_cast(user); + return session->record(static_cast(slot), stream); +} + +modalities::Status Sm110BackendSession::set_prompt_length(int prompt_tokens) { + return workspace_.set_fixed_prompt_length(prompt_tokens); +} + +int Sm110BackendSession::replay(GraphKind kind) const { + return graphs_.replay(static_cast(kind)); +} + +modalities::Status Sm110BackendSession::synchronize() const { + return graphs_.synchronize(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/backends/sm120/native_bf16_forward.cpp b/cpp/models/pi05/src/backends/sm120/native_bf16_forward.cpp new file mode 100644 index 00000000..5c5f6094 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_bf16_forward.cpp @@ -0,0 +1,853 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/native_bf16_forward.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +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 && + buffer->shape == std::vector(shape); +} + +bool shape_is(const NativeAttentionBuffer* buffer, + std::initializer_list shape) { + return buffer && buffer->dtype == NativeAttentionDType::kBf16 && + buffer->shape == std::vector(shape); +} + +#ifdef FLASHRT_CPP_WITH_FA2 +bool shape_is(const NativeDeviceWeight* weight, + std::initializer_list shape) { + return weight && weight->dtype == NativeWeightDType::kBf16 && + weight->shape == std::vector(shape); +} +#endif + +} // namespace + +modalities::Status NativeBf16Forward::encoder_qkv( + int layer, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || layer < 0 || layer >= 18) { + return invalid("native encoder QKV owner is invalid"); + } + const int sequence = workspace->encoder_sequence(); + if (sequence <= 0) { + return invalid("native encoder sequence is invalid"); + } + const NativeWorkspaceBuffer* x = workspace->find("encoder_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("encoder_x_norm"); + const NativeWorkspaceBuffer* qkv = workspace->find("encoder_QKV"); + const NativeWorkspaceBuffer* rms = workspace->find("encoder_rms_ones"); + const NativeWorkspaceBuffer* rope = + workspace->find("encoder_rope_weights"); + const NativeAttentionBuffer* query = attention->find("attn_enc_Q"); + const NativeAttentionBuffer* cache = attention->find("attn_enc_K"); + const NativeAttentionBuffer* value_cache = attention->find("attn_enc_V"); + const 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}) || + !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 || + !value_cache || value_cache->dtype != NativeAttentionDType::kBf16 || + value_cache->shape != cache->shape || + !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); + void* value = attention->encoder_v_layer_dptr(layer); + if (!key || !value) { + return invalid("native encoder QKV cache layer is invalid"); + } + 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), + frt_buffer_dptr(query->buffer), key, value, sequence, 2048, 256, 256, + 256, stream); +} + +#ifdef FLASHRT_CPP_WITH_FA2 +modalities::Status NativeBf16Forward::vision_layer( + int layer, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver || + !attention_driver->status().ok_status() || layer < 0 || layer >= 27) { + return invalid("native vision layer owner is invalid"); + } + const int sequence = workspace->vision_sequence(); + const int num_views = sequence / 256; + const NativeWorkspaceBuffer* x = workspace->find("vision_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("vision_x_norm"); + const NativeWorkspaceBuffer* qkv = workspace->find("vision_QKV"); + const NativeWorkspaceBuffer* hidden = workspace->find("vision_hidden"); + const NativeAttentionBuffer* query = attention->find("attn_vis_Q"); + const NativeAttentionBuffer* key = attention->find("attn_vis_K"); + const NativeAttentionBuffer* value = attention->find("attn_vis_V"); + const std::string suffix = std::to_string(layer); + const 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_bias = + weights.find("vision_attn_o_b_" + suffix); + const NativeDeviceWeight* up_bias = + weights.find("vision_ffn_up_b_" + 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}) || + !linear_->weight_shape_is(weights, qkv_name, {1152, 3456}) || + !shape_is(qkv_bias, {3456}) || + !linear_->weight_shape_is(weights, output_name, {1152, 1152}) || + !shape_is(output_bias, {1152}) || + !linear_->weight_shape_is(weights, up_name, {1152, 4304}) || + !shape_is(up_bias, {4304}) || + !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 = 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), + 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 = 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( + 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 = 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), + 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 = 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), + 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 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 || + 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}) || + !linear_->weight_shape_is(weights, projector_name, {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 = 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), + frt_buffer_dptr(projector_bias->buffer), encoder_sequence, 2048, + stream); +} + +modalities::Status NativeBf16Forward::encoder_layer( + int layer, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + modalities::Status st = + encoder_qkv(layer, weights, workspace, attention, stream); + if (!st.ok_status() || layer == 17) return st; + if (!attention_driver || !attention_driver->status().ok_status()) { + return invalid("native encoder attention driver is invalid"); + } + const int sequence = workspace->encoder_sequence(); + const NativeWorkspaceBuffer* x = workspace->find("encoder_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("encoder_x_norm"); + const NativeWorkspaceBuffer* gate = + workspace->find("encoder_gate_merged"); + const NativeWorkspaceBuffer* hidden = workspace->find("encoder_hidden"); + const NativeWorkspaceBuffer* rms = workspace->find("encoder_rms_ones"); + const 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}) || !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 = 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); + 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; + 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 = 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), + static_cast(sequence) * 2048, stream); +} + +modalities::Status NativeBf16Forward::encoder( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + for (int layer = 0; layer < 18; ++layer) { + modalities::Status st = encoder_layer( + layer, weights, workspace, attention, attention_driver, stream); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} + +modalities::Status NativeBf16Forward::decoder_layer( + int layer, + int step, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver || + !attention_driver->status().ok_status() || layer < 0 || layer >= 18) { + return invalid("native decoder layer owner is invalid"); + } + const NativeWorkspaceBuffer* x = workspace->find("decoder_x"); + const NativeWorkspaceBuffer* x_norm = workspace->find("x_normed_buf"); + const NativeWorkspaceBuffer* gate = workspace->find("gate_buf"); + const NativeWorkspaceBuffer* qkv = workspace->find("decoder_QKV"); + const NativeWorkspaceBuffer* hidden = workspace->find("decoder_hidden"); + const NativeWorkspaceBuffer* gate_projection = + workspace->find("decoder_gate_merged"); + const NativeWorkspaceBuffer* rms = workspace->find("decoder_rms_ones"); + const NativeWorkspaceBuffer* rope = + workspace->find("decoder_rope_weights"); + const NativeWorkspaceBuffer* style_attn = + workspace->find("decoder_style_attn"); + const NativeWorkspaceBuffer* style_ffn = + workspace->find("decoder_style_ffn"); + if (!x || x->shape.size() != 2) { + return invalid("native decoder workspace is invalid"); + } + const int sequence = static_cast(x->shape[0]); + const NativeAttentionBuffer* query = attention->find("attn_dec_Q"); + const NativeAttentionBuffer* devpos = attention->find("attn_dec_devpos"); + const std::string suffix = std::to_string(layer); + const 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}) || + !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}) || + !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 = + (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; + 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), + 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 = 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), + 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; + 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 = 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), + frt_buffer_dptr(gate->buffer), + static_cast(sequence) * 1024, stream); +} + +modalities::Status NativeBf16Forward::diffusion_step( + int step, + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!driver_ || !workspace || !attention || !attention_driver || step < 0) { + return invalid("native diffusion step owner is invalid"); + } + const NativeWorkspaceBuffer* noise = workspace->find("diffusion_noise"); + const NativeWorkspaceBuffer* x = workspace->find("decoder_x"); + const NativeWorkspaceBuffer* action = + workspace->find("decoder_action_buf"); + const NativeWorkspaceBuffer* x_norm = workspace->find("x_normed_buf"); + const NativeWorkspaceBuffer* gate = workspace->find("gate_buf"); + const NativeWorkspaceBuffer* rms = workspace->find("decoder_rms_ones"); + const NativeWorkspaceBuffer* style = + workspace->find("decoder_style_final"); + if (!noise || noise->shape.size() != 2) { + return invalid("native diffusion workspace is invalid"); + } + const int sequence = static_cast(noise->shape[0]); + const NativeDeviceWeight* input_weight = + weights.find("decoder_action_in_proj_w"); + const NativeDeviceWeight* input_bias = + weights.find("decoder_action_in_proj_b"); + const NativeDeviceWeight* output_weight = + weights.find("decoder_action_out_proj_w"); + const NativeDeviceWeight* output_bias = + weights.find("decoder_action_out_proj_b"); + if (sequence <= 0 || + !shape_is(noise, {static_cast(sequence), 32}) || + !shape_is(x, {static_cast(sequence), 1024}) || + !shape_is(action, {static_cast(sequence), 32}) || + !shape_is(x_norm, {static_cast(sequence), 1024}) || + !shape_is(gate, {static_cast(sequence), 1024}) || + !shape_is(rms, {1024}) || !style || + style->dtype != modalities::DType::kBFloat16 || + style->shape.size() != 3 || + style->shape[0] <= static_cast(step) || + style->shape[1] != static_cast(sequence) || + style->shape[2] != 3072 || + !shape_is(input_weight, {32, 1024}) || + !shape_is(input_bias, {1024}) || + !shape_is(output_weight, {1024, 32}) || + !shape_is(output_bias, {32})) { + return invalid("native diffusion buffers or weights are invalid"); + } + modalities::Status st = driver_->bf16_nn( + frt_buffer_dptr(noise->buffer), frt_buffer_dptr(input_weight->buffer), + frt_buffer_dptr(x->buffer), sequence, 1024, 32, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(input_bias->buffer), + sequence, 1024, stream); + if (!st.ok_status()) return st; + for (int layer = 0; layer < 18; ++layer) { + st = decoder_layer(layer, step, weights, workspace, attention, + attention_driver, stream); + if (!st.ok_status()) return st; + } + const std::size_t style_offset = + static_cast(step) * sequence * 3072 * + sizeof(std::uint16_t); + const auto* final_style = + static_cast(frt_buffer_dptr(style->buffer)) + + style_offset; + st = driver_->ada_rms_norm_style_bf16( + frt_buffer_dptr(x->buffer), frt_buffer_dptr(rms->buffer), final_style, + frt_buffer_dptr(x_norm->buffer), frt_buffer_dptr(gate->buffer), + sequence, 1024, 1e-6f, stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(x_norm->buffer), + frt_buffer_dptr(output_weight->buffer), frt_buffer_dptr(action->buffer), + sequence, 32, 1024, stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(action->buffer), frt_buffer_dptr(output_bias->buffer), + sequence, 32, stream); + if (!st.ok_status()) return st; + return driver_->residual_add_bf16( + frt_buffer_dptr(noise->buffer), frt_buffer_dptr(action->buffer), + static_cast(sequence) * 32, stream); +} + +modalities::Status NativeBf16Forward::diffusion( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + NativeRtxAttentionWorkspace* attention, + const NativeRtxAttentionDriver* attention_driver, + std::uintptr_t stream) const { + if (!workspace) return invalid("native diffusion workspace is invalid"); + const NativeWorkspaceBuffer* style = + workspace->find("decoder_style_final"); + if (!style || style->shape.size() != 3 || !style->shape[0] || + style->shape[0] > static_cast(INT_MAX)) { + return invalid("native diffusion step count is invalid"); + } + const int steps = static_cast(style->shape[0]); + for (int step = 0; step < steps; ++step) { + modalities::Status st = diffusion_step( + step, weights, workspace, attention, attention_driver, stream); + if (!st.ok_status()) return st; + } + return modalities::Status::ok(); +} +#endif + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/backends/sm120/native_kernel_driver.cu b/cpp/models/pi05/src/backends/sm120/native_kernel_driver.cu new file mode 100644 index 00000000..c02c38f8 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_kernel_driver.cu @@ -0,0 +1,589 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/native_kernel_driver.h" + +#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 +#include + +#include +#include + +void add_bias_bf16(__nv_bfloat16* x, const __nv_bfloat16* b, + int rows, int columns, cudaStream_t stream); + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +__global__ void native_silu_bf16_kernel(__nv_bfloat16* values, + std::size_t elements) { + const std::size_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < elements) { + const float value = __bfloat162float(values[index]); + values[index] = + __float2bfloat16(value / (1.0f + expf(-value))); + } +} + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const std::string& message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +modalities::Status launch_status() { + const cudaError_t rc = cudaGetLastError(); + return rc == cudaSuccess ? modalities::Status::ok() + : backend(cudaGetErrorString(rc)); +} + +} // namespace + +struct NativeKernelDriver::Impl { + GemmRunner gemm; +}; + +NativeKernelDriver::NativeKernelDriver() noexcept { + try { + impl_.reset(new Impl()); + } catch (const std::exception& e) { + error_ = e.what(); + } catch (...) { + error_ = "native kernel driver initialization failed"; + } +} + +NativeKernelDriver::~NativeKernelDriver() = default; + +modalities::Status NativeKernelDriver::status() const { + return impl_ ? modalities::Status::ok() : backend(error_); +} + +modalities::Status NativeKernelDriver::bf16_nn( + void* a, + void* b, + void* output, + int m, + int n, + int k, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!a || !b || !output || m <= 0 || n <= 0 || k <= 0) { + return invalid("native BF16 GEMM arguments are invalid"); + } + try { + impl_->gemm.bf16_nn(a, b, output, m, n, k, + reinterpret_cast(stream)); + return modalities::Status::ok(); + } catch (const std::exception& e) { + return backend(e.what()); + } catch (...) { + return backend("native BF16 GEMM launch failed"); + } +} + +modalities::Status NativeKernelDriver::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, + int rows, + int columns, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !bias || rows <= 0 || columns <= 0) { + return invalid("native BF16 bias arguments are invalid"); + } + ::add_bias_bf16(static_cast<__nv_bfloat16*>(values), + static_cast(bias), rows, columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::silu_bf16( + void* values, + std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !elements) { + return invalid("native BF16 SiLU arguments are invalid"); + } + native_silu_bf16_kernel<<<(elements + 255) / 256, 256, 0, + reinterpret_cast(stream)>>>( + static_cast<__nv_bfloat16*>(values), elements); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gelu_bf16( + void* values, std::size_t elements, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !elements || (elements & 1) || + elements > static_cast(INT_MAX)) { + return invalid("native BF16 GELU arguments are invalid"); + } + ::gelu_inplace(static_cast<__nv_bfloat16*>(values), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::gate_gelu_bf16( + const void* gate, const void* up, void* output, std::size_t elements, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!gate || !up || !output || !elements || + elements > static_cast(INT_MAX)) { + return invalid("native BF16 gated GELU arguments are invalid"); + } + ::gate_silu_mul(static_cast(gate), + static_cast(up), + static_cast<__nv_bfloat16*>(output), + static_cast(elements), + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::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 { + 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::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 { + 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::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, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!qkv || !query || !key || !value || rows <= 0 || query_columns <= 0 || + key_columns <= 0 || value_columns <= 0) { + return invalid("native BF16 QKV split arguments are invalid"); + } + ::qkv_split(static_cast(qkv), + static_cast<__nv_bfloat16*>(query), + static_cast<__nv_bfloat16*>(key), + static_cast<__nv_bfloat16*>(value), rows, query_columns, + key_columns, value_columns, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::qkv_split_rope_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + int rows, int query_columns, int key_columns, int value_columns, + int head_dimension, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!qkv || !rope || !query || !key || !value || rows <= 0 || + query_columns <= 0 || key_columns <= 0 || value_columns <= 0 || + head_dimension <= 0 || (head_dimension & 1) || + query_columns % head_dimension || key_columns % head_dimension) { + return invalid("native BF16 QKV RoPE arguments are invalid"); + } + ::qkv_split_rope( + static_cast(qkv), + static_cast(rope), + static_cast<__nv_bfloat16*>(query), + static_cast<__nv_bfloat16*>(key), + static_cast<__nv_bfloat16*>(value), rows, query_columns, key_columns, + value_columns, head_dimension, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::qkv_split_rope_devpos_bf16( + const void* qkv, const void* rope, void* query, void* key, void* value, + const void* device_position, int rows, int query_columns, int key_columns, + int value_columns, int head_dimension, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!qkv || !rope || !query || !key || !value || !device_position || + rows <= 0 || query_columns <= 0 || key_columns <= 0 || + value_columns <= 0 || head_dimension <= 0 || (head_dimension & 1) || + query_columns % head_dimension || key_columns % head_dimension) { + return invalid("native BF16 QKV devpos arguments are invalid"); + } + ::qkv_split_rope_devpos( + static_cast(qkv), + static_cast(rope), + static_cast<__nv_bfloat16*>(query), + static_cast<__nv_bfloat16*>(key), + static_cast<__nv_bfloat16*>(value), + static_cast(device_position), rows, query_columns, + key_columns, value_columns, head_dimension, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::patch_im2col_16bit( + const void* images, void* patches, int num_views, + std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!images || !patches || num_views <= 0) { + return invalid("native patch im2col arguments are invalid"); + } + ::patch_im2col(static_cast(images), + static_cast<__half*>(patches), num_views, + reinterpret_cast(stream)); + return launch_status(); +} + +modalities::Status NativeKernelDriver::avg_pool_vision_bf16( + const void* values, void* output, int num_views, int height, int width, + int columns, int pool_factor, std::uintptr_t stream) const { + if (!impl_) return backend(error_); + if (!values || !output || num_views <= 0 || height <= 0 || width <= 0 || + columns <= 0 || pool_factor <= 0 || height % pool_factor || + width % pool_factor) { + return invalid("native vision pooling arguments are invalid"); + } + ::avg_pool_vision_tokens( + static_cast(values), + static_cast<__nv_bfloat16*>(output), num_views, height, width, columns, + pool_factor, reinterpret_cast(stream)); + return launch_status(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/backends/sm120/native_rtx_attention.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_attention.cpp new file mode 100644 index 00000000..40e2d812 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_attention.cpp @@ -0,0 +1,205 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_attention.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +std::size_t dtype_size(NativeAttentionDType dtype) { + switch (dtype) { + case NativeAttentionDType::kBf16: return sizeof(std::uint16_t); + case NativeAttentionDType::kFloat32: return sizeof(float); + case NativeAttentionDType::kInt32: return sizeof(std::int32_t); + } + return 0; +} + +bool element_count(std::initializer_list shape, + std::size_t* out) { + std::size_t count = 1; + for (std::uint64_t dim : shape) { + if (!dim || dim > std::numeric_limits::max() || + count > std::numeric_limits::max() / + static_cast(dim)) { + return false; + } + count *= static_cast(dim); + } + if (out) *out = count; + return true; +} + +std::uint64_t round_up_128(std::uint64_t value) { + return ((value + 127) / 128) * 128; +} + +} // namespace + +modalities::Status NativeRtxAttentionWorkspace::add( + const std::string& name, + std::initializer_list shape, + NativeAttentionDType dtype) { + if (!ctx_ || name.empty() || buffers_.find(name) != buffers_.end()) { + return invalid("native attention buffer definition is invalid"); + } + std::size_t elements = 0; + const std::size_t width = dtype_size(dtype); + if (!width || !element_count(shape, &elements) || + elements > std::numeric_limits::max() / width) { + return invalid("native attention buffer shape is invalid"); + } + const std::size_t bytes = elements * width; + frt_buffer buffer = frt_buffer_alloc(ctx_, name.c_str(), bytes); + if (!buffer) return backend("native attention allocation failed"); + buffers_.emplace(name, NativeAttentionBuffer{ + buffer, std::vector(shape), + dtype}); + allocated_bytes_ += bytes; + return modalities::Status::ok(); +} + +modalities::Status NativeRtxAttentionWorkspace::allocate( + const NativeRtxAttentionConfig& config) { + if (!ctx_ || !buffers_.empty() || config.num_views < 1 || + config.num_views > 3 || config.encoder_sequence <= 0 || + config.encoder_vision_sequence <= 0 || + config.encoder_vision_sequence > config.encoder_sequence || + config.chunk_size <= 0 || config.encoder_layers != 18) { + return invalid("Pi0.5 RTX attention configuration is invalid"); + } + num_views_ = config.num_views; + encoder_sequence_ = config.encoder_sequence; + encoder_vision_sequence_ = config.encoder_vision_sequence; + chunk_size_ = config.chunk_size; + encoder_layers_ = config.encoder_layers; + const std::uint64_t nv = static_cast(num_views_); + const std::uint64_t es = static_cast(encoder_sequence_); + const std::uint64_t ds = static_cast(chunk_size_); + const std::uint64_t layers = static_cast(encoder_layers_); + const std::uint64_t total_kv = es + ds; + encoder_splits_ = std::min(128, (encoder_sequence_ + 63) / 64); + decoder_splits_ = + std::min(128, (encoder_sequence_ + chunk_size_ + 63) / 64); + kv_layer_stride_bytes_ = + static_cast(total_kv) * 256 * sizeof(std::uint16_t); + modalities::Status st; +#define FRT_ADD(...) \ + do { \ + st = add(__VA_ARGS__); \ + if (!st.ok_status()) return st; \ + } while (false) + FRT_ADD("attn_vis_Q", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_vis_K", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_vis_V", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_Q", {es, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_K", {layers, total_kv, 1, 256}, + NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_V", {layers, total_kv, 1, 256}, + NativeAttentionDType::kBf16); + FRT_ADD("attn_dec_Q", {ds, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_seqused", {1}, NativeAttentionDType::kInt32); + FRT_ADD("attn_dec_seqused", {1}, NativeAttentionDType::kInt32); + FRT_ADD("attn_dec_devpos", {1}, NativeAttentionDType::kInt32); + + FRT_ADD("attn_vis_O", {nv, 256, 16, 72}, NativeAttentionDType::kBf16); + FRT_ADD("attn_vis_lse", {nv, 16, 256}, NativeAttentionDType::kFloat32); + FRT_ADD("attn_vis_lse_accum", {2, nv, 16, 256}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_vis_o_accum", {2, nv, 16, 256, 96}, + NativeAttentionDType::kFloat32); + + FRT_ADD("attn_enc_O", {1, es, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_enc_lse", {1, 8, round_up_128(es)}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_enc_lse_accum", + {static_cast(encoder_splits_), 1, 8, es}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_enc_o_accum", + {static_cast(encoder_splits_), 1, 8, es, 256}, + NativeAttentionDType::kFloat32); + + FRT_ADD("attn_dec_O", {1, ds, 8, 256}, NativeAttentionDType::kBf16); + FRT_ADD("attn_dec_lse", {1, 8, round_up_128(ds)}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_dec_lse_accum", + {static_cast(decoder_splits_), 1, 8, ds}, + NativeAttentionDType::kFloat32); + FRT_ADD("attn_dec_o_accum", + {static_cast(decoder_splits_), 1, 8, ds, 256}, + NativeAttentionDType::kFloat32); +#undef FRT_ADD + const char* prompt_length_names[] = { + "attn_enc_seqused", "attn_dec_seqused", "attn_dec_devpos"}; + for (int i = 0; i < 3; ++i) { + const NativeAttentionBuffer* target = find(prompt_length_names[i]); + if (!target) return invalid("prompt length buffer was not allocated"); + prompt_length_buffers_[i] = target->buffer; + } + return set_fixed_prompt_length(0); +} + +modalities::Status NativeRtxAttentionWorkspace::set_fixed_prompt_length( + int prompt_tokens) { + const int max_prompt = encoder_sequence_ - encoder_vision_sequence_; + if (prompt_tokens < 0 || prompt_tokens > max_prompt || buffers_.empty()) { + return invalid("Pi0.5 fixed attention prompt length is invalid"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "fixed attention update requires the CUDA build"); +#else + const std::int32_t valid = encoder_vision_sequence_ + prompt_tokens; + const std::int32_t values[] = {valid, valid + chunk_size_, valid}; + for (int i = 0; i < 3; ++i) { + if (!prompt_length_buffers_[i] || + cudaMemcpy(frt_buffer_dptr(prompt_length_buffers_[i]), &values[i], + sizeof(values[i]), cudaMemcpyHostToDevice) != + cudaSuccess) { + return backend("fixed attention length upload failed"); + } + } + return modalities::Status::ok(); +#endif +} + +const NativeAttentionBuffer* NativeRtxAttentionWorkspace::find( + const std::string& name) const { + const auto it = buffers_.find(name); + return it == buffers_.end() ? nullptr : &it->second; +} + +void* NativeRtxAttentionWorkspace::encoder_k_layer_dptr(int layer) const { + const NativeAttentionBuffer* cache = find("attn_enc_K"); + if (!cache || layer < 0 || layer >= encoder_layers_) return nullptr; + return static_cast(frt_buffer_dptr(cache->buffer)) + + static_cast(layer) * kv_layer_stride_bytes_; +} + +void* NativeRtxAttentionWorkspace::encoder_v_layer_dptr(int layer) const { + const NativeAttentionBuffer* cache = find("attn_enc_V"); + if (!cache || layer < 0 || layer >= encoder_layers_) return nullptr; + return static_cast(frt_buffer_dptr(cache->buffer)) + + static_cast(layer) * kv_layer_stride_bytes_; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/backends/sm120/native_rtx_attention_driver.cu b/cpp/models/pi05/src/backends/sm120/native_rtx_attention_driver.cu new file mode 100644 index 00000000..7aaca702 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_attention_driver.cu @@ -0,0 +1,213 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/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); +} + +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( + 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, + inverse_sqrt(72), 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, + inverse_sqrt(256), 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, inverse_sqrt(256), + num_sms_, reinterpret_cast(stream)); + return launch_status(); +} + +void* NativeRtxAttentionDriver::vision_output() const { + const NativeAttentionBuffer* output = find("attn_vis_O"); + return output ? frt_buffer_dptr(output->buffer) : nullptr; +} + +void* NativeRtxAttentionDriver::encoder_output() const { + const NativeAttentionBuffer* output = find("attn_enc_O"); + return output ? frt_buffer_dptr(output->buffer) : nullptr; +} + +void* NativeRtxAttentionDriver::decoder_output() const { + const NativeAttentionBuffer* output = find("attn_dec_O"); + return output ? frt_buffer_dptr(output->buffer) : nullptr; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/backends/sm120/native_rtx_autotune.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_autotune.cpp new file mode 100644 index 00000000..eaf19fc2 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_autotune.cpp @@ -0,0 +1,80 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/native_rtx_calibration_session.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp new file mode 100644 index 00000000..a41e493d --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_calibration_session.cpp @@ -0,0 +1,449 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_calibration_session.h" + +#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/session.h" +#include "flashrt/cpp/models/pi05/model/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; + }); + + 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 = BackendPrecision::kFp8E4M3; + modalities::Status st; + graph = Sm120BackendSession::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/backends/sm120/native_rtx_linear.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_linear.cpp new file mode 100644 index 00000000..c09ebd27 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_linear.cpp @@ -0,0 +1,212 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/native_rtx_weight_packer.cpp b/cpp/models/pi05/src/backends/sm120/native_rtx_weight_packer.cpp new file mode 100644 index 00000000..a08e8c48 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_rtx_weight_packer.cpp @@ -0,0 +1,149 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/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/backends/sm120/native_style_precompute.cu b/cpp/models/pi05/src/backends/sm120/native_style_precompute.cu new file mode 100644 index 00000000..003a07da --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/native_style_precompute.cu @@ -0,0 +1,197 @@ +#include "flashrt/cpp/models/pi05/backends/sm120/native_style_precompute.h" + +#include + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +modalities::Status backend(const char* message) { + return modalities::Status::error(modalities::StatusCode::kBackend, + message); +} + +bool weight_shape(const NativeDeviceWeightStore& weights, + const std::string& name, + std::initializer_list shape, + const NativeDeviceWeight** out) { + const NativeDeviceWeight* weight = weights.find(name); + if (!weight || weight->dtype != NativeWeightDType::kBf16 || + weight->shape != std::vector(shape)) { + return false; + } + if (out) *out = weight; + return true; +} + +void* offset(void* base, std::size_t elements) { + return static_cast(base) + + elements * sizeof(std::uint16_t); +} + +} // namespace + +modalities::Status NativeStylePrecomputer::run( + const NativeDeviceWeightStore& weights, + NativeWorkspace* workspace, + std::uintptr_t stream) const { + if (!driver_ || !driver_->status().ok_status() || !workspace) { + return invalid("native style precomputer is invalid"); + } + const NativeWorkspaceBuffer* time_output = + workspace->find("decoder_time_emb"); + const NativeWorkspaceBuffer* style_attn = + workspace->find("decoder_style_attn"); + const NativeWorkspaceBuffer* style_ffn = + workspace->find("decoder_style_ffn"); + const NativeWorkspaceBuffer* style_final = + workspace->find("decoder_style_final"); + const NativeWorkspaceBuffer* scratch_a = workspace->find("decoder_x"); + const NativeWorkspaceBuffer* scratch_b = workspace->find("x_normed_buf"); + if (!time_output || !style_attn || !style_ffn || !style_final || + !scratch_a || !scratch_b || time_output->shape.size() != 3 || + style_attn->shape.size() != 4 || style_ffn->shape != style_attn->shape || + style_final->shape.size() != 3) { + return invalid("native style workspace layout is invalid"); + } + const int steps = static_cast(time_output->shape[0]); + const int chunk = static_cast(time_output->shape[1]); + if (time_output->shape[2] != 1024 || style_attn->shape[0] != steps || + style_attn->shape[1] != 18 || style_attn->shape[2] != chunk || + style_attn->shape[3] != 3072 || + style_final->shape != + std::vector( + {static_cast(steps), + static_cast(chunk), 3072})) { + return invalid("native style workspace shape is invalid"); + } + + const NativeDeviceWeight* time_source = nullptr; + const NativeDeviceWeight* time_in_w = nullptr; + const NativeDeviceWeight* time_in_b = nullptr; + const NativeDeviceWeight* time_out_w = nullptr; + const NativeDeviceWeight* time_out_b = nullptr; + const NativeDeviceWeight* final_w = nullptr; + const NativeDeviceWeight* final_b = nullptr; + if (!weight_shape(weights, "decoder_time_embeds", + {static_cast(steps), 1024}, + &time_source) || + !weight_shape(weights, "decoder_time_mlp_in_w", {1024, 1024}, &time_in_w) || + !weight_shape(weights, "decoder_time_mlp_in_b", {1024}, &time_in_b) || + !weight_shape(weights, "decoder_time_mlp_out_w", {1024, 1024}, &time_out_w) || + !weight_shape(weights, "decoder_time_mlp_out_b", {1024}, &time_out_b) || + !weight_shape(weights, "decoder_final_norm_mod_w", {1024, 3072}, &final_w) || + !weight_shape(weights, "decoder_final_norm_mod_b", {3072}, &final_b)) { + return invalid("native style global weights are incomplete"); + } + const cudaStream_t cuda_stream = reinterpret_cast(stream); + const std::uintptr_t native_stream = stream; + for (int step = 0; step < steps; ++step) { + void* time_row = offset(frt_buffer_dptr(time_source->buffer), + static_cast(step) * 1024); + modalities::Status st = driver_->bf16_nn( + time_row, frt_buffer_dptr(time_in_w->buffer), + frt_buffer_dptr(scratch_a->buffer), 1, 1024, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(scratch_a->buffer), + frt_buffer_dptr(time_in_b->buffer), 1, 1024, native_stream); + if (!st.ok_status()) return st; + st = driver_->silu_bf16(frt_buffer_dptr(scratch_a->buffer), 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn( + frt_buffer_dptr(scratch_a->buffer), + frt_buffer_dptr(time_out_w->buffer), + frt_buffer_dptr(scratch_b->buffer), 1, 1024, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16( + frt_buffer_dptr(scratch_b->buffer), + frt_buffer_dptr(time_out_b->buffer), 1, 1024, native_stream); + if (!st.ok_status()) return st; + st = driver_->silu_bf16(frt_buffer_dptr(scratch_b->buffer), 1024, + native_stream); + if (!st.ok_status()) return st; + + void* expanded = offset( + frt_buffer_dptr(time_output->buffer), + static_cast(step) * chunk * 1024); + for (int row = 0; row < chunk; ++row) { + const cudaError_t rc = cudaMemcpyAsync( + offset(expanded, static_cast(row) * 1024), + frt_buffer_dptr(scratch_b->buffer), + 1024 * sizeof(std::uint16_t), cudaMemcpyDeviceToDevice, + cuda_stream); + if (rc != cudaSuccess) return backend("time style expansion failed"); + } + + for (int layer = 0; layer < 18; ++layer) { + const std::string suffix = std::to_string(layer); + const NativeDeviceWeight* attn_w = nullptr; + const NativeDeviceWeight* attn_b = nullptr; + const NativeDeviceWeight* ffn_w = nullptr; + const NativeDeviceWeight* ffn_b = nullptr; + if (!weight_shape(weights, "decoder_pre_attn_norm_mod_w_" + suffix, + {1024, 3072}, &attn_w) || + !weight_shape(weights, "decoder_pre_attn_norm_mod_b_" + suffix, + {3072}, &attn_b) || + !weight_shape(weights, "decoder_pre_ffn_norm_mod_w_" + suffix, + {1024, 3072}, &ffn_w) || + !weight_shape(weights, "decoder_pre_ffn_norm_mod_b_" + suffix, + {3072}, &ffn_b)) { + return invalid("native style layer weights are incomplete"); + } + const std::size_t style_offset = + (static_cast(step) * 18 + layer) * chunk * 3072; + void* attn_target = + offset(frt_buffer_dptr(style_attn->buffer), style_offset); + void* ffn_target = + offset(frt_buffer_dptr(style_ffn->buffer), style_offset); + st = driver_->bf16_nn(expanded, frt_buffer_dptr(attn_w->buffer), + attn_target, chunk, 3072, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16(attn_target, + frt_buffer_dptr(attn_b->buffer), + chunk, 3072, native_stream); + if (!st.ok_status()) return st; + st = driver_->bf16_nn(expanded, frt_buffer_dptr(ffn_w->buffer), + ffn_target, chunk, 3072, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16(ffn_target, + frt_buffer_dptr(ffn_b->buffer), + chunk, 3072, native_stream); + if (!st.ok_status()) return st; + } + void* final_target = offset( + frt_buffer_dptr(style_final->buffer), + static_cast(step) * chunk * 3072); + st = driver_->bf16_nn(expanded, frt_buffer_dptr(final_w->buffer), + final_target, chunk, 3072, 1024, + native_stream); + if (!st.ok_status()) return st; + st = driver_->add_bias_bf16(final_target, + frt_buffer_dptr(final_b->buffer), + chunk, 3072, native_stream); + if (!st.ok_status()) return st; + } + const cudaError_t rc = cudaStreamSynchronize(cuda_stream); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("native style precompute synchronization failed"); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/backends/sm120/session.cpp b/cpp/models/pi05/src/backends/sm120/session.cpp new file mode 100644 index 00000000..d2db1f05 --- /dev/null +++ b/cpp/models/pi05/src/backends/sm120/session.cpp @@ -0,0 +1,390 @@ +#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" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_autotune.h" +#include "flashrt/cpp/models/pi05/backends/sm120/native_rtx_weight_packer.h" + +#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 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 + +Sm120BackendSession::Sm120BackendSession( + frt_ctx ctx, + const BackendConfig& config, + NativeRtxLinearMode linear_mode) + : graphs_(ctx, static_cast(GraphKind::kCount)), + config_(config), + weights_(ctx), + workspace_(ctx), + attention_(ctx), + linear_(&driver_, linear_mode), + forward_(&driver_, &linear_) {} + +Sm120BackendSession::~Sm120BackendSession() = default; + +std::unique_ptr Sm120BackendSession::create( + const std::string& checkpoint_path, + const BackendConfig& config, + modalities::Status* status) { + 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 || + 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"); + return nullptr; + } + frt_ctx ctx = frt_ctx_create(); + if (!ctx) { + if (status) *status = backend("native graph context creation failed"); + return nullptr; + } + std::unique_ptr session( + new (std::nothrow) Sm120BackendSession( + ctx, config, NativeRtxLinearMode::kBf16)); + if (!session) { + frt_ctx_destroy(ctx); + if (status) *status = backend("SM120 backend session allocation failed"); + return 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 session; +} + +std::unique_ptr Sm120BackendSession::create( + const std::string& checkpoint_path, + const BackendConfig& config, + const NativeCalibrationArtifact& calibration, + modalities::Status* status) { + if (config.precision != BackendPrecision::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 session( + new (std::nothrow) Sm120BackendSession( + ctx, config, NativeRtxLinearMode::kFp8Static)); + if (!session) { + frt_ctx_destroy(ctx); + if (status) *status = backend("SM120 backend session allocation failed"); + return nullptr; + } + 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 session; +} + +std::unique_ptr Sm120BackendSession::create_calibration( + const std::string& checkpoint_path, + const BackendConfig& config, + modalities::Status* status) { + if (config.precision != BackendPrecision::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 session( + new (std::nothrow) Sm120BackendSession( + ctx, config, NativeRtxLinearMode::kFp8Dynamic)); + if (!session) { + frt_ctx_destroy(ctx); + if (status) *status = backend("SM120 backend session allocation failed"); + return 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 session; +} + +modalities::Status Sm120BackendSession::initialize( + const std::string& checkpoint_path, + const NativeCalibrationArtifact* calibration) { + const bool fp8 = linear_.fp8(); + if ((calibration != nullptr) != linear_.static_fp8() || + (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; + 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; + 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; + 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 = fp8 ? NativeWorkspaceFlavor::kRtxFp8 + : NativeWorkspaceFlavor::kBf16; + st = workspace_.allocate(workspace_config); + if (!st.ok_status()) return st; + if (linear_.static_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; + } 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; + + 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; + if (fp8) { + st = autotune_native_rtx_fp8( + weights_, &workspace_, linear_, config_.num_views, + config_.chunk_size); + if (!st.ok_status()) return st; + report("fp8_autotune"); + } + report("workspace_style"); + + st = resolve_backend_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"}) { + 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"); + } + report("input_init"); + + st = capture_backend_graph( + &graphs_, + 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_backend_graph( + &graphs_, + 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_backend_graph( + &graphs_, + GraphKind::kContext, workspace_, + {"observation_images_normalized", "prompt_embedding", "encoder_x"}, + record_graph, this); + if (!st.ok_status()) return st; + report("capture"); + + st = graphs_.create_replay_stream(); + if (!st.ok_status()) return st; + 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(); +} + +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( + 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 st; +} + +modalities::Status Sm120BackendSession::record_action(void* stream) { + return forward_.diffusion(weights_, &workspace_, &attention_, + attention_driver_.get(), + reinterpret_cast(stream)); +} + +modalities::Status Sm120BackendSession::record(GraphKind kind, + void* stream) { + 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 Sm120BackendSession::record_graph( + void* user, std::size_t slot, void* stream) { + auto* session = static_cast(user); + return session->record(static_cast(slot), stream); +} + +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 Sm120BackendSession::replay(GraphKind kind) const { + return graphs_.replay(static_cast(kind)); +} + +modalities::Status Sm120BackendSession::synchronize() const { + return graphs_.synchronize(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/config_map.h b/cpp/models/pi05/src/config_map.h deleted file mode 100644 index 3c085e28..00000000 --- a/cpp/models/pi05/src/config_map.h +++ /dev/null @@ -1,101 +0,0 @@ -/* config_map.h — shared internal helpers for the Pi0.5 C faces (c_api and - * model_runtime): C config -> RuntimeConfig mapping and status/enum - * translation. Not installed. */ -#ifndef FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H -#define FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H - -#include "flashrt/cpp/models/pi05/c_api.h" -#include "flashrt/cpp/models/pi05/runtime.h" - -#include - -namespace flashrt { -namespace models { -namespace pi05 { -namespace cface { - -inline int status_code(const modalities::Status& st) { - using Code = modalities::StatusCode; - switch (st.code) { - case Code::kOk: return 0; - case Code::kInvalidArgument: return -1; - case Code::kNotFound: return -2; - case Code::kUnsupported: return -3; - case Code::kShapeMismatch: return -4; - case Code::kInsufficientStorage: return -5; - case Code::kBackend: return -6; - } - return -127; -} - -inline modalities::PixelFormat pixel_format(int value) { - using modalities::PixelFormat; - switch (value) { - case FRT_PI05_PIXEL_BGR8: return PixelFormat::kBGR8; - case FRT_PI05_PIXEL_RGBA8: return PixelFormat::kRGBA8; - case FRT_PI05_PIXEL_BGRA8: return PixelFormat::kBGRA8; - case FRT_PI05_PIXEL_GRAY8: return PixelFormat::kGRAY8; - case FRT_PI05_PIXEL_RGB8: - default: return PixelFormat::kRGB8; - } -} - -inline modalities::DType dtype(int value) { - using modalities::DType; - switch (value) { - case FRT_PI05_DTYPE_FLOAT16: return DType::kFloat16; - case FRT_PI05_DTYPE_FLOAT32: return DType::kFloat32; - case FRT_PI05_DTYPE_BFLOAT16: - case FRT_PI05_DTYPE_DEFAULT: - default: return DType::kBFloat16; - } -} - -inline bool has_field(const frt_pi05_runtime_config* in, std::size_t offset, - std::size_t bytes) { - return in && in->struct_size >= offset + bytes; -} - -inline RuntimeConfig make_config(const frt_pi05_runtime_config* in) { - RuntimeConfig cfg; - if (!in) return cfg; - if (in->num_views > 0) cfg.num_views = in->num_views; - if (in->chunk > 0) cfg.chunk = in->chunk; - if (in->model_action_dim > 0) cfg.model_action_dim = in->model_action_dim; - if (in->robot_action_dim > 0) cfg.robot_action_dim = in->robot_action_dim; - if (in->action_mean && in->n_action_mean) { - cfg.action_mean.assign(in->action_mean, - in->action_mean + in->n_action_mean); - } - if (in->action_stddev && in->n_action_stddev) { - cfg.action_stddev.assign(in->action_stddev, - in->action_stddev + in->n_action_stddev); - } - if (in->graph_name) cfg.graph_name = in->graph_name; - if (in->image_buffer_name) cfg.image_buffer_name = in->image_buffer_name; - if (in->action_buffer_name) cfg.action_buffer_name = in->action_buffer_name; - if (has_field(in, offsetof(frt_pi05_runtime_config, image_dtype), - sizeof(in->image_dtype))) { - cfg.image_dtype = dtype(in->image_dtype); - } - if (has_field(in, offsetof(frt_pi05_runtime_config, action_dtype), - sizeof(in->action_dtype))) { - cfg.action_dtype = dtype(in->action_dtype); - } - if (has_field(in, offsetof(frt_pi05_runtime_config, max_frame_width), - sizeof(in->max_frame_width)) && in->max_frame_width > 0) { - cfg.max_frame_width = in->max_frame_width; - } - if (has_field(in, offsetof(frt_pi05_runtime_config, max_frame_height), - sizeof(in->max_frame_height)) && in->max_frame_height > 0) { - cfg.max_frame_height = in->max_frame_height; - } - return cfg; -} - -} // namespace cface -} // namespace pi05 -} // namespace models -} // namespace flashrt - -#endif // FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H diff --git a/cpp/models/pi05/src/io.cpp b/cpp/models/pi05/src/io.cpp deleted file mode 100644 index fbe95220..00000000 --- a/cpp/models/pi05/src/io.cpp +++ /dev/null @@ -1,42 +0,0 @@ -#include "flashrt/cpp/models/pi05/io.h" - -namespace flashrt { -namespace models { -namespace pi05 { - -RuntimeIo::RuntimeIo(int num_views, - modalities::TensorView image_input, - modalities::TensorView action_output, - std::vector action_mean, - std::vector action_stddev, - void* stream, - int chunk, - int model_action_dim, - int robot_action_dim, - modalities::DType image_dtype, - modalities::VisionStaging* staging) - : image_input_(image_input), - action_output_(action_output), - stream_(stream), - staging_(staging), - vision_spec_(vision_preprocess_spec(num_views)), - action_spec_(action_postprocess_spec(action_mean, action_stddev, chunk, - model_action_dim, robot_action_dim)) { - vision_spec_.output_dtype = image_dtype; -} - -modalities::Status RuntimeIo::prepare_vision( - const std::vector& frames) const { - return modalities::preprocess_vision(vision_spec_, frames, image_input_, - stream_, staging_); -} - -modalities::Status RuntimeIo::read_actions( - std::vector* robot_actions) const { - return modalities::postprocess_action(action_spec_, action_output_, - robot_actions, stream_); -} - -} // namespace pi05 -} // namespace models -} // namespace flashrt diff --git a/cpp/models/pi05/src/model/io.cpp b/cpp/models/pi05/src/model/io.cpp new file mode 100644 index 00000000..5ad26d3a --- /dev/null +++ b/cpp/models/pi05/src/model/io.cpp @@ -0,0 +1,92 @@ +#include "flashrt/cpp/models/pi05/model/io.h" + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status validate_pi05_frame_contract( + const modalities::VisionFrame& frame, bool strict_rgb8) { + if (strict_rgb8 && frame.format != modalities::PixelFormat::kRGB8) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "Pi05 image input must be RGB8"); + } + if (frame.image.dtype != modalities::DType::kUInt8 || + frame.image.layout != modalities::Layout::kHWC) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "Pi05 image input must be u8 HWC"); + } + std::uint64_t channels = 3; + if (frame.format == modalities::PixelFormat::kRGBA8 || + frame.format == modalities::PixelFormat::kBGRA8) { + channels = 4; + } else if (frame.format == modalities::PixelFormat::kGRAY8) { + channels = 1; + } + if (frame.width <= 0 || frame.height <= 0 || + frame.image.shape.rank != 3 || + frame.image.shape.dims[0] != static_cast(frame.height) || + frame.image.shape.dims[1] != static_cast(frame.width) || + frame.image.shape.dims[2] != channels) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "Pi05 image shape must match HWC dimensions"); + } + if (frame.image.place != modalities::MemoryPlace::kHost && + frame.image.place != modalities::MemoryPlace::kHostPinned) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "Pi05 image input must be host memory"); + } + return modalities::Status::ok(); +} + +} // namespace + +RuntimeIo::RuntimeIo(int num_views, + modalities::TensorView image_input, + modalities::TensorView action_output, + std::vector action_mean, + std::vector action_stddev, + void* stream, + int chunk, + int model_action_dim, + int robot_action_dim, + modalities::DType image_dtype, + modalities::VisionStaging* staging, + modalities::ActionStaging* action_staging, + bool strict_rgb8) + : image_input_(image_input), + action_output_(action_output), + stream_(stream), + staging_(staging), + action_staging_(action_staging), + strict_rgb8_(strict_rgb8), + vision_spec_(vision_preprocess_spec(num_views)), + action_spec_(action_postprocess_spec(action_mean, action_stddev, chunk, + model_action_dim, robot_action_dim)) { + vision_spec_.output_dtype = image_dtype; +} + +modalities::Status RuntimeIo::prepare_vision( + const std::vector& frames) const { + for (const auto& frame : frames) { + auto st = validate_pi05_frame_contract(frame, strict_rgb8_); + if (!st.ok_status()) return st; + } + return modalities::preprocess_vision(vision_spec_, frames, image_input_, + stream_, staging_); +} + +modalities::Status RuntimeIo::read_actions( + std::vector* robot_actions) const { + return modalities::postprocess_action(action_spec_, action_output_, + robot_actions, stream_, + action_staging_); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/model/prompt_embed.cpp b/cpp/models/pi05/src/model/prompt_embed.cpp new file mode 100644 index 00000000..1f8fdaae --- /dev/null +++ b/cpp/models/pi05/src/model/prompt_embed.cpp @@ -0,0 +1,171 @@ +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" + +#include "flashrt/cpp/models/pi05/model/prompt_format.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status validate_output_capacity( + const PromptEmbeddingSpec& spec, + const modalities::TensorView& output) { + if (!spec.vocab_size || !spec.hidden_dim || !spec.max_tokens) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "invalid prompt embedding dimensions"); + } + if (!output.data) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt_embedding has null data"); + } + if (output.place != modalities::MemoryPlace::kHost && + output.place != modalities::MemoryPlace::kHostPinned && + output.place != modalities::MemoryPlace::kDevice) { + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "prompt_embedding memory place is unsupported"); + } + if (output.layout != modalities::Layout::kFlat || + output.shape.rank != 2 || + output.shape.dims[0] != spec.max_tokens || + output.shape.dims[1] != spec.hidden_dim) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "prompt_embedding shape mismatch"); + } + const std::uint64_t need = + spec.max_tokens * spec.hidden_dim * modalities::dtype_size(output.dtype); + if (output.bytes < need) { + return modalities::Status::error( + modalities::StatusCode::kInsufficientStorage, + "prompt_embedding storage is too small"); + } + return modalities::Status::ok(); +} + +modalities::Status zero_prompt_output(const modalities::TensorView& output, + void* stream) { + if (output.place == modalities::MemoryPlace::kHost || + output.place == modalities::MemoryPlace::kHostPinned) { + std::memset(output.data, 0, static_cast(output.bytes)); + return modalities::Status::ok(); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)stream; + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "device prompt zeroing requires the CUDA build"); +#else + cudaStream_t cuda_stream = reinterpret_cast(stream); + cudaError_t rc = cudaMemsetAsync(output.data, 0, output.bytes, + cuda_stream); + if (rc == cudaSuccess) rc = cudaStreamSynchronize(cuda_stream); + if (rc != cudaSuccess) { + return modalities::Status::error( + modalities::StatusCode::kBackend, + std::string("cuda prompt zeroing failed: ") + + cudaGetErrorString(rc)); + } + return modalities::Status::ok(); +#endif +} + +} // namespace + +modalities::Status embed_prompt( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len, + void* stream, + modalities::TextEmbeddingStaging* staging, + std::string* formatted_workspace) { + if (!token_ids || !prompt_len) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt embedding outputs are null"); + } + token_ids->clear(); + *prompt_len = 0; + auto st = validate_output_capacity(spec, output); + if (!st.ok_status()) return st; + if (!tokenizer.loaded()) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "SentencePiece model is not loaded"); + } + + modalities::SentencePieceEncodeOptions options; + options.add_bos = true; + options.max_tokens = spec.max_tokens; + if (state) { + std::string local; + std::string* formatted = formatted_workspace ? formatted_workspace + : &local; + format_state_prompt_into(prompt, state, n_state, formatted); + st = tokenizer.encode(*formatted, options, token_ids); + } else { + st = tokenizer.encode(prompt, options, token_ids); + if (st.ok_status() && spec.no_state_suffix_token_id >= 0) { + token_ids->push_back(spec.no_state_suffix_token_id); + } + } + if (!st.ok_status()) return st; + if (token_ids->size() > spec.max_tokens) { + return modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "prompt token count exceeds max_tokens"); + } + + if (spec.zero_pad_output) { + st = zero_prompt_output(output, stream); + if (!st.ok_status()) return st; + } + modalities::TensorView prefix = output; + prefix.shape = modalities::Shape{static_cast( + token_ids->size()), + spec.hidden_dim}; + prefix.bytes = static_cast(token_ids->size()) * + spec.hidden_dim * modalities::dtype_size(output.dtype); + + modalities::EmbeddingGatherSpec gather{spec.vocab_size, spec.hidden_dim, + spec.scale}; + st = modalities::gather_token_embeddings( + gather, token_ids->data(), token_ids->size(), embedding_table, prefix, + stream, staging); + if (!st.ok_status()) return st; + *prompt_len = static_cast(token_ids->size()); + return modalities::Status::ok(); +} + +modalities::Status embed_prompt_cpu( + modalities::SentencePieceTokenizer& tokenizer, + const PromptEmbeddingSpec& spec, + const std::string& prompt, + const float* state, + std::uint64_t n_state, + modalities::TensorView embedding_table, + modalities::TensorView output, + std::vector* token_ids, + std::uint64_t* prompt_len) { + return embed_prompt(tokenizer, spec, prompt, state, n_state, + embedding_table, output, token_ids, prompt_len); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/model/prompt_format.cpp b/cpp/models/pi05/src/model/prompt_format.cpp new file mode 100644 index 00000000..7f0cf64f --- /dev/null +++ b/cpp/models/pi05/src/model/prompt_format.cpp @@ -0,0 +1,95 @@ +#include "flashrt/cpp/models/pi05/model/prompt_format.h" + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +std::vector make_openpi_bins() { + std::vector bins; + bins.reserve(256); + for (int i = 0; i < 256; ++i) { + bins.push_back(-1.0f + static_cast(i) * (2.0f / 256.0f)); + } + return bins; +} + +bool ascii_space(char c) { + return std::isspace(static_cast(c)) != 0; +} + +} // namespace + +std::vector discretize_state_prompt_bins( + const float* state, std::uint64_t n) { + static const std::vector bins = make_openpi_bins(); + std::vector out; + out.reserve(static_cast(n)); + for (std::uint64_t i = 0; i < n; ++i) { + const auto it = std::upper_bound(bins.begin(), bins.end(), state[i]); + out.push_back(static_cast(it - bins.begin()) - 1); + } + return out; +} + +std::string clean_task_prompt(const std::string& prompt) { + auto begin = prompt.begin(); + auto end = prompt.end(); + while (begin != end && ascii_space(*begin)) ++begin; + while (begin != end && ascii_space(*(end - 1))) --end; + + std::string cleaned(begin, end); + for (char& c : cleaned) { + if (c == '_' || c == '\n') c = ' '; + } + return cleaned; +} + +std::string format_state_prompt(const std::string& prompt, + const float* state, + std::uint64_t n_state) { + std::string out; + out.reserve(prompt.size() + static_cast(n_state) * 5 + 32); + format_state_prompt_into(prompt, state, n_state, &out); + return out; +} + +void format_state_prompt_into(const std::string& prompt, + const float* state, + std::uint64_t n_state, + std::string* out) { + if (!out) return; + out->clear(); + auto begin = prompt.begin(); + auto end = prompt.end(); + while (begin != end && ascii_space(*begin)) ++begin; + while (begin != end && ascii_space(*(end - 1))) --end; + + if (state) out->append("Task: "); + for (auto it = begin; it != end; ++it) { + out->push_back(*it == '_' || *it == '\n' ? ' ' : *it); + } + if (!state) return; + + static const std::vector bins = make_openpi_bins(); + out->append(", State: "); + char number[24]; + for (std::uint64_t i = 0; i < n_state; ++i) { + if (i) out->push_back(' '); + const auto bin = static_cast( + std::upper_bound(bins.begin(), bins.end(), state[i]) - + bins.begin()) - + 1; + const auto result = std::to_chars(number, number + sizeof(number), bin); + out->append(number, result.ptr); + } + out->append(";\nAction: "); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/runtime.cpp b/cpp/models/pi05/src/model/runtime.cpp similarity index 54% rename from cpp/models/pi05/src/runtime.cpp rename to cpp/models/pi05/src/model/runtime.cpp index 37b48b2c..7faba9d5 100644 --- a/cpp/models/pi05/src/runtime.cpp +++ b/cpp/models/pi05/src/model/runtime.cpp @@ -1,6 +1,8 @@ -#include "flashrt/cpp/models/pi05/runtime.h" +#include "flashrt/cpp/models/pi05/model/runtime.h" +#include #include +#include namespace flashrt { namespace models { @@ -70,6 +72,8 @@ Runtime::Runtime(const frt_runtime_export_v1* exp, RuntimeConfig config) } Runtime::~Runtime() { + modalities::text_embedding_staging_destroy(&prompt_embedding_staging_); + modalities::action_staging_destroy(&action_staging_); modalities::vision_staging_destroy(&staging_); release_export(); } @@ -166,18 +170,85 @@ modalities::Status Runtime::bind() { staging = &staging_; } + modalities::ActionStaging* action_staging = nullptr; + if (action.place == modalities::MemoryPlace::kDevice) { + modalities::ActionPostprocessSpec action_spec = action_postprocess_spec( + config_.action_mean, config_.action_stddev, config_.chunk, + config_.model_action_dim, config_.robot_action_dim); + modalities::Status st = modalities::action_staging_create( + &action_staging_, + modalities::required_action_output_bytes(action_spec, + action.dtype)); + if (!st.ok_status()) return st; + action_staging = &action_staging_; + } + io_ = RuntimeIo(config_.num_views, image, action, config_.action_mean, config_.action_stddev, find_native_stream(exp_, stream_id_), config_.chunk, config_.model_action_dim, - config_.robot_action_dim, config_.image_dtype, staging); - return modalities::Status::ok(); + config_.robot_action_dim, config_.image_dtype, staging, + action_staging, config_.strict_rgb8); + return bind_prompt_staging(); } int Runtime::set_prompt(const char* text) { - /* The adopted-export path assumes prompt/token embedding was prepared by - * the producer before capture/export. A native Pi0.5 producer will replace - * this with tokenizer + prompt-region binding without changing Nexus. */ - return (text == nullptr || text[0] == '\0') ? 0 : -1; + return set_prompt_state(text, nullptr, 0); +} + +int Runtime::set_prompt_state(const char* text, const float* state, + std::uint64_t n_state) { + if (!prompt_staging_enabled_) { + return (text == nullptr || text[0] == '\0') ? 0 : -1; + } + if (!text) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt text is null"); + return -1; + } + const std::size_t text_bytes = std::strlen(text); + if (text_bytes > max_task_prompt_bytes_) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "prompt text exceeds the configured hot-path capacity"); + return -1; + } + task_prompt_workspace_.assign(text, text_bytes); + const float* state_for_prompt = state; + if (state && state_normalization_enabled()) { + if (n_state != config_.state_q01.size()) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kShapeMismatch, + "state dimension does not match norm stats"); + return -1; + } + for (std::uint64_t i = 0; i < n_state; ++i) { + const float lo = config_.state_q01[i]; + const float hi = config_.state_q99[i]; + normalized_state_[i] = + ((state[i] - lo) / (hi - lo + 1e-6f)) * 2.0f - 1.0f; + } + state_for_prompt = normalized_state_.data(); + } + prompt_status_ = embed_prompt( + prompt_tokenizer_, prompt_spec_, task_prompt_workspace_, + state_for_prompt, n_state, + prompt_embedding_table_, prompt_embedding_output_, &prompt_token_ids_, + ¤t_prompt_len_, find_native_stream(exp_, stream_id_), + prompt_embedding_output_.place == modalities::MemoryPlace::kDevice + ? &prompt_embedding_staging_ + : nullptr, + &formatted_prompt_workspace_); + if (prompt_status_.ok_status() && config_.prompt_length_update_fn) { + const int rc = config_.prompt_length_update_fn( + config_.prompt_length_update_user, current_prompt_len_); + if (rc != 0) { + prompt_status_ = modalities::Status::error( + modalities::StatusCode::kBackend, + "prompt length device update failed"); + } + } + return prompt_status_.ok_status() ? 0 : -1; } modalities::Status Runtime::prepare_vision( @@ -203,6 +274,70 @@ int Runtime::default_replay(frt_graph graph, frt_shape_key key, return frt_graph_replay(graph, key, stream_id); } +modalities::Status Runtime::bind_prompt_staging() { + const bool any = + !config_.prompt_tokenizer_model_path.empty() || + config_.prompt_embedding_table.data || + config_.prompt_embedding_output.data || + config_.prompt_vocab_size || config_.prompt_hidden_dim || + config_.prompt_max_tokens; + if (!any) { + prompt_status_ = modalities::Status::ok(); + return modalities::Status::ok(); + } + if (config_.prompt_tokenizer_model_path.empty() || + !config_.prompt_embedding_table.data || + !config_.prompt_embedding_output.data || + !config_.prompt_vocab_size || !config_.prompt_hidden_dim || + !config_.prompt_max_tokens) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "incomplete Pi05 prompt staging config"); + } + prompt_status_ = + prompt_tokenizer_.load_model(config_.prompt_tokenizer_model_path); + if (!prompt_status_.ok_status()) return prompt_status_; + + prompt_embedding_table_ = config_.prompt_embedding_table; + prompt_embedding_output_ = config_.prompt_embedding_output; + prompt_spec_.vocab_size = config_.prompt_vocab_size; + prompt_spec_.hidden_dim = config_.prompt_hidden_dim; + prompt_spec_.max_tokens = config_.prompt_max_tokens; + prompt_spec_.scale = config_.prompt_embedding_scale > 0.0f + ? config_.prompt_embedding_scale + : std::sqrt(static_cast( + config_.prompt_hidden_dim)); + if (config_.state_q01.size() > + (std::numeric_limits::max() - 32ull) / 5ull) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "state workspace capacity overflows size_t"); + } + const std::size_t state_bytes = config_.state_q01.size() * 5ull + 32ull; + if (config_.prompt_max_tokens > + (std::numeric_limits::max() - state_bytes) / 8ull) { + return modalities::Status::error( + modalities::StatusCode::kInvalidArgument, + "prompt workspace capacity overflows size_t"); + } + const std::size_t max_prompt_bytes = + static_cast(config_.prompt_max_tokens * 8ull); + max_task_prompt_bytes_ = max_prompt_bytes; + task_prompt_workspace_.reserve(max_prompt_bytes); + formatted_prompt_workspace_.reserve(max_prompt_bytes + state_bytes); + prompt_token_ids_.reserve(static_cast( + config_.prompt_max_tokens + 1ull)); + normalized_state_.resize(config_.state_q01.size()); + prompt_tokenizer_.reserve(config_.prompt_max_tokens); + if (prompt_embedding_output_.place == modalities::MemoryPlace::kDevice) { + prompt_status_ = modalities::text_embedding_staging_create( + &prompt_embedding_staging_, config_.prompt_max_tokens); + if (!prompt_status_.ok_status()) return prompt_status_; + } + prompt_staging_enabled_ = true; + return modalities::Status::ok(); +} + } // namespace pi05 } // namespace models } // namespace flashrt diff --git a/cpp/models/pi05/src/spec.cpp b/cpp/models/pi05/src/model/spec.cpp similarity index 89% rename from cpp/models/pi05/src/spec.cpp rename to cpp/models/pi05/src/model/spec.cpp index 4d3efea2..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 @@ -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/models/pi05/src/c_api.cpp b/cpp/models/pi05/src/service/c_api.cpp similarity index 57% rename from cpp/models/pi05/src/c_api.cpp rename to cpp/models/pi05/src/service/c_api.cpp index 4582f8aa..ea696b2b 100644 --- a/cpp/models/pi05/src/c_api.cpp +++ b/cpp/models/pi05/src/service/c_api.cpp @@ -4,23 +4,30 @@ #include #include +#include #include #include #include #include #include +#include #include struct frt_pi05_runtime_s { std::unique_ptr runtime; std::string last_error; + std::vector vision_frames; + std::vector vision_seen; + std::vector action_values; }; namespace { using flashrt::models::pi05::cface::make_config; +using flashrt::models::pi05::cface::pixel_channels; using flashrt::models::pi05::cface::pixel_format; using flashrt::models::pi05::cface::status_code; +using flashrt::models::pi05::cface::valid_pixel_format; } // namespace @@ -38,8 +45,10 @@ extern "C" int frt_pi05_runtime_create( auto* h = new (std::nothrow) frt_pi05_runtime_s(); if (!h) return -5; try { + auto runtime_config = make_config(config); + runtime_config.strict_rgb8 = false; h->runtime.reset( - new flashrt::models::pi05::Runtime(exp, make_config(config))); + new flashrt::models::pi05::Runtime(exp, std::move(runtime_config))); } catch (const std::exception& e) { h->last_error = e.what(); delete h; @@ -54,6 +63,14 @@ extern "C" int frt_pi05_runtime_create( delete h; return rc; } + const auto& manifest = h->runtime->manifest(); + h->vision_frames.resize(manifest.vision.view_order.size()); + h->vision_seen.resize(manifest.vision.view_order.size()); + for (std::size_t i = 0; i < h->vision_frames.size(); ++i) { + h->vision_frames[i].name = manifest.vision.view_order[i]; + } + h->action_values.resize(static_cast( + manifest.action.chunk * manifest.action.robot_dim)); *out = h; return 0; } @@ -66,7 +83,32 @@ extern "C" int frt_pi05_runtime_set_prompt(frt_pi05_runtime* h, const char* text) { if (!h || !h->runtime) return -1; int rc = h->runtime->set_prompt(text); - if (rc != 0) h->last_error = "prompt updates are not supported by adopted-export Pi05 runtime"; + if (rc != 0) { + const auto& st = h->runtime->prompt_status(); + h->last_error = st.message.empty() + ? "prompt updates are not supported by this Pi05 runtime" + : st.message; + } else { + h->last_error.clear(); + } + return rc; +} + +extern "C" int frt_pi05_runtime_set_prompt_state( + frt_pi05_runtime* h, + const char* text, + const float* state, + uint64_t n_state) { + if (!h || !h->runtime || (!state && n_state)) return -1; + int rc = h->runtime->set_prompt_state(text, state, n_state); + if (rc != 0) { + const auto& st = h->runtime->prompt_status(); + h->last_error = st.message.empty() + ? "prompt updates are not supported by this Pi05 runtime" + : st.message; + } else { + h->last_error.clear(); + } return rc; } @@ -75,8 +117,11 @@ extern "C" int frt_pi05_runtime_prepare_vision( const frt_pi05_vision_frame* frames, uint64_t n_frames) { if (!h || !h->runtime || (!frames && n_frames)) return -1; - std::vector v; - v.reserve(static_cast(n_frames)); + if (n_frames != h->vision_frames.size()) { + h->last_error = "Pi05 vision frame count does not match the runtime"; + return -4; + } + std::fill(h->vision_seen.begin(), h->vision_seen.end(), 0); for (uint64_t i = 0; i < n_frames; ++i) { const frt_pi05_vision_frame& in = frames[i]; if (in.struct_size < sizeof(frt_pi05_vision_frame) || @@ -84,8 +129,23 @@ extern "C" int frt_pi05_runtime_prepare_vision( h->last_error = "invalid Pi05 vision frame"; return -1; } - flashrt::modalities::VisionFrame out; - out.name = in.name; + if (!valid_pixel_format(in.pixel_format)) { + h->last_error = "Pi05 vision pixel format is invalid"; + return -4; + } + std::size_t slot = h->vision_frames.size(); + for (std::size_t j = 0; j < h->vision_frames.size(); ++j) { + if (h->vision_frames[j].name == in.name) { + slot = j; + break; + } + } + if (slot == h->vision_frames.size() || h->vision_seen[slot]) { + h->last_error = "Pi05 vision frame name is unknown or duplicated"; + return -4; + } + h->vision_seen[slot] = 1; + auto& out = h->vision_frames[slot]; out.image.data = const_cast(in.data); out.image.bytes = in.bytes; out.image.dtype = flashrt::modalities::DType::kUInt8; @@ -94,15 +154,14 @@ extern "C" int frt_pi05_runtime_prepare_vision( out.image.shape = flashrt::modalities::Shape{ static_cast(std::max(0, in.height)), static_cast(std::max(0, in.width)), - 3}; + pixel_channels(in.pixel_format)}; out.format = pixel_format(in.pixel_format); out.width = in.width; out.height = in.height; out.stride_bytes = in.stride_bytes; out.timestamp_ns = in.timestamp_ns; - v.push_back(std::move(out)); } - auto st = h->runtime->prepare_vision(v); + auto st = h->runtime->prepare_vision(h->vision_frames); if (!st.ok_status()) { h->last_error = st.message; return status_code(st); @@ -123,19 +182,19 @@ extern "C" int frt_pi05_runtime_read_actions(frt_pi05_runtime* h, uint64_t out_capacity, uint64_t* n_written) { if (!h || !h->runtime || !out_actions) return -1; - std::vector actions; - auto st = h->runtime->read_actions(&actions); + auto st = h->runtime->read_actions(&h->action_values); if (!st.ok_status()) { h->last_error = st.message; return status_code(st); } - if (out_capacity < actions.size()) { + if (out_capacity < h->action_values.size()) { h->last_error = "action output buffer is too small"; - if (n_written) *n_written = actions.size(); + if (n_written) *n_written = h->action_values.size(); return -5; } - std::memcpy(out_actions, actions.data(), actions.size() * sizeof(float)); - if (n_written) *n_written = actions.size(); + std::memcpy(out_actions, h->action_values.data(), + h->action_values.size() * sizeof(float)); + if (n_written) *n_written = h->action_values.size(); h->last_error.clear(); return 0; } diff --git a/cpp/models/pi05/src/service/config_map.h b/cpp/models/pi05/src/service/config_map.h new file mode 100644 index 00000000..fb78d803 --- /dev/null +++ b/cpp/models/pi05/src/service/config_map.h @@ -0,0 +1,215 @@ +/* config_map.h — shared internal helpers for the Pi0.5 C faces (c_api and + * model_runtime): C config -> RuntimeConfig mapping and status/enum + * translation. Not installed. */ +#ifndef FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H +#define FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H + +#include "flashrt/cpp/models/pi05/c_api.h" +#include "flashrt/cpp/models/pi05/model/runtime.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace cface { + +inline int status_code(const modalities::Status& st) { + using Code = modalities::StatusCode; + switch (st.code) { + case Code::kOk: return 0; + case Code::kInvalidArgument: return -1; + case Code::kNotFound: return -2; + case Code::kUnsupported: return -3; + case Code::kShapeMismatch: return -4; + case Code::kInsufficientStorage: return -5; + case Code::kBackend: return -6; + } + return -127; +} + +inline modalities::PixelFormat pixel_format(int value) { + using modalities::PixelFormat; + switch (value) { + case FRT_PI05_PIXEL_BGR8: return PixelFormat::kBGR8; + case FRT_PI05_PIXEL_RGBA8: return PixelFormat::kRGBA8; + case FRT_PI05_PIXEL_BGRA8: return PixelFormat::kBGRA8; + case FRT_PI05_PIXEL_GRAY8: return PixelFormat::kGRAY8; + case FRT_PI05_PIXEL_RGB8: + default: return PixelFormat::kRGB8; + } +} + +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) { + case FRT_PI05_DTYPE_FLOAT16: return DType::kFloat16; + case FRT_PI05_DTYPE_FLOAT32: return DType::kFloat32; + case FRT_PI05_DTYPE_BFLOAT16: + case FRT_PI05_DTYPE_DEFAULT: + default: return DType::kBFloat16; + } +} + +inline bool has_field(const frt_pi05_runtime_config* in, std::size_t offset, + std::size_t bytes) { + return in && in->struct_size >= offset + bytes; +} + +inline RuntimeConfig make_config(const frt_pi05_runtime_config* in) { + RuntimeConfig cfg; + if (!in) return cfg; + if (in->num_views > 0) cfg.num_views = in->num_views; + if (in->chunk > 0) cfg.chunk = in->chunk; + if (in->model_action_dim > 0) cfg.model_action_dim = in->model_action_dim; + if (in->robot_action_dim > 0) cfg.robot_action_dim = in->robot_action_dim; + if (in->action_mean && in->n_action_mean) { + cfg.action_mean.assign(in->action_mean, + in->action_mean + in->n_action_mean); + } + if (in->action_stddev && in->n_action_stddev) { + cfg.action_stddev.assign(in->action_stddev, + in->action_stddev + in->n_action_stddev); + } + if (in->graph_name) cfg.graph_name = in->graph_name; + if (in->image_buffer_name) cfg.image_buffer_name = in->image_buffer_name; + if (in->action_buffer_name) cfg.action_buffer_name = in->action_buffer_name; + if (has_field(in, offsetof(frt_pi05_runtime_config, image_dtype), + sizeof(in->image_dtype))) { + cfg.image_dtype = dtype(in->image_dtype); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, action_dtype), + sizeof(in->action_dtype))) { + cfg.action_dtype = dtype(in->action_dtype); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, max_frame_width), + sizeof(in->max_frame_width)) && in->max_frame_width > 0) { + cfg.max_frame_width = in->max_frame_width; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, max_frame_height), + sizeof(in->max_frame_height)) && in->max_frame_height > 0) { + cfg.max_frame_height = in->max_frame_height; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_tokenizer_model_path), + sizeof(in->prompt_tokenizer_model_path)) && + in->prompt_tokenizer_model_path) { + cfg.prompt_tokenizer_model_path = in->prompt_tokenizer_model_path; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_table_data), + sizeof(in->prompt_embedding_table_data)) && + in->prompt_embedding_table_data) { + cfg.prompt_embedding_table.data = + const_cast(in->prompt_embedding_table_data); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_table_bytes), + sizeof(in->prompt_embedding_table_bytes))) { + cfg.prompt_embedding_table.bytes = in->prompt_embedding_table_bytes; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_table_dtype), + sizeof(in->prompt_embedding_table_dtype))) { + cfg.prompt_embedding_table.dtype = + dtype(in->prompt_embedding_table_dtype); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_vocab_size), + sizeof(in->prompt_embedding_vocab_size))) { + cfg.prompt_vocab_size = in->prompt_embedding_vocab_size; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_hidden_dim), + sizeof(in->prompt_embedding_hidden_dim))) { + cfg.prompt_hidden_dim = in->prompt_embedding_hidden_dim; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, prompt_embedding_data), + sizeof(in->prompt_embedding_data)) && + in->prompt_embedding_data) { + cfg.prompt_embedding_output.data = in->prompt_embedding_data; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, prompt_embedding_bytes), + sizeof(in->prompt_embedding_bytes))) { + cfg.prompt_embedding_output.bytes = in->prompt_embedding_bytes; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, prompt_embedding_dtype), + sizeof(in->prompt_embedding_dtype))) { + cfg.prompt_embedding_output.dtype = dtype(in->prompt_embedding_dtype); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, max_prompt_tokens), + sizeof(in->max_prompt_tokens))) { + cfg.prompt_max_tokens = in->max_prompt_tokens; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_scale), + sizeof(in->prompt_embedding_scale))) { + cfg.prompt_embedding_scale = in->prompt_embedding_scale; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, state_q01), + sizeof(in->state_q01)) && + has_field(in, offsetof(frt_pi05_runtime_config, n_state_q01), + sizeof(in->n_state_q01)) && + in->state_q01 && in->n_state_q01) { + cfg.state_q01.assign(in->state_q01, in->state_q01 + in->n_state_q01); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, state_q99), + sizeof(in->state_q99)) && + has_field(in, offsetof(frt_pi05_runtime_config, n_state_q99), + sizeof(in->n_state_q99)) && + in->state_q99 && in->n_state_q99) { + cfg.state_q99.assign(in->state_q99, in->state_q99 + in->n_state_q99); + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_length_update), + sizeof(in->prompt_length_update))) { + cfg.prompt_length_update_fn = in->prompt_length_update; + } + if (has_field(in, offsetof(frt_pi05_runtime_config, + prompt_length_update_user), + sizeof(in->prompt_length_update_user))) { + cfg.prompt_length_update_user = in->prompt_length_update_user; + } + const bool prompt_on_device = + has_field(in, offsetof(frt_pi05_runtime_config, + prompt_embedding_on_device), + sizeof(in->prompt_embedding_on_device)) && + in->prompt_embedding_on_device != 0; + if (cfg.prompt_vocab_size && cfg.prompt_hidden_dim) { + cfg.prompt_embedding_table.place = + prompt_on_device ? modalities::MemoryPlace::kDevice + : modalities::MemoryPlace::kHost; + cfg.prompt_embedding_table.layout = modalities::Layout::kFlat; + cfg.prompt_embedding_table.shape = + modalities::Shape{cfg.prompt_vocab_size, cfg.prompt_hidden_dim}; + } + if (cfg.prompt_max_tokens && cfg.prompt_hidden_dim) { + cfg.prompt_embedding_output.place = + prompt_on_device ? modalities::MemoryPlace::kDevice + : modalities::MemoryPlace::kHost; + cfg.prompt_embedding_output.layout = modalities::Layout::kFlat; + cfg.prompt_embedding_output.shape = + modalities::Shape{cfg.prompt_max_tokens, cfg.prompt_hidden_dim}; + } + return cfg; +} + +} // namespace cface +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_CONFIG_MAP_H diff --git a/cpp/models/pi05/src/model_runtime.cpp b/cpp/models/pi05/src/service/model_runtime.cpp similarity index 71% rename from cpp/models/pi05/src/model_runtime.cpp rename to cpp/models/pi05/src/service/model_runtime.cpp index b91988fd..d2730124 100644 --- a/cpp/models/pi05/src/model_runtime.cpp +++ b/cpp/models/pi05/src/service/model_runtime.cpp @@ -31,6 +31,15 @@ struct Adapter { uint32_t images_port = kPortImages; uint32_t noise_port = kPortNoise; uint32_t actions_port = kPortActions; + uint32_t prompt_port = kNoPort; + uint32_t state_port = kNoPort; + bool has_prompt_text = false; + bool has_state = false; + std::size_t prompt_text_limit = 0; + std::string prompt_text; + std::vector state_values; + std::vector vision_frames; + std::vector action_values; int64_t image_shape[4] = {0, 0, 0, 3}; int64_t noise_shape[2] = {0, 0}; @@ -119,18 +128,21 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, } const auto* views = static_cast(data); const uint64_t n = bytes / sizeof(frt_image_view); - std::vector frames; - frames.reserve(n); + if (n != a->vision_frames.size()) { + a->last_error = "image view count does not match the runtime"; + return -4; + } for (uint64_t i = 0; i < n; ++i) { const frt_image_view& in = views[i]; if (in.struct_size < sizeof(frt_image_view) || !in.data) { a->last_error = "invalid image view"; return -1; } - flashrt::modalities::VisionFrame f; - /* generic views carry no names: positional, declared order */ - f.name = i < a->view_order.size() ? a->view_order[i] - : "view" + std::to_string(i); + if (in.pixel_format > FRT_RT_PIXEL_GRAY8) { + a->last_error = "image pixel format is invalid"; + return -4; + } + auto& f = a->vision_frames[static_cast(i)]; f.image.data = const_cast(in.data); f.image.bytes = in.bytes; f.image.dtype = flashrt::modalities::DType::kUInt8; @@ -144,9 +156,8 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, f.height = in.height; f.stride_bytes = in.stride_bytes; f.timestamp_ns = in.timestamp_ns; - frames.push_back(std::move(f)); } - Status st = a->runtime->prepare_vision(frames); + Status st = a->runtime->prepare_vision(a->vision_frames); if (!st.ok_status()) { a->last_error = st.message; return status_code(st); @@ -159,6 +170,69 @@ int set_input(void* self, uint32_t port, const void* data, uint64_t bytes, "noise is a SWAP port: write its buffer window directly"; return -3; } + if (port == a->prompt_port) { + if (!data && bytes) { + a->last_error = "prompt payload is null"; + return -1; + } + if (bytes > a->prompt_text_limit) { + a->last_error = "prompt payload exceeds the hot-path capacity"; + return -4; + } + const char* begin = static_cast(data); + if (bytes) { + a->prompt_text.assign(begin, begin + bytes); + } else { + a->prompt_text.clear(); + } + a->has_prompt_text = true; + const int rc = a->has_state + ? a->runtime->set_prompt_state( + a->prompt_text.c_str(), + a->state_values.data(), + a->state_values.size()) + : a->runtime->set_prompt(a->prompt_text.c_str()); + if (rc != 0) { + const auto& st = a->runtime->prompt_status(); + a->last_error = st.message.empty() + ? "prompt staging is not configured" + : st.message; + return status_code(st); + } + a->last_error.clear(); + return 0; + } + if (port == a->state_port) { + if (!data || !bytes || bytes % sizeof(float)) { + a->last_error = "state payload must be f32[]"; + return -1; + } + const uint64_t n = bytes / sizeof(float); + if (n != a->state_values.size()) { + a->last_error = "state dimension does not match the runtime"; + return -4; + } + const auto* values = static_cast(data); + std::memcpy(a->state_values.data(), values, + static_cast(bytes)); + a->has_state = true; + if (!a->has_prompt_text) { + a->last_error.clear(); + return 0; + } + const int rc = a->runtime->set_prompt_state( + a->prompt_text.c_str(), a->state_values.data(), + a->state_values.size()); + if (rc != 0) { + const auto& st = a->runtime->prompt_status(); + a->last_error = st.message.empty() + ? "state staging failed" + : st.message; + return status_code(st); + } + a->last_error.clear(); + return 0; + } a->last_error = "unknown or non-input port"; return -1; } @@ -172,29 +246,40 @@ int get_output(void* self, uint32_t port, void* out, uint64_t capacity, a->last_error = "unknown or non-output port"; return -1; } - std::vector actions; - Status st = a->runtime->read_actions(&actions); + Status st = a->runtime->read_actions(&a->action_values); if (!st.ok_status()) { a->last_error = st.message; return status_code(st); } - const uint64_t need = actions.size() * sizeof(float); + const uint64_t need = a->action_values.size() * sizeof(float); if (written) *written = need; if (capacity < need) { a->last_error = "action output buffer is too small"; return -5; } - std::memcpy(out, actions.data(), need); + std::memcpy(out, a->action_values.data(), need); a->last_error.clear(); return 0; } int prepare(void* self, uint32_t graph, frt_shape_key key) { - (void)graph; - (void)key; auto* a = static_cast(self); - if (a) a->last_error = "adopted-export Pi05 runtime has fixed variants"; - return -3; + if (!a) return -1; + if (!a->source_model || !a->source_model->exp) { + a->last_error = "Pi05 fixed graph variants are captured at setup"; + return -3; + } + const frt_runtime_export_v1* exp = a->source_model->exp; + if (graph >= exp->n_graphs) { + a->last_error = "Pi05 prepare graph index is out of range"; + return -2; + } + if (!frt_graph_has_variant(exp->graphs[graph].handle, key)) { + a->last_error = "Pi05 fixed graph variant was not captured"; + return -2; + } + a->last_error.clear(); + return 0; } int step(void* self) { @@ -306,6 +391,12 @@ extern "C" int frt_pi05_model_runtime_create( const auto& manifest = a->runtime->manifest(); a->view_order = manifest.vision.view_order; + a->vision_frames.resize(a->view_order.size()); + for (std::size_t i = 0; i < a->view_order.size(); ++i) { + a->vision_frames[i].name = a->view_order[i]; + } + a->action_values.resize(static_cast( + manifest.action.chunk * manifest.action.robot_dim)); a->image_shape[0] = static_cast(a->view_order.size()); a->image_shape[1] = manifest.vision.target_height; a->image_shape[2] = manifest.vision.target_width; @@ -335,11 +426,12 @@ extern "C" int frt_pi05_model_runtime_create( FRT_RT_PORT_SWAP, 0, a->noise_shape, 2, 0, action_buf ? action_buf->handle : nullptr, 0, action_buf ? action_buf->bytes : 0}; - ports[kPortActions] = {"actions", FRT_RT_MOD_ACTION, io_dtype, + ports[kPortActions] = {"actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, a->action_shape, 2, 0, - action_buf ? action_buf->handle : nullptr, 0, - action_buf ? action_buf->bytes : 0}; + nullptr, 0, + static_cast(manifest.action.chunk) * + manifest.action.robot_dim * sizeof(float)}; const std::string graph_name = config && config->graph_name ? config->graph_name : "infer"; @@ -393,10 +485,14 @@ extern "C" int frt_pi05_model_runtime_create_over( const uint32_t images = find_port_index(model, "images"); const uint32_t noise = find_port_index(model, "noise"); const uint32_t actions = find_port_index(model, "actions"); + const uint32_t actions_raw = find_port_index(model, "actions_raw"); + const uint32_t prompt = find_port_index(model, "prompt"); + const uint32_t state = find_port_index(model, "state"); if (!compatible_port(model, images, FRT_RT_MOD_IMAGE, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED) || !compatible_port(model, actions, FRT_RT_MOD_ACTION, FRT_RT_PORT_OUT, - FRT_RT_PORT_STAGED)) { + FRT_RT_PORT_STAGED) || + model->ports[actions].dtype != FRT_RT_DTYPE_F32) { return -2; } if (noise != kNoPort && @@ -404,6 +500,21 @@ extern "C" int frt_pi05_model_runtime_create_over( FRT_RT_PORT_SWAP)) { return -2; } + if (actions_raw != kNoPort && + !compatible_port(model, actions_raw, FRT_RT_MOD_TENSOR, + FRT_RT_PORT_OUT, FRT_RT_PORT_SWAP)) { + return -2; + } + if (prompt != kNoPort && + !compatible_port(model, prompt, FRT_RT_MOD_TEXT, FRT_RT_PORT_IN, + FRT_RT_PORT_STAGED)) { + return -2; + } + if (state != kNoPort && + !compatible_port(model, state, FRT_RT_MOD_STATE, FRT_RT_PORT_IN, + FRT_RT_PORT_STAGED)) { + return -2; + } auto cfg = flashrt::models::pi05::cface::make_config(config); if (model->n_stages) { @@ -412,7 +523,8 @@ extern "C" int frt_pi05_model_runtime_create_over( cfg.graph_name = graph->name; } cfg.image_input_override = tensor_from_port(model->ports[images]); - cfg.action_output_override = tensor_from_port(model->ports[actions]); + cfg.action_output_override = tensor_from_port( + model->ports[actions_raw != kNoPort ? actions_raw : actions]); auto a = std::unique_ptr(new (std::nothrow) Adapter()); if (!a) return -5; @@ -420,11 +532,36 @@ extern "C" int frt_pi05_model_runtime_create_over( flashrt::models::pi05::Runtime(model->exp, cfg)); if (!a->runtime) return -5; if (!a->runtime->ok()) return status_code(a->runtime->status()); + if (prompt != kNoPort && !a->runtime->prompt_staging_enabled()) { + return -2; + } + if (state != kNoPort && + (!a->runtime->prompt_staging_enabled() || + !a->runtime->state_normalization_enabled())) { + return -2; + } a->source_model = model; a->images_port = images; a->noise_port = noise; a->actions_port = actions; + a->prompt_port = prompt; + a->state_port = state; a->view_order = a->runtime->manifest().vision.view_order; + a->vision_frames.resize(a->view_order.size()); + for (std::size_t i = 0; i < a->view_order.size(); ++i) { + a->vision_frames[i].name = a->view_order[i]; + } + const auto& action = a->runtime->manifest().action; + a->action_values.resize(static_cast(action.chunk * + action.robot_dim)); + if (cfg.prompt_max_tokens) { + a->prompt_text_limit = static_cast( + cfg.prompt_max_tokens * 8ull); + a->prompt_text.reserve(a->prompt_text_limit); + } + if (state != kNoPort) { + a->state_values.resize(cfg.state_q01.size()); + } frt_model_runtime_verbs verbs{}; verbs.struct_size = sizeof(verbs); diff --git a/cpp/models/pi05/src/service/native_calibration_c_api.cpp b/cpp/models/pi05/src/service/native_calibration_c_api.cpp new file mode 100644 index 00000000..6757b7ac --- /dev/null +++ b/cpp/models/pi05/src/service/native_calibration_c_api.cpp @@ -0,0 +1,289 @@ +#include "flashrt/cpp/models/pi05/c_api.h" + +#include "config_map.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/backend/native_calibration_session.h" +#if defined(FLASHRT_CPP_WITH_FA2) +#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/backends/sm110/native_thor_calibration_session.h" +#endif +#include +#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_HAS_NATIVE_CALIBRATION) + std::unique_ptr impl; +#endif + std::string last_error; + std::vector view_names; + std::vector frames; +}; + +namespace { + +#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; + 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_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); + 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::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; + 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); + + 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; + 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); + } + 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 a native FP8 backend 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_HAS_NATIVE_CALIBRATION) + 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_HAS_NATIVE_CALIBRATION) + 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_HAS_NATIVE_CALIBRATION) + 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/service/native_model_runtime.cpp b/cpp/models/pi05/src/service/native_model_runtime.cpp new file mode 100644 index 00000000..6f64ea26 --- /dev/null +++ b/cpp/models/pi05/src/service/native_model_runtime.cpp @@ -0,0 +1,516 @@ +#include "native_open_internal.h" + +#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/support/native_calibration.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/session.h" +#endif +#if defined(FLASHRT_CPP_WITH_THOR_FP8) +#include "flashrt/cpp/models/pi05/backends/sm110/session.h" +#endif + +#include + +#include +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +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); + if (!graph || prompt_len > static_cast(INT_MAX)) return -1; + return cface::status_code( + graph->set_prompt_length(static_cast(prompt_len))); +} + +bool add_identity(frt_runtime_builder builder, const char* key, + const std::string& value) { + return frt_runtime_builder_add_identity(builder, key, value.c_str()) == 0; +} + +int unpublished_set_input(void*, uint32_t, const void*, uint64_t, int) { + return -3; +} +int unpublished_get_output(void*, uint32_t, void*, uint64_t, uint64_t*, int) { + return -3; +} + +frt_model_runtime_verbs unpublished_verbs() { + frt_model_runtime_verbs verbs{}; + verbs.struct_size = sizeof(verbs); + verbs.set_input = unpublished_set_input; + verbs.get_output = unpublished_get_output; + return verbs; +} + +int fail_builder(frt_runtime_builder builder, std::string* error, + const char* message) { + frt_model_runtime_verbs discard_verbs = unpublished_verbs(); + frt_model_runtime_v1* discarded = frt_runtime_builder_finish_model( + builder, &discard_verbs, nullptr, nullptr, nullptr, nullptr); + if (discarded) discarded->release(discarded->owner); + if (error) *error = message; + return -6; +} + +} // namespace + +int build_native_model_runtime(const NativeOpenConfig& config, + frt_model_runtime_v1** out, + std::string* error) { + if (!out) return -1; + *out = nullptr; + int device = 0; + cudaDeviceProp properties{}; + cudaError_t cuda_rc = cudaGetDevice(&device); + if (cuda_rc == cudaSuccess) { + cuda_rc = cudaGetDeviceProperties(&properties, device); + } + if (cuda_rc != cudaSuccess) { + if (error) *error = cudaGetErrorString(cuda_rc); + return -6; + } + 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 = config.calibration_path.empty() + ? Precision::kBf16 + : Precision::kFp8E4M3Fn; + } 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.major == 12) && + properties.minor == 0)) { + if (error) *error = "Pi0.5 native FP8 requires SM110 or SM120"; + 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 && 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; + 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.tokenizer_model_path, &tokenizer_sha256, + &hash_error)) { + if (error) *error = hash_error; + return -2; + } + + NativeCalibrationArtifact calibration; + std::string calibration_sha256; + if (precision == Precision::kFp8E4M3Fn) { + 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.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 || + 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; + } + } + + 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 + ? BackendPrecision::kFp8E4M3 + : BackendPrecision::kBf16; + modalities::Status st; + std::unique_ptr graph; + 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 = rtx_fp8 + ? Sm120BackendSession::create( + config.checkpoint_path, graph_config, calibration, + &st) + : Sm120BackendSession::create( + config.checkpoint_path, graph_config, &st); +#endif + } else if (thor_fp8) { +#if defined(FLASHRT_CPP_WITH_THOR_FP8) + graph = Sm110BackendSession::create( + config.checkpoint_path, graph_config, calibration, &st); +#endif + } + if (!graph) { + 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; + } + if (precision == Precision::kFp8E4M3Fn && + calibration.weights_sha256 != weights_sha256.digest) { + if (error) *error = "Pi0.5 calibration checkpoint digest mismatch"; + return -2; + } + + const BackendArtifacts& 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) { + 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, backend_graph_name(GraphKind::kInfer), + graph->graph(GraphKind::kInfer), 0, keys, 1, + graph->stream_id()) == 0 && + frt_runtime_builder_add_graph( + builder, backend_graph_name(GraphKind::kDecodeOnly), + graph->graph(GraphKind::kDecodeOnly), 0, keys, 1, + graph->stream_id()) == 0 && + frt_runtime_builder_add_graph( + 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, + 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"); + + 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") && + add_identity(builder, "pipeline", pipeline_id) && + add_identity(builder, "hardware", hardware_id) && + 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") && + 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", + 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", + std::to_string(kModelActionDim)) && + add_identity(builder, "robot_action_dim", + std::to_string(config.action_q01.size())); + if (ok && 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\"," + << "\"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"); + } + + const int64_t prompt_shape[] = {-1}; + const int64_t state_shape[] = {config.state_dim}; + 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 = + 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, + 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, 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, 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 && + frt_runtime_builder_add_port( + builder, "actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, + action_shape, 2, 0, nullptr, 0, action_bytes) == 0 && + frt_runtime_builder_add_port( + builder, "actions_raw", FRT_RT_MOD_TENSOR, 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; + 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"); + + 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_backend_session); + if (!base) { + delete raw_graph; + if (error) *error = "native integrated runtime finish failed"; + return -6; + } + + std::vector action_mean(config.action_q01.size()); + std::vector action_stddev(config.action_q01.size()); + for (std::size_t i = 0; i < action_mean.size(); ++i) { + action_stddev[i] = + (config.action_q99[i] - config.action_q01[i] + 1e-6f) * 0.5f; + action_mean[i] = config.action_q01[i] + action_stddev[i]; + } + frt_pi05_runtime_config runtime_config{}; + runtime_config.struct_size = sizeof(runtime_config); + runtime_config.num_views = config.num_views; + runtime_config.chunk = config.chunk; + runtime_config.model_action_dim = 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(); + 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"; + 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 = runtime_dtype; + runtime_config.prompt_embedding_vocab_size = embedding->shape[0]; + 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(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(); + 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 graph backend and SentencePiece are unavailable"; + } + return -3; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif diff --git a/cpp/models/pi05/src/service/native_open.cpp b/cpp/models/pi05/src/service/native_open.cpp new file mode 100644 index 00000000..d73050a6 --- /dev/null +++ b/cpp/models/pi05/src/service/native_open.cpp @@ -0,0 +1,606 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/loader/safetensors.h" +#include "flashrt/cpp/models/pi05/support/native_weights.h" +#include "flashrt/cpp/native/config_object.h" +#include "native_open_internal.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +thread_local std::string g_last_error; + +bool path_exists(const std::string& path) { + struct stat st {}; + return !path.empty() && ::stat(path.c_str(), &st) == 0; +} + +bool regular_file_exists(const std::string& path) { + struct stat st {}; + return !path.empty() && ::stat(path.c_str(), &st) == 0 && + S_ISREG(st.st_mode); +} + +std::string join_path(const std::string& dir, const char* leaf) { + if (dir.empty() || dir.back() == '/') return dir + leaf; + return dir + "/" + leaf; +} + +std::string quoted_key(const std::string& key) { + std::string out = "\""; + for (char c : key) { + if (c == '"' || c == '\\') out.push_back('\\'); + out.push_back(c); + } + out.push_back('"'); + return out; +} + +bool object_for_key(const std::string& json, + const std::string& key, + std::string* object) { + const std::string q = quoted_key(key); + size_t pos = json.find(q); + while (pos != std::string::npos) { + size_t p = pos + q.size(); + while (p < json.size() && + std::isspace(static_cast(json[p]))) { + ++p; + } + if (p < json.size() && json[p] == ':') { + ++p; + while (p < json.size() && + std::isspace(static_cast(json[p]))) { + ++p; + } + if (p < json.size() && json[p] == '{') { + int depth = 0; + bool in_string = false; + bool escaped = false; + for (size_t i = p; i < json.size(); ++i) { + const char c = json[i]; + if (in_string) { + if (escaped) { + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + in_string = false; + } + continue; + } + if (c == '"') { + in_string = true; + } else if (c == '{') { + ++depth; + } else if (c == '}') { + --depth; + if (depth == 0) { + if (object) *object = json.substr(p, i - p + 1); + return true; + } + } + } + } + } + pos = json.find(q, pos + 1); + } + return false; +} + +bool parse_f64_array_property(const std::string& object, + const char* name, + std::vector* out) { + const std::string q = quoted_key(name); + size_t p = object.find(q); + if (p == std::string::npos) return false; + p += q.size(); + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p >= object.size() || object[p++] != ':') return false; + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p >= object.size() || object[p++] != '[') return false; + std::vector values; + while (p < object.size()) { + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p < object.size() && object[p] == ']') { + ++p; + if (out) *out = std::move(values); + return true; + } + errno = 0; + char* end = nullptr; + const double value = std::strtod(object.c_str() + p, &end); + if (errno || end == object.c_str() + p) return false; + values.push_back(value); + p = static_cast(end - object.c_str()); + while (p < object.size() && + std::isspace(static_cast(object[p]))) ++p; + if (p < object.size() && object[p] == ',') { + ++p; + continue; + } + if (p < object.size() && object[p] == ']') continue; + return false; + } + return false; +} + +bool read_safetensors_f32_vector(const std::string& path, + const char* key, + std::vector* out) { + if (!out) return false; + flashrt::loader::SafetensorsFile file; + if (!file.open(path)) { + g_last_error = file.error(); + return false; + } + const auto* tensor = file.find(key); + if (!tensor || tensor->dtype != "F32" || tensor->shape.size() != 1) { + g_last_error = "safetensors F32 vector metadata is invalid"; + return false; + } + const uint64_t n = tensor->shape[0]; + if (n > (1ull << 20)) { + g_last_error = "safetensors vector is too large"; + return false; + } + std::vector values(static_cast(n)); + std::memcpy(values.data(), file.data(*tensor), + static_cast(tensor->bytes)); + *out = std::move(values); + return true; +} + +bool sane_quantile_pair(const std::vector& q01, + const std::vector& q99) { + if (q01.empty() || q01.size() != q99.size()) return false; + for (size_t i = 0; i < q01.size(); ++i) { + if (!std::isfinite(q01[i]) || !std::isfinite(q99[i]) || + q99[i] <= q01[i]) { + return false; + } + } + return true; +} + +bool sane_quantile_pair(const std::vector& q01, + const std::vector& q99) { + if (q01.empty() || q01.size() != q99.size()) return false; + for (size_t i = 0; i < q01.size(); ++i) { + if (!std::isfinite(q01[i]) || !std::isfinite(q99[i]) || + q99[i] <= q01[i]) { + return false; + } + } + return true; +} + +bool read_text_file(const std::string& path, std::string* out) { + if (!out) return false; + std::ifstream f(path); + if (!f) return false; + out->assign((std::istreambuf_iterator(f)), + std::istreambuf_iterator()); + return f.good() || f.eof(); +} + +std::string dirname(const std::string& path) { + const size_t p = path.find_last_of('/'); + if (p == std::string::npos) return "."; + if (p == 0) return "/"; + return path.substr(0, p); +} + +bool norm_block_values(const std::string& json, + const char* block_name, + std::vector* q01_out, + std::vector* q99_out) { + std::string block; + if (!object_for_key(json, block_name, &block)) return false; + std::vector q01; + std::vector q99; + if (!parse_f64_array_property(block, "q01", &q01) || + !parse_f64_array_property(block, "q99", &q99) || + !sane_quantile_pair(q01, q99)) { + return false; + } + if (q01_out) q01_out->assign(q01.begin(), q01.end()); + if (q99_out) q99_out->assign(q99.begin(), q99.end()); + return true; +} + +bool validate_norm_stats_file(const std::string& path, + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* config) { + std::string json; + if (!read_text_file(path, &json)) return false; + std::vector action_q01; + std::vector action_q99; + std::vector state_q01; + std::vector state_q99; + if (!norm_block_values(json, "actions", &action_q01, &action_q99) || + !norm_block_values(json, "state", &state_q01, &state_q99)) { + g_last_error = "norm_stats.json is missing actions/state q01/q99"; + return false; + } + if (action_q01.empty() || action_q01.size() > 32) { + g_last_error = "norm_stats action dimension is invalid"; + return false; + } + if (state_q01.size() != static_cast(state_dim)) { + g_last_error = "norm_stats state dimension does not match config"; + return false; + } + if (config) { + config->state_q01 = std::move(state_q01); + config->state_q99 = std::move(state_q99); + config->action_q01 = std::move(action_q01); + config->action_q99 = std::move(action_q99); + } + g_last_error.clear(); + return true; +} + +bool has_prefix(const std::string& s, const char* prefix) { + const size_t n = std::strlen(prefix); + return s.size() >= n && s.compare(0, n, prefix) == 0; +} + +bool has_suffix(const std::string& s, const char* suffix) { + const size_t n = std::strlen(suffix); + return s.size() >= n && s.compare(s.size() - n, n, suffix) == 0; +} + +std::string find_child(const std::string& dir, + const char* prefix, + const char* suffix) { + DIR* d = ::opendir(dir.c_str()); + if (!d) return ""; + std::string found; + while (dirent* ent = ::readdir(d)) { + const std::string name = ent->d_name; + if (has_prefix(name, prefix) && has_suffix(name, suffix)) { + found = join_path(dir, name.c_str()); + break; + } + } + ::closedir(d); + return found; +} + +bool validate_lerobot_policy_norm_stats(const std::string& checkpoint_path, + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* + config) { + const std::string pre = find_child( + checkpoint_path, "policy_preprocessor_step_", + "_normalizer_processor.safetensors"); + const std::string post = find_child( + checkpoint_path, "policy_postprocessor_step_", + "_unnormalizer_processor.safetensors"); + if (pre.empty() || post.empty()) return false; + + std::vector state_q01; + std::vector state_q99; + std::vector action_q01; + std::vector action_q99; + if (!read_safetensors_f32_vector(pre, "observation.state.q01", + &state_q01) || + !read_safetensors_f32_vector(pre, "observation.state.q99", + &state_q99) || + !read_safetensors_f32_vector(post, "action.q01", &action_q01) || + !read_safetensors_f32_vector(post, "action.q99", &action_q99)) { + g_last_error = + "lerobot policy stats are missing action/state q01/q99"; + return false; + } + if (state_q01.size() != static_cast(state_dim) || + !sane_quantile_pair(state_q01, state_q99)) { + g_last_error = + "lerobot policy state dimension does not match config"; + return false; + } + if (action_q01.size() > 32 || + !sane_quantile_pair(action_q01, action_q99)) { + g_last_error = "lerobot policy action dimension is invalid"; + return false; + } + if (config) { + config->state_q01 = std::move(state_q01); + config->state_q99 = std::move(state_q99); + config->action_q01 = std::move(action_q01); + config->action_q99 = std::move(action_q99); + } + g_last_error.clear(); + return true; +} + +bool validate_norm_stats(const std::string& checkpoint_path, + int64_t state_dim, + flashrt::models::pi05::NativeOpenConfig* config) { + const std::string parent = dirname(checkpoint_path); + const std::string candidates[] = { + join_path(checkpoint_path, + "assets/physical-intelligence/libero/norm_stats.json"), + join_path(checkpoint_path, "assets/droid/norm_stats.json"), + join_path(checkpoint_path, "norm_stats.json"), + join_path(parent, + "pi05_libero/assets/physical-intelligence/libero/" + "norm_stats.json"), + join_path(parent, "pi05_droid/assets/droid/norm_stats.json"), + join_path(parent, "pi05_droid_pytorch/assets/droid/norm_stats.json"), + }; + bool saw_malformed = false; + std::string malformed_error; + for (const std::string& path : candidates) { + if (!regular_file_exists(path)) continue; + if (validate_norm_stats_file(path, state_dim, config)) return true; + saw_malformed = true; + malformed_error = g_last_error; + } + if (validate_lerobot_policy_norm_stats(checkpoint_path, state_dim, + config)) { + return true; + } + g_last_error = saw_malformed + ? malformed_error + : "norm_stats.json not found for Pi0.5 native_v2"; + return false; +} + +bool validate_pi05_safetensors(const std::string& checkpoint_path) { + const std::string path = join_path(checkpoint_path, "model.safetensors"); + if (!regular_file_exists(path)) { + g_last_error = "checkpoint_path must contain model.safetensors"; + return false; + } + flashrt::loader::SafetensorsFile file; + if (!file.open(path)) { + g_last_error = file.error(); + return false; + } + + for (const auto& req : + flashrt::models::pi05::native_tensor_requirements()) { + std::string key = req.key; + const flashrt::loader::SafetensorInfo* meta = file.find(key); + if (!meta) { + key = std::string("model.") + req.key; + meta = file.find(key); + if (!meta) { + g_last_error = std::string("model.safetensors is missing ") + + req.key; + return false; + } + } + if (meta->dtype != "BF16" && meta->dtype != "F16" && + meta->dtype != "F32") { + g_last_error = std::string("Pi0.5 tensor dtype is unsupported: ") + + req.key; + return false; + } + if (meta->shape != req.shape) { + g_last_error = std::string("Pi0.5 tensor shape mismatch: ") + + req.key; + return false; + } + } + g_last_error.clear(); + return true; +} + +int validate_config( + const char* config_json, + flashrt::models::pi05::NativeOpenConfig* parsed) { + if (!config_json) { + g_last_error = "config_json is null"; + return -1; + } + flashrt::native::ConfigObject config_object; + if (!config_object.parse(config_json)) { + g_last_error = config_object.error(); + return -1; + } + + std::string io; + std::string checkpoint_path; + std::string tokenizer_model_path; + std::string state_prompt_mode; + std::string precision; + std::string calibration_path; + std::string stage_plan; + 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") { + 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 (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 (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; + } + if (!regular_file_exists(tokenizer_model_path)) { + 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; + } + + int64_t max_prompt_tokens = 0; + int64_t state_dim = 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; + int64_t max_frame_height = 720; + 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) { + g_last_error = "max_prompt_tokens must be in [1, INT_MAX]"; + return -1; + } + if (state_dim <= 0 || state_dim > INT_MAX) { + g_last_error = "state_dim must be in [1, INT_MAX]"; + return -1; + } + if (num_views < 1 || num_views > 3) { + g_last_error = "num_views must be in [1, 3]"; + return -1; + } + if (chunk <= 0 || chunk > INT_MAX) { + g_last_error = "chunk must be in [1, INT_MAX]"; + return -1; + } + if (num_steps <= 0 || num_steps > INT_MAX) { + g_last_error = "num_steps must be in [1, INT_MAX]"; + return -1; + } + if (vision_pool_factor != 1 && vision_pool_factor != 2 && + vision_pool_factor != 4) { + g_last_error = "vision_pool_factor must be one of 1, 2, or 4"; + return -1; + } + 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.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); + 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); + config.max_frame_height = static_cast(max_frame_height); + if (!validate_norm_stats(checkpoint_path, state_dim, &config)) { + return -2; + } + + if (parsed) *parsed = std::move(config); + g_last_error.clear(); + return 0; +} + +} // 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) try { + if (!out) { + g_last_error = "out is null"; + return -1; + } + *out = nullptr; + flashrt::models::pi05::NativeOpenConfig 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_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 " + "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() { + return g_last_error.c_str(); +} diff --git a/cpp/models/pi05/src/service/native_open_internal.h b/cpp/models/pi05/src/service/native_open_internal.h new file mode 100644 index 00000000..60fff664 --- /dev/null +++ b/cpp/models/pi05/src/service/native_open_internal.h @@ -0,0 +1,45 @@ +#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; + 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; + 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); + +} // namespace pi05 +} // namespace models +} // namespace flashrt + +#endif // FLASHRT_CPP_MODELS_PI05_NATIVE_OPEN_INTERNAL_H diff --git a/cpp/models/pi05/src/support/native_calibration.cpp b/cpp/models/pi05/src/support/native_calibration.cpp new file mode 100644 index 00000000..e32b96a2 --- /dev/null +++ b/cpp/models/pi05/src/support/native_calibration.cpp @@ -0,0 +1,425 @@ +#include "flashrt/cpp/models/pi05/support/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 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 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) { + 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", artifact.activation_dtype == kBfloat16 + ? kSchemaV2 + : kSchemaV1}, + {"model", kModel}, + {"precision", kPrecision}, + {"tensor_dtype", artifact.activation_dtype}, + {"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 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 = + 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 << '}'; + 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\":[" << vision_bytes << ',' + << vision_bytes + encoder_bytes + << "]},\"decoder_scales\":{\"dtype\":\"F32\",\"shape\":[" + << artifact.decoder_scales.size() << "],\"data_offsets\":[" + << vision_bytes + encoder_bytes << ',' + << vision_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.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 || + 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"); + } + 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) || + !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.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(), + 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(); + 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, "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"); + 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) || + (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"); + } + 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/support/native_device_weights.cpp b/cpp/models/pi05/src/support/native_device_weights.cpp new file mode 100644 index 00000000..6c0d4e28 --- /dev/null +++ b/cpp/models/pi05/src/support/native_device_weights.cpp @@ -0,0 +1,162 @@ +#include "flashrt/cpp/models/pi05/support/native_device_weights.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +bool element_count(const std::vector& shape, + std::size_t* out) { + std::size_t count = 1; + for (std::uint64_t dim : shape) { + if (dim > std::numeric_limits::max() || + (dim && count > std::numeric_limits::max() / + static_cast(dim))) { + return false; + } + count *= static_cast(dim); + } + if (out) *out = count; + return true; +} + +std::size_t element_bytes(NativeWeightDType dtype) { + switch (dtype) { + case NativeWeightDType::kBf16: return sizeof(std::uint16_t); + case NativeWeightDType::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); + } + 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( + 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"); + } + std::size_t elements = 0; + const std::size_t width = element_bytes(dtype); + if (!data || !width || !element_count(shape, &elements) || + elements > std::numeric_limits::max() / width || + elements * width != bytes) { + return invalid("device weight shape does not match typed payload"); + } + if (!bytes) return invalid("device weight payload is empty"); + +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + (void)data; + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "device weight upload requires the CUDA build"); +#else + frt_buffer buffer = frt_buffer_alloc(ctx_, name.c_str(), bytes); + if (!buffer) { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + cudaMemGetInfo(&free_bytes, &total_bytes); + std::ostringstream message; + message << "device weight allocation failed: " << name + << " requested=" << bytes << " free=" << free_bytes + << " total=" << total_bytes; + return modalities::Status::error(modalities::StatusCode::kBackend, + message.str()); + } + const cudaError_t rc = cudaMemcpy(frt_buffer_dptr(buffer), + data, bytes, + cudaMemcpyHostToDevice); + if (rc != cudaSuccess) { + return modalities::Status::error( + modalities::StatusCode::kBackend, + std::string("device weight upload failed: ") + + cudaGetErrorString(rc)); + } + weights_.emplace(name, NativeDeviceWeight{buffer, shape, dtype}); + return modalities::Status::ok(); +#endif +} + +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); + return it == weights_.end() ? nullptr : &it->second; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/support/native_quantization.cu b/cpp/models/pi05/src/support/native_quantization.cu new file mode 100644 index 00000000..cb4eba3b --- /dev/null +++ b/cpp/models/pi05/src/support/native_quantization.cu @@ -0,0 +1,167 @@ +#include "flashrt/cpp/models/pi05/support/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; +} + +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( + 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); + // 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] * 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) { + 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/support/native_quantization_unavailable.cpp b/cpp/models/pi05/src/support/native_quantization_unavailable.cpp new file mode 100644 index 00000000..cec2fb4f --- /dev/null +++ b/cpp/models/pi05/src/support/native_quantization_unavailable.cpp @@ -0,0 +1,40 @@ +#include "flashrt/cpp/models/pi05/support/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_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*) { + return unavailable(); +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/support/native_rope.cu b/cpp/models/pi05/src/support/native_rope.cu new file mode 100644 index 00000000..e9a926c0 --- /dev/null +++ b/cpp/models/pi05/src/support/native_rope.cu @@ -0,0 +1,62 @@ +#include "flashrt/cpp/models/pi05/support/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/support/native_rope_unavailable.cpp b/cpp/models/pi05/src/support/native_rope_unavailable.cpp new file mode 100644 index 00000000..ec43576e --- /dev/null +++ b/cpp/models/pi05/src/support/native_rope_unavailable.cpp @@ -0,0 +1,16 @@ +#include "flashrt/cpp/models/pi05/support/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/support/native_weight_materializer.cpp b/cpp/models/pi05/src/support/native_weight_materializer.cpp new file mode 100644 index 00000000..21dd765f --- /dev/null +++ b/cpp/models/pi05/src/support/native_weight_materializer.cpp @@ -0,0 +1,439 @@ +#include "flashrt/cpp/models/pi05/support/native_weight_materializer.h" + +#include +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +std::string encoder_prefix(int layer) { + return "paligemma_with_expert.paligemma.model.language_model.layers." + + std::to_string(layer); +} + +std::string decoder_prefix(int layer) { + return "paligemma_with_expert.gemma_expert.model.layers." + + std::to_string(layer); +} + +const std::string& vision_prefix() { + static const std::string prefix = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; + return prefix; +} + +std::string layer_name(const char* stem, int layer) { + return std::string(stem) + std::to_string(layer); +} + +} // namespace + +modalities::Status NativeWeightMaterializer::load( + const std::string& key, + NativeFloatTensor* out) { + return load_native_float_tensor(source_, key, out); +} + +modalities::Status NativeWeightMaterializer::upload( + const std::string& name, + const NativeFloatTensor& tensor) { + if (!destination_) return invalid("native weight destination is null"); + NativeBf16Tensor bf16; + modalities::Status st = native_to_bf16(tensor, &bf16); + if (!st.ok_status()) return st; + return destination_->upload(name, bf16); +} + +modalities::Status NativeWeightMaterializer::upload_rounded_transpose( + const std::string& source_key, + const std::string& destination_name) { + 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_source_to_bf16(source, true, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); +} + +modalities::Status NativeWeightMaterializer::upload_rounded_copy( + const std::string& source_key, + const std::string& destination_name) { + 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_source_to_bf16(source, false, &converted); + if (!st.ok_status()) return st; + 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) { + 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_source_fold_rms_columns_transpose(source, norm, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); +} + +modalities::Status NativeWeightMaterializer::upload_rounded_scaled( + const std::string& source_key, + const std::string& destination_name, + float scale, + bool transpose) { + 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_source_round_scale_to_bf16( + source, scale, transpose, &converted); + if (!st.ok_status()) return st; + return destination_->upload(destination_name, converted); +} + +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; + + 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 = 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; + st = native_source_qkv_to_bf16(q, k, v, 8, 1, &norm, &qkv); + if (!st.ok_status()) return st; + st = destination_->upload(layer_name("encoder_attn_qkv_w_", layer), qkv); + if (!st.ok_status()) return st; + + st = upload_rounded_transpose( + prefix + ".self_attn.o_proj.weight", + layer_name("encoder_attn_o_w_", layer)); + if (!st.ok_status()) return st; + + st = load(prefix + ".post_attention_layernorm.weight", &norm); + if (!st.ok_status()) return st; + st = upload_folded_transpose( + prefix + ".mlp.gate_proj.weight", norm, + layer_name("encoder_ffn_gate_w_", layer)); + if (!st.ok_status()) return st; + st = upload_folded_transpose( + prefix + ".mlp.up_proj.weight", norm, + layer_name("encoder_ffn_up_w_", layer)); + if (!st.ok_status()) return st; + return upload_rounded_transpose( + prefix + ".mlp.down_proj.weight", + layer_name("encoder_ffn_down_w_", layer)); +} + +modalities::Status NativeWeightMaterializer::materialize_decoder_layer( + int layer, + bool merge_gate_up) { + if (layer < 0 || layer >= 18 || !destination_) { + return invalid("Pi0.5 decoder layer index is invalid"); + } + const std::string prefix = decoder_prefix(layer); + 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 = 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; + st = native_source_qkv_to_bf16(q, k, v, 8, 1, nullptr, &qkv); + if (!st.ok_status()) return st; + st = destination_->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; + + st = upload_rounded_transpose( + prefix + ".mlp.gate_proj.weight", + layer_name("decoder_ffn_gate_w_", layer)); + if (!st.ok_status()) return st; + 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) { + 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 = destination_->upload( + layer_name("decoder_ffn_gate_up_w_", layer), gate_up); + if (!st.ok_status()) return st; + } + st = upload_rounded_transpose( + prefix + ".mlp.down_proj.weight", + layer_name("decoder_ffn_down_w_", layer)); + if (!st.ok_status()) return st; + + st = upload_rounded_transpose( + prefix + ".input_layernorm.dense.weight", + layer_name("decoder_pre_attn_norm_mod_w_", layer)); + if (!st.ok_status()) return st; + st = upload_rounded_copy( + prefix + ".input_layernorm.dense.bias", + layer_name("decoder_pre_attn_norm_mod_b_", layer)); + if (!st.ok_status()) return st; + st = upload_rounded_transpose( + prefix + ".post_attention_layernorm.dense.weight", + layer_name("decoder_pre_ffn_norm_mod_w_", layer)); + if (!st.ok_status()) return st; + return upload_rounded_copy( + prefix + ".post_attention_layernorm.dense.bias", + layer_name("decoder_pre_ffn_norm_mod_b_", layer)); +} + +modalities::Status NativeWeightMaterializer::materialize_vision_layer( + int layer) { + if (layer < 0 || layer >= 27 || !destination_) { + return invalid("Pi0.5 vision layer index is invalid"); + } + const std::string prefix = vision_prefix() + ".encoder.layers." + + std::to_string(layer); + 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 = 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; + st = native_source_qkv_to_bf16(q, k, v, 0, 0, nullptr, &qkv); + if (!st.ok_status()) return st; + st = destination_->upload(layer_name("vision_attn_qkv_w_", layer), qkv); + 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_bf16({&q, &k, &v}, &qkv); + if (!st.ok_status()) return st; + st = destination_->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(); + 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_source_patch_oihw_to_hwio_bf16(patch, &permuted); + if (!st.ok_status()) return st; + 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"); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".embeddings.position_embedding.weight", + "vision_position_embedding"); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".post_layernorm.weight", + "vision_final_norm_w"); + if (!st.ok_status()) return st; + st = upload_rounded_copy(prefix + ".post_layernorm.bias", + "vision_final_norm_b"); + if (!st.ok_status()) return st; + + const std::string projector = + "paligemma_with_expert.paligemma.model.multi_modal_projector.linear"; + st = upload_rounded_transpose(projector + ".weight", + "encoder_multi_modal_projector_w"); + if (!st.ok_status()) return st; + return upload_rounded_copy(projector + ".bias", + "encoder_multi_modal_projector_b"); +} + +modalities::Status NativeWeightMaterializer::materialize_decoder_globals( + int num_steps) { + if (!destination_ || num_steps <= 0) { + return invalid("Pi0.5 decoder global configuration is invalid"); + } + const struct { + const char* source; + const char* destination; + bool transpose; + } entries[] = { + {"paligemma_with_expert.gemma_expert.model.norm.dense.weight", + "decoder_final_norm_mod_w", true}, + {"paligemma_with_expert.gemma_expert.model.norm.dense.bias", + "decoder_final_norm_mod_b", false}, + {"time_mlp_in.weight", "decoder_time_mlp_in_w", true}, + {"time_mlp_in.bias", "decoder_time_mlp_in_b", false}, + {"time_mlp_out.weight", "decoder_time_mlp_out_w", true}, + {"time_mlp_out.bias", "decoder_time_mlp_out_b", false}, + {"action_in_proj.weight", "decoder_action_in_proj_w", true}, + {"action_in_proj.bias", "decoder_action_in_proj_b", false}, + }; + for (const auto& entry : entries) { + const modalities::Status st = + entry.transpose + ? upload_rounded_transpose(entry.source, entry.destination) + : upload_rounded_copy(entry.source, entry.destination); + if (!st.ok_status()) return st; + } + + NativeFloatTensor time_embeddings; + modalities::Status st = + native_pi05_time_embeddings(num_steps, 1024, &time_embeddings); + if (!st.ok_status()) return st; + st = upload("decoder_time_embeds", time_embeddings); + if (!st.ok_status()) return st; + + const float step_scale = -1.0f / static_cast(num_steps); + st = upload_rounded_scaled( + "action_out_proj.weight", "decoder_action_out_proj_w", step_scale, + true); + if (!st.ok_status()) return st; + return upload_rounded_scaled( + "action_out_proj.bias", "decoder_action_out_proj_b", step_scale, + false); +} + +modalities::Status NativeWeightMaterializer::materialize_embedding() { + if (!destination_) return invalid("native weight destination is null"); + return upload_rounded_copy( + "paligemma_with_expert.paligemma.lm_head.weight", + "embedding_weight"); +} + +modalities::Status NativeWeightMaterializer::materialize_all( + const NativeMaterializationOptions& options) { + if (!destination_ || options.num_steps <= 0) { + return invalid("Pi0.5 materialization options are invalid"); + } + 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; + report("globals"); + st = materialize_embedding(); + report("embedding"); + return st; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/support/native_weight_ops.cpp b/cpp/models/pi05/src/support/native_weight_ops.cpp new file mode 100644 index 00000000..aff9888e --- /dev/null +++ b/cpp/models/pi05/src/support/native_weight_ops.cpp @@ -0,0 +1,1204 @@ +#include "flashrt/cpp/models/pi05/support/native_weight_ops.h" + +#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); +} + +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(); +} + +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({32, 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(); +} + +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) { + const loader::SafetensorInfo* tensor = file.find(key); + if (!tensor) tensor = file.find(std::string("model.") + key); + 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]); + } + std::uint16_t f16(std::size_t index) const { + return modalities::float_to_float16(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]; } + std::uint16_t f16(std::size_t index) const { + return modalities::float_to_float16((*this)[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])); + } + std::uint16_t f16(std::size_t index) const { return 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]); + } + std::uint16_t f16(std::size_t index) const { + return modalities::float_to_float16((*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]); + } + std::uint16_t f16(std::size_t index) const { + return IsBf16 ? modalities::float_to_float16((*this)[index]) + : bits(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_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, + 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, + NativeFloatTensor* out) { + 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(source.shape, &count)) { + return invalid("native tensor shape overflows size_t"); + } + + NativeFloatTensor loaded; + loaded.shape = source.shape; + loaded.values.resize(count); + 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] = reader[i]; + } + }); + return modalities::Status::ok(); + }); + if (!st.ok_status()) return st; + *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()); + 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_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) { + if (!out || !valid_tensor(input)) { + return invalid("invalid BF16 round-trip input"); + } + NativeFloatTensor rounded = input; + 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(); +} + +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()); + 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(); +} + +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]); + 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(); +} + +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) { + const std::uint64_t input_rows = input->shape[0]; + const std::uint64_t output_offset = row_offset; + 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); + 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); + 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(); +} + +modalities::Status native_concat_vectors( + const std::vector& inputs, + NativeFloatTensor* out) { + if (!out || inputs.empty()) return invalid("vector concat has no inputs"); + std::size_t total = 0; + for (const NativeFloatTensor* input : inputs) { + if (!input || !valid_tensor(*input) || input->shape.size() != 1 || + input->values.size() > + std::numeric_limits::max() - total) { + return invalid("vector concat tensors have incompatible shapes"); + } + total += input->values.size(); + } + NativeFloatTensor joined; + joined.shape = {static_cast(total)}; + joined.values.reserve(total); + for (const NativeFloatTensor* input : inputs) { + joined.values.insert(joined.values.end(), input->values.begin(), + input->values.end()); + } + *out = std::move(joined); + return modalities::Status::ok(); +} + +modalities::Status native_scale(const NativeFloatTensor& input, + float scale, + NativeFloatTensor* out) { + if (!out || !valid_tensor(input)) return invalid("invalid scale input"); + NativeFloatTensor scaled = input; + 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(); +} + +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(); +} + +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/models/pi05/src/support/native_weights.cpp b/cpp/models/pi05/src/support/native_weights.cpp new file mode 100644 index 00000000..bf684832 --- /dev/null +++ b/cpp/models/pi05/src/support/native_weights.cpp @@ -0,0 +1,115 @@ +#include "flashrt/cpp/models/pi05/support/native_weights.h" + +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +using Requirement = NativeTensorRequirement; + +void add(std::vector* out, const std::string& key, + std::initializer_list shape) { + out->push_back(Requirement{key, shape}); +} + +std::vector build_requirements() { + std::vector out; + out.reserve(820); + + const std::string vision = + "paligemma_with_expert.paligemma.model.vision_tower.vision_model"; + add(&out, vision + ".embeddings.patch_embedding.weight", {1152, 3, 14, 14}); + add(&out, vision + ".embeddings.patch_embedding.bias", {1152}); + add(&out, vision + ".embeddings.position_embedding.weight", {256, 1152}); + for (int layer = 0; layer < 27; ++layer) { + const std::string p = vision + ".encoder.layers." + + std::to_string(layer); + for (const char* projection : {"q_proj", "k_proj", "v_proj", + "out_proj"}) { + add(&out, p + ".self_attn." + projection + ".weight", + {1152, 1152}); + add(&out, p + ".self_attn." + projection + ".bias", {1152}); + } + add(&out, p + ".mlp.fc1.weight", {4304, 1152}); + add(&out, p + ".mlp.fc1.bias", {4304}); + add(&out, p + ".mlp.fc2.weight", {1152, 4304}); + add(&out, p + ".mlp.fc2.bias", {1152}); + for (const char* norm : {"layer_norm1", "layer_norm2"}) { + add(&out, p + "." + norm + ".weight", {1152}); + add(&out, p + "." + norm + ".bias", {1152}); + } + } + add(&out, vision + ".post_layernorm.weight", {1152}); + add(&out, vision + ".post_layernorm.bias", {1152}); + + const std::string projector = + "paligemma_with_expert.paligemma.model.multi_modal_projector.linear"; + add(&out, projector + ".weight", {2048, 1152}); + add(&out, projector + ".bias", {2048}); + + const std::string encoder = + "paligemma_with_expert.paligemma.model.language_model.layers."; + for (int layer = 0; layer < 18; ++layer) { + const std::string p = encoder + std::to_string(layer); + add(&out, p + ".self_attn.q_proj.weight", {2048, 2048}); + add(&out, p + ".self_attn.k_proj.weight", {256, 2048}); + add(&out, p + ".self_attn.v_proj.weight", {256, 2048}); + add(&out, p + ".self_attn.o_proj.weight", {2048, 2048}); + add(&out, p + ".mlp.gate_proj.weight", {16384, 2048}); + add(&out, p + ".mlp.up_proj.weight", {16384, 2048}); + add(&out, p + ".mlp.down_proj.weight", {2048, 16384}); + add(&out, p + ".input_layernorm.weight", {2048}); + add(&out, p + ".post_attention_layernorm.weight", {2048}); + } + add(&out, "paligemma_with_expert.paligemma.model.language_model.norm.weight", + {2048}); + add(&out, "paligemma_with_expert.paligemma.lm_head.weight", + {257152, 2048}); + + const std::string decoder = + "paligemma_with_expert.gemma_expert.model.layers."; + for (int layer = 0; layer < 18; ++layer) { + const std::string p = decoder + std::to_string(layer); + add(&out, p + ".self_attn.q_proj.weight", {2048, 1024}); + add(&out, p + ".self_attn.k_proj.weight", {256, 1024}); + add(&out, p + ".self_attn.v_proj.weight", {256, 1024}); + add(&out, p + ".self_attn.o_proj.weight", {1024, 2048}); + add(&out, p + ".mlp.gate_proj.weight", {4096, 1024}); + add(&out, p + ".mlp.up_proj.weight", {4096, 1024}); + add(&out, p + ".mlp.down_proj.weight", {1024, 4096}); + for (const char* norm : {"input_layernorm", "post_attention_layernorm"}) { + add(&out, p + "." + norm + ".dense.weight", {3072, 1024}); + add(&out, p + "." + norm + ".dense.bias", {3072}); + } + } + add(&out, "paligemma_with_expert.gemma_expert.model.norm.dense.weight", + {3072, 1024}); + add(&out, "paligemma_with_expert.gemma_expert.model.norm.dense.bias", + {3072}); + add(&out, "paligemma_with_expert.gemma_expert.lm_head.weight", + {257152, 1024}); + + add(&out, "action_in_proj.weight", {1024, 32}); + add(&out, "action_in_proj.bias", {1024}); + add(&out, "action_out_proj.weight", {32, 1024}); + add(&out, "action_out_proj.bias", {32}); + add(&out, "time_mlp_in.weight", {1024, 1024}); + add(&out, "time_mlp_in.bias", {1024}); + add(&out, "time_mlp_out.weight", {1024, 1024}); + add(&out, "time_mlp_out.bias", {1024}); + return out; +} + +} // namespace + +const std::vector& native_tensor_requirements() { + static const std::vector requirements = + build_requirements(); + return requirements; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/src/support/native_workspace.cpp b/cpp/models/pi05/src/support/native_workspace.cpp new file mode 100644 index 00000000..2d80f06d --- /dev/null +++ b/cpp/models/pi05/src/support/native_workspace.cpp @@ -0,0 +1,550 @@ +#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 +#endif + +#include +#include +#include + +namespace flashrt { +namespace models { +namespace pi05 { +namespace { + +modalities::Status invalid(const char* message) { + return modalities::Status::error(modalities::StatusCode::kInvalidArgument, + message); +} + +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"); + 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); + if (rc != cudaSuccess) return backend("native RMS upload failed"); + } + return modalities::Status::ok(); +#endif +} + +modalities::Status NativeWorkspace::initialize_rope() { +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "native RoPE initialization requires the CUDA build"); +#else + const int max_positions = encoder_sequence_ + chunk_size_; + rope_table_.resize(static_cast(max_positions) * 256); + 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) { + 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 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::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()) { + return invalid("Pi0.5 decoder RoPE prompt length is invalid"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "decoder RoPE update requires the CUDA build"); +#else + if (!decoder_rope_buffer_) + return invalid("decoder RoPE buffer was not allocated"); + 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; + const std::size_t elements = + static_cast(chunk_size_) * 256; + if (start > rope_table_.size() || + elements > rope_table_.size() - start) { + return invalid("decoder RoPE slice exceeds the generated table"); + } + const cudaError_t rc = cudaMemcpy( + frt_buffer_dptr(decoder_rope_buffer_), rope_table_.data() + start, + elements * sizeof(std::uint16_t), cudaMemcpyHostToDevice); + return rc == cudaSuccess + ? modalities::Status::ok() + : backend("decoder RoPE upload failed"); +#endif +} + +modalities::Status NativeWorkspace::expand_vision_position_embedding( + const NativeDeviceWeightStore& weights) { + const NativeDeviceWeight* source = + weights.find("vision_position_embedding"); + const NativeWorkspaceBuffer* destination = + find("vision_pos_embed_expanded"); + 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"); + } +#ifndef FLASHRT_CPP_WITH_CUDA_STAGING + return modalities::Status::error( + modalities::StatusCode::kUnsupported, + "position embedding expansion requires the CUDA build"); +#else + const std::size_t view_bytes = 256 * 1152 * sizeof(std::uint16_t); + if (frt_buffer_bytes(destination->buffer) != + static_cast(num_views_) * view_bytes) { + return invalid("expanded position embedding buffer size is invalid"); + } + for (int view = 0; view < num_views_; ++view) { + auto* target = static_cast( + frt_buffer_dptr(destination->buffer)) + + static_cast(view) * view_bytes; + const cudaError_t rc = cudaMemcpy( + target, frt_buffer_dptr(source->buffer), view_bytes, + cudaMemcpyDeviceToDevice); + if (rc != cudaSuccess) { + return backend("vision position embedding expansion failed"); + } + } + return modalities::Status::ok(); +#endif +} + +modalities::Status NativeWorkspace::allocate( + const NativeWorkspaceConfig& config) { + if (!ctx_ || !buffers_.empty() || config.num_views < 1 || + config.num_views > 3 || config.max_prompt_tokens <= 0 || + config.max_prompt_tokens > std::numeric_limits::max() - 768 || + config.chunk_size <= 0 || config.num_steps <= 0 || + config.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.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 = + 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; + num_steps_ = config.num_steps; + flavor_ = config.flavor; + 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) + 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); + FRT_ADD("vision_x_norm", {vs, 1152}, modalities::DType::kBFloat16); + if (config.vision_pool_factor == 1) { + st = add_alias("vision_x_pooled", "vision_x", {vs_enc, 1152}); + if (!st.ok_status()) return st; + } else { + FRT_ADD("vision_x_pooled", {vs_enc, 1152}, + modalities::DType::kBFloat16); + } + FRT_ADD("vision_QKV", {vs, 3456}, modalities::DType::kBFloat16); + FRT_ADD("vision_hidden", {vs, 4304}, modalities::DType::kBFloat16); + FRT_ADD("vision_pos_embed_expanded", {vs, 1152}, + modalities::DType::kBFloat16); + FRT_ADD("vision_patches", {vs, 588}, modalities::DType::kBFloat16); + + FRT_ADD("encoder_rope_weights", {es, 256}, + modalities::DType::kBFloat16); + FRT_ADD("prompt_embedding", + {static_cast(max_prompt_tokens_), 2048}, + modalities::DType::kBFloat16); + FRT_ADD("encoder_x", {es, 2048}, modalities::DType::kBFloat16); + FRT_ADD("encoder_x_norm", {es, 2048}, modalities::DType::kBFloat16); + FRT_ADD("encoder_QKV", {es, 2560}, modalities::DType::kBFloat16); + FRT_ADD("encoder_hidden", {es, 16384}, modalities::DType::kBFloat16); + FRT_ADD("encoder_gate_merged", {es, 32768}, + modalities::DType::kBFloat16); + FRT_ADD("encoder_gate_buf", {es, 16384}, + modalities::DType::kBFloat16); + FRT_ADD("encoder_rms_ones", {2048}, modalities::DType::kBFloat16); + + FRT_ADD("decoder_rope_weights", {ds, 256}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_x", {ds, 1024}, modalities::DType::kBFloat16); + FRT_ADD("decoder_action_buf", {ds, 32}, modalities::DType::kBFloat16); + FRT_ADD("decoder_time_emb", {steps, ds, 1024}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_style_attn", {steps, 18, ds, 3072}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_style_ffn", {steps, 18, ds, 3072}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_style_final", {steps, ds, 3072}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_QKV", {ds, 2560}, modalities::DType::kBFloat16); + FRT_ADD("decoder_hidden", {ds, 4096}, modalities::DType::kBFloat16); + FRT_ADD("decoder_gate_merged", {ds, 8192}, + modalities::DType::kBFloat16); + FRT_ADD("decoder_gate_buf", {ds, 4096}, + modalities::DType::kBFloat16); + FRT_ADD("diffusion_noise", {ds, 32}, modalities::DType::kBFloat16); + FRT_ADD("rtc_prev_action_chunk", {ds, 32}, + modalities::DType::kBFloat16); + FRT_ADD("rtc_prefix_weights", {ds}, modalities::DType::kFloat32); + FRT_ADD("rtc_guidance_weight", {1}, modalities::DType::kFloat32); + FRT_ADD("x_normed_buf", {ds, 1024}, modalities::DType::kBFloat16); + FRT_ADD("gate_buf", {ds, 1024}, modalities::DType::kBFloat16); + FRT_ADD("decoder_rms_ones", {1024}, modalities::DType::kBFloat16); + 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"); + decoder_rope_buffer_ = decoder->buffer; + st = initialize_rms_ones(); + if (!st.ok_status()) return st; + return initialize_rope(); +} + +const NativeWorkspaceBuffer* NativeWorkspace::find( + const std::string& name) const { + const auto it = buffers_.find(name); + return it == buffers_.end() ? nullptr : &it->second; +} + +} // namespace pi05 +} // namespace models +} // namespace flashrt diff --git a/cpp/models/pi05/tests/CMakeLists.txt b/cpp/models/pi05/tests/CMakeLists.txt new file mode 100644 index 00000000..23f49cb8 --- /dev/null +++ b/cpp/models/pi05/tests/CMakeLists.txt @@ -0,0 +1,91 @@ +# 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") +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_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(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_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) + 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_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/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/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/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/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 new file mode 100644 index 00000000..b1cbc4b6 --- /dev/null +++ b/cpp/tests/CMakeLists.txt @@ -0,0 +1,35 @@ +# 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_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) + +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) diff --git a/cpp/tests/data/pi05_native_v2_schema.records b/cpp/tests/data/pi05_native_v2_schema.records new file mode 100644 index 00000000..de6fb986 --- /dev/null +++ b/cpp/tests/data/pi05_native_v2_schema.records @@ -0,0 +1,8 @@ +region:0:rollout_boundary:0:640:3 +port:0:prompt:2:0:0:0:1:1:-1:-1:0:0 +port:1:state:3:1:0:0:1:1:8:-1:0:0 +port:2:images:1:3:2:0:1:1:2,224,224,3:0:0:602112 +port:3:noise:0:3:0:0:0:0:10,32:1:0:640 +port:4:actions:4:1:0:1:1:0:10,7:-1:0:280 +port:5:actions_raw:0:3:0:1:0:0:10,32:1:0:640 +stage:0:0: diff --git a/cpp/tests/gate_pi05_hot_allocator.py b/cpp/tests/gate_pi05_hot_allocator.py new file mode 100644 index 00000000..8d809049 --- /dev/null +++ b/cpp/tests/gate_pi05_hot_allocator.py @@ -0,0 +1,86 @@ +"""Reject CUDA allocation APIs in a replay-only Pi0.5 Nsight trace.""" + +from __future__ import annotations + +import argparse +from collections import Counter +import csv +from pathlib import Path + + +FORBIDDEN = ( + "cudaMalloc", + "cudaFree", + "cudaHostAlloc", + "cudaHostRegister", + "cudaHostUnregister", + "cudaMemPoolCreate", + "cudaMemPoolDestroy", + "cuMemAlloc", + "cuMemFree", + "cuMemHostAlloc", + "cuMemHostRegister", + "cuMemHostUnregister", + "cuMemCreate", + "cuMemMap", + "cuMemUnmap", + "cuMemAddressReserve", + "cuMemAddressFree", +) + + +def _read_rows(path: Path) -> list[dict[str, str]]: + lines = path.read_text(encoding="utf-8").splitlines(keepends=True) + try: + header = next( + index for index, line in enumerate(lines) + if line.startswith("Start (ns),") + ) + except StopIteration as exc: + raise ValueError(f"{path}: cuda_api_trace CSV header is missing") from exc + rows = list(csv.DictReader(lines[header:])) + if not rows: + raise ValueError(f"{path}: CUDA API trace is empty") + return rows + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--trace", type=Path, required=True) + parser.add_argument("--expected-replays", type=int, default=1000) + args = parser.parse_args() + if args.expected_replays <= 0: + parser.error("--expected-replays must be positive") + + rows = _read_rows(args.trace) + failed = [row for row in rows if int(row["Result"]) != 0] + if failed: + raise AssertionError(f"CUDA API calls failed: {failed[:4]}") + names = Counter(row["Name"] for row in rows) + graph_launches = sum( + count for name, count in names.items() + if name.startswith("cudaGraphLaunch") or name.startswith("cuGraphLaunch") + ) + if graph_launches != args.expected_replays: + raise AssertionError( + f"graph replay count differs: actual={graph_launches} " + f"expected={args.expected_replays}" + ) + allocator_calls = { + name: count for name, count in names.items() + if any(marker in name for marker in FORBIDDEN) + } + if allocator_calls: + raise AssertionError(f"hot path called CUDA allocators: {allocator_calls}") + print({ + "ok": True, + "graph_replays": graph_launches, + "allocator_calls": 0, + "cuda_api_calls": sum(names.values()), + }) + print("PASS Pi0.5 hot replay performed no CUDA allocation calls") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/gate_pi05_model_runtime_export.py b/cpp/tests/gate_pi05_model_runtime_export.py index 7af2a09c..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: - - PYTHONPATH=.:./exec/build-container:./runtime/build-container \ - 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 +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. @@ -19,6 +16,7 @@ import argparse import ctypes +import os import statistics import sys import time @@ -34,6 +32,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 +205,16 @@ def _cos(a: np.ndarray, b: np.ndarray, dtype: torch.dtype) -> float: return float(np.dot(af, bf) / (na * nb)) +def _cos_f32(a: np.ndarray, b: np.ndarray) -> float: + af = np.asarray(a, dtype=np.float64).reshape(-1) + bf = np.asarray(b, dtype=np.float64).reshape(-1) + na = float(np.linalg.norm(af)) + nb = float(np.linalg.norm(bf)) + if na == 0.0 or nb == 0.0: + return float("nan") + return float(np.dot(af, bf) / (na * nb)) + + def _make_images(num_views: int, seed: int) -> list[np.ndarray]: rng = np.random.default_rng(seed) return [ @@ -331,6 +346,13 @@ def _python_replay(pipe, pl, obs, start_noise: np.ndarray) -> np.ndarray: return _read(pl.input_noise_buf) +def _python_replay_staged(pl, start_noise: np.ndarray) -> np.ndarray: + """Replay without touching the already-staged image buffer.""" + _upload_bytes(pl.input_noise_buf, start_noise) + pl.forward() + return _read(pl.input_noise_buf) + + def _python_replay_actions(pipe, pl, obs, start_noise: np.ndarray, dtype: torch.dtype) -> np.ndarray: raw = _python_replay(pipe, pl, obs, start_noise) @@ -395,6 +417,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 +495,28 @@ def main() -> None: assert m_split.n_stages == 2, m_split.n_stages assert m_rtc.n_ports == 5, m_rtc.n_ports assert m_rtc.n_stages == 2, m_rtc.n_stages + for runtime in (mr_full, mr_split, mr_rtc): + action_port = next( + port for port in runtime.ports() + if port["name"] == "actions" + ) + assert action_port["dtype"] == 1, action_port + assert action_port["update"] == 1, action_port + assert action_port["buffer"] == 0, action_port + assert action_port["bytes"] == ( + int(pl.chunk_size) * LIBERO_ACTION_DIM * 4 + ), action_port + rtc_raw_port = next( + port for port in mr_rtc.ports() + if port["name"] == "actions_raw" + ) + rtc_noise_port = next( + port for port in mr_rtc.ports() + if port["name"] == "noise" + ) + assert rtc_raw_port["dtype"] == rtc_noise_port["dtype"], rtc_raw_port + assert rtc_raw_port["update"] == 0, rtc_raw_port + assert rtc_raw_port["buffer"] != 0, rtc_raw_port assert mr_full.fingerprint != mr_split.fingerprint assert mr_rtc.fingerprint not in ( mr_full.fingerprint, mr_split.fingerprint) @@ -484,6 +531,18 @@ def main() -> None: _raw_to_float(cxx_img, torch_dtype)))) start_noise = _seed_noise(pipe, pl, args.seed + 1009, torch_dtype) + + # Mechanism lane: hold the exact image/noise bytes fixed and compare + # the Python graph entry with the model-runtime graph entry. Numerical + # differences from the two image preprocessors do not belong here. + _upload_bytes(pl.input_images_buf, cxx_img) + mechanism_py_raw = _python_replay_staged(pl, start_noise) + _upload_bytes(pl.input_images_buf, cxx_img) + _upload_bytes(pl.input_noise_buf, start_noise) + mechanism_full_actions = _model_step_get_actions( + m_full, int(pl.chunk_size)) + mechanism_full_raw = _read(pl.input_noise_buf) + py_raw = _python_replay(pipe, pl, obs, start_noise) _upload_bytes(pl.input_noise_buf, start_noise) @@ -517,21 +576,41 @@ def main() -> None: py_actions = unnormalize_actions(py_raw_f, pipe.norm_stats)[ :, :LIBERO_ACTION_DIM].astype(np.float32) + mechanism_raw_exact = bool(np.array_equal( + mechanism_py_raw, mechanism_full_raw)) + mechanism_raw_max = float(np.max(np.abs( + _raw_to_float(mechanism_py_raw, torch_dtype) - + _raw_to_float(mechanism_full_raw, torch_dtype)))) + mechanism_py_actions = unnormalize_actions( + _raw_to_float(mechanism_py_raw, torch_dtype).reshape( + int(pl.chunk_size), 32), + pipe.norm_stats, + )[:, :LIBERO_ACTION_DIM].astype(np.float32) + mechanism_action_max = float(np.max(np.abs( + mechanism_py_actions - mechanism_full_actions))) + mechanism_action_close = bool(np.allclose( + mechanism_py_actions, mechanism_full_actions, + rtol=1e-6, atol=1e-6, + )) + raw_exact = bool(np.array_equal(py_raw, full_raw)) raw_cos = _cos(py_raw, full_raw, torch_dtype) raw_max = float(np.max(np.abs( _raw_to_float(py_raw, torch_dtype) - _raw_to_float(full_raw, torch_dtype)))) act_max = float(np.max(np.abs(py_actions - full_actions))) - act_ok = bool(np.allclose(py_actions, full_actions, rtol=1e-4, atol=1e-3)) + act_cos = _cos_f32(py_actions, full_actions) + act_close = bool(np.allclose( + py_actions, full_actions, rtol=1e-4, atol=1e-3 + )) + act_ok = act_cos >= 0.999 and act_close split_raw_exact = bool(np.array_equal(full_raw, split_raw)) split_raw_cos = _cos(full_raw, split_raw, torch_dtype) split_raw_max = float(np.max(np.abs( _raw_to_float(full_raw, torch_dtype) - _raw_to_float(split_raw, torch_dtype)))) + split_act_exact = bool(np.array_equal(full_actions, split_actions)) split_act_max = float(np.max(np.abs(full_actions - split_actions))) - split_act_ok = bool(np.allclose( - full_actions, split_actions, rtol=1e-4, atol=1e-3)) print("\n===== REAL PI0.5 MODEL-RUNTIME EXPORT GATE =====") print(f"full fingerprint : 0x{mr_full.fingerprint:016x}") @@ -541,20 +620,49 @@ def main() -> None: print(f"split runtime : ports={m_split.n_ports} stages={m_split.n_stages}") print(f"rtc runtime : ports={m_rtc.n_ports} stages={m_rtc.n_stages} prefix={prefix_len}") print(f"image buffer exact : {img_exact} cos={img_cos:.8f} max_abs={img_max:.6g}") + print( + "same-bytes mechanism raw: " + f"exact={mechanism_raw_exact} max_abs={mechanism_raw_max:.6g}" + ) + print( + "same-bytes action stage : " + f"allclose={mechanism_action_close} " + f"max_abs={mechanism_action_max:.6g}" + ) print(f"py vs full raw exact : {raw_exact} cos={raw_cos:.8f} max_abs={raw_max:.6g}") - print(f"py vs full action : {act_ok} max_abs={act_max:.6g}") + print( + f"py vs full action : accepted={act_ok} " + f"cos={act_cos:.8f} allclose={act_close} max_abs={act_max:.6g}" + ) print(f"full vs split raw exact: {split_raw_exact} cos={split_raw_cos:.8f} max_abs={split_raw_max:.6g}") - print(f"full vs split action : {split_act_ok} max_abs={split_act_max:.6g}") + print( + f"full vs split action : exact={split_act_exact} " + f"max_abs={split_act_max:.6g}" + ) print(f"rtc prefix exact : {rtc_prefix_exact} max_abs={rtc_prefix_max:.6g}") assert img_cos >= 0.999, f"image preprocess cosine too low: {img_cos}" - assert raw_cos >= 0.999, f"raw replay cosine too low: {raw_cos}" - assert act_ok, f"robot actions differ: max_abs={act_max}" + assert mechanism_raw_exact, ( + "same image/noise bytes must replay bit-exactly through Python " + f"and model-runtime entries; max_abs={mechanism_raw_max:.6g}" + ) + assert mechanism_action_close, ( + "logical action staging differs for identical raw bytes: " + f"max_abs={mechanism_action_max:.6g}" + ) + assert raw_exact, ( + "Python and model-runtime replay must be bit-exact after image " + f"staging; cos={raw_cos:.8f} max_abs={raw_max:.6g}" + ) + assert act_ok, ( + f"robot actions differ: cos={act_cos:.8f} max_abs={act_max}" + ) assert split_raw_exact, ( "split replay must be bit-exact against full replay; " f"cos={split_raw_cos:.8f} max_abs={split_raw_max:.6g}") - assert split_act_ok, ( - f"split robot actions differ: max_abs={split_act_max}") + assert split_act_exact, ( + f"split robot actions are not bit-exact: max_abs={split_act_max}" + ) assert rtc_prefix_exact, ( "RTC-prefix action graph did not preserve prev_action_chunk " f"prefix; max_abs={rtc_prefix_max:.6g}") diff --git a/cpp/tests/gate_pi05_native_schema_parity.py b/cpp/tests/gate_pi05_native_schema_parity.py new file mode 100644 index 00000000..845cf698 --- /dev/null +++ b/cpp/tests/gate_pi05_native_schema_parity.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Compare Python and C++ native-v2 port/stage/region declarations.""" + +from __future__ import annotations + +import argparse +import difflib +import os +from pathlib import Path +import subprocess +import sys +import tempfile + +import numpy as np + + +ROOT = Path(__file__).resolve().parents[2] +GOLDEN = Path(__file__).with_name("data") / "pi05_native_v2_schema.records" +sys.path.insert(0, str(ROOT)) +configured_build = os.environ.get("FLASHRT_BUILD_DIR") +if not configured_build: + raise SystemExit("Set FLASHRT_BUILD_DIR to the Python producer CMake build") +for subdir in ("exec", "runtime"): + sys.path.insert(0, str(Path(configured_build).resolve() / subdir)) + +import _flashrt_runtime as runtime_abi # noqa: E402 +import flash_rt # noqa: E402 + + +def canonical_records(identity: str) -> list[str]: + prefixes = ("region:", "port:", "stage:") + return [line for line in identity.splitlines() + if line.startswith(prefixes)] + + +def assert_records(label: str, actual: list[str], expected: list[str]) -> None: + if actual == expected: + return + diff = "\n".join(difflib.unified_diff( + expected, actual, fromfile="golden", tofile=label, lineterm="", + )) + raise AssertionError(f"{label} schema mismatch:\n{diff}") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--native-probe", type=Path, required=True) + args = parser.parse_args() + for name in ("checkpoint", "tokenizer", "native_probe"): + path = getattr(args, name).resolve() + if not path.exists(): + parser.error(f"--{name.replace('_', '-')} does not exist: {path}") + setattr(args, name, path) + golden_records = GOLDEN.read_text(encoding="utf-8").splitlines() + expected_ports = sum(line.startswith("port:") for line in golden_records) + expected_stages = sum(line.startswith("stage:") for line in golden_records) + expected_regions = sum(line.startswith("region:") for line in golden_records) + + rng = np.random.default_rng(20260710) + images = [ + np.ascontiguousarray( + rng.integers(0, 256, size=(224, 224, 3), dtype=np.uint8) + ) + for _ in range(2) + ] + state = np.linspace(-0.25, 0.25, 8, dtype=np.float32) + model = flash_rt.load_model( + str(args.checkpoint), framework="torch", config="pi05", + hardware="auto", num_views=2, num_steps=10, cache_frames=1, + use_fp8=True, use_fp16=False, state_prompt_mode="fixed", + ) + model.predict(images, prompt="pick up the red block", state=state) + producer = model._pipe.pipeline.export_model_runtime( + identity={"gate": "native_v2_schema_parity"}, + stage_plan="full", io="native_v2", + ) + try: + counts = dict(runtime_abi.export_counts(producer.export_ptr)) + if len(producer.ports()) != expected_ports or \ + len(producer.stages()) != expected_stages or \ + counts.get("capsule_regions") != expected_regions: + raise RuntimeError( + f"unexpected Python native-v2 counts: ports=" + f"{len(producer.ports())} stages={len(producer.stages())} " + f"regions={counts.get('capsule_regions')}" + ) + python_records = canonical_records(producer.identity) + assert_records("python-native-v2", python_records, golden_records) + with tempfile.TemporaryDirectory(prefix="pi05_schema_parity_") as tmp: + native_path = Path(tmp) / "native.schema" + env = dict(os.environ) + env["FLASHRT_SCHEMA_OUTPUT"] = str(native_path) + env["FLASHRT_SCHEMA_ONLY"] = "1" + subprocess.run( + [str(args.native_probe), str(args.checkpoint), + str(args.tokenizer)], + check=True, env=env, + ) + native_records = native_path.read_text().splitlines() + + assert_records("cpp-native-v2", native_records, golden_records) + + print("\n===== PI0.5 NATIVE-V2 SCHEMA PARITY =====") + print(f"records : {len(python_records)}") + print(f"ports/stage : {expected_ports} / {expected_stages}") + print(f"regions : {expected_regions}") + print("PASS - Python and C++ native-v2 schemas match the golden records") + return 0 + finally: + producer.release() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/tests/pi05_native_dlopen_probe.cpp b/cpp/tests/pi05_native_dlopen_probe.cpp new file mode 100644 index 00000000..88d19964 --- /dev/null +++ b/cpp/tests/pi05_native_dlopen_probe.cpp @@ -0,0 +1,82 @@ +#include "flashrt/model_runtime.h" + +#include + +#include +#include +#include +#include + +namespace { + +std::string json_string(const std::string& value) { + std::string output = "\""; + for (char c : value) { + if (c == '\\' || c == '"') output.push_back('\\'); + output.push_back(c); + } + output.push_back('"'); + return output; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 5) { + std::cerr << "usage: pi05_native_dlopen_probe SO CHECKPOINT " + "TOKENIZER CYCLES\n"; + return 2; + } + const int cycles = std::atoi(argv[4]); + if (cycles <= 0) return 2; + std::ostringstream config; + config << "{\"io\":\"native_v2\",\"checkpoint_path\":" + << json_string(argv[2]) << ",\"tokenizer_model_path\":" + << json_string(argv[3]) + << ",\"state_prompt_mode\":\"fixed\"," + "\"max_prompt_tokens\":200,\"state_dim\":8," + "\"num_views\":2,\"chunk\":10,\"num_steps\":10," + "\"vision_pool_factor\":1}"; + for (int cycle = 0; cycle < cycles; ++cycle) { + void* library = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL); + if (!library) { + std::cerr << "dlopen failed: " << dlerror() << '\n'; + return 1; + } + auto open = reinterpret_cast( + dlsym(library, FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL)); + auto last_error = reinterpret_cast( + dlsym(library, "frt_pi05_native_open_last_error")); + if (!open || !last_error) { + std::cerr << "native factory symbols are missing\n"; + dlclose(library); + return 1; + } + frt_model_runtime_v1* model = nullptr; + const int rc = open(config.str().c_str(), &model); + if (rc != 0 || !model) { + std::cerr << "native open failed: rc=" << rc << " error=" + << last_error() << '\n'; + dlclose(library); + return 1; + } + if (model->abi_version != FRT_MODEL_RUNTIME_ABI_VERSION || + model->struct_size < sizeof(*model) || !model->retain || + !model->release) { + std::cerr << "native model ABI is invalid\n"; + if (model->release) model->release(model->owner); + dlclose(library); + return 1; + } + model->retain(model->owner); + model->release(model->owner); + model->release(model->owner); + if (dlclose(library) != 0) { + std::cerr << "dlclose failed: " << dlerror() << '\n'; + return 1; + } + std::cout << "cycle " << (cycle + 1) << " released\n"; + } + std::cout << "PASS native model dlopen/release/dlclose lifecycle\n"; + return 0; +} diff --git a/cpp/tests/pi05_native_fp8_calibration_probe.cpp b/cpp/tests/pi05_native_fp8_calibration_probe.cpp new file mode 100644 index 00000000..74ed6c98 --- /dev/null +++ b/cpp/tests/pi05_native_fp8_calibration_probe.cpp @@ -0,0 +1,443 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/modalities/types.h" +#include "flashrt/cpp/models/pi05/c_api.h" +#include "flashrt/cpp/models/pi05/support/native_calibration.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +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, + int max_prompt_tokens, + 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\":" + << max_prompt_tokens << ",\"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 && argc != 8) { + std::cerr << "usage: pi05_native_fp8_calibration_probe CHECKPOINT " + "TOKENIZER " + "ARTIFACT SAMPLES VIEWS [RAW_ACTION_OUTPUT " + "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"; + return 2; + } + const int num_views = std::atoi(argv[5]); + 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, max_prompt_tokens); + 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()); + 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", "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]; + 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 > 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( + (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 || + artifact.hardware != hardware || + artifact.activation_dtype != (rtx ? "bfloat16" : "float16")) { + 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 || + 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, 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) { + 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 && 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") && + std::strstr(exp->identity, + hardware_identity.c_str()) && + std::strstr(exp->identity, "calibration_sha256=") && + 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; + } + 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_payload(noise.size()); + for (std::size_t i = 0; i < noise.size(); ++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_payload.data(), + noise_payload.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; + } + 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_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 || + 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::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; + } + } + 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"; + return 0; +} diff --git a/cpp/tests/pi05_native_open_probe.cpp b/cpp/tests/pi05_native_open_probe.cpp new file mode 100644 index 00000000..68ceaf50 --- /dev/null +++ b/cpp/tests/pi05_native_open_probe.cpp @@ -0,0 +1,545 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/modalities/types.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +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; +} + +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) { + if (argc != 3 && argc != 4) { + std::cerr << "usage: pi05_native_open_probe CHECKPOINT TOKENIZER " + "[CALIBRATION]\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); + } + bool profile_service_loop = + std::getenv("FLASHRT_PROFILE_SERVICE_LOOP") != nullptr; + 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) { + 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); + } + 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) { + char* end = nullptr; + hot_state_p99_limit_us = std::strtod(hot_limit_env, &end); + if (!end || *end != '\0' || !std::isfinite(hot_state_p99_limit_us) || + hot_state_p99_limit_us <= 0.0) { + std::cerr << "FLASHRT_HOT_STATE_P99_US must be positive\n"; + return 2; + } + } + std::ostringstream json; + json << "{\"io\":\"native_v2\",\"checkpoint_path\":\"" + << argv[1] << "\",\"tokenizer_model_path\":\"" << argv[2] + << "\",\"state_prompt_mode\":\"fixed\"," + << "\"max_prompt_tokens\":" << 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] << '"'; + } + 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) { + std::cerr << "native open failed: rc=" << open_rc << " error=" + << frt_pi05_native_open_last_error() << '\n'; + return 1; + } + const char* port_names[] = { + "prompt", "state", "images", "noise", "actions", "actions_raw"}; + const frt_runtime_export_v1* exp = model->exp; + int active_device = 0; + cudaDeviceProp active_properties{}; + const bool device_identity_ok = + cudaGetDevice(&active_device) == cudaSuccess && + cudaGetDeviceProperties(&active_properties, active_device) == + cudaSuccess; + const std::string hardware_id = device_identity_ok + ? "sm" + std::to_string(active_properties.major * 10 + + active_properties.minor) + : std::string(); + const std::string hardware_identity = "hardware=" + hardware_id; + const std::string hardware_manifest = + "\"hardware\":\"" + hardware_id + "\""; + 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 && + exp->struct_size == sizeof(frt_runtime_export_v1) && + 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") && + 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 == (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; + } + ok = ok && + std::strcmp(exp->capsule_regions[0].name, "rollout_boundary") == 0 && + model->ports[0].modality == FRT_RT_MOD_TEXT && + model->ports[0].update == FRT_RT_PORT_STAGED && + model->ports[1].modality == FRT_RT_MOD_STATE && + model->ports[2].modality == FRT_RT_MOD_IMAGE && + model->ports[3].update == FRT_RT_PORT_SWAP && + model->ports[4].direction == FRT_RT_PORT_OUT && + model->ports[4].dtype == FRT_RT_DTYPE_F32 && + model->ports[4].update == FRT_RT_PORT_STAGED && + model->ports[4].buffer == nullptr && + model->ports[4].bytes == 10 * 7 * sizeof(float) && + model->ports[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 && + model->ports[5].bytes == model->ports[3].bytes; + if (!ok) { + std::cerr << "native schema validation failed\n"; + model->release(model->owner); + return 1; + } + const char* schema_output = std::getenv("FLASHRT_SCHEMA_OUTPUT"); + if (schema_output && schema_output[0] != '\0') { + std::ofstream output(schema_output); + std::istringstream identity(exp->identity ? exp->identity : ""); + std::string line; + while (std::getline(identity, line)) { + if (line.compare(0, 7, "region:") == 0 || + line.compare(0, 5, "port:") == 0 || + line.compare(0, 6, "stage:") == 0) { + output << line << '\n'; + } + } + if (!output) { + std::cerr << "native schema output failed\n"; + model->release(model->owner); + return 1; + } + if (std::getenv("FLASHRT_SCHEMA_ONLY")) { + model->release(model->owner); + std::cout << "PASS native schema export\n"; + return 0; + } + } + if (model->verbs.prepare(model->self, 0, 0) != 0 || + model->verbs.prepare(model->self, 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); + return 1; + } + + const std::string prompt = "pick up the black bowl"; + float state[8] = {0.1f, -0.2f, 0.3f, -0.4f, + 0.5f, -0.6f, 0.7f, -0.8f}; + if (model->verbs.set_input(model->self, 0, prompt.data(), prompt.size(), + -1) != 0 || + model->verbs.set_input(model->self, 1, state, sizeof(state), -1) != 0) { + std::cerr << "native prompt/state staging failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + if (hot_state_updates) { + constexpr int kWarmUpdates = 20; + std::vector hot_state_latencies; + hot_state_latencies.reserve(hot_state_updates); + for (int update = -kWarmUpdates; update < hot_state_updates; ++update) { + for (int dim = 0; dim < 8; ++dim) { + state[dim] = std::sin( + static_cast((update + kWarmUpdates) * 8 + dim) * + 0.017f); + } + const auto begin = std::chrono::steady_clock::now(); + const int rc = model->verbs.set_input( + model->self, 1, state, sizeof(state), -1); + const auto end = std::chrono::steady_clock::now(); + if (rc != 0) { + std::cerr << "native hot state update failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + if (update >= 0) { + hot_state_latencies.push_back( + std::chrono::duration(end - begin) + .count()); + } + } + std::sort(hot_state_latencies.begin(), hot_state_latencies.end()); + const std::size_t p99_index = + (hot_state_latencies.size() * 99 + 99) / 100 - 1; + const double p50 = hot_state_latencies[hot_state_latencies.size() / 2]; + const double p99 = hot_state_latencies[p99_index]; + const double maximum = hot_state_latencies.back(); + std::cout << "hot state updates: n=" << hot_state_latencies.size() + << " p50_us=" << p50 << " p99_us=" << p99 + << " max_us=" << maximum << '\n'; + if ((hot_state_p99_limit_us > 0.0 && + p99 > hot_state_p99_limit_us) || + !all_graph_variants_stable(exp)) { + std::cerr << "native hot state update gate failed\n"; + model->release(model->owner); + return 1; + } + } + 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); + } + } + 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 = 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.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); + 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) { + 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; + float actions[10 * 7]{}; + std::uint64_t written = 0; + if (!noise || model->verbs.set_input(model->self, 3, host_noise.data(), + host_noise.size() * 2, -1) != -3 || + cudaMemcpy(frt_buffer_dptr(noise), host_noise.data(), + host_noise.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice) != cudaSuccess || + (profile_range && cudaProfilerStart() != cudaSuccess)) { + std::cerr << "native step failed: " + << model->verbs.last_error(model->self) << '\n'; + model->release(model->owner); + return 1; + } + int step_rc = 0; + cudaError_t upload_rc = cudaSuccess; + 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( + 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.data(), + views.size() * sizeof(frt_image_view), -1) != 0) { + step_rc = -1; + break; + } + } + if (replay != 0 || profile_service_loop) { + upload_rc = cudaMemcpy( + frt_buffer_dptr(noise), host_noise.data(), + host_noise.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice); + if (upload_rc != cudaSuccess) break; + } + step_rc = model->verbs.step(model->self); + if (step_rc != 0) break; + if (profile_service_loop && + (model->verbs.get_output( + model->self, 4, actions, sizeof(actions), &written, -1) != 0 || + written != sizeof(actions))) { + step_rc = -1; + break; + } + 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 || + 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)) { + 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; + } + } + 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); + std::cout << "PASS native open_v1 full lifecycle\n"; + return 0; +} diff --git a/cpp/tests/test_device_staging.cpp b/cpp/tests/test_device_staging.cpp index 0c673ebf..3c89b24f 100644 --- a/cpp/tests/test_device_staging.cpp +++ b/cpp/tests/test_device_staging.cpp @@ -1,24 +1,33 @@ #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 +#include +#include #include #include #include #include +#include #include using flashrt::modalities::DType; +using flashrt::modalities::ActionStaging; using flashrt::modalities::Layout; using flashrt::modalities::MemoryPlace; using flashrt::modalities::PixelFormat; using flashrt::modalities::Shape; using flashrt::modalities::TensorView; +using flashrt::modalities::EmbeddingGatherSpec; +using flashrt::modalities::TextEmbeddingStaging; using flashrt::modalities::VisionFrame; using flashrt::modalities::bfloat16_to_float; using flashrt::modalities::float_to_bfloat16; +using flashrt::modalities::gather_token_embeddings; +using flashrt::modalities::gather_token_embeddings_cpu; using flashrt::modalities::postprocess_action; using flashrt::modalities::preprocess_vision_cpu; using flashrt::modalities::preprocess_vision; @@ -37,7 +46,25 @@ 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() { + 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); @@ -59,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); @@ -110,6 +137,105 @@ void test_vision_h2d_staging() { cudaFree(device); } +void test_vision_resize_matrix() { + struct Case { int width; int height; int padding; }; + const std::array cases{{ + {1, 1, 0}, {3, 2, 5}, {17, 19, 1}, {63, 47, 7}, + {224, 224, 0}, {321, 181, 3}, {181, 321, 9}, + }}; + std::uint64_t max_frame_bytes = 0; + for (const auto& item : cases) { + max_frame_bytes = std::max( + max_frame_bytes, + static_cast(item.width * 3 + item.padding) * + static_cast(item.height)); + } + + const auto spec = flashrt::models::pi05::vision_preprocess_spec(1); + const std::uint64_t output_bytes = required_vision_output_bytes(spec); + void* device = nullptr; + assert(cudaMalloc(&device, output_bytes) == cudaSuccess); + flashrt::modalities::VisionStaging pool; + auto st = flashrt::modalities::vision_staging_create( + &pool, 1, max_frame_bytes); + assert(st.ok_status()); + std::vector actual(output_bytes / 2); + std::vector expected(output_bytes / 2); + std::uint32_t matrix_max_ulp = 0; + float matrix_max_abs = 0.0f; + Case worst_case{}; + std::size_t worst_index = 0; + std::uint16_t worst_actual = 0; + std::uint16_t worst_expected = 0; + + for (const auto& item : cases) { + const int stride = item.width * 3 + item.padding; + std::vector pixels( + static_cast(stride) * item.height, 0xa5); + for (int y = 0; y < item.height; ++y) { + for (int x = 0; x < item.width; ++x) { + for (int c = 0; c < 3; ++c) { + pixels[static_cast(y) * stride + x * 3 + c] = + static_cast( + (x * 13 + y * 17 + c * 71) & 0xff); + } + } + } + VisionFrame frame; + frame.name = "image"; + frame.image = { + pixels.data(), pixels.size(), DType::kUInt8, MemoryPlace::kHost, + Layout::kHWC, + Shape{static_cast(item.height), + static_cast(item.width), 3}}; + frame.format = PixelFormat::kRGB8; + frame.width = item.width; + frame.height = item.height; + frame.stride_bytes = stride; + TensorView device_output{ + device, output_bytes, DType::kBFloat16, MemoryPlace::kDevice, + Layout::kNHWC, Shape{1, 224, 224, 3}}; + st = preprocess_vision(spec, {frame}, device_output, nullptr, &pool); + assert(st.ok_status()); + assert(cudaMemcpy(actual.data(), device, output_bytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + TensorView host_output{ + expected.data(), output_bytes, DType::kBFloat16, + MemoryPlace::kHost, Layout::kNHWC, Shape{1, 224, 224, 3}}; + st = preprocess_vision_cpu(spec, {frame}, host_output); + assert(st.ok_status()); + for (std::size_t i = 0; i < actual.size(); ++i) { + const std::uint32_t ulp = + bf16_ulp_distance(actual[i], expected[i]); + const float absolute = std::fabs( + bfloat16_to_float(actual[i]) - + bfloat16_to_float(expected[i])); + matrix_max_abs = std::max(matrix_max_abs, absolute); + if (ulp > matrix_max_ulp) { + matrix_max_ulp = ulp; + worst_case = item; + worst_index = i; + worst_actual = actual[i]; + worst_expected = expected[i]; + } + } + } + if (matrix_max_ulp > 1) { + std::cerr << "vision resize max_ulp=" << matrix_max_ulp + << " max_abs=" << matrix_max_abs + << " size=" << worst_case.width << 'x' + << worst_case.height << " index=" << worst_index << '\n'; + std::cerr << "vision resize values actual=" + << bfloat16_to_float(worst_actual) << " expected=" + << bfloat16_to_float(worst_expected) << '\n'; + } + std::cout << "vision resize matrix max BF16 ULP: " + << matrix_max_ulp << '\n'; + assert(matrix_max_ulp <= 1); + flashrt::modalities::vision_staging_destroy(&pool); + cudaFree(device); +} + void test_action_d2h_staging() { auto spec = flashrt::models::pi05::action_postprocess_spec( {10.0f, 20.0f, 30.0f}, {2.0f, 3.0f, 4.0f}, @@ -135,9 +261,83 @@ void test_action_d2h_staging() { assert(std::fabs(actions[0] - 12.0f) < 0.01f); assert(std::fabs(actions[1] - 17.0f) < 0.01f); assert(std::fabs(actions[2] - 34.0f) < 0.01f); + ActionStaging staging; + st = flashrt::modalities::action_staging_create(&staging, bytes); + assert(st.ok_status() && staging.host_pinned && staging.bytes == bytes); + const std::size_t action_capacity = actions.capacity(); + for (int round = 0; round < 1000; ++round) { + st = postprocess_action(spec, src, &actions, nullptr, &staging); + assert(st.ok_status()); + assert(actions.capacity() == action_capacity); + } + ActionStaging too_small; + st = flashrt::modalities::action_staging_create(&too_small, bytes - 1); + assert(st.ok_status()); + st = postprocess_action(spec, src, &actions, nullptr, &too_small); + assert(!st.ok_status()); + assert(st.code == flashrt::modalities::StatusCode::kInsufficientStorage); + flashrt::modalities::action_staging_destroy(&too_small); + flashrt::modalities::action_staging_destroy(&staging); + assert(staging.host_pinned == nullptr && staging.bytes == 0); cudaFree(device); } +void test_text_embedding_device_gather() { + const std::vector table = { + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, + 9.0f, 10.0f, 11.0f, 12.0f, + }; + const std::int32_t ids[] = {2, 0}; + std::vector ref(2 * 4, 0.0f); + TensorView host_table{const_cast(table.data()), + static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{3, 4}}; + TensorView host_out{ref.data(), static_cast(ref.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + EmbeddingGatherSpec spec{3, 4, 2.0f}; + auto st = gather_token_embeddings_cpu(spec, ids, 2, host_table, host_out); + assert(st.ok_status()); + + void* d_table = nullptr; + void* d_out = nullptr; + assert(cudaMalloc(&d_table, table.size() * sizeof(float)) == cudaSuccess); + assert(cudaMalloc(&d_out, ref.size() * sizeof(float)) == cudaSuccess); + assert(cudaMemcpy(d_table, table.data(), table.size() * sizeof(float), + cudaMemcpyHostToDevice) == cudaSuccess); + TensorView device_table{d_table, + static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kDevice, + Layout::kFlat, Shape{3, 4}}; + TensorView device_out{d_out, static_cast(ref.size() * 4), + DType::kFloat32, MemoryPlace::kDevice, + Layout::kFlat, Shape{2, 4}}; + + TextEmbeddingStaging staging; + st = flashrt::modalities::text_embedding_staging_create(&staging, 2); + assert(st.ok_status()); + std::vector got(ref.size(), 0.0f); + for (int round = 0; round < 3; ++round) { + assert(cudaMemset(d_out, 0, ref.size() * sizeof(float)) == cudaSuccess); + st = gather_token_embeddings(spec, ids, 2, device_table, device_out, + nullptr, &staging); + assert(st.ok_status()); + assert(cudaMemcpy(got.data(), d_out, got.size() * sizeof(float), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(got == ref); + } + st = gather_token_embeddings(spec, ids, 3, device_table, device_out, + nullptr, &staging); + assert(!st.ok_status()); + assert(st.code == flashrt::modalities::StatusCode::kInsufficientStorage); + flashrt::modalities::text_embedding_staging_destroy(&staging); + assert(staging.device_token_ids == nullptr); + cudaFree(d_out); + cudaFree(d_table); +} + } // namespace int main() { @@ -146,7 +346,9 @@ int main() { return 0; } test_vision_h2d_staging(); + test_vision_resize_matrix(); test_action_d2h_staging(); + test_text_embedding_device_gather(); std::cout << "PASS - CUDA modality kernels/staging\n"; return 0; } diff --git a/cpp/tests/test_modalities.cpp b/cpp/tests/test_modalities.cpp index 117f11c7..5cc44045 100644 --- a/cpp/tests/test_modalities.cpp +++ b/cpp/tests/test_modalities.cpp @@ -1,12 +1,13 @@ #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 #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); @@ -32,6 +58,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 +98,7 @@ void test_pi05_vision_spec_and_preprocess() { assert(std::fabs(first_r - (-1.0f)) < 0.01f); assert(std::fabs(first_g - (127.0f / 127.5f - 1.0f)) < 0.01f); assert(std::fabs(first_b - 1.0f) < 0.01f); + assert(out[1] == float_to_bfloat16(127.0f / 127.5f - 1.0f)); } void test_view_order_guard() { @@ -117,7 +146,7 @@ void test_action_postprocess() { assert(std::fabs(out[1] - 17.0f) < 0.01f); assert(std::fabs(out[2] - 34.0f) < 0.01f); assert(std::fabs(out[3] - 11.0f) < 0.01f); - assert(std::fabs(out[4] - 24.5f) < 0.01f); + assert(std::fabs(out[4] - 23.0f) < 0.01f); assert(std::fabs(out[5] - 26.0f) < 0.01f); } @@ -158,11 +187,39 @@ void test_pi05_runtime_io_adapter() { auto st = io.prepare_vision({image}); assert(st.ok_status()); + VisionFrame invalid = image; + invalid.format = PixelFormat::kBGR8; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.image.dtype = DType::kFloat32; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.image.layout = Layout::kCHW; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.image.shape = Shape{2, 3, 3}; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.stride_bytes = -1; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + invalid = image; + invalid.stride_bytes = 5; + st = io.prepare_vision({invalid}); + assert(st.code == StatusCode::kShapeMismatch); + st = io.prepare_vision({}); + assert(st.code == StatusCode::kShapeMismatch); + st = io.prepare_vision({image, image}); + assert(st.code == StatusCode::kShapeMismatch); std::vector actions; st = io.read_actions(&actions); assert(st.ok_status()); assert(actions.size() == 3); - assert(std::fabs(actions[0] - 21.0f) < 0.01f); + assert(std::fabs(actions[0] - 11.0f) < 0.01f); assert(std::fabs(actions[1] - 22.0f) < 0.01f); assert(std::fabs(actions[2] - 33.0f) < 0.01f); } @@ -170,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_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; +} 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; +} diff --git a/cpp/tests/test_pi05_c_api.cpp b/cpp/tests/test_pi05_c_api.cpp index ef3b10ba..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); @@ -144,6 +149,33 @@ int main() { frame.pixel_format = FRT_PI05_PIXEL_RGB8; rc = frt_pi05_runtime_prepare_vision(rt, &frame, 1); assert(rc == 0); + std::vector rgb_staged(image_bytes / 2); + assert(cudaMemcpy(rgb_staged.data(), frt_buffer_dptr(image), image_bytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + frt_pi05_vision_frame invalid = frame; + invalid.pixel_format = 999; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == -4); + assert(std::strstr(frt_pi05_runtime_last_error(rt), "pixel format")); + const std::uint8_t bgr[] = { + 255, 127, 0, 0, 127, 255, + 30, 20, 10, 60, 50, 40, + }; + invalid = frame; + invalid.data = bgr; + invalid.bytes = sizeof(bgr); + invalid.pixel_format = FRT_PI05_PIXEL_BGR8; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == 0); + std::vector bgr_staged(image_bytes / 2); + assert(cudaMemcpy(bgr_staged.data(), frt_buffer_dptr(image), image_bytes, + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(bgr_staged == rgb_staged); + invalid = frame; + invalid.stride_bytes = 5; + rc = frt_pi05_runtime_prepare_vision(rt, &invalid, 1); + assert(rc == -4); + assert(std::strstr(frt_pi05_runtime_last_error(rt), "stride")); float out[3] = {}; uint64_t n_written = 0; diff --git a/cpp/tests/test_pi05_model_runtime.cpp b/cpp/tests/test_pi05_model_runtime.cpp index e250486e..20555acb 100644 --- a/cpp/tests/test_pi05_model_runtime.cpp +++ b/cpp/tests/test_pi05_model_runtime.cpp @@ -11,7 +11,9 @@ #include +#include #include +#include #include #include #include @@ -61,6 +63,13 @@ bool has_cuda_device() { return n > 0; } +int producer_set_input(void*, uint32_t, const void*, uint64_t, int) { + return 0; +} +int producer_get_output(void*, uint32_t, void*, uint64_t, uint64_t*, int) { + return 0; +} + } // namespace int main() { @@ -153,6 +162,12 @@ int main() { m->ports[1].update == FRT_RT_PORT_SWAP && m->ports[1].buffer == action, "noise SWAP port exposes the device window"); + CHECK(std::strcmp(m->ports[2].name, "actions") == 0 && + m->ports[2].dtype == FRT_RT_DTYPE_F32 && + m->ports[2].update == FRT_RT_PORT_STAGED && + m->ports[2].buffer == nullptr && + m->ports[2].bytes == 3 * sizeof(float), + "actions port declares the logical F32 payload"); CHECK(m->stages[0].graph == 0, "stage resolves the infer graph"); /* staged image input lands in the device buffer */ @@ -167,6 +182,29 @@ int main() { view.height = 2; CHECK(m->verbs.set_input(m->self, 0, &view, sizeof(view), -1) == 0, "set_input(images) accepts frt_image_view[]"); + frt_image_view bgr_view = view; + bgr_view.pixel_format = FRT_RT_PIXEL_BGR8; + CHECK(m->verbs.set_input(m->self, 0, &bgr_view, sizeof(bgr_view), -1) + == -4, + "set_input(images) rejects non-RGB image formats"); + frt_image_view invalid_format = view; + invalid_format.pixel_format = 999; + CHECK(m->verbs.set_input(m->self, 0, &invalid_format, + sizeof(invalid_format), -1) == -4, + "set_input(images) rejects unknown pixel formats"); + CHECK(std::strstr(m->verbs.last_error(m->self), "pixel format") != nullptr, + "unknown image format reports a readable error"); + frt_image_view two_views[2] = {view, view}; + CHECK(m->verbs.set_input(m->self, 0, two_views, sizeof(two_views), -1) + == -4, + "set_input(images) rejects the wrong view count"); + frt_image_view bad_stride = view; + bad_stride.stride_bytes = 5; + CHECK(m->verbs.set_input(m->self, 0, &bad_stride, sizeof(bad_stride), -1) + == -4, + "set_input(images) rejects a short row stride"); + CHECK(std::strstr(m->verbs.last_error(m->self), "stride") != nullptr, + "invalid image stride reports a readable error"); std::vector img_host(image_bytes / 2); cudaMemcpy(img_host.data(), frt_buffer_dptr(image), image_bytes, cudaMemcpyDeviceToHost); @@ -242,25 +280,42 @@ int main() { ports[1].bytes = action_bytes; ports[2].name = "actions"; ports[2].modality = FRT_RT_MOD_ACTION; - ports[2].dtype = FRT_RT_DTYPE_BF16; + ports[2].dtype = FRT_RT_DTYPE_F32; ports[2].layout = FRT_RT_LAYOUT_FLAT; ports[2].direction = FRT_RT_PORT_OUT; ports[2].update = FRT_RT_PORT_STAGED; ports[2].shape = action_shape; ports[2].rank = 2; - ports[2].buffer = action; - ports[2].bytes = action_bytes; + ports[2].bytes = 3 * sizeof(float); uint32_t after_action[1] = {0}; frt_runtime_stage_desc stages[2]{}; stages[0].graph = 0; stages[1].graph = 1; stages[1].after = after_action; stages[1].n_after = 1; + frt_model_runtime_verbs producer_verbs{}; + producer_verbs.struct_size = sizeof(producer_verbs); + producer_verbs.set_input = producer_set_input; + producer_verbs.get_output = producer_get_output; frt_model_runtime_v1* producer = frt_model_runtime_wrap( - &exp, ports, 3, stages, 2, nullptr, nullptr, nullptr, nullptr); + &exp, ports, 3, stages, 2, &producer_verbs, nullptr, nullptr, nullptr); CHECK(producer != nullptr, "producer model declaration for create_over"); + frt_runtime_port_desc wrong_action_ports[3] = {}; + for (int i = 0; i < 3; ++i) wrong_action_ports[i] = ports[i]; + wrong_action_ports[2].dtype = FRT_RT_DTYPE_BF16; + frt_model_runtime_v1* wrong_action_producer = frt_model_runtime_wrap( + &exp, wrong_action_ports, 3, stages, 2, &producer_verbs, nullptr, + nullptr, nullptr); + frt_model_runtime_v1* wrong_action_over = nullptr; + CHECK(wrong_action_producer && + frt_pi05_model_runtime_create_over( + wrong_action_producer, &cfg, &wrong_action_over) == -2 && + wrong_action_over == nullptr, + "create_over rejects a non-F32 logical action port"); + wrong_action_producer->release(wrong_action_producer->owner); + frt_model_runtime_v1* over = nullptr; CHECK(frt_pi05_model_runtime_create_over(producer, &cfg, &over) == 0 && over, @@ -272,7 +327,7 @@ int main() { over->stages[1].after[0] == 0, "create_over preserves a producer-declared two-stage DAG"); producer->release(producer->owner); - CHECK(owner.release == retains, + CHECK(owner.release < owner.retain, "producer declaration stays alive through the override"); CHECK(over->verbs.set_input(over->self, 0, &view, sizeof(view), -1) == 0, @@ -290,6 +345,145 @@ int main() { CHECK(owner.release == owner.retain, "create_over releases its native runtime and inherited producer"); + /* A producer may declare prompt only when the native runtime can serve it. */ + const int64_t prompt_shape[1] = {-1}; + frt_runtime_port_desc prompt_ports[4] = {}; + for (int i = 0; i < 3; ++i) prompt_ports[i] = ports[i]; + prompt_ports[3].name = "prompt"; + prompt_ports[3].modality = FRT_RT_MOD_TEXT; + prompt_ports[3].dtype = FRT_RT_DTYPE_U8; + prompt_ports[3].layout = FRT_RT_LAYOUT_FLAT; + prompt_ports[3].direction = FRT_RT_PORT_IN; + prompt_ports[3].update = FRT_RT_PORT_STAGED; + prompt_ports[3].shape = prompt_shape; + prompt_ports[3].rank = 1; + frt_model_runtime_v1* prompt_producer = frt_model_runtime_wrap( + &exp, prompt_ports, 4, stages, 2, &producer_verbs, nullptr, nullptr, + nullptr); + CHECK(prompt_producer != nullptr, + "producer declaration with prompt port"); + frt_model_runtime_v1* prompt_over = nullptr; + CHECK(frt_pi05_model_runtime_create_over(prompt_producer, &cfg, + &prompt_over) == -2 && + prompt_over == nullptr, + "prompt port is refused without prompt staging config"); + + frt_runtime_port_desc state_ports[5] = {}; + for (int i = 0; i < 4; ++i) state_ports[i] = prompt_ports[i]; + const int64_t state_shape[1] = {3}; + state_ports[4].name = "state"; + state_ports[4].modality = FRT_RT_MOD_STATE; + state_ports[4].dtype = FRT_RT_DTYPE_F32; + state_ports[4].layout = FRT_RT_LAYOUT_FLAT; + state_ports[4].direction = FRT_RT_PORT_IN; + state_ports[4].update = FRT_RT_PORT_STAGED; + state_ports[4].required = 1; + state_ports[4].shape = state_shape; + state_ports[4].rank = 1; + frt_model_runtime_v1* state_producer = frt_model_runtime_wrap( + &exp, state_ports, 5, stages, 2, &producer_verbs, nullptr, nullptr, + nullptr); + CHECK(state_producer != nullptr, + "producer declaration with prompt and state ports"); + frt_model_runtime_v1* state_over = nullptr; + CHECK(frt_pi05_model_runtime_create_over(state_producer, &cfg, + &state_over) == -2 && + state_over == nullptr, + "state port is refused without state normalization config"); + +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const char* tokenizer = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + if (tokenizer && tokenizer[0] != '\0') { + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + std::vector prompt_out(max_tokens * hidden, 9.0f); + frt_pi05_runtime_config prompt_cfg = cfg; + prompt_cfg.prompt_tokenizer_model_path = tokenizer; + prompt_cfg.prompt_embedding_table_data = table.data(); + prompt_cfg.prompt_embedding_table_bytes = table.size() * sizeof(float); + prompt_cfg.prompt_embedding_table_dtype = FRT_PI05_DTYPE_FLOAT32; + prompt_cfg.prompt_embedding_vocab_size = vocab; + prompt_cfg.prompt_embedding_hidden_dim = hidden; + prompt_cfg.prompt_embedding_data = prompt_out.data(); + prompt_cfg.prompt_embedding_bytes = + prompt_out.size() * sizeof(float); + prompt_cfg.prompt_embedding_dtype = FRT_PI05_DTYPE_FLOAT32; + prompt_cfg.max_prompt_tokens = max_tokens; + prompt_cfg.prompt_embedding_scale = 0.5f; + const float state_q01[3] = {0.0f, 0.0f, 0.0f}; + const float state_q99[3] = {2.0f, 2.0f, 2.0f}; + prompt_cfg.state_q01 = state_q01; + prompt_cfg.n_state_q01 = 3; + prompt_cfg.state_q99 = state_q99; + prompt_cfg.n_state_q99 = 3; + + CHECK(frt_pi05_model_runtime_create_over(prompt_producer, + &prompt_cfg, + &prompt_over) == 0 && + prompt_over, + "prompt port accepted with prompt staging config"); + const char prompt_text[] = "pick up cube"; + CHECK(prompt_over->verbs.set_input( + prompt_over->self, 3, prompt_text, + sizeof(prompt_text) - 1, -1) == 0, + "set_input(prompt) writes staged embeddings"); + CHECK(std::fabs(prompt_out[0] - 1.0f) < 0.001f && + std::fabs(prompt_out[1] + 1.0f) < 0.001f, + "prompt staging wrote the BOS embedding row"); + prompt_over->release(prompt_over->owner); + + std::fill(prompt_out.begin(), prompt_out.end(), 9.0f); + CHECK(frt_pi05_model_runtime_create_over(state_producer, + &prompt_cfg, + &state_over) == 0 && + state_over, + "state port accepted with prompt staging and norm stats"); + const float raw_state[3] = {1.0f, 2.0f, 0.0f}; + CHECK(state_over->verbs.set_input( + state_over->self, 4, raw_state, sizeof(raw_state), -1) == 0, + "set_input(state) accepts f32 state before prompt"); + CHECK(state_over->verbs.set_input( + state_over->self, 3, prompt_text, + sizeof(prompt_text) - 1, -1) == 0, + "set_input(prompt) renders cached state"); + CHECK(std::fabs(prompt_out[0] - 1.0f) < 0.001f && + std::fabs(prompt_out[1] + 1.0f) < 0.001f, + "state prompt staging wrote embeddings"); + const std::size_t variants_before = frt_graph_variant_count(graph); + for (int tick = 0; tick < 1000; ++tick) { + const float changing_state[3] = { + static_cast(tick % 3), 2.0f, 0.0f}; + CHECK(state_over->verbs.set_input( + state_over->self, 4, changing_state, + sizeof(changing_state), -1) == 0, + "state hot update remains available"); + } + CHECK(frt_graph_variant_count(graph) == variants_before, + "state hot updates do not recapture graph variants"); + const float wrong_state[2] = {0.0f, 0.0f}; + CHECK(state_over->verbs.set_input( + state_over->self, 4, wrong_state, sizeof(wrong_state), -1) == + -4, + "state hot update rejects dimension changes"); + const std::string oversized_prompt(max_tokens * 8 + 1, 'x'); + CHECK(state_over->verbs.set_input( + state_over->self, 3, oversized_prompt.data(), + oversized_prompt.size(), -1) == -4, + "prompt hot update rejects capacity growth"); + state_over->release(state_over->owner); + } else { + std::printf("SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"); + } +#endif + state_producer->release(state_producer->owner); + prompt_producer->release(prompt_producer->owner); + frt_graph_destroy(graph); frt_ctx_destroy(ctx); std::printf(g_fail ? "\n== PI05 MODEL RUNTIME FAILED ==\n" diff --git a/cpp/tests/test_pi05_native_calibration.cpp b/cpp/tests/test_pi05_native_calibration.cpp new file mode 100644 index 00000000..3e9f6875 --- /dev/null +++ b/cpp/tests/test_pi05_native_calibration.cpp @@ -0,0 +1,160 @@ +#include "flashrt/cpp/models/pi05/support/native_calibration.h" +#include "flashrt/cpp/models/pi05/backend/native_calibration_session.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::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( + {{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(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"); + return 0; +} diff --git a/cpp/tests/test_pi05_native_open.cpp b/cpp/tests/test_pi05_native_open.cpp new file mode 100644 index 00000000..758d3720 --- /dev/null +++ b/cpp/tests/test_pi05_native_open.cpp @@ -0,0 +1,355 @@ +#include "flashrt/model_runtime.h" +#include "flashrt/cpp/models/pi05/support/native_weights.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +extern "C" const char* frt_pi05_native_open_last_error(); + +namespace { + +std::string make_temp_dir() { + char tmpl[] = "/tmp/frt_pi05_native_open_XXXXXX"; + char* path = ::mkdtemp(tmpl); + assert(path); + return path; +} + +void write_file(const std::string& path) { + std::ofstream f(path, std::ios::binary); + f << "x"; + assert(f.good()); +} + +void write_raw_safetensors(const std::string& path, + const std::string& header, + const std::string& payload) { + std::ofstream f(path, std::ios::binary); + uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char b = static_cast((n >> (8 * i)) & 0xffu); + f.write(&b, 1); + } + f.write(header.data(), static_cast(header.size())); + f.write(payload.data(), static_cast(payload.size())); + assert(f.good()); +} + +void write_raw_safetensors(const std::string& path, + const std::string& header, + uint64_t payload_bytes) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + uint64_t n = header.size(); + for (int i = 0; i < 8; ++i) { + const char b = static_cast((n >> (8 * i)) & 0xffu); + f.write(&b, 1); + } + f.write(header.data(), static_cast(header.size())); + if (payload_bytes) { + f.seekp(static_cast(payload_bytes - 1), std::ios::cur); + const char zero = 0; + f.write(&zero, 1); + } + assert(f.good()); +} + +void append_f32(std::string* out, float value) { + char bytes[sizeof(float)]; + std::memcpy(bytes, &value, sizeof(value)); + out->append(bytes, sizeof(bytes)); +} + +void write_safetensors(const std::string& path, + const std::string& omit = "") { + const auto& entries = + flashrt::models::pi05::native_tensor_requirements(); + std::string header = "{"; + uint64_t offset = 0; + bool first = true; + for (size_t i = 0; i < entries.size(); ++i) { + const auto& e = entries[i]; + if (e.key == omit) continue; + if (!first) header += ","; + first = false; + header += "\""; + header += e.key; + header += "\":{\"dtype\":\"F32\",\"shape\":["; + uint64_t elements = 1; + for (size_t dim = 0; dim < e.shape.size(); ++dim) { + if (dim) header += ","; + header += std::to_string(e.shape[dim]); + elements *= e.shape[dim]; + } + header += "]"; + header += ",\"data_offsets\":["; + header += std::to_string(offset); + header += ","; + offset += elements * sizeof(float); + header += std::to_string(offset); + header += "]}"; + } + header += ",\"__metadata__\":{\"format\":\"pt\"}}"; + write_raw_safetensors(path, header, offset); +} + +void write_bad_safetensors(const std::string& path) { + const uint64_t bytes = 1001ull * 2048ull * 2ull; + write_raw_safetensors( + path, + "{\"paligemma_with_expert.paligemma.lm_head.weight\":{" + "\"dtype\":\"BF16\",\"shape\":[1001,2048]," + "\"data_offsets\":[0," + std::to_string(bytes) + "]}}", + 1024); +} + +void write_lerobot_policy_stats(const std::string& root, bool valid = true) { + std::string state_payload; + for (int i = 0; i < 8; ++i) append_f32(&state_payload, 0.0f); + for (int i = 0; i < 8; ++i) append_f32(&state_payload, valid ? 1.0f : 0.0f); + write_raw_safetensors( + root + "/policy_preprocessor_step_2_normalizer_processor.safetensors", + "{\"observation.state.q01\":{\"dtype\":\"F32\",\"shape\":[8]," + "\"data_offsets\":[0,32]}," + "\"observation.state.q99\":{\"dtype\":\"F32\",\"shape\":[8]," + "\"data_offsets\":[32,64]}}", + state_payload); + std::string action_payload; + for (int i = 0; i < 7; ++i) append_f32(&action_payload, 0.0f); + for (int i = 0; i < 7; ++i) append_f32(&action_payload, valid ? 1.0f : 0.0f); + write_raw_safetensors( + root + "/policy_postprocessor_step_0_unnormalizer_processor.safetensors", + "{\"action.q01\":{\"dtype\":\"F32\",\"shape\":[7]," + "\"data_offsets\":[0,28]}," + "\"action.q99\":{\"dtype\":\"F32\",\"shape\":[7]," + "\"data_offsets\":[28,56]}}", + action_payload); +} + +void write_norm_stats(const std::string& path, bool valid = true) { + std::ofstream f(path); + f << "{" + << "\"norm_stats\":{" + << "\"state\":{\"q01\":[0,0,0,0,0,0,0,0]," + << (valid ? "\"q99\":[1,1,1,1,1,1,1,1]}," + : "\"q99\":[0,0,0,0,0,0,0,0]},") + << "\"actions\":{\"q01\":[0,0,0,0,0,0,0]," + << (valid ? "\"q99\":[1,1,1,1,1,1,1]}" + : "\"q99\":[0,0,0,0,0,0,0]}") + << "}}"; + assert(f.good()); +} + +std::string config(const std::string& ckpt, + const std::string& tokenizer, + const char* extra = "") { + return std::string("{") + + "\"io\":\"native_v2\"," + + "\"checkpoint_path\":\"" + ckpt + "\"," + + "\"tokenizer_model_path\":\"" + tokenizer + "\"," + + "\"state_prompt_mode\":\"fixed\"," + + "\"max_prompt_tokens\":200," + + "\"state_dim\":8," + + "\"num_views\":2," + + "\"chunk\":10" + + extra + "}"; +} + +} // namespace + +int main() { + const auto& inventory = + flashrt::models::pi05::native_tensor_requirements(); + assert(inventory.size() == 812); + const auto has_key = [&inventory](const char* key) { + return std::any_of(inventory.begin(), inventory.end(), + [key](const auto& item) { return item.key == key; }); + }; + assert(has_key( + "paligemma_with_expert.paligemma.model.vision_tower.vision_model." + "encoder.layers.26.mlp.fc2.weight")); + assert(has_key( + "paligemma_with_expert.paligemma.model.language_model.layers.17." + "mlp.down_proj.weight")); + assert(has_key( + "paligemma_with_expert.gemma_expert.model.layers.17." + "post_attention_layernorm.dense.bias")); + + frt_model_runtime_v1* out = reinterpret_cast(0x1); + int rc = frt_model_runtime_open_v1(nullptr, &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "null")); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1("{", &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "JSON")); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + "{\"io\":\"native\",\"checkpoint_path\":\"/tmp\"," + "\"tokenizer_model_path\":\"/tmp/x\"," + "\"state_prompt_mode\":\"fixed\"," + "\"max_prompt_tokens\":200,\"state_dim\":1}", + &out); + assert(rc == -1); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "native_v2")); + + const std::string root = make_temp_dir(); + const std::string tokenizer = root + "/tokenizer.model"; + write_file(tokenizer); + + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, ",\"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, ",\"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, + ",\"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); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), + "model.safetensors")); + + write_bad_safetensors(root + "/model.safetensors"); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "byte range")); + + const std::string missing_key = + flashrt::models::pi05::native_tensor_requirements().back().key; + write_safetensors(root + "/model.safetensors", missing_key); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), + missing_key.c_str())); + + write_safetensors(root + "/model.safetensors"); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(root, tokenizer, ",\"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); + 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); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "norm_stats")); + + write_norm_stats(root + "/norm_stats.json", false); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(config(root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "q01/q99")); + + write_norm_stats(root + "/norm_stats.json"); +#ifndef FLASHRT_CPP_PI05_NATIVE_OPEN_ENABLED + const std::string good = config(root, tokenizer); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1(good.c_str(), &out); + assert(rc == -3); + assert(out == nullptr); + assert(std::strstr(frt_pi05_native_open_last_error(), "validated")); +#endif + + const std::string invalid_prompt_capacity = + std::string("{") + + "\"io\":\"native_v2\"," + + "\"checkpoint_path\":\"" + root + "\"," + + "\"tokenizer_model_path\":\"" + tokenizer + "\"," + + "\"state_prompt_mode\":\"fixed\"," + + "\"max_prompt_tokens\":0," + + "\"state_dim\":8}"; + 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")); + + const std::string lerobot_root = make_temp_dir(); + write_safetensors(lerobot_root + "/model.safetensors"); + write_lerobot_policy_stats(lerobot_root, false); + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(lerobot_root, tokenizer).c_str(), &out); + assert(rc == -2); + assert(out == nullptr); + + write_lerobot_policy_stats(lerobot_root); +#ifndef FLASHRT_CPP_PI05_NATIVE_OPEN_ENABLED + out = reinterpret_cast(0x1); + rc = frt_model_runtime_open_v1( + config(lerobot_root, tokenizer).c_str(), &out); + assert(rc == -3); + assert(out == nullptr); +#endif + + ::unlink((lerobot_root + "/model.safetensors").c_str()); + ::unlink((lerobot_root + + "/policy_preprocessor_step_2_normalizer_processor.safetensors") + .c_str()); + ::unlink((lerobot_root + + "/policy_postprocessor_step_0_unnormalizer_processor.safetensors") + .c_str()); + ::rmdir(lerobot_root.c_str()); + + ::unlink(tokenizer.c_str()); + ::unlink((root + "/model.safetensors").c_str()); + ::unlink((root + "/norm_stats.json").c_str()); + ::rmdir(root.c_str()); + std::printf("PASS - Pi05 native open scaffold\n"); + return 0; +} diff --git a/cpp/tests/test_pi05_prompt_embed.cpp b/cpp/tests/test_pi05_prompt_embed.cpp new file mode 100644 index 00000000..c625a615 --- /dev/null +++ b/cpp/tests/test_pi05_prompt_embed.cpp @@ -0,0 +1,200 @@ +#include "flashrt/cpp/models/pi05/model/prompt_embed.h" + +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING +#include +#endif + +#include +#include +#include +#include +#include +#include +#include + +using flashrt::modalities::DType; +using flashrt::modalities::Layout; +using flashrt::modalities::MemoryPlace; +using flashrt::modalities::SentencePieceTokenizer; +using flashrt::modalities::Shape; +using flashrt::modalities::StatusCode; +using flashrt::modalities::TensorView; +using flashrt::modalities::TextEmbeddingStaging; +using flashrt::models::pi05::PromptEmbeddingSpec; +using flashrt::models::pi05::embed_prompt; +using flashrt::models::pi05::embed_prompt_cpu; + +namespace { + +std::string tokenizer_model_path() { + const char* env = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + return env ? std::string(env) : std::string(); +} + +bool has_cuda_device() { +#ifdef FLASHRT_CPP_WITH_CUDA_STAGING + int n = 0; + cudaError_t rc = cudaGetDeviceCount(&n); + if (rc != cudaSuccess) { + cudaGetLastError(); + return false; + } + return n > 0; +#else + return false; +#endif +} + +void test_requires_loaded_tokenizer() { + SentencePieceTokenizer tokenizer; + std::vector table(8, 0.0f); + std::vector out(8, 0.0f); + std::vector ids; + std::uint64_t prompt_len = 0; + + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + PromptEmbeddingSpec spec{2, 4, 2, 1.0f}; + auto st = embed_prompt_cpu(tokenizer, spec, "pick", nullptr, 0, src, dst, + &ids, &prompt_len); + assert(!st.ok_status()); + assert(st.code == StatusCode::kInvalidArgument); +} + +void test_paligemma_prompt_embedding_when_configured() { +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const std::string path = tokenizer_model_path(); + if (path.empty()) { + std::cout << "SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"; + return; + } + SentencePieceTokenizer tokenizer; + auto st = tokenizer.load_model(path); + assert(st.ok_status()); + + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + std::vector out(max_tokens * hidden, 7.0f); + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{vocab, hidden}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{max_tokens, hidden}}; + + const float state[] = {0.0f, 1.0f, -1.0f}; + PromptEmbeddingSpec spec{vocab, hidden, max_tokens, 0.5f}; + std::vector ids; + ids.reserve(max_tokens + 1); + tokenizer.reserve(max_tokens); + std::string formatted; + formatted.reserve(512); + std::uint64_t prompt_len = 0; + st = embed_prompt(tokenizer, spec, "pick_up_cube", state, 3, src, dst, + &ids, &prompt_len, nullptr, nullptr, &formatted); + assert(st.ok_status()); + const std::vector expected_ids = { + 2, 7071, 235292, 4788, 908, 28660, 235269, 3040, 235292, + 235248, 235274, 235284, 235321, 235248, 235284, 235308, + 235308, 235248, 235276, 235289, 108, 4022, 235292, 235248, + }; + assert(ids == expected_ids); + assert(prompt_len == expected_ids.size()); + for (std::uint64_t i = 0; i < prompt_len; ++i) { + const float id = static_cast(expected_ids[i]); + assert(std::fabs(out[i * hidden + 0] - id * 0.5f) < 0.001f); + assert(std::fabs(out[i * hidden + 1] + id * 0.5f) < 0.001f); + } + for (std::uint64_t i = prompt_len * hidden; i < out.size(); ++i) { + assert(out[i] == 0.0f); + } + const std::size_t id_capacity = ids.capacity(); + const std::size_t formatted_capacity = formatted.capacity(); + const std::uint64_t tokenizer_capacity = tokenizer.workspace_capacity(); + for (int round = 0; round < 1000; ++round) { + st = embed_prompt(tokenizer, spec, "pick_up_cube", state, 3, src, dst, + &ids, &prompt_len, nullptr, nullptr, &formatted); + assert(st.ok_status()); + assert(ids.capacity() == id_capacity); + assert(formatted.capacity() == formatted_capacity); + assert(tokenizer.workspace_capacity() == tokenizer_capacity); + } +#endif +} + +void test_paligemma_prompt_embedding_device_when_configured() { +#if defined(FLASHRT_CPP_HAS_SENTENCEPIECE) && defined(FLASHRT_CPP_WITH_CUDA_STAGING) + const std::string path = tokenizer_model_path(); + if (path.empty() || !has_cuda_device()) { + std::cout << "SKIP - tokenizer or CUDA device not available\n"; + return; + } + SentencePieceTokenizer tokenizer; + auto st = tokenizer.load_model(path); + assert(st.ok_status()); + + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + void* d_table = nullptr; + void* d_out = nullptr; + assert(cudaMalloc(&d_table, table.size() * sizeof(float)) == cudaSuccess); + assert(cudaMalloc(&d_out, max_tokens * hidden * sizeof(float)) == + cudaSuccess); + assert(cudaMemcpy(d_table, table.data(), table.size() * sizeof(float), + cudaMemcpyHostToDevice) == cudaSuccess); + TensorView src{d_table, static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kDevice, Layout::kFlat, + Shape{vocab, hidden}}; + TensorView dst{d_out, + static_cast(max_tokens * hidden * 4), + DType::kFloat32, MemoryPlace::kDevice, Layout::kFlat, + Shape{max_tokens, hidden}}; + TextEmbeddingStaging staging; + st = flashrt::modalities::text_embedding_staging_create(&staging, + max_tokens); + assert(st.ok_status()); + std::vector ids; + std::uint64_t prompt_len = 0; + PromptEmbeddingSpec spec{vocab, hidden, max_tokens, 0.5f}; + st = embed_prompt(tokenizer, spec, "pick up cube", nullptr, 0, src, dst, + &ids, &prompt_len, nullptr, &staging); + assert(st.ok_status()); + std::vector out(max_tokens * hidden, 1.0f); + assert(cudaMemcpy(out.data(), d_out, out.size() * sizeof(float), + cudaMemcpyDeviceToHost) == cudaSuccess); + assert(prompt_len == ids.size()); + assert(ids[0] == 2); + assert(std::fabs(out[0] - 1.0f) < 0.001f); + assert(std::fabs(out[1] + 1.0f) < 0.001f); + assert(out[prompt_len * hidden] == 0.0f); + flashrt::modalities::text_embedding_staging_destroy(&staging); + cudaFree(d_out); + cudaFree(d_table); +#endif +} + +} // namespace + +int main() { + test_requires_loaded_tokenizer(); + test_paligemma_prompt_embedding_when_configured(); + test_paligemma_prompt_embedding_device_when_configured(); + std::cout << "PASS - Pi05 prompt embedding\n"; + return 0; +} diff --git a/cpp/tests/test_pi05_prompt_format.cpp b/cpp/tests/test_pi05_prompt_format.cpp new file mode 100644 index 00000000..f7431909 --- /dev/null +++ b/cpp/tests/test_pi05_prompt_format.cpp @@ -0,0 +1,73 @@ +#include "flashrt/cpp/models/pi05/model/prompt_format.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +void test_discretize_matches_python_reference() { + const float state[] = {-1.0f, 0.0f, 1.0f, 2.0f, -2.0f}; + const auto bins = flashrt::models::pi05::discretize_state_prompt_bins( + state, 5); + const std::vector expected = {0, 128, 255, 255, -1}; + assert(bins == expected); +} + +void test_prompt_format_matches_python_reference() { + const float state[] = {-1.0f, 0.0f, 1.0f, 2.0f, -2.0f}; + const std::string out = flashrt::models::pi05::format_state_prompt( + "pick_up\nred", state, 5); + assert(out == + "Task: pick up red, State: 0 128 255 255 -1;\nAction: "); + std::string workspace; + workspace.reserve(128); + const std::size_t capacity = workspace.capacity(); + for (int i = 0; i < 1000; ++i) { + flashrt::models::pi05::format_state_prompt_into( + "pick_up\nred", state, 5, &workspace); + assert(workspace == out); + assert(workspace.capacity() == capacity); + } +} + +void test_prompt_without_state_keeps_text_only_format() { + const std::string out = flashrt::models::pi05::format_state_prompt( + " pick_up\nred ", nullptr, 0); + assert(out == "pick up red"); +} + +void test_boundary_values() { + const float eps = 1.0f / 1024.0f; + const float state[] = { + -1.0f - eps, + -1.0f, + -1.0f + eps, + 1.0f - eps, + 1.0f, + std::numeric_limits::quiet_NaN(), + }; + const auto bins = flashrt::models::pi05::discretize_state_prompt_bins( + state, 6); + assert(bins[0] == -1); + assert(bins[1] == 0); + assert(bins[2] == 0); + assert(bins[3] == 255); + assert(bins[4] == 255); + assert(bins[5] == 255); +} + +} // namespace + +int main() { + test_discretize_matches_python_reference(); + test_prompt_format_matches_python_reference(); + test_prompt_without_state_keeps_text_only_format(); + test_boundary_values(); + std::cout << "PASS - Pi05 prompt formatter\n"; + return 0; +} diff --git a/cpp/tests/test_pi05_runtime.cpp b/cpp/tests/test_pi05_runtime.cpp index 607c58dd..5e2a1e18 100644 --- a/cpp/tests/test_pi05_runtime.cpp +++ b/cpp/tests/test_pi05_runtime.cpp @@ -1,7 +1,8 @@ -#include "flashrt/cpp/models/pi05/runtime.h" +#include "flashrt/cpp/models/pi05/model/runtime.h" #include #include +#include #include #include #include @@ -111,6 +112,7 @@ void test_adopted_export_runtime_flow() { assert(runtime.export_runtime() == &exp); assert(runtime.manifest().vision.view_order.size() == 1); assert(runtime.manifest().graphs.infer == "infer"); + assert(runtime.set_prompt("pick up the cube") != 0); const std::uint8_t image_rgb[] = { 0, 127, 255, 255, 127, 0, @@ -127,6 +129,11 @@ void test_adopted_export_runtime_flow() { auto st = runtime.prepare_vision({image}); assert(st.ok_status()); + VisionFrame bgr = image; + bgr.format = PixelFormat::kBGR8; + st = runtime.prepare_vision({bgr}); + assert(!st.ok_status()); + assert(st.code == flashrt::modalities::StatusCode::kShapeMismatch); assert(runtime.replay_tick() == 0); assert(probe.calls == 1); @@ -142,10 +149,76 @@ void test_adopted_export_runtime_flow() { assert(owner.release == 1); } +void test_prompt_staging_when_configured() { +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const char* tokenizer = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + if (!tokenizer || tokenizer[0] == '\0') { + std::cout << "SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"; + return; + } + + Owner owner; + frt_runtime_graph_desc graph{}; + graph.name = "infer"; + graph.handle = reinterpret_cast(0x2000); + graph.default_key = 9; + graph.stream_id = 5; + auto exp = make_export(&owner, &graph); + + constexpr std::uint64_t vocab = 257152; + constexpr std::uint64_t hidden = 2; + constexpr std::uint64_t max_tokens = 32; + std::vector table(vocab * hidden); + for (std::uint64_t i = 0; i < vocab; ++i) { + table[i * hidden + 0] = static_cast(i); + table[i * hidden + 1] = -static_cast(i); + } + std::vector prompt(max_tokens * hidden, 3.0f); + std::vector image_input(1 * 224 * 224 * 3); + std::vector action_model(4, 0.0f); + + flashrt::models::pi05::RuntimeConfig cfg; + cfg.num_views = 1; + cfg.chunk = 1; + cfg.model_action_dim = 4; + cfg.robot_action_dim = 3; + cfg.image_input_override = TensorView{ + image_input.data(), static_cast(image_input.size() * 2), + DType::kBFloat16, MemoryPlace::kHost, Layout::kNHWC, + Shape{1, 224, 224, 3}}; + cfg.action_output_override = TensorView{ + action_model.data(), static_cast(action_model.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, Shape{1, 4}}; + cfg.prompt_tokenizer_model_path = tokenizer; + cfg.prompt_vocab_size = vocab; + cfg.prompt_hidden_dim = hidden; + cfg.prompt_max_tokens = max_tokens; + cfg.prompt_embedding_scale = 0.5f; + cfg.prompt_embedding_table = TensorView{ + table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{vocab, hidden}}; + cfg.prompt_embedding_output = TensorView{ + prompt.data(), static_cast(prompt.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{max_tokens, hidden}}; + + flashrt::models::pi05::Runtime runtime(&exp, cfg); + assert(runtime.ok()); + const float state[] = {0.0f, 1.0f, -1.0f}; + assert(runtime.set_prompt_state("pick_up_cube", state, 3) == 0); + assert(runtime.current_prompt_len() == 24); + assert(std::fabs(prompt[0] - 1.0f) < 0.001f); + assert(std::fabs(prompt[1] + 1.0f) < 0.001f); + assert(prompt[24 * hidden] == 0.0f); +#endif +} + } // namespace int main() { test_adopted_export_runtime_flow(); + test_prompt_staging_when_configured(); std::cout << "PASS - Pi05 C++ runtime flow\n"; return 0; } diff --git a/cpp/tests/test_safetensors_loader.cpp b/cpp/tests/test_safetensors_loader.cpp new file mode 100644 index 00000000..7d5c23a5 --- /dev/null +++ b/cpp/tests/test_safetensors_loader.cpp @@ -0,0 +1,135 @@ +#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); + assert(file.metadata().size() == 1); + assert(file.metadata().at("format") == "pt"); + 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(moved.metadata().at("format") == "pt"); + 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, + "{\"__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]}}", + std::string(2, '\0')); + assert(!file.open(invalid)); + assert(file.error().find("exceeds") != std::string::npos); + + assert(::unlink(invalid.c_str()) == 0); + std::printf("PASS - safetensors mmap loader\n"); + return 0; +} diff --git a/cpp/tests/test_sha256.cpp b/cpp/tests/test_sha256.cpp new file mode 100644 index 00000000..323dca13 --- /dev/null +++ b/cpp/tests/test_sha256.cpp @@ -0,0 +1,30 @@ +#include "flashrt/cpp/loader/sha256.h" + +#include +#include +#include +#include +#include + +int main() { + char path[] = "/tmp/flashrt_sha256_XXXXXX"; + const int fd = ::mkstemp(path); + assert(fd >= 0); + ::close(fd); + { + std::ofstream file(path, std::ios::binary | std::ios::trunc); + file << "abc"; + assert(file.good()); + } + std::string digest; + std::string error; + assert(flashrt::loader::sha256_file(path, &digest, &error)); + assert(digest == + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert(error.empty()); + ::unlink(path); + assert(!flashrt::loader::sha256_file(path, &digest, &error)); + assert(!error.empty()); + std::printf("PASS - SHA-256 file hashing\n"); + return 0; +} diff --git a/cpp/tests/test_text_modalities.cpp b/cpp/tests/test_text_modalities.cpp new file mode 100644 index 00000000..94645572 --- /dev/null +++ b/cpp/tests/test_text_modalities.cpp @@ -0,0 +1,92 @@ +#include "flashrt/cpp/modalities/text.h" + +#include +#include +#include +#include +#include + +using flashrt::modalities::DType; +using flashrt::modalities::EmbeddingGatherSpec; +using flashrt::modalities::Layout; +using flashrt::modalities::MemoryPlace; +using flashrt::modalities::Shape; +using flashrt::modalities::StatusCode; +using flashrt::modalities::TensorView; +using flashrt::modalities::bfloat16_to_float; +using flashrt::modalities::float_to_bfloat16; +using flashrt::modalities::gather_token_embeddings_cpu; + +namespace { + +void test_f32_embedding_gather() { + std::vector table = { + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, + 9.0f, 10.0f, 11.0f, 12.0f, + }; + std::int32_t ids[] = {2, 0}; + std::vector out(2 * 4, 0.0f); + + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{3, 4}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 4}}; + EmbeddingGatherSpec spec{3, 4, 2.0f}; + auto st = gather_token_embeddings_cpu(spec, ids, 2, src, dst); + assert(st.ok_status()); + const std::vector expected = { + 18.0f, 20.0f, 22.0f, 24.0f, + 2.0f, 4.0f, 6.0f, 8.0f, + }; + assert(out == expected); +} + +void test_bf16_embedding_gather() { + std::vector table = { + float_to_bfloat16(0.5f), float_to_bfloat16(-1.0f), + float_to_bfloat16(2.0f), float_to_bfloat16(3.0f), + }; + std::int32_t ids[] = {1}; + std::vector out(2); + + TensorView src{table.data(), static_cast(table.size() * 2), + DType::kBFloat16, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 2}}; + TensorView dst{out.data(), static_cast(out.size() * 2), + DType::kBFloat16, MemoryPlace::kHost, Layout::kFlat, + Shape{1, 2}}; + EmbeddingGatherSpec spec{2, 2, 1.5f}; + auto st = gather_token_embeddings_cpu(spec, ids, 1, src, dst); + assert(st.ok_status()); + assert(std::fabs(bfloat16_to_float(out[0]) - 3.0f) < 0.01f); + assert(std::fabs(bfloat16_to_float(out[1]) - 4.5f) < 0.01f); +} + +void test_invalid_token_rejected() { + std::vector table(4, 0.0f); + std::vector out(2, 0.0f); + std::int32_t ids[] = {2}; + TensorView src{table.data(), static_cast(table.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{2, 2}}; + TensorView dst{out.data(), static_cast(out.size() * 4), + DType::kFloat32, MemoryPlace::kHost, Layout::kFlat, + Shape{1, 2}}; + EmbeddingGatherSpec spec{2, 2, 1.0f}; + auto st = gather_token_embeddings_cpu(spec, ids, 1, src, dst); + assert(!st.ok_status()); + assert(st.code == StatusCode::kInvalidArgument); +} + +} // namespace + +int main() { + test_f32_embedding_gather(); + test_bf16_embedding_gather(); + test_invalid_token_rejected(); + std::cout << "PASS - text modality contracts\n"; + return 0; +} diff --git a/cpp/tests/test_text_tokenizer.cpp b/cpp/tests/test_text_tokenizer.cpp new file mode 100644 index 00000000..1f6d7cb1 --- /dev/null +++ b/cpp/tests/test_text_tokenizer.cpp @@ -0,0 +1,81 @@ +#include "flashrt/cpp/modalities/tokenizer.h" + +#include +#include +#include +#include +#include +#include + +using flashrt::modalities::SentencePieceEncodeOptions; +using flashrt::modalities::SentencePieceTokenizer; +using flashrt::modalities::StatusCode; + +namespace { + +std::string tokenizer_model_path() { + const char* env = std::getenv("FLASH_RT_PALIGEMMA_TOKENIZER"); + return env ? std::string(env) : std::string(); +} + +void test_unavailable_build_reports_unsupported() { + SentencePieceTokenizer tokenizer; +#ifndef FLASHRT_CPP_HAS_SENTENCEPIECE + auto st = tokenizer.load_model("missing.model"); + assert(!st.ok_status()); + assert(st.code == StatusCode::kUnsupported); +#else + (void)tokenizer; +#endif +} + +void test_paligemma_token_exact_when_configured() { +#ifdef FLASHRT_CPP_HAS_SENTENCEPIECE + const std::string path = tokenizer_model_path(); + if (path.empty()) { + std::cout << "SKIP - FLASH_RT_PALIGEMMA_TOKENIZER not set\n"; + return; + } + SentencePieceTokenizer tokenizer; + auto st = tokenizer.load_model(path); + assert(st.ok_status()); + assert(tokenizer.loaded()); + assert(tokenizer.vocab_size() == 257152); + assert(tokenizer.bos_id() == 2); + assert(tokenizer.eos_id() == 1); + assert(tokenizer.unk_id() == 3); + assert(tokenizer.pad_id() == 0); + + std::vector ids; + SentencePieceEncodeOptions options; + options.add_bos = true; + st = tokenizer.encode( + "Task: pick up cube, State: 0 128 255;\nAction: ", + options, &ids); + assert(st.ok_status()); + const std::vector expected = { + 2, 7071, 235292, 4788, 908, 28660, 235269, 3040, 235292, + 235248, 235276, 235248, 235274, 235284, 235321, 235248, + 235284, 235308, 235308, 235289, 108, 4022, 235292, 235248, + }; + assert(ids == expected); + + options.max_tokens = expected.size() + 2; + options.pad_to_max_tokens = true; + st = tokenizer.encode( + "Task: pick up cube, State: 0 128 255;\nAction: ", + options, &ids); + assert(st.ok_status()); + assert(ids.size() == expected.size() + 2); + assert(ids[ids.size() - 1] == 0); +#endif +} + +} // namespace + +int main() { + test_unavailable_build_reports_unsupported(); + test_paligemma_token_exact_when_configured(); + std::cout << "PASS - text tokenizer contracts\n"; + return 0; +} diff --git a/csrc/attention/fa2_wrapper.cu b/csrc/attention/fa2_wrapper.cu index dd043327..cc8bfcaa 100644 --- a/csrc/attention/fa2_wrapper.cu +++ b/csrc/attention/fa2_wrapper.cu @@ -24,6 +24,7 @@ #include "flash_attn_2_src/flash_attn/namespace_config.h" #include "flash_attn_2_src/flash_attn/flash.h" +#include "attention/fa2_wrapper.h" namespace FLASH_NAMESPACE { template diff --git a/csrc/attention/fa2_wrapper.h b/csrc/attention/fa2_wrapper.h new file mode 100644 index 00000000..d07de0b0 --- /dev/null +++ b/csrc/attention/fa2_wrapper.h @@ -0,0 +1,73 @@ +#ifndef FLASHRT_ATTENTION_FA2_WRAPPER_H +#define FLASHRT_ATTENTION_FA2_WRAPPER_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +void fvk_attention_fa2_fwd_fp16( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_seqused( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_seqused_splitkv( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_causal( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +#ifdef __cplusplus +} +#endif + +#endif // FLASHRT_ATTENTION_FA2_WRAPPER_H diff --git a/csrc/attention/fa2_wrapper_causal.cu b/csrc/attention/fa2_wrapper_causal.cu index f651b194..fda63187 100644 --- a/csrc/attention/fa2_wrapper_causal.cu +++ b/csrc/attention/fa2_wrapper_causal.cu @@ -25,6 +25,7 @@ #include "flash_attn_2_src/flash_attn/namespace_config.h" #include "flash_attn_2_src/flash_attn/flash.h" +#include "attention/fa2_wrapper.h" namespace FLASH_NAMESPACE { template @@ -180,20 +181,17 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal( int o_batch_stride, int o_row_stride, int o_head_stride, float softmax_scale, int num_sms, cudaStream_t stream) { - if ((head_dim != 128) -#ifdef FA2_HAS_HDIM_256 - && head_dim != 256 + bool supported = false; +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_128) + supported = supported || head_dim == 128; #endif - ) { -#ifdef FA2_HAS_HDIM_256 - fprintf(stderr, - "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. " - "Only head_dim=128 and 256 are currently instantiated.\n", head_dim); -#else +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_256) + supported = supported || head_dim == 256; +#endif + if (!supported) { fprintf(stderr, "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. " - "Only head_dim=128 is currently instantiated.\n", head_dim); -#endif + "Enable its FA2_HDIMS entry and rebuild.\n", head_dim); std::abort(); } @@ -208,26 +206,36 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal( o_batch_stride, o_row_stride, o_head_stride, softmax_scale); - int num_splits = setup_splitkv_causal(params, softmax_lse_accum_ptr, o_accum_ptr, - num_sms, seqlen_q, seqlen_k, - head_dim, batch, num_heads_q); - if (head_dim == 128 && num_splits > 1) { - FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(params, stream); - } else if (head_dim == 128) { - FLASH_NAMESPACE::run_mha_fwd_(params, stream); - } -#ifdef FA2_HAS_HDIM_256 - else if (num_splits > 1) { - FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(params, stream); - } else { - FLASH_NAMESPACE::run_mha_fwd_(params, stream); - } -#else - else { - fprintf(stderr, - "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built " - "(hdim=256 disabled at compile time).\n", head_dim); - std::abort(); - } + int num_splits = setup_splitkv_causal( + params, softmax_lse_accum_ptr, o_accum_ptr, num_sms, seqlen_q, + seqlen_k, head_dim, batch, num_heads_q); + switch (head_dim) { +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_128) + case 128: + if (num_splits > 1) { + FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch< + cutlass::bfloat16_t, 128, true>(params, stream); + } else { + FLASH_NAMESPACE::run_mha_fwd_< + cutlass::bfloat16_t, 128, true>(params, stream); + } + return; +#endif +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_256) + case 256: + if (num_splits > 1) { + FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch< + cutlass::bfloat16_t, 256, true>(params, stream); + } else { + FLASH_NAMESPACE::run_mha_fwd_< + cutlass::bfloat16_t, 256, true>(params, stream); + } + return; #endif + default: + fprintf(stderr, + "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built " + "in this FA2 matrix.\n", head_dim); + std::abort(); + } } diff --git a/csrc/fa2_bindings.cpp b/csrc/fa2_bindings.cpp index 41d13d04..75360388 100644 --- a/csrc/fa2_bindings.cpp +++ b/csrc/fa2_bindings.cpp @@ -22,79 +22,14 @@ #include #include +#include "attention/fa2_wrapper.h" + namespace py = pybind11; static cudaStream_t to_stream(uintptr_t s) { return reinterpret_cast(s); } -// Forward declarations (definitions in csrc/attention/fa2_wrapper.cu). -extern "C" void fvk_attention_fa2_fwd_fp16( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -extern "C" void fvk_attention_fa2_fwd_bf16( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// seqused_k variant — definition in csrc/attention/fa2_wrapper.cu. Reads the -// per-batch K length from device memory so one captured graph serves any pos. -extern "C" void fvk_attention_fa2_fwd_bf16_seqused( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// seqused_k + split-KV variant (experimental; caller pre-inits lse_accum=-inf). -extern "C" void fvk_attention_fa2_fwd_bf16_seqused_splitkv( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// Causal variant — definition in csrc/attention/fa2_wrapper_causal.cu. -// Currently only (bf16, head_dim=128) is built. Used by Qwen3-8B -// prefill (S=N causal self-attention). -extern "C" void fvk_attention_fa2_fwd_bf16_causal( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - - // Shared docstring. pybind::def's doc arg takes a single string; we want the // same text for both fwd_fp16 and fwd_bf16 so deduplicate via static const. static const char* kDocstring = R"(FlashAttention-2 fwd (vendored). GQA-capable cross-attention. diff --git a/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); 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 ca3515cc..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 @@ -53,25 +52,36 @@ 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 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 @@ -104,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: @@ -123,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 e53e5856..e91399ad 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,21 +82,27 @@ 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 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. -## 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/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 new file mode 100644 index 00000000..f96e84c7 --- /dev/null +++ b/docs/mindon_pi05_integration.md @@ -0,0 +1,263 @@ +# Mindon Pi0.5 Integration Guide + +This guide describes how a Mindon C++ host should integrate Pi0.5 through the +existing FlashRT runtime/model-runtime contracts. It is a deployment guide, not +a new ABI. + +## Layer Ownership + +FlashRT owns: + +- checkpoint loading and graph capture; +- ports, stages, streams, buffers, capsule regions, identity, and fingerprint; +- model-specific IO semantics: tokenizer, prompt formatting, image + preprocess, state normalization/discretization, and action postprocess; +- `set_input`, `get_output`, `prepare`, and `step` producer verbs. + +Nexus owns: + +- adoption of `frt_runtime_export_v1` / `frt_model_runtime_v1`; +- capsule snapshot/restore/fork/move over declared regions; +- stage scheduling and interaction modes; +- embedded and transport adapters that map external payloads to declared + ports. + +Mindon owns: + +- the application/control loop; +- camera/state/prompt transport into the adopted ports; +- action publication and deadline policy. + +Nexus should not learn Pi0.5 tokenizer, tensor layout, normalization, or +action schema rules. If a host needs richer state, the producer must export a +richer model-runtime face. + +## Recommended Lanes + +### Lane A: Available Now + +Use a resident Python setup producer, then run the hot loop in C++. + +Flow: + +1. Start a process that embeds CPython or calls a Python setup function. +2. Load the Pi0.5 checkpoint through the FlashRT Python frontend. +3. Capture graphs and call `pipeline.export_model_runtime(io="native", ...)`. +4. Pass the returned `frt_model_runtime_v1*` to the C++ host. +5. Adopt it into Nexus with `flashrt_adopt_model_runtime`. +6. Warm any declared graph variants. +7. Drive `images`, `noise`, and `actions` through the C++ hot loop. + +In Lane A, prompt is setup-time. The current adopted-export face does not +export hot `prompt` or `state` ports. A request may repeat the setup prompt for +bookkeeping, but it cannot change the model prompt dynamically. + +### Lane B: Adopted Prompt/State Staging + +Use the same setup/adopt path as Lane A, but the producer additionally exports +real hot ports: + +- `prompt: TEXT/STAGED` +- `state: STATE/STAGED` + +The C++ host updates these ports with `cap_model_set_input` or the embedded +session equivalent. The producer formats, tokenizes, embeds, and writes the +fixed prompt window. Nexus remains unchanged. + +### Lane C: Native Producer + +Load a native FlashRT shared object and call: + +```c +int frt_model_runtime_open_v1(const char* config_json, + frt_model_runtime_v1** out); +``` + +The returned struct must expose the same public model-runtime contract as the +Python setup producer. The host and Nexus adoption code must not change when +switching between Lane A and Lane C. + +The current C++ shared object implements this symbol as a complete native-v2 +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 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. + +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. + +The native loader maps the checkpoint read-only and directly emits final BF16 +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 + +For same-process control loops, prefer Nexus embedded/session APIs over HTTP. +The high-level loop is: + +``` +producer setup -> frt_model_runtime_v1 +adopt -> cap_model_runtime +open embedded session +for each control tick: + update declared input ports + tick or fire stages + read declared output ports +optional: + snapshot / restore named capsules +``` + +The C++ loop should discover ports by name and then rely on the declared port +shape, dtype, direction, and update class. It should not hard-code `(10, 7)`, +graph names, or internal buffer names. + +## Port Update Rules + +For `SWAP` ports: + +- write the declared buffer window directly through the capsule/backend copy + mechanism; +- do not call `set_input`; +- verify byte count against `port.bytes`. + +For `STAGED` ports: + +- call the producer verb through `cap_model_set_input` or + `nexus_embedded_set_input`; +- pass bytes in the payload convention declared by `frt_model_runtime_v1`; +- expect shape/status errors for invalid input. + +For `SETUP` ports: + +- never update them inside a control tick. + +## Mapping Existing Mindon Calls + +| Mindon call | Integration point | +|---|---| +| `Prepare` | warm phase, producer `prepare(graph, key)` | +| `Warmup` | host policy: `prepare` plus warm ticks | +| `Infer` | `cap_model_tick`, `nexus_embedded_tick`, or explicit stage firing | +| `Sync` | host/backend stream sync or embedded session synchronization | +| `GetOutput` | `cap_model_get_output` / `nexus_embedded_get_output` | + +Do not introduce a second runtime API beside `frt_model_runtime_v1`. The +existing verbs already carry these phases. + +## Prompt and State + +Pi0.5 state is rendered into the language prompt. It is not an independent +model tensor. The producer path is: + +``` +raw proprioception -> normalize -> 256-bin discretize -> prompt string +-> token ids -> embedding gather -> encoder_x prompt window +``` + +Lane A still requires a setup-time producer refresh for prompt/state changes. +Lanes B and C accept task text through `prompt` and raw proprioception through +`state`. The producer owns all formatting and normalization details. + +## Image Input + +Mindon should pass camera frames as `frt_image_view[]` to the `images` +`IMAGE/STAGED` port, or through the matching Nexus embedded input. Frames are +matched by declared position, not by runtime graph names. + +The current Pi0.5 native producer stages host pixels into the +`observation_images_normalized` device tensor and normalizes to `[-1, 1]`. +Pass `u8` `RGB8` frames in HWC layout. BGR/RGBA/GRAY, CHW, and non-`u8` inputs +are rejected instead of silently reinterpreted. Use the producer documentation +in `docs/pi05_io_contract.md` for accepted formats and shape rules. + +## Action Output + +Read the `actions` port shape to determine chunk length and action dimension. +The output is the host-visible robot action chunk after producer postprocess. +For raw model action state, use `actions_raw` when the producer exports it. In +the Pi0.5 `native_v2` face this is a raw `TENSOR/SWAP` alias of +`diffusion_noise` with shape `(chunk, 32)`. + +## Capsule Boundaries + +Capsules snapshot exactly the producer-declared regions, in declared order. +Mindon should treat capsule contents as opaque bytes. A fingerprint mismatch +on restore is a deployment mismatch and must fail loudly. + +The native-v2 producer currently declares only `rollout_boundary`. Prompt +embedding, attention lengths, RoPE, and CPU prompt/state caches are not a +capsule region because partial restoration would be incorrect. If a future +face makes the entire prompt context restorable, its added ordered regions and +new fingerprint will intentionally reject old capsules. + +## Configuration Sketch + +Lane A setup in a Python producer plugin should export: + +```python +model = pipeline.export_model_runtime( + identity={"deployment": "mindon-pi05"}, + stage_plan="full", + io="native", +) +``` + +A split or RTC deployment may choose another producer-registered stage plan, +but the C++ host still sees only the adopted stage array. + +Lane C opens the native producer with a setup config such as: + +```json +{ + "io": "native_v2", + "checkpoint_path": "/models/pi05", + "tokenizer_model_path": "/models/paligemma/tokenizer.model", + "state_prompt_mode": "fixed", + "max_prompt_tokens": 200, + "state_dim": 8, + "num_views": 2, + "chunk": 10, + "num_steps": 10, + "vision_pool_factor": 1 +} +``` + +## Acceptance Checklist + +- The host discovers ports and shapes from `cap_model_runtime`. +- `images` updates use `STAGED`; `noise` updates use `SWAP`. +- `actions` capacity is computed from the declared output shape and dtype. +- The warm phase finishes before the first control tick. +- The hot loop performs no graph capture, allocation, or graph rebinding. +- Prompt/state/image/action staging capacities are fixed at setup; oversized + payloads fail instead of growing a hot-path workspace. +- Snapshot/restore is tested within one live capture. +- Nexus core code remains unchanged for model-specific semantics. diff --git a/docs/model_runtime_api.md b/docs/model_runtime_api.md index 7730bfc4..2d39b275 100644 --- a/docs/model_runtime_api.md +++ b/docs/model_runtime_api.md @@ -19,6 +19,11 @@ windows are `TENSOR`. A `STAGED` declaration is a promise the port accepts hot updates — a producer that cannot deliver that declares `SETUP` or omits the port, never advertise-and-refuse. +The Python `build_model_runtime()` helper enforces this mechanically: STAGED +inputs require `set_input`, and STAGED outputs require `get_output`. Internal +declaration-only handoffs may bypass that check only while constructing a +native verb overlay; the declaration is not an adoptable runtime. + ## Payload conventions (STAGED `set_input`) | modality | `data` points at | `bytes` | @@ -30,7 +35,7 @@ port, never advertise-and-refuse. ## Descriptors `frt_runtime_port_desc` — one dynamic input/output: -`name`, `modality`, `dtype` (device-side tensor), `layout`, `direction`, +`name`, `modality`, `dtype` (logical payload/target tensor), `layout`, `direction`, `update`, `required`, `shape[rank]` (−1 = bucket-variable), `cadence_hint_hz` (advisory only), and the SWAP window `buffer`/`offset`/ `bytes` (null buffer = staged-only). Strings/arrays are owned by the runtime @@ -58,8 +63,10 @@ frt_model_runtime_v1 { } ``` -**Verbs** (`frt_model_runtime_verbs`; every entry is always callable — absent -producer verbs are filled with unsupported stubs returning `-3`): +**Verbs** (`frt_model_runtime_verbs`; every entry on a successfully constructed +object is callable). Construction rejects a STAGED input without `set_input` +and a STAGED output without `get_output`. Other absent verbs are filled with +unsupported stubs returning `-3`: | verb | phase | semantics | |---|---|---| @@ -129,6 +136,11 @@ The override retains `producer_model`, so inherited port/stage pointers remain valid even if the original producer reference is released first. Deployment identity is unchanged. +The integrated, adapter, and verb-override C paths enforce the same STAGED +verb-presence rule before retaining owners or consuming builders. An internal +intermediate may exist while one factory assembles an override, but it must not +escape that factory or be independently adoptable. + **Native factory (symbol convention)** — a model-runtime `.so` exports `FRT_MODEL_RUNTIME_OPEN_V1_SYMBOL`: `int frt_model_runtime_open_v1(const char* config_json, frt_model_runtime_v1** out)`. @@ -266,3 +278,6 @@ PYTHONPATH=.:./exec/build:./runtime/build \ The consumer side (adoption, hot-input contract, real-model tick) is validated in the FlashRT-Nexus repository. + +For producer implementation and review rules, see +[`native_model_runtime_producer.md`](native_model_runtime_producer.md). diff --git a/docs/native_model_runtime_producer.md b/docs/native_model_runtime_producer.md new file mode 100644 index 00000000..0d9fd5d5 --- /dev/null +++ b/docs/native_model_runtime_producer.md @@ -0,0 +1,116 @@ +# 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. + +## 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 +structural boundary but may have different graph counts, internal buffers, +workspace, identities, and synchronization implementations. Validate each +producer's invariants independently. Compare only a deliberately shared +semantic face through checked-in canonical records. + +FlashRT supplies one exec implementation per process. Heterogeneous backend +instances enter above it through capsule backend vtables; do not add a runtime +backend registry or backend-kind ABI field. + +## Hot path + +Setup allocates storage, resolves names, loads weights, captures or adopts +graphs, and prepares variants. A control tick may update SWAP windows, execute +STAGED verbs, fire stages, and read output. It must not allocate device memory, +recapture, rebind graph pointers, or grow capacity. Oversized payloads fail. + +Measure CUDA allocator APIs over the complete service iteration. Host +allocation claims require a host allocation counter scoped to the component; +CUDA traces cannot prove host allocator behavior. + +## Schema workflow + +For a face implemented by more than one producer: + +1. Check in canonical `region:`, `port:`, and `stage:` records. +2. Generate records from every producer independently. +3. Diff each producer against the golden records. +4. Derive expected counts from the records instead of repeating constants. +5. Treat a golden update as a public contract review with fingerprint and + restore-compatibility analysis. + +Do not require implementation-private graph, buffer, manifest, or identity +records to match across backends. + +## Pull request evidence + +- C++ runtime tests in CUDA-off and affected hardware builds. +- Python runtime contract tests when the Python producer is supported. +- Producer-local lifecycle, schema, negative-input, and hot-loop gates. +- Numerical evidence appropriate to the boundary: bit-exact for identical + graph/input bytes; a documented fixed tolerance for genuinely different + backend math. +- Consumer adoption tests when descriptor or enum mapping changes. +- Migration notes for payload, fingerprint, packaging, or compatibility + changes. + +Use placeholders such as `` and `` in public commands. +Do not publish local paths, host/container names, credentials, environment +dumps, internal URLs, or proprietary dataset/checkpoint identifiers. diff --git a/docs/pi05_cpp_runtime_migration.md b/docs/pi05_cpp_runtime_migration.md new file mode 100644 index 00000000..5a0ff9ea --- /dev/null +++ b/docs/pi05_cpp_runtime_migration.md @@ -0,0 +1,60 @@ +# 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 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 + +A published runtime with STAGED input or output ports must install the matching +`set_input` or `get_output` verb. Declaration-only native handoff objects are +internal overlay inputs and are marked as such; they are not independently +adoptable runtimes. + +Port, stage, binding-window, stream-placement, and capsule-region changes alter +the runtime fingerprint. Capsules produced under an older fingerprint must be +regenerated; rejecting their restore is required behavior. + +## Native FA2 dependency + +The Python FA2 adapter and the Python-free `libflashrt_fa2_raw.so` are one +install unit. Native producers link the raw library and Python producers reach +the same symbols through the adapter. Deployment packages must install both in +the same directory so their `$ORIGIN` runtime lookup remains relocatable. diff --git a/docs/pi05_io_contract.md b/docs/pi05_io_contract.md new file mode 100644 index 00000000..cac778c3 --- /dev/null +++ b/docs/pi05_io_contract.md @@ -0,0 +1,729 @@ +# 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. + +## 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 +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 | +|---|---|---|---|---| +| `images` | `IMAGE/STAGED` | input | device tensor dtype, `NHWC`, `(num_views, 224, 224, 3)` | `observation_images_normalized` | +| `noise` | `TENSOR/SWAP` | input | device tensor dtype, `FLAT`, `(chunk_length, 32)` | `diffusion_noise` | +| `actions` | `ACTION/STAGED` | output | host `f32`, `FLAT`, `(chunk_length, robot_action_dim)` | staged-only; no raw window | + +Current source of truth: + +- Export declaration: `flash_rt/models/pi05/runtime_export.py`, + `export_model_runtime(..., io="native")` +- Native verb implementation: `cpp/models/pi05/src/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 +`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 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. + +## Complete Native V2 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 | +|---|---|---|---|---| +| `prompt` | `TEXT/STAGED` | input | UTF-8 bytes, `FLAT`, variable length | staged by C++ runtime | +| `state` | `STATE/STAGED` | input | host `f32`, `FLAT`, `(state_dim,)` | staged by C++ runtime | +| `images` | `IMAGE/STAGED` | input | device tensor dtype, `NHWC`, `(num_views, 224, 224, 3)` | `observation_images_normalized` | +| `noise` | `TENSOR/SWAP` | input | device tensor dtype, `FLAT`, `(chunk_length, 32)` | `diffusion_noise` | +| `actions` | `ACTION/STAGED` | output | host `f32`, `FLAT`, `(chunk_length, robot_action_dim)` | staged-only; no raw window | +| `actions_raw` | `TENSOR/SWAP` | output | device tensor dtype, `FLAT`, `(chunk_length, 32)` | `diffusion_noise` | + +For Pi0.5, proprioceptive state is not an independent model tensor. It is +normalized, discretized into OpenPI-compatible 256-bin state tokens, rendered +into the prompt text, tokenized, embedded, and written into the language rows +of `encoder_x`. Therefore prompt and state updates are one producer-owned text +staging path. + +`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 +of a restorable boundary. + +## STAGED Payloads + +The payload conventions are inherited from `frt_model_runtime_v1`. + +| modality | `set_input` data | bytes | +|---|---|---| +| `IMAGE` / `DEPTH` | `frt_image_view[]` | `n_frames * sizeof(frt_image_view)` | +| `TEXT` | UTF-8 bytes | byte length | +| `TENSOR` / `STATE` / `ACTION` / `AUDIO` | raw bytes per the port dtype and shape | byte length | + +For `IMAGE`, frames are matched positionally to the producer-declared camera +view order. The Pi0.5 view order is: + +1. `image` +2. `wrist_image` +3. `wrist_image_right` + +Deployments with fewer views export a shorter `num_views` and use the prefix +of that view order. + +## Image Input + +The current native image input accepts host `frt_image_view[]` and stages the +data into the device `observation_images_normalized` buffer before replay. + +Producer-owned preprocessing: + +- view count is checked against the exported `images` port shape; +- frame payloads are host `u8` pixels in `RGB8`/`HWC` layout; +- target tensor is `(num_views, 224, 224, 3)`; +- output layout is `NHWC`; +- output dtype is the exported tensor dtype: 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 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 +non-`u8` frames are not silently converted at the Pi0.5 contract boundary. If a +deployment supports more pixel formats, the supported set must be documented by +the producer and tested against the CPU reference path. + +Compatibility note: the legacy `frt_pi05_runtime_prepare_vision` C API keeps +accepting its explicit `BGR8`, `RGBA8`, `BGRA8`, and `GRAY8` enum values and +converts them through the shared modality implementation. This compatibility +does not expand the `frt_model_runtime_v1` native face: its `IMAGE/STAGED` +deployment contract remains strict `RGB8/u8/HWC`. + +## Noise Input + +`noise` is a `TENSOR/SWAP` port. The host writes its raw bytes directly into +the `diffusion_noise` window, usually by `cap_swap` after Nexus adoption or by +the equivalent runtime/backend copy mechanism. Calling `set_input` on this +port is unsupported by design: SWAP means the device window is the interface. + +Shape is `(chunk_length, 32)`. `chunk_length` is declared by the producer and +must be read from the port shape; host code must not assume `(10, 32)`. + +## Action Output + +`actions` is the host-visible robot action chunk after producer-owned +postprocess. Its declared dtype is F32 because that is the payload returned by +`get_output`; it does not expose the backend's diffusion window as backing. + +The logical output shape is: + +``` +(chunk_length, robot_action_dim) +``` + +For LIBERO-style Pi0.5 deployments, `robot_action_dim` is typically 7. Other +deployments may export a different fixed robot action dimension. Consumers and +schedulers must read the declared port shape instead of hard-coding `(10, 7)`. + +The internal model output remains `(chunk_length, 32)` in `diffusion_noise`. +The native `actions` STAGED output slices the robot dimensions and applies the +deployment action normalization statistics. With q01/q99 stats, the affine +parameters are: + +``` +mean = (q01 + q99) / 2 +stddev = (q99 - q01) / 2 +``` + +The C++ postprocess path clamps normalized action values to the configured +domain before applying the affine transform. Any raw `(chunk_length, 32)` face +must be exported as a separate `TENSOR/SWAP` output. The Pi0.5 `native_v2` +face declares this as `actions_raw`; RTC stage plans also use the same port +name. Nexus must treat it as a declared raw byte window, not model internals. + +## Lifecycle Mapping + +Mindon-style lifecycle names map to the existing ABI. Do not add a parallel +API family for the same phases. + +| requested name | existing contract | phase | +|---|---|---| +| `Prepare` | `prepare(graph, key)` | warm only | +| `Warmup` | host policy: call `prepare` for needed variants, then run warm ticks | warm | +| `Infer` | `step()` sugar or host-scheduled stage replay | hot | +| `Sync` | host/backend stream synchronization | hot or drain | +| `GetOutput` | `get_output(port, out, capacity, &written, stream)` | hot | + +`prepare` is the only place a shape-bucket miss may capture or materialize a +variant. A hot tick must not recapture, allocate, or rebind graph pointers. +The Pi0.5 C++ face fixes its vision frames, action D2H staging, task/formatted +prompt strings, tokenizer ids, and normalized-state storage during setup. +Payloads that would grow those workspaces return a shape/capacity error; there +is no larger-allocation fallback in the hot path. These workspace changes do +not alter the port schema or deployment fingerprint. + +## Identity and Capsule Regions + +The following changes are deployment identity changes: + +- adding/removing/reordering ports; +- changing a port modality, dtype, layout, direction, update class, required + flag, shape, bound buffer index, offset, or byte window; +- changing graph names or default stream placement; +- changing the stage DAG; +- adding/removing/reordering capsule regions; +- changing a region name, buffer, offset, byte length, or flags. + +The following are not deployment identity changes: + +- editing `manifest_json`; +- changing `cadence_hint_hz`. + +Prompt/state staging does not by itself make prompt context a capsule region. +A restorable prompt context would have to include the embedding, attention +lengths, decoder position/RoPE, and the CPU semantic cache used by later +independent prompt/state updates. The current face can rebuild those values +from its declared inputs, so it does not advertise partial prompt restoration. +Region layout and order are fingerprinted; adding a complete prompt region in +a later face will intentionally invalidate old capsules. + +## Current Integration Lanes + +There are three supported integration lanes: + +- Lane A, current: Python setup/capture/export stays resident in the process; + the hot loop adopts `frt_model_runtime_v1` and runs through C++/Nexus. +- Lane B: an adopted setup producer exposes real hot `prompt`/`state` STAGED + ports and the C++ overlay owns their transforms. +- Lane C, current on 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, SentencePiece, and the backend +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. 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 +verify the complete 812-tensor Pi0.5 inventory: all 27 vision layers, all 18 +language encoder layers, all 18 action-expert layers, embeddings/final norms, +projectors, action projections, and time MLP. It also verifies action/state q01/q99 +dimensions from either openpi `norm_stats.json` or LeRobot policy +normalizer/unnormalizer safetensors. Safetensors tensor byte ranges must match +dtype/shape, and normalization quantiles must be finite ordered pairs. Builds +without 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, +`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. + +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 +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 +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, 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. + +On a representative Thor SM110 Release run with the same checkpoint size, +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 +capsule regions. Upload is complete before capture; duplicate logical names or +typed shape/payload mismatches fail setup. The same store carries BF16, FP8 +E4M3, INT8, and FP32 scale buffers without introducing a model-level state +object. Destroying the context releases the device weights after graphs and +plans, preserving the exec ownership order. + +The composed materializer covers language-encoder, action-expert, and vision +weights. Encoder layers emit the five pipeline groups (`attn_qkv`, `attn_o`, +`ffn_gate`, `ffn_up`, and `ffn_down`) with FP32 RMS folds. Decoder layers emit +those groups plus the four AdaRMS modulation tensors and the optional merged +gate/up buffer used by the FP16 path. Vision setup emits patch/position/final +norm and multimodal-projector globals plus the twelve per-layer attention, +FFN, and normalization buffers. Decoder globals include final AdaRMS +modulation, time MLP, generated time embeddings, and action projections. The +action output projection is pre-scaled by `-1/num_steps` after source BF16 +rounding; 5-step and 10-step schedules are byte-exact with the PyTorch +producer. The prompt embedding table is materialized separately to keep its +approximately 1 GiB allocation explicit. These paths have been exercised +against the two supported real checkpoint layouts. The checkpoint inventory +also validates the language final norm and expert LM head even though the +current Pi0.5 pipeline does not consume them. The native producer materializes +the full BF16 store before capture and keeps it under the graph context lifetime. + +Native setup quantization reproduces the PyTorch producer's per-tensor FP8 +E4M3 weights in either `kn` or `nk` layout and per-output-channel INT8 weights +in `[N,K]` layout. FP8 scalar descales and INT8 channel scales are FP32 device +buffers. Real-checkpoint gates compare both quantized bytes and scale bytes; +the precision choice remains producer setup policy and does not alter ports, +regions, or the exec mechanism. + +The setup packer derives low-precision buffers from the already uploaded BF16 +fallback, so both paths share exactly the same transformed source bytes. It +stores packed weights under `fp8.*` or `int8.*` names and their typed FP32 +scales under the matching `.scale` names in the same context-owned store. + +Full BF16 assembly has one ordered setup path: vision globals and 27 layers, +18 language-encoder layers, 18 action-expert layers, decoder globals, then the +prompt embedding table. With merged decoder gate/up buffers enabled this owns +613 logical device buffers. Assembly options make `num_steps`, merged gate/up, +and the large embedding allocation explicit; they are producer configuration, +not ABI fields. + +Full FP8 packing follows the producer's exact site inventory: four GEMM +weights for each vision layer plus the projector, four for each encoder layer, +and four for each decoder layer. Encoder gate/up columns are merged during +setup. INT8 packing remains independently selectable for vision, encoder, and +decoder and preserves their existing four/five/five weights-per-layer policy. + +The native kernel layer is CPython-independent and links the existing +`GemmRunner` implementation directly. Setup warms required BF16 GEMM shapes, +captures the complete `infer` graph through `frt_graph_capture`, and exports +exactly one shape-key variant (`0`). + +The native core workspace maps every vision, encoder, decoder, style, action, +RTC, and reusable scratch allocation to a context-owned `frt_buffer`. There is +no model-level State object. With vision pooling disabled, `vision_x_pooled` +is an explicit alias of `vision_x` (34 logical names, 33 allocations); pooled +deployments allocate it separately. Buffer shapes are fixed from `num_views`, +`max_prompt_tokens`, `chunk_size`, `num_steps`, and `vision_pool_factor` before +capture, and BF16 RMS-one constants, attention backend storage, and generated +decoder style contents are initialized during setup. + +Native RoPE setup uses the same float64 frequency/phase computation and BF16 +interleaved `[cos, sin]` layout as the Python producer. Encoder and +prompt-relative decoder slices are byte-exact against NumPy/ml_dtypes for +pooled and unpooled configurations. Decoder slice updates reuse one stable +buffer across prompt lengths; vision position embeddings are expanded per view +with setup-side D2D copies from the typed weight store. + +Decoder time/style precompute is also native setup work. It consumes the +generated time embeddings, time-MLP weights, 18 layers of AdaRMS modulation, +and final modulation from the typed store; it reuses existing workspace +buffers as scratch and writes the four persistent style buffers without a +temporary device allocation. The GEMM, explicit BF16 bias round-trip, and +float-SiLU sequence is BF16 bit-exact with the PyTorch producer on both +supported checkpoint layouts. + +The native kernel driver also owns the BF16 forward primitives used around +GEMM and attention: RMS/Layer/AdaRMS normalization, residual and gated +residual updates, GELU/gated GELU, QKV split with fixed or device-position +RoPE, patch im2col, and vision pooling. These are direct typed calls to the +existing CUDA implementations, with CPU-reference and captured-replay gates; +they do not route through pybind or introduce a second kernel implementation. + +The native vision graph composes patch im2col/embedding, all 27 SigLIP layers, +per-view FA2, optional fixed-factor spatial pooling, final LayerNorm, and the +1152-to-2048 multimodal projector. Position embedding expansion remains setup +work. With inputs restored before each of 100 replays, the graph keeps one +variant; final SigLIP and projected encoder tokens reach cosine 0.9999 or +better against the layer-by-layer PyTorch reference on both supported +checkpoint layouts. + +The first composed BF16 forward segment is the encoder QKV path: +RMSNorm, the folded QKV projection, RoPE split, and writes into the selected +layer of the shared K/V cache. Layer 17 is also the complete final encoder +layer behavior because the producer intentionally stops after populating its +cache. Its outputs are bit-exact (`cos=1`, `max=0`) against the PyTorch +checkpoint path for both OpenPI and LeRobot layouts, and the segment captures +and replays with one graph variant. + +Encoder layers 0-16 extend that segment through fixed-shape FA2, output +projection, residual/RMS normalization, the separate gate/up projections, +gated GELU, down projection, and the final residual update. A captured layer 0 +replayed 100 times remains a single variant and reaches cosine 0.999992 versus +the original PyTorch path on both checkpoint layouts. Layer 17 keeps the +intentional cache-only early exit described above. + +The native encoder composes all 18 layers into one captured graph while +preserving that final cache-only behavior. Restoring the input before each of +100 replays produces one graph variant. On both OpenPI and LeRobot checkpoint +layouts, the final encoder state and layer-17 Q/K/V each reach cosine 0.9999 or +better against the layer-by-layer PyTorch reference. This composition owns no +state object: activations and K/V remain context-backed buffers. + +The native decoder composes one BF16 AdaRMS/cross-attention/FFN layer, one +flow-matching update, and the complete 10-step diffusion graph. Decoder K/V is +appended at the device-side fixed-prompt position in the encoder cache; style +and noise remain context-backed buffers. Full 10-step captures replay 100 times +with one variant on both checkpoint layouts. Independent first and final +schedule steps reach cosine 0.9999 or better against PyTorch; the accumulated +endpoint gate remains part of the real-episode end-to-end validation because +synthetic random K/V amplifies SDPA-versus-FA2 rounding across steps. + +The native graph owner now assembles the completed segments into one `infer` +capture: prompt copy, vision, encoder, then diffusion. Prompt embeddings live +in a separate persistent buffer because `encoder_x` is an in-place residual +stream; each replay captures a D2D copy into its language window. Both +checkpoint layouts complete 100 full replays with one variant, bit-identical +outputs for restored inputs, and a constant workspace allocation count. The +persistent prompt source, not the overwritten encoder rows, is the primary +prompt-context capsule candidate. + +RTX attention owns a separate context-backed buffer set rather than borrowing +Torch tensors: SigLIP Q/K/V, encoder Q and 18-layer shared K/V cache, decoder +Q, fixed-shape `seqused/devpos` int32 values, FA2 outputs/LSE, and split-KV +accumulators. Layer K/V pointers are stable offsets into one cache allocation. +Updating a fixed prompt length writes the same three scalar buffers without +allocation or rebinding. The Python-free attention driver calls the vendored +FA2 raw C entries directly for SigLIP, fixed-shape encoder `seqused`, and +decoder `seqused` split-KV. Its graph gate changes the prompt length after +capture, replays 100 times with one variant, and verifies the new device-side +valid length is observed. `flash_rt_fa2` remains a thin Python adapter over the +same `libflashrt_fa2_raw` kernel owner. + +The native builder publishes `infer`, `context`, and `decode_only` graphs and +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 +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 +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 +order without a second lifecycle owner. + +CUDA graph execs are process-local objects. They are not serialized as a +portable artifact. The native producer therefore loads assets and captures the +graph in the replay process. + +## Validation + +The minimum regression set for this contract is: + +``` +PYTHONPATH=.:./exec/build:./runtime/build python runtime/tests/test_runtime_export.py +PYTHONPATH=.:./exec/build:./runtime/build python runtime/tests/test_model_runtime_py.py +./runtime/build/test_model_runtime +ctest --test-dir cpp/build --output-on-failure +``` + +Detailed real-checkpoint oracles are maintained in the +[MindOn validation suite](https://github.com/LiangSu8899/MindOn-dev/tree/main/validation/pi05_cpp). +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: + +``` +python /oracles/gate_pi05_native_weight_ops.py \ + --checkpoint \ + --probe /pi05_native_weight_probe +``` + +``` +FLASHRT_BUILD_DIR= \ + python cpp/tests/gate_pi05_model_runtime_export.py \ + --lib /libflashrt_cpp_pi05_c.so ... +python cpp/tests/gate_pi05_c_api_export.py ... +``` + +The Python-produced overlay gate uses the FA2-enabled, SentencePiece-off SM120 +build because Python already owns prompt/tokenizer setup. Native-v2 factory +gates use a SentencePiece-enabled build with either SM120 FA2 or SM110 Thor FP8. +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 +producer must not retain the declarations if any required verb is unavailable. + +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 variant per captured graph, prompt/state/image staging, +direct SWAP noise input, finite action output, and retain/release teardown. + +SM110 and SM120 FP8 calibration/runtime math are gated separately against a +Torch producer built from the same source revision: + +``` +python /oracles/gate_pi05_native_thor_fp8.py \ + --probe /pi05_native_fp8_calibration_probe \ + --checkpoint \ + --tokenizer \ + --artifact \ + --samples 1 --views 1 + +python /oracles/gate_pi05_native_calibration.py \ + --probe /pi05_native_fp8_calibration_probe \ + --checkpoint \ + --tokenizer \ + --artifact \ + --samples 3 --views 3 +``` + +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 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 +different): + +``` +FLASHRT_BUILD_DIR= \ + python cpp/tests/gate_pi05_native_schema_parity.py \ + --checkpoint \ + --tokenizer \ + --native-probe /pi05_native_open_probe +``` + +Both producers are compared independently against +`cpp/tests/data/pi05_native_v2_schema.records`. Update that golden file only +for an intentional public-face change, and review the resulting fingerprint +and capsule compatibility impact in the same change. + +The native formatter and tokenizer must also remain token-exact over real +prompt/state traffic: + +``` +python /oracles/gate_pi05_tokenizer_corpus.py \ + --dataset \ + --checkpoint \ + --tokenizer \ + --probe /pi05_tokenizer_corpus_probe \ + --count 10000 +``` + +This gate normalizes every recorded state with the checkpoint q01/q99 values, +renders the full state prompt through the native formatter, and compares every +valid token ID with OpenPI `PaligemmaTokenizer`. The 10,000-record reference +run used a lightweight oracle whose tokenizer and formatter logic is kept +line-equivalent with upstream OpenPI; it covered 20 token lengths from 43 +through 62 with zero mismatches. This source-level oracle is distinct from the +official-environment end-to-end gate below. + +The real-episode numerical gate compares against the official OpenPI PyTorch +`PI0Pytorch.sample_actions` path, not another native intermediate: + +``` +python /oracles/gate_pi05_native_e2e.py \ + --checkpoint \ + --tokenizer \ + --dataset \ + --probe /pi05_native_e2e_probe \ + --episode 0 --frame 0 +``` + +The gate rounds the shared initial noise to the exported BF16 contract before +both runs. Raw and robot action outputs must each reach cosine 0.9999 against +the official FP32 residual path. Separately, the STAGED `actions` bytes must +match q01/q99 postprocess recomputed from the native BF16 `actions_raw` window +at `rtol=atol=1e-6`; this keeps numerical precision and IO semantics as two +independent acceptance checks. 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 +remain outside the CUDA profiler range: + +``` +FLASHRT_PROFILE_RANGE=1 nsys profile --trace=cuda \ + --cuda-graph-trace=node \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + /pi05_native_open_probe \ + + +nsys profile --trace=cuda --cuda-graph-trace=node \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + python /perf/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 /perf/gate_pi05_kernel_sequence.py \ + --native .csv --python .csv +``` + +The comparator rejects unknown kernels and requires equal raw event counts. +Its explicit equivalence list covers only selected GEMM kernel variants, +GEMM workspace-init versus split-K reduction helpers, `add_bias` versus the +equivalent `bias_res` form, and the two negative-infinity fill symbols. On the +reference RTX 5090 SM120 run both traces contained 3,576 raw events and their +3,172 logical-kernel sequences were exactly equal. + +The separate hot allocator gate profiles 1,000 complete service iterations +without tracing individual graph nodes. Each measured iteration updates prompt, +state, image, and noise inputs, launches one graph replay, and reads the logical +action output: + +``` +FLASHRT_PROFILE_REPLAYS=1000 FLASHRT_PROFILE_SERVICE_LOOP=1 \ +nsys profile --trace=cuda \ + --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o \ + /pi05_native_open_probe \ + [calibration.safetensors] +nsys stats --report cuda_api_trace --format csv \ + .nsys-rep > .csv +python cpp/tests/gate_pi05_hot_allocator.py \ + --trace .csv --expected-replays 1000 +``` + +The gate requires exactly 1,000 `cudaGraphLaunch` calls and rejects CUDA/driver +device allocation, host registration, mempool creation, virtual-memory map, +and corresponding release APIs. The probe independently requires one graph +variant after the final replay. Omitting `FLASHRT_PROFILE_SERVICE_LOOP` retains +the replay-only diagnostic mode for kernel-sequence comparison. +This trace proves the absence of CUDA/driver allocation APIs across the full +service iteration. Host allocation claims are scoped to components with an +explicit allocation-counter test, such as exec graph-cache LRU maintenance; +the trace does not infer host allocator behavior from CUDA API events. + +The same native probe can gate the complete hot state staging chain +(normalization, formatting, tokenization, embedding gather, and prompt-length +device update): + +``` +FLASHRT_HOT_STATE_UPDATES=1000 FLASHRT_HOT_STATE_P99_US=1000 \ + /pi05_native_open_probe \ + [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. 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, +releases the final model reference, and only then unloads the producer: + +``` +/pi05_native_dlopen_probe \ + /libflashrt_cpp_pi05_c.so \ + 1 +``` + +For the ASAN build, instrument the producer, runtime, exec library, and probe, +not only the executable. CUDA needs `protect_shadow_gap=0` in this environment +to avoid an address-space collision with ASAN's default shadow gap: + +``` +ASAN_OPTIONS=detect_leaks=1:halt_on_error=1:protect_shadow_gap=0 \ + /pi05_native_dlopen_probe \ + /libflashrt_cpp_pi05_c.so \ + 1 +``` diff --git a/docs/pi05_thor_native_fp8.md b/docs/pi05_thor_native_fp8.md new file mode 100644 index 00000000..c6af8aeb --- /dev/null +++ b/docs/pi05_thor_native_fp8.md @@ -0,0 +1,439 @@ +# Pi0.5 Native C++ FP8 + +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 +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 | +| 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. 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 + +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" +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. + +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=`. 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 + +Calibration and runtime open consume the same model configuration. Calibration +does not need `calibration_path`; every FP8 runtime open requires it. + +```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. + +| 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 +`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-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 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 +artifact until an identity input or calibration policy changes. + +## Artifact Contract + +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; +- 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-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 `infer`, +`decode_only`, and `context` graphs. The selected stage plan is either one +`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 | +|---|---|---|---|---| +| `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 +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 +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 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 + +Run CPU/CUDA-off tests in addition to SM110 Release and Debug builds: + +```bash +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 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 \ + --probe "$BUILD_DIR/pi05_native_fp8_calibration_probe" \ + --checkpoint "$CHECKPOINT_DIR" \ + --tokenizer "$TOKENIZER_MODEL" \ + --artifact "$CALIBRATION_FILE" \ + --samples 1 --views 1 + +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 3 +``` + +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 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. 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: + +```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 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 | 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 +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. + +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. diff --git a/docs/pr_review_checklist.md b/docs/pr_review_checklist.md index 2a7d78dd..25e39b24 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: @@ -440,6 +454,21 @@ 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. +- 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: @@ -450,6 +479,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: @@ -483,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 @@ -630,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. @@ -688,3 +736,32 @@ 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. +- [ ] 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. diff --git a/docs/runtime_contract.md b/docs/runtime_contract.md index f4420b16..b8bdccf0 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,22 @@ 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 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 +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 diff --git a/exec/CMakeLists.txt b/exec/CMakeLists.txt index 2bcc975e..eaa6837a 100644 --- a/exec/CMakeLists.txt +++ b/exec/CMakeLists.txt @@ -33,6 +33,15 @@ target_include_directories(flashrt_exec ${CMAKE_CURRENT_SOURCE_DIR}/backend) target_link_libraries(flashrt_exec PUBLIC CUDA::cudart) +if(BUILD_TESTING) + add_executable(test_graph_lru_alloc tests/test_graph_lru_alloc.cpp) + target_include_directories(test_graph_lru_alloc PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/src) + target_link_libraries(test_graph_lru_alloc PRIVATE flashrt_exec) + add_test(NAME graph_lru_alloc COMMAND test_graph_lru_alloc) +endif() + # --- pybind dev module (optional) --- find_package(Python3 COMPONENTS Interpreter Development.Module QUIET) find_package(pybind11 CONFIG QUIET) diff --git a/exec/src/graph.cpp b/exec/src/graph.cpp index 100c67a3..b5a776fc 100644 --- a/exec/src/graph.cpp +++ b/exec/src/graph.cpp @@ -4,7 +4,10 @@ void frt_graph_s::touch(frt_shape_key key) { for (auto it = lru.begin(); it != lru.end(); ++it) { - if (*it == key) { lru.erase(it); break; } + if (*it == key) { + lru.splice(lru.end(), lru, it); + return; + } } lru.push_back(key); // back = most recently used } diff --git a/exec/tests/test_graph_lru_alloc.cpp b/exec/tests/test_graph_lru_alloc.cpp new file mode 100644 index 00000000..f3f572e5 --- /dev/null +++ b/exec/tests/test_graph_lru_alloc.cpp @@ -0,0 +1,46 @@ +#include "internal.h" + +#include +#include +#include +#include + +namespace { + +std::atomic count_allocations{false}; +std::atomic allocation_count{0}; + +} // namespace + +void* operator new(std::size_t bytes) { + if (count_allocations.load(std::memory_order_relaxed)) { + allocation_count.fetch_add(1, std::memory_order_relaxed); + } + if (void* p = std::malloc(bytes)) return p; + throw std::bad_alloc(); +} + +void operator delete(void* p) noexcept { std::free(p); } +void operator delete(void* p, std::size_t) noexcept { std::free(p); } + +int main() { + frt_graph_s graph; + graph.lru = {1, 2, 3}; + + allocation_count.store(0, std::memory_order_relaxed); + count_allocations.store(true, std::memory_order_relaxed); + for (int i = 0; i < 1000; ++i) graph.touch((i % 3) + 1); + count_allocations.store(false, std::memory_order_relaxed); + + assert(allocation_count.load(std::memory_order_relaxed) == 0); + assert(graph.lru.size() == 3); + assert(graph.lru.back() == 1); + + allocation_count.store(0, std::memory_order_relaxed); + count_allocations.store(true, std::memory_order_relaxed); + graph.touch(4); + count_allocations.store(false, std::memory_order_relaxed); + assert(allocation_count.load(std::memory_order_relaxed) > 0); + assert(graph.lru.back() == 4); + return 0; +} diff --git a/flash_rt/frontends/torch/pi05_rtx.py b/flash_rt/frontends/torch/pi05_rtx.py index 5edf211d..6601feb6 100644 --- a/flash_rt/frontends/torch/pi05_rtx.py +++ b/flash_rt/frontends/torch/pi05_rtx.py @@ -1071,6 +1071,7 @@ def _set_prompt_fixed(self, prompt_len: int) -> None: vision_pool_factor=self._vision_pool_factor, vision_num_layers=self._vision_num_layers, fixed_shape=True, + norm_stats=self.norm_stats, **self._pipeline_precision_kwargs()) if self._fixed_pipeline.use_int8_vision_static: self._fixed_pipeline.vis_int8_static_calibrated = False @@ -1121,6 +1122,7 @@ def _set_prompt_per_length(self, state, prompt_len: int) -> None: num_steps=self._num_steps, vision_pool_factor=self._vision_pool_factor, vision_num_layers=self._vision_num_layers, + norm_stats=self.norm_stats, **self._pipeline_precision_kwargs()) self._prompt_pipeline_cache[prompt_len] = self.pipeline # Static INT8 vision scales are per-pipeline-instance. diff --git a/flash_rt/frontends/torch/pi05_rtx_fp16.py b/flash_rt/frontends/torch/pi05_rtx_fp16.py index f0e78266..3d9b6121 100644 --- a/flash_rt/frontends/torch/pi05_rtx_fp16.py +++ b/flash_rt/frontends/torch/pi05_rtx_fp16.py @@ -989,6 +989,7 @@ def set_prompt(self, prompt_text: str, state=None) -> None: num_steps=self._num_steps, vision_pool_factor=self._vision_pool_factor, vision_num_layers=self._vision_num_layers, + norm_stats=self.norm_stats, **self._pipeline_precision_kwargs()) self._prompt_pipeline_cache[prompt_len] = self.pipeline # Static INT8 vision scales are per-pipeline-instance. diff --git a/flash_rt/models/pi05/pipeline_rtx.py b/flash_rt/models/pi05/pipeline_rtx.py index 01be1102..4c43c5ac 100644 --- a/flash_rt/models/pi05/pipeline_rtx.py +++ b/flash_rt/models/pi05/pipeline_rtx.py @@ -137,6 +137,7 @@ class Pi05Pipeline: use_fp8_decoder: Enable FP8 on decoder branch (else BF16). use_int8_decoder: Enable experimental decoder-only INT8 GEMMs. num_steps: Diffusion denoise steps (default 10). + norm_stats: Producer metadata for logical state/action IO contracts. Expected weights dict keys: Vision BF16: @@ -185,11 +186,13 @@ def __init__(self, gemm, fvk, attn_backend, weights, *, vision_pool_factor: int = 1, vision_num_layers: int = VIS_L, num_steps: int = NUM_STEPS_DEFAULT, - fixed_shape: bool = False): + fixed_shape: bool = False, + norm_stats=None): self.gemm = gemm self.fvk = fvk self.attn = attn_backend self.weights = weights + self.norm_stats = norm_stats or {} # Fixed-shape state-prompt mode: one captured graph at the MAX prompt # length serves every length via seqused masking + devpos K/V append. diff --git a/flash_rt/models/pi05/pipeline_rtx_fp16.py b/flash_rt/models/pi05/pipeline_rtx_fp16.py index 09b92584..59065c41 100644 --- a/flash_rt/models/pi05/pipeline_rtx_fp16.py +++ b/flash_rt/models/pi05/pipeline_rtx_fp16.py @@ -186,6 +186,7 @@ class Pi05PipelineFP16: use_fp8_decoder: Enable FP8 on decoder branch (else BF16). use_int8_decoder: Enable experimental decoder-only INT8 GEMMs. num_steps: Diffusion denoise steps (default 10). + norm_stats: Producer metadata for logical state/action IO contracts. Expected weights dict keys: Vision BF16: @@ -233,7 +234,8 @@ def __init__(self, gemm, fvk, attn_backend, weights, *, use_int8_vision_static: bool = False, vision_pool_factor: int = 1, vision_num_layers: int = VIS_L, - num_steps: int = NUM_STEPS_DEFAULT): + num_steps: int = NUM_STEPS_DEFAULT, + norm_stats=None): if use_fp8 or use_fp8_decoder: raise ValueError("Pi05PipelineFP16 supports only use_fp8=False") if use_int8_decoder or use_int8_encoder or use_int8_vision or use_int8_vision_static: @@ -244,6 +246,7 @@ def __init__(self, gemm, fvk, attn_backend, weights, *, self.fvk = _Fp16KernelProxy(fvk) self.attn = attn_backend self.weights = weights + self.norm_stats = norm_stats or {} self.num_views = int(num_views) self.max_prompt_len = int(max_prompt_len) diff --git a/flash_rt/models/pi05/runtime_export.py b/flash_rt/models/pi05/runtime_export.py index 9ecb55fa..ae2cf05c 100644 --- a/flash_rt/models/pi05/runtime_export.py +++ b/flash_rt/models/pi05/runtime_export.py @@ -9,6 +9,8 @@ from __future__ import annotations +import hashlib + def exec_enable(pl) -> None: """Create the exec ctx/graphs for a captured pipeline and adopt any @@ -84,13 +86,25 @@ def export_model_runtime(pl, identity=None, extra_regions=None, graphs: images/actions are STAGED, noise is SWAP, and the C++ runtime supplies the verbs through ``frt_model_runtime_override_verbs``. + ``io="native_v2"`` extends that face with prompt/state STAGED ports for + fixed state-prompt deployments. The declaration is intended to be consumed + through the C++ verb overlay; port/window/region changes are part of the + export identity and therefore intentionally change the fingerprint. + Prompt staging (text -> embeds) stays with the frontend / the native tokenizer producer. ``stage_plan`` defaults to the full infer graph; an explicit StagePlan or registered plan name may select already-captured graphs from this export. ``stage_plan_kwargs`` are passed only to registered plan factories, for deployment-specific graph cuts. """ - parts = _parts(pl, identity, extra_regions) + identity_for_parts = identity + if io == "native_v2": + _require_native_v2_ready(pl) + identity_for_parts = { + **{str(k): str(v) for k, v in (identity or {}).items()}, + "io": "native_v2", + } + parts = _parts(pl, identity_for_parts, extra_regions) from flash_rt.runtime import export as _rt from flash_rt.subgraphs.pi05 import stage_plans as _pi05_stage_plans # noqa: F401 from flash_rt.subgraphs.stage_plan import resolve_stage_plan @@ -138,7 +152,7 @@ def export_model_runtime(pl, identity=None, extra_regions=None, "in", "swap", shape=(1,), buffer=wrap["rtc_guidance_weight"]), ]) - elif io == "native": + elif io in ("native", "native_v2"): ports = [ _rt.PortSpec("images", "image", tensor_dtype, "nhwc", "in", "staged", required=True, shape=(num_views, 224, 224, 3), @@ -146,10 +160,23 @@ def export_model_runtime(pl, identity=None, extra_regions=None, buffer=wrap["observation_images_normalized"]), _rt.PortSpec("noise", "tensor", tensor_dtype, "flat", "in", "swap", shape=(chunk, 32), buffer=wrap["diffusion_noise"]), - _rt.PortSpec("actions", "action", tensor_dtype, "flat", "out", + _rt.PortSpec("actions", "action", "f32", "flat", "out", "staged", shape=(chunk, robot_action_dim), - buffer=wrap["diffusion_noise"]), + nbytes=chunk * robot_action_dim * 4), ] + if io == "native_v2": + state_dim = _state_dim(pl) + ports = [ + _rt.PortSpec("prompt", "text", "u8", "flat", "in", "staged", + required=True, shape=(-1,)), + _rt.PortSpec("state", "state", "f32", "flat", "in", "staged", + required=True, shape=(state_dim,)), + ] + ports + if not (uses_rtc_prefix or uses_rtc_vjp): + ports.append( + _rt.PortSpec("actions_raw", "tensor", tensor_dtype, + "flat", "out", "swap", shape=(chunk, 32), + buffer=wrap["diffusion_noise"])) if uses_rtc_prefix or uses_rtc_vjp: ports.extend([ _rt.PortSpec("prev_action_chunk", "tensor", tensor_dtype, @@ -193,6 +220,16 @@ def step(): return rc return rc + manifest_extra = {"stage_plan": plan.manifest(), "io": io} + declaration_only = io in ("native", "native_v2") + if declaration_only: + manifest_extra["declaration_only"] = True + if io == "native_v2": + manifest_extra["prompt"] = { + "state_prompt_mode": "fixed", + "max_prompt_len": int(getattr(pl, "max_prompt_len", 0) or 0), + "state_dim": _state_dim(pl), + } return _rt.build_model_runtime( pl._exec_ctx, streams=parts["streams"], @@ -202,9 +239,10 @@ def step(): ports=ports, stages=stages, identity=parts["identity"], - manifest_extra={"stage_plan": plan.manifest(), "io": io}, + manifest_extra=manifest_extra, owner=parts["owner"], step=step, + _allow_incomplete_staged=declaration_only, ) @@ -230,6 +268,37 @@ def _robot_action_dim(pl): return int(LIBERO_ACTION_DIM) +def _state_dim(pl): + """Raw proprioception dimension exposed by native_v2 STATE/STAGED.""" + try: + return int(len(pl.norm_stats["state"]["q01"])) + except Exception as e: + raise ValueError( + "Pi05 native_v2 requires norm_stats['state']['q01']") from e + + +def _tokenizer_sha256() -> str: + from flash_rt.utils.paligemma_tokenizer import ( + resolve_paligemma_tokenizer_path, + ) + h = hashlib.sha256() + with open(resolve_paligemma_tokenizer_path(), "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def _require_native_v2_ready(pl) -> None: + mode = getattr(pl, "_state_prompt_mode", None) + fixed = bool(getattr(pl, "_fixed_shape", False)) + if mode != "fixed" and not fixed: + raise ValueError( + "Pi05 native_v2 requires state_prompt_mode='fixed'") + if int(getattr(pl, "max_prompt_len", 0) or 0) < 200: + raise ValueError("Pi05 native_v2 requires max_prompt_len >= 200") + _state_dim(pl) + + def _parts(pl, identity, extra_regions): """Shared assembly for the plain export and the model-runtime export.""" if getattr(pl, "_graph", None) is None: @@ -296,6 +365,10 @@ def _parts(pl, identity, extra_regions): "robot_action_dim": str(_robot_action_dim(pl)), } ident.update({str(k): str(v) for k, v in (identity or {}).items()}) + if ident.get("io") == "native_v2": + ident["state_prompt_mode"] = "fixed" + ident["state_dim"] = str(_state_dim(pl)) + ident["tokenizer_sha256"] = _tokenizer_sha256() return { "wrap": wrap, diff --git a/flash_rt/runtime/export.py b/flash_rt/runtime/export.py index 05437dc0..8906f7f4 100644 --- a/flash_rt/runtime/export.py +++ b/flash_rt/runtime/export.py @@ -145,7 +145,7 @@ class PortSpec: name: str modality: int | str # MODALITY name or value - dtype: int | str = "bf16" # device-side tensor dtype + dtype: int | str = "bf16" # logical payload/target tensor dtype layout: int | str = "flat" direction: int | str = "in" update: int | str = "swap" @@ -276,6 +276,7 @@ def build_model_runtime( get_output=None, prepare=None, step=None, + _allow_incomplete_staged=False, ) -> ModelRuntime: """Assemble an ``frt_model_runtime_v1``: an export plus the dynamic-IO contract (ports, stage DAG, optional verb callables) under one identity — @@ -289,6 +290,22 @@ def build_model_runtime( them from any thread. SWAP ports need no callable — hosts write the declared buffer window directly. """ + staged_inputs = any( + _enum(UPDATE, port.update) == UPDATE["staged"] and + _enum(DIRECTION, port.direction) == DIRECTION["in"] + for port in ports + ) + staged_outputs = any( + _enum(UPDATE, port.update) == UPDATE["staged"] and + _enum(DIRECTION, port.direction) == DIRECTION["out"] + for port in ports + ) + if not _allow_incomplete_staged: + if staged_inputs and set_input is None: + raise ValueError("STAGED input ports require set_input") + if staged_outputs and get_output is None: + raise ValueError("STAGED output ports require get_output") + b, anchor, manifest_json = _assemble( ctx, streams=streams, graphs=graphs, buffers=buffers, regions=regions, ports=ports, stages=stages, identity=identity, diff --git a/runtime/include/flashrt/model_runtime.h b/runtime/include/flashrt/model_runtime.h index 51dbe689..bf14d025 100644 --- a/runtime/include/flashrt/model_runtime.h +++ b/runtime/include/flashrt/model_runtime.h @@ -131,7 +131,7 @@ typedef struct frt_image_view { typedef struct frt_runtime_port_desc { const char* name; /* "images", "prompt", "state", "actions" */ uint32_t modality; /* frt_rt_modality */ - uint32_t dtype; /* frt_rt_dtype (of the DEVICE-side tensor) */ + uint32_t dtype; /* frt_rt_dtype of logical payload/target */ uint32_t layout; /* frt_rt_layout */ uint32_t direction; /* frt_rt_port_direction */ uint32_t update; /* frt_rt_port_update */ @@ -247,8 +247,10 @@ int frt_runtime_builder_add_stage(frt_runtime_builder, uint32_t graph, /* Like frt_runtime_builder_finish, but returns the model runtime whose * `exp` is the internally-built export (one object, one refcount). `verbs` - * is copied; entries may be null (the runtime then reports them - * unsupported). Consumes the builder. */ + * is copied; entries may be null only when no matching STAGED declaration + * requires them (other missing verbs report unsupported). A STAGED input + * requires set_input and a STAGED output requires get_output. On validation + * failure the builder is not consumed. Consumes the builder on success. */ frt_model_runtime_v1* frt_runtime_builder_finish_model( frt_runtime_builder, const frt_model_runtime_verbs* verbs, void* verbs_self, @@ -262,7 +264,8 @@ frt_model_runtime_v1* frt_runtime_builder_finish_model( /* producer builds both. Descriptor arrays are copied. The wrapper */ /* takes one export reference and calls `wrapper_release(wrapper_owner)`*/ /* exactly once when its refcount hits zero (use it to destroy the */ -/* producer instance behind `verbs_self`). */ +/* producer instance behind `verbs_self`). STAGED declarations require */ +/* matching input/output verbs, as on construction path 1. */ /* ------------------------------------------------------------------ */ frt_model_runtime_v1* frt_model_runtime_wrap( const frt_runtime_export_v1* exp, @@ -278,7 +281,8 @@ frt_model_runtime_v1* frt_model_runtime_wrap( /* a native runtime owns hot-path transforms. The override retains `in` */ /* so all inherited descriptor pointers stay valid; consumers release */ /* only the returned object. `retain_owner`/`release_owner` manage the */ -/* native verb object, called once at construction/destruction. */ +/* native verb object, called once at construction/destruction. The new */ +/* verbs must satisfy every inherited STAGED input/output declaration. */ /* ------------------------------------------------------------------ */ frt_model_runtime_v1* frt_model_runtime_override_verbs( const frt_model_runtime_v1* in, diff --git a/runtime/src/model_runtime.cpp b/runtime/src/model_runtime.cpp index 9e82e706..e7b284d7 100644 --- a/runtime/src/model_runtime.cpp +++ b/runtime/src/model_runtime.cpp @@ -21,6 +21,21 @@ bool valid_port_args(const char* name, uint32_t direction, uint32_t update, return true; } +bool staged_verbs_present(const frt_runtime_port_desc* ports, uint64_t n_ports, + const frt_model_runtime_verbs* verbs) { + bool needs_input = false; + bool needs_output = false; + for (uint64_t i = 0; i < n_ports; ++i) { + if (ports[i].update != FRT_RT_PORT_STAGED) continue; + needs_input |= ports[i].direction == FRT_RT_PORT_IN; + needs_output |= ports[i].direction == FRT_RT_PORT_OUT; + } + const bool complete = + verbs && verbs->struct_size >= sizeof(frt_model_runtime_verbs); + return (!needs_input || (complete && verbs->set_input)) && + (!needs_output || (complete && verbs->get_output)); +} + /* Default stubs for verbs a producer does not provide: report unsupported * (-3) instead of leaving null function pointers for consumers to crash on. */ int stub_set_input(void*, uint32_t, const void*, uint64_t, int) { return -3; } @@ -108,6 +123,8 @@ extern "C" frt_model_runtime_v1* frt_runtime_builder_finish_model( void (*release_owner)(void*)) { if (!b) return nullptr; Holder* h = b->h; + if (!staged_verbs_present(h->ports.data(), h->ports.size(), verbs)) + return nullptr; frt_rt::finish_export_into(h, b, owner, retain_owner, release_owner); frt_model_runtime_v1& m = h->model; @@ -172,6 +189,7 @@ extern "C" frt_model_runtime_v1* frt_model_runtime_wrap( if (!valid_port_args(ports[i].name, ports[i].direction, ports[i].update, ports[i].shape, ports[i].rank)) return nullptr; + if (!staged_verbs_present(ports, n_ports, verbs)) return nullptr; for (uint64_t i = 0; i < n_stages; ++i) { if (stages[i].graph >= exp->n_graphs) return nullptr; if (stages[i].n_after && !stages[i].after) return nullptr; @@ -255,6 +273,7 @@ extern "C" frt_model_runtime_v1* frt_model_runtime_override_verbs( void* owner, void (*retain_owner)(void*), void (*release_owner)(void*)) { if (!valid_model_runtime(in)) return nullptr; + if (!staged_verbs_present(in->ports, in->n_ports, verbs)) return nullptr; auto* o = new VerbOverride(); o->base = in; diff --git a/runtime/tests/test_model_runtime.cpp b/runtime/tests/test_model_runtime.cpp index 0c9b29a5..a9362b33 100644 --- a/runtime/tests/test_model_runtime.cpp +++ b/runtime/tests/test_model_runtime.cpp @@ -56,16 +56,30 @@ static const int64_t IMG_SHAPE[4] = {3, 224, 224, 3}; static const int64_t ACT_SHAPE[2] = {50, 32}; static void add_ports_and_stages(frt_runtime_builder b, int64_t img_h = 224, - uint64_t act_bytes = 3200) { + uint64_t act_bytes = 3200, + uint32_t cadence_hz = 30, + bool add_prompt_state = false) { const int64_t img[4] = {3, img_h, 224, 3}; frt_runtime_builder_add_port(b, "images", FRT_RT_MOD_IMAGE, FRT_RT_DTYPE_BF16, FRT_RT_LAYOUT_NHWC, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, - img, 4, 30, nullptr, 0, 0); + img, 4, cadence_hz, nullptr, 0, 0); frt_runtime_builder_add_port(b, "actions", FRT_RT_MOD_ACTION, FRT_RT_DTYPE_BF16, FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_OUT, FRT_RT_PORT_STAGED, 0, ACT_SHAPE, 2, 0, FAKE_B0, 0, act_bytes); + if (add_prompt_state) { + const int64_t text_shape[1] = {-1}; + const int64_t state_shape[1] = {8}; + frt_runtime_builder_add_port( + b, "prompt", FRT_RT_MOD_TEXT, FRT_RT_DTYPE_U8, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, + text_shape, 1, 0, nullptr, 0, 0); + frt_runtime_builder_add_port( + b, "state", FRT_RT_MOD_STATE, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_STAGED, 1, + state_shape, 1, 0, nullptr, 0, 0); + } const uint32_t after1[1] = {0}; frt_runtime_builder_add_stage(b, 0, nullptr, 0); frt_runtime_builder_add_stage(b, 1, after1, 1); @@ -87,18 +101,32 @@ int main() { add_ports_and_stages(b); CHECK(frt_runtime_builder_finish(b, nullptr, nullptr, nullptr) == nullptr, "plain finish refuses a builder that declared ports/stages"); - /* the builder survives that refusal — finish_model consumes it */ + CHECK(frt_runtime_builder_finish_model( + b, nullptr, nullptr, nullptr, nullptr, nullptr) == nullptr, + "finish_model rejects STAGED ports without verbs"); + /* The builder survives both refusals; valid verbs consume it. */ + frt_model_runtime_verbs staged_verbs{}; + staged_verbs.struct_size = sizeof(staged_verbs); + staged_verbs.set_input = v_set_input; + staged_verbs.get_output = v_get_output; frt_model_runtime_v1* m = frt_runtime_builder_finish_model( - b, nullptr, nullptr, nullptr, nullptr, nullptr); - CHECK(m != nullptr, "finish_model after refused finish"); - /* absent producer verbs become unsupported stubs, never null */ - CHECK(m->verbs.set_input && m->verbs.step && m->verbs.last_error, - "null verbs are stubbed"); - CHECK(m->verbs.set_input(m->self, 0, nullptr, 0, -1) == -3 && - m->verbs.step(m->self) == -3 && - m->verbs.last_error(m->self)[0] != '\0', - "stubs report unsupported (-3) with an explanation"); + b, &staged_verbs, nullptr, nullptr, nullptr, nullptr); + CHECK(m != nullptr, "finish_model after refused finishes"); m->release(m->owner); + + frt_runtime_builder sb = make_builder(); + const int64_t shape[1] = {16}; + CHECK(frt_runtime_builder_add_port( + sb, "setup", FRT_RT_MOD_TENSOR, FRT_RT_DTYPE_F32, + FRT_RT_LAYOUT_FLAT, FRT_RT_PORT_IN, FRT_RT_PORT_SETUP, 0, + shape, 1, 0, nullptr, 0, 0) == 0, + "add non-staged port"); + frt_model_runtime_v1* sm = frt_runtime_builder_finish_model( + sb, nullptr, nullptr, nullptr, nullptr, nullptr); + CHECK(sm && sm->verbs.step(sm->self) == -3 && + sm->verbs.last_error(sm->self)[0] != '\0', + "missing non-staged verbs retain unsupported stubs"); + sm->release(sm->owner); } /* --- integrated build: struct, identity, fingerprint, verbs --- */ @@ -179,6 +207,39 @@ int main() { "graph stream change changes the fingerprint"); m4->release(m4->owner); } + /* prompt/state port additions define a new native face */ + { + frt_runtime_builder b5 = make_builder(); + add_ports_and_stages(b5, 224, 3200, 30, true); + frt_model_runtime_v1* m5 = frt_runtime_builder_finish_model( + b5, &verbs, &vlog, nullptr, nullptr, nullptr); + CHECK(m5 && m5->exp->fingerprint != m->exp->fingerprint, + "adding prompt/state ports changes the fingerprint"); + m5->release(m5->owner); + } + /* cadence is a scheduling hint, not deployment identity */ + { + frt_runtime_builder b6 = make_builder(); + add_ports_and_stages(b6, 224, 3200, 60); + frt_model_runtime_v1* m6 = frt_runtime_builder_finish_model( + b6, &verbs, &vlog, nullptr, nullptr, nullptr); + CHECK(m6 && m6->exp->fingerprint == m->exp->fingerprint, + "cadence hint does not change the fingerprint"); + m6->release(m6->owner); + } + /* manifest is discovery metadata, not deployment identity */ + { + frt_runtime_builder b7 = make_builder(); + add_ports_and_stages(b7); + CHECK(frt_runtime_builder_set_manifest( + b7, "{\"note\":\"discovery-only\"}") == 0, + "set manifest metadata"); + frt_model_runtime_v1* m7 = frt_runtime_builder_finish_model( + b7, &verbs, &vlog, nullptr, nullptr, nullptr); + CHECK(m7 && m7->exp->fingerprint == m->exp->fingerprint, + "manifest does not change the fingerprint"); + m7->release(m7->owner); + } /* verbs plumb through self */ m->verbs.set_input(m->self, 0, nullptr, 0, -1); @@ -224,6 +285,9 @@ int main() { CHECK(frt_model_runtime_wrap(exp, ports, 1, &bad, 1, &verbs, &vlog, nullptr, nullptr) == nullptr, "wrap rejects a stage over a missing graph"); + CHECK(frt_model_runtime_wrap(exp, ports, 1, stages, 1, nullptr, + nullptr, nullptr, nullptr) == nullptr, + "wrap rejects STAGED input without set_input"); frt_model_runtime_v1* wm = frt_model_runtime_wrap( exp, ports, 1, stages, 1, &verbs, &vlog, &wrapper_freed, @@ -262,6 +326,14 @@ int main() { native_verbs.step = v_step; native_verbs.last_error = v_last_error; + frt_model_runtime_verbs incomplete_verbs = native_verbs; + incomplete_verbs.get_output = nullptr; + CHECK(frt_model_runtime_override_verbs( + base, &incomplete_verbs, &native_vlog, &native_owner, + owner_retain, owner_release) == nullptr && + native_owner.retains == 0, + "override rejects missing STAGED output verb without retain"); + frt_model_runtime_v1* over = frt_model_runtime_override_verbs( base, &native_verbs, &native_vlog, &native_owner, owner_retain, owner_release); diff --git a/runtime/tests/test_model_runtime_py.py b/runtime/tests/test_model_runtime_py.py index 39dde2e7..8d982ab0 100644 --- a/runtime/tests/test_model_runtime_py.py +++ b/runtime/tests/test_model_runtime_py.py @@ -95,7 +95,10 @@ def record(stream): def build(setup, img_h=224, verbs=None): ctx, sid, src, dst, g = setup - verbs = verbs or {} + verbs = verbs or { + "set_input": lambda port, payload, stream: 0, + "get_output": lambda port, stream: b"", + } return build_model_runtime( ctx, streams=[StreamSpec("main", sid)], @@ -142,9 +145,48 @@ def build_split(setup): stages=plan.to_stage_specs(export_mod), identity={"model": "trivial", "quant": "none"}, manifest_extra={"stage_plan": plan.manifest()}, + set_input=lambda port, payload, stream: 0, + get_output=lambda port, stream: b"", ) +def check_staged_verb_guards(setup): + ctx, sid, src, dst, g = setup + common = dict( + ctx=ctx, + streams=[StreamSpec("main", sid)], + graphs=[GraphSpec("infer", g, 0, (0,))], + buffers=[BufferSpec("src", src, "input"), + BufferSpec("dst", dst, "output")], + stages=[StageSpec("infer")], + identity={"model": "verb_guard"}, + ) + try: + build_model_runtime( + **common, + ports=[PortSpec("input", "state", "f32", "flat", "in", + "staged", shape=(1,))], + ) + except ValueError as exc: + input_rejected = "require set_input" in str(exc) + else: + input_rejected = False + check("STAGED input without set_input is rejected", input_rejected) + + try: + build_model_runtime( + **common, + ports=[PortSpec("output", "action", "f32", "flat", "out", + "staged", shape=(1,))], + set_input=lambda port, payload, stream: 0, + ) + except ValueError as exc: + output_rejected = "require get_output" in str(exc) + else: + output_rejected = False + check("STAGED output without get_output is rejected", output_rejected) + + def check_stage_plan_registry(): register_stage_plan( "unit_chain", @@ -334,6 +376,9 @@ def py_step(): check_stage_plan_registry() check_vjp_guided_port_lowering(setup) + print("== staged verb declarations ==") + check_staged_verb_guards(setup) + print("== verbs through the C function pointers ==") rc = rt.model_set_input(mr.ptr, 1, b"\xAA\xBB", -1) check("set_input reaches the Python callable", 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())