diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..3614161 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,26 @@ +--- +name: Bug report +about: Something built or ran, and did the wrong thing +labels: bug +--- + +**What happened** + +**Expected** + +**Repro** + +``` +# command, including the model and backend +``` + +**Environment** +- vla.cpp commit: +- Backend: CPU / CUDA / Metal / SYCL +- OS and compiler: +- GPU and driver (if relevant): +- Model and checkpoint: + +**Output** + +Paste the shortest decisive part of the log, not the whole run. diff --git a/.github/ISSUE_TEMPLATE/model_request.md b/.github/ISSUE_TEMPLATE/model_request.md new file mode 100644 index 0000000..b7c8a51 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/model_request.md @@ -0,0 +1,20 @@ +--- +name: Architecture request +about: Ask for a VLA policy that vla.cpp does not run yet +labels: enhancement +--- + +**Policy** + +Name, paper or repo link, and the reference implementation. + +**Checkpoint** + +Where the weights live and under what license. + +**Why it is worth adding** + +Benchmark numbers, or what it does that the supported archs do not. + +CONTRIBUTING.md has the six-site walkthrough if you want to send the port +yourself. diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..14ebff8 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,12 @@ +## What + +## Why + +## Verified + +- [ ] Builds clean under `-Wall -Wextra` (first-party code) +- [ ] `ctest` passes +- [ ] Numeric output unchanged (`vla_predict_check` diff), or the change is + meant to move it and a LIBERO sweep is below + +Archs and backends tested: diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 1838000..782e2ef 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -17,10 +17,14 @@ jobs: runs-on: ubuntu-24.04 steps: - uses: actions/checkout@v4 - - name: pixel-shuffle channel order + # Compiled directly: pure, no llama.cpp or protobuf/zmq needed. + - name: pure unit tests run: | - g++ -std=c++17 -Isrc -Wall -Wextra tests/test_vision_common.cpp -o /tmp/test_vision_common - /tmp/test_vision_common + for t in test_vision_common test_rope_conventions; do + g++ -std=c++17 -Isrc -Wall -Wextra -fsanitize=address,undefined \ + -fno-omit-frame-pointer "tests/$t.cpp" -o "/tmp/$t" + "/tmp/$t" + done py-tooling: runs-on: ubuntu-24.04 @@ -31,6 +35,8 @@ jobs: python-version: '3.11' - name: converter remap run: python tests/py/test_converters.py + - name: binding struct parity + run: python tests/py/test_bindings.py build-gate: runs-on: ubuntu-24.04 @@ -45,8 +51,11 @@ jobs: - uses: actions/cache@v4 with: path: build/_deps - key: llama-b9866-${{ runner.os }} - - name: build vla-server + vlm-server + vla-cli (CPU, -Wall -Wextra) + key: llama-b10331-${{ runner.os }} + # Everything, not a target list: ctest registers tests this job must build, + # and a named list goes stale the next time one is added. + - name: build + ctest (CPU, -Wall -Wextra) run: | - cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=OFF - cmake --build build -j"$(nproc)" --target vla-server vlm-server vla-cli + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=OFF -DVLA_BUILD_TESTS=ON + cmake --build build -j"$(nproc)" + ctest --test-dir build --output-on-failure diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..7919ed3 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,144 @@ +# Tagged binaries and a container image. build.yml already compiles all of this +# on every push; this is the same work with the artifacts kept. +name: release + +on: + push: + tags: ['v*'] + workflow_dispatch: + inputs: + tag: + description: Tag to build (dry run, nothing is published) + required: true + +permissions: + contents: write + packages: write + +env: + BINARIES: vla-server vlm-server vla-cli vla-bench + +jobs: + linux: + runs-on: ${{ matrix.runner }} + strategy: + fail-fast: false + matrix: + include: + - name: linux-x86_64-cpu + cmake: -DGGML_CUDA=OFF + cuda: false + runner: ubuntu-24.04 + - name: linux-x86_64-cuda + cmake: -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=75;86;89;120 + cuda: true + runner: ubuntu-24.04 + # Jetson and other aarch64 boards. Native arm64 runner, CPU only: the + # hosted images carry no CUDA for arm64, so a Jetson GPU build still + # has to happen on the device. + - name: linux-aarch64-cpu + cmake: -DGGML_CUDA=OFF + cuda: false + runner: ubuntu-24.04-arm + steps: + - uses: actions/checkout@v4 + + - name: deps + run: | + sudo apt-get update -qq + sudo apt-get install -y -qq --no-install-recommends \ + build-essential cmake git ca-certificates pkg-config \ + libzmq3-dev cppzmq-dev libprotobuf-dev protobuf-compiler + + - name: cuda toolkit + if: matrix.cuda + run: | + wget -q https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb + sudo dpkg -i cuda-keyring_1.1-1_all.deb + sudo apt-get update -qq + sudo apt-get install -y -qq --no-install-recommends cuda-toolkit-12-6 + echo "/usr/local/cuda/bin" >> "$GITHUB_PATH" + + - name: build + run: | + cmake -B build -DCMAKE_BUILD_TYPE=Release ${{ matrix.cmake }} + cmake --build build -j"$(nproc)" --target $BINARIES vla + + - name: package + run: | + out="vla.cpp-${{ github.ref_name }}-${{ matrix.name }}" + mkdir -p "$out" + for b in $BINARIES; do cp "build/$b" "$out/"; done + cp build/libvla.so "$out/" + cp include/vla.h LICENSE.md README.md "$out/" + # vla-cli --text runs this; VLA_TOKENIZE_SCRIPT points at it. + mkdir -p "$out/scripts" && cp scripts/tokenize_prompt.py "$out/scripts/" + tar -czf "$out.tar.gz" "$out" + + - uses: actions/upload-artifact@v4 + with: + name: ${{ matrix.name }} + path: '*.tar.gz' + + macos: + runs-on: macos-14 + steps: + - uses: actions/checkout@v4 + + - name: deps + run: brew install cmake zeromq cppzmq protobuf + + - name: build + run: | + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_METAL=ON + cmake --build build -j"$(sysctl -n hw.ncpu)" --target $BINARIES vla + + - name: package + run: | + out="vla.cpp-${{ github.ref_name }}-macos-arm64-metal" + mkdir -p "$out" + for b in $BINARIES; do cp "build/$b" "$out/"; done + cp build/libvla.dylib "$out/" + cp include/vla.h LICENSE.md README.md "$out/" + mkdir -p "$out/scripts" && cp scripts/tokenize_prompt.py "$out/scripts/" + # Metal needs the shader library next to the binary. + find build -name 'default.metallib' -exec cp {} "$out/" \; + tar -czf "$out.tar.gz" "$out" + + - uses: actions/upload-artifact@v4 + with: + name: macos-arm64-metal + path: '*.tar.gz' + + docker: + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v4 + - uses: docker/setup-buildx-action@v3 + - uses: docker/login-action@v3 + if: startsWith(github.ref, 'refs/tags/') + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - uses: docker/build-push-action@v6 + with: + context: . + push: ${{ startsWith(github.ref, 'refs/tags/') }} + tags: | + ghcr.io/${{ github.repository }}:${{ github.ref_name }} + ghcr.io/${{ github.repository }}:latest + cache-from: type=gha + cache-to: type=gha,mode=max + + publish: + needs: [linux, macos, docker] + if: startsWith(github.ref, 'refs/tags/') + runs-on: ubuntu-24.04 + steps: + - uses: actions/download-artifact@v4 + with: { path: dist, merge-multiple: true } + - uses: softprops/action-gh-release@v2 + with: + files: dist/*.tar.gz + generate_release_notes: true diff --git a/CHANGELOG.md b/CHANGELOG.md index a043125..26c52a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,28 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://keepachangelog.com). +## [0.2.0] - 2026-08-09 + +### Added +- SYCL backend for Intel GPUs (Arc, Flex, Data Center Max, Xe iGPU). `VLA_DEVICE` picks the ordinal on CUDA and SYCL alike. See `docs/backend/sycl.md`. +- Stable C ABI (`include/vla.h`, `libvla`) and Python bindings over it (`bindings/python`). +- Four more architectures: π0.5, VLA-Adapter, OpenVLA-OFT and VLA-JEPA. +- `vla-bench` for engine-only latency, and `-hf user/repo[:file.gguf]` to fetch a checkpoint on first use. +- `vla-cli --text`, tokenized by `scripts/tokenize_prompt.py` with the tokenizer the architecture was trained on. +- Release workflow publishing Linux x86-64 (CPU and CUDA), Linux aarch64 (CPU), macOS Metal and a GHCR image. + +### Changed +- One shared backend ladder (`src/backend.h`) instead of a copy per arch. CMake rejects two accelerators in one build directory. +- Shared headers for the Qwen3-VL tower, the DINOv2+SigLIP dual tower, the DiT time embeddings, the causal mask and CHW image preprocessing. +- `vla::graph_cache` keeps the compute graph across `predict` calls in nine architectures, not just GR00T N1.7. Output is unchanged. +- llama.cpp pinned at b10331. GR00T N1.5 and N1.6 shift by up to 4.6e-4 on actions peaking near 0.87, from an upstream ggml kernel change in the SigLIP tower they share. The other nine architectures are bit-identical. + +### Fixed +- Reject checkpoint geometry that contradicts itself before it sizes a buffer, in smolvla, bitvla, gr00tn1d6, vla_adapter and the Qwen3-VL position resample. +- A peer that stalls mid-message no longer parks either server. +- Treat a missing state vector as zeros in every architecture rather than dereferencing it. +- Build every registered test before `ctest`, so the four that were never built stop reporting as not run. + ## [0.1.1] - 2026-07-04 ### Added @@ -40,5 +62,6 @@ expert + dataset stats), CPU or CUDA, no external mmproj and no patch to llama.c - llama.cpp is fetched + pinned via CMake `FetchContent` (tag `b9866`); bumping is a one-line `GIT_TAG` change. Removed the `patches/` fetch script. +[0.2.0]: https://github.com/VinRobotics/vla.cpp/releases/tag/v0.2.0 [0.1.1]: https://github.com/VinRobotics/vla.cpp/releases/tag/v0.1.1 [0.1.0]: https://github.com/VinRobotics/vla.cpp/releases/tag/v0.1.0 diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e757e1..50bc0a6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -11,6 +11,22 @@ if(NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE Release CACHE STRING "" FORCE) endif() +# src/backend.h compiles in exactly one accelerator, so two GGML_* flags would +# give ggml both backends and vla_core only one. Reject before FetchContent. +set(_vla_accel "") +foreach(_flag GGML_CUDA GGML_SYCL GGML_METAL) + if(${_flag}) + list(APPEND _vla_accel ${_flag}) + endif() +endforeach() +list(LENGTH _vla_accel _vla_accel_n) +if(_vla_accel_n GREATER 1) + string(REPLACE ";" ", " _vla_accel_str "${_vla_accel}") + message(FATAL_ERROR + "Enable one accelerator backend at a time; got ${_vla_accel_str}. " + "Configure a separate build directory per backend.") +endif() + set(LLAMA_BUILD_COMMON ON CACHE BOOL "" FORCE) set(LLAMA_BUILD_TOOLS ON CACHE BOOL "" FORCE) set(LLAMA_BUILD_SERVER OFF CACHE BOOL "" FORCE) @@ -20,7 +36,7 @@ set(LLAMA_BUILD_TESTS OFF CACHE BOOL "" FORCE) include(FetchContent) FetchContent_Declare(llama GIT_REPOSITORY https://github.com/ggml-org/llama.cpp - GIT_TAG b9866 + GIT_TAG b10331 GIT_SHALLOW TRUE ) FetchContent_MakeAvailable(llama) @@ -88,7 +104,23 @@ if(GGML_CUDA) target_include_directories(vla_core PUBLIC ${CUDAToolkit_INCLUDE_DIRS}) endif() -if(GGML_METAL AND NOT GGML_CUDA) +# Intel GPUs (Arc / Flex / Data Center Max / Xe iGPU) through oneAPI SYCL. +# ggml's SYCL sources only compile under the oneAPI DPC++ driver, and +# CMAKE_CXX_COMPILER is global, so our targets are built by icpx too. Fail +# loudly here rather than let ggml die deep in a kernel compile. +if(GGML_SYCL AND NOT GGML_CUDA) + if(NOT CMAKE_CXX_COMPILER_ID STREQUAL "IntelLLVM") + message(FATAL_ERROR + "GGML_SYCL=ON needs the oneAPI DPC++ compiler, but CMAKE_CXX_COMPILER is " + "'${CMAKE_CXX_COMPILER_ID}'. Source /opt/intel/oneapi/setvars.sh and configure a " + "fresh build dir with -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx. " + "See docs/backend/sycl.md.") + endif() + target_compile_definitions(vla_core PUBLIC GGML_USE_SYCL) +endif() + +# Backend precedence matches the ladder in src/backend.h: CUDA, then SYCL, then Metal. +if(GGML_METAL AND NOT GGML_CUDA AND NOT GGML_SYCL) target_compile_definitions(vla_core PUBLIC GGML_USE_METAL) endif() @@ -169,6 +201,7 @@ add_executable(vlm-server ) target_include_directories(vlm-server PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/proto-gen + ${llama_SOURCE_DIR}/vendor/stb ${CPPZMQ_INCLUDE_DIR} ) target_link_libraries(vlm-server PRIVATE @@ -177,6 +210,20 @@ target_link_libraries(vlm-server PRIVATE PkgConfig::ZeroMQ ) +# Stable C ABI. Shared so bindings can dlopen it; visibility hidden so only the +# vla_* symbols are exported and llama/ggml stay internal. +add_library(vla SHARED src/vla_c_api.cpp) +target_include_directories(vla + PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src +) +target_link_libraries(vla PRIVATE vla_core) +set_target_properties(vla PROPERTIES + CXX_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON + PUBLIC_HEADER ${CMAKE_CURRENT_SOURCE_DIR}/include/vla.h +) + # One-shot inference CLI: image + tokens -> action, no server or simulator. add_executable(vla-cli src/serving/vla-cli.cpp @@ -185,9 +232,17 @@ target_include_directories(vla-cli PRIVATE ${llama_SOURCE_DIR}/vendor/stb ) target_link_libraries(vla-cli PRIVATE vla_core) +# --text shells out to the tokenizer script; VLA_TOKENIZE_SCRIPT overrides it. +target_compile_definitions(vla-cli PRIVATE VLA_SOURCE_DIR="${CMAKE_CURRENT_SOURCE_DIR}") + +# Latency for one checkpoint, emits the README table rows. +add_executable(vla-bench + src/serving/vla-bench.cpp +) +target_link_libraries(vla-bench PRIVATE vla_core) # --- First-party build hygiene (never applied to the vendored llama.cpp subtree) -- -set(VLA_FIRST_PARTY_TARGETS vla_core vlm_core vla-server vlm-server vla-cli) +set(VLA_FIRST_PARTY_TARGETS vla_core vlm_core vla vla-server vlm-server vla-cli vla-bench) foreach(tgt IN LISTS VLA_FIRST_PARTY_TARGETS) # Warn on our own C++ only; nvcc device code keeps its own diagnostics. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..2ce5d34 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,83 @@ +# Contributing + +## Build and test + +```bash +cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=OFF -DVLA_BUILD_TESTS=ON +cmake --build build -j"$(nproc)" +ctest --test-dir build --output-on-failure +``` + +First-party code must compile clean under `-Wall -Wextra`. Warnings from +`build/_deps` are upstream and not your problem. + +## Proving a change is numerically neutral + +Refactors, performance work and dependency bumps must not move the output. +`vla_predict_check` feeds fixed images, tokens, state and noise, then prints the +action chunk: + +```bash +VLA_IMG_SIZE=224 ./build/tests/vla_predict_check model.gguf "" 1 > before.txt +# ... make the change, rebuild ... +VLA_IMG_SIZE=224 ./build/tests/vla_predict_check model.gguf "" 1 > after.txt +diff before.txt after.txt +``` + +Any difference is a bug unless the change is meant to alter numerics, in which +case say so in the commit message and back it with a LIBERO sweep. + +`VLA_IMG_SIZE` must match the model or `predict` returns empty: 512 for +SmolVLA, 448 for Evo-1, 256 for GR00T N1.7 and VLA-JEPA, 224 for the rest. Other +knobs: `VLA_BENCH_ITERS` (timing), `VLA_TIMING=phase`, `VLA_EXTRA_TOKEN` / +`VLA_EXTRA_COUNT` (VLA-JEPA needs its `` tokens), `VLA_N_THREADS`, +`VLA_DEVICE`. + +Checkpoints are at [huggingface.co/vrfai](https://huggingface.co/vrfai), or let +the binaries fetch them: + +```bash +./build/vla-cli -hf vrfai/smolvla-libero-gguf --image assets/front.jpg --tokens 1,100,2 +``` + +## Adding an architecture + +Six sites, all mechanical. `smolvla` is the reference for a two-file (mmproj + +ckpt) model, `bitvla` for a vision-baked one. + +1. `src/arch.h` - add to `enum class Arch`. +2. `src/arch.h` - declare `_create(mmproj_path, ckpt_path, config_path)`. +3. `src/model.cpp` - add `.architecture` to the `try_str` list in + `detect_arch_gguf`. +4. `src/model.cpp` - map the string to the enum in the same function. +5. `src/model.cpp` - add a `case` to the `model_load` switch. +6. `CMakeLists.txt` - add `src/models/.cpp` to `vla_core`. + +Then write `src/models/.cpp`. Before adding a helper, check +`src/models/`: `gguf_reader.h` (tensor and KV reads), `vision_common.h` +(preprocessing, pixel shuffle), `dual_tower.h` (DINOv2 + SigLIP), +`qwen3vl_vit.h` (Qwen3-VL tower), `dit_common.h` (DiT time embeddings), +`scratch_ctx.h` (compute context reuse), `backend.h` (accelerator selection). + +Your loader must fail rather than return a half-built model: check every tensor +lookup, and check `real_*_dim <= max_*_dim` (`config_is_sane` in `src/model.cpp` +does this for all archs). + +A converter goes in `scripts/convert__to_gguf.py`, and its tensor-name +remap should be covered by `tests/py/test_converters.py`. + +## Ports of upstream quirks + +Some references do surprising things, and we match them because the weights were +trained that way. Three so far: VLA-Adapter's RoPE pairs a half-split frequency +table with an interleaved rotation, OpenVLA-OFT's LM attention is bidirectional, +and BitVLA's is too. Each carries a comment naming the reference file and lines. + +If you find something that looks wrong, diff against the reference before +changing it, and leave a comment with the line numbers so the next reader does +not have to. + +## Commits and pull requests + +One logical change per commit, imperative subject, no trailing period. Say what +you verified: which archs, which backends, whether the numeric output moved. diff --git a/Dockerfile b/Dockerfile index c399342..f703649 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,6 +1,6 @@ # vla-server, CPU or CUDA. cmake fetches llama.cpp at build time. -# GPU sm_120 (default): docker build -t vla-cpp . -# GPU other arch: --build-arg CUDA_ARCH=89 (89=RTX40 90=H100 87=Orin; sm_120 needs CUDA>=12.8) +# GPU sm_89 (default): docker build -t vla-cpp . +# GPU other arch: --build-arg CUDA_ARCH=120 (86=RTX30 90=H100 87=Orin 120=RTX50; sm_120 needs CUDA>=12.8) # older card: --build-arg BASE_IMAGE=nvidia/cuda:12.4.1-devel-ubuntu24.04 --build-arg CUDA_ARCH=86 # CPU: --build-arg BACKEND=cpu --build-arg BASE_IMAGE=ubuntu:24.04 -t vla-cpp-cpu # run: docker run --gpus all -p5555:5555 -v $PWD/models:/models vla-cpp --bind tcp://*:5555 /models/M.gguf @@ -18,7 +18,7 @@ WORKDIR /src COPY . . ARG BACKEND=cuda -ARG CUDA_ARCH=120 +ARG CUDA_ARCH=89 # nvcc can segfault on the flash-attn kernels under high -j; lower JOBS if so. ARG JOBS= # CUDA: -devel ships only a libcuda stub (real driver injected at runtime), so @@ -26,7 +26,7 @@ ARG JOBS= RUN set -eux; \ if [ "$BACKEND" = "cuda" ]; then \ export LIBRARY_PATH=/usr/local/cuda/lib64/stubs; \ - cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=ON -DGGML_CUDA_GRAPHS=ON \ + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=ON \ -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCH}" \ -DCMAKE_SHARED_LINKER_FLAGS="-lcuda" -DCMAKE_EXE_LINKER_FLAGS="-lcuda"; \ else \ diff --git a/README.md b/README.md index 58f4924..5d58cb0 100644 --- a/README.md +++ b/README.md @@ -11,8 +11,8 @@ A C++ inference engine for **Vision-Language-Action (VLA) models**, built on [`llama.cpp`](https://github.com/ggml-org/llama.cpp). It runs the open VLA policies - SmolVLA, π0, BitVLA, Evo-1, GR00T N1.5/1.6/1.7 and more - under one runtime, each packaged as a single self-contained GGUF that needs no Python or -PyTorch at inference time. The binaries drive robots on **CPU**, **Apple Silicon**, or -**CUDA**, from consumer GPUs down to Jetson-class boards. +PyTorch at inference time. The binaries drive robots on **CPU**, **Apple Silicon**, **CUDA** - +from consumer GPUs down to Jetson-class boards - or **Intel GPUs** via SYCL. [**Learn vla.cpp**](https://fai-modelopt-tech.github.io/learn-vla-cpp/) walks through the engine design and how each policy is implemented on ggml. @@ -24,11 +24,13 @@ PyTorch at inference time. The binaries drive robots on **CPU**, **Apple Silicon - CMake ≥ 3.22 - A C++17 compiler (GCC 11+ or Clang 14+) -- CUDA 12.x (optional - required only for GPU builds) -- `libzmq3-dev`, `libprotobuf-dev`, `protobuf-compiler` +- CUDA 12.x (optional - required only for CUDA GPU builds) +- Intel oneAPI 2025.x + GPU compute runtime (optional - only for Intel GPU + builds, see [docs/backend/sycl.md](docs/backend/sycl.md)) +- `libzmq3-dev`, `cppzmq-dev`, `libprotobuf-dev`, `protobuf-compiler` ```bash -sudo apt-get install -y libzmq3-dev libprotobuf-dev protobuf-compiler +sudo apt-get install -y libzmq3-dev cppzmq-dev libprotobuf-dev protobuf-compiler ``` ### From source @@ -54,12 +56,26 @@ cmake --build build -j$(nproc) # CUDA build (set CMAKE_CUDA_ARCHITECTURES for your GPU): cmake -B build \ -DGGML_CUDA=ON \ - -DGGML_CUDA_GRAPHS=ON \ -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_CUDA_ARCHITECTURES=$CUDA_ARCHITECTURE cmake --build build -j$(nproc) ``` +```bash +# Intel GPU build (Arc / Flex / Max / Xe iGPU). ggml's SYCL sources need the +# oneAPI DPC++ driver, so the whole project is compiled by icpx: +source /opt/intel/oneapi/setvars.sh +cmake -B build \ + -DGGML_SYCL=ON \ + -DCMAKE_C_COMPILER=icx \ + -DCMAKE_CXX_COMPILER=icpx \ + -DCMAKE_BUILD_TYPE=Release +cmake --build build -j$(nproc) +``` + +The driver and oneAPI setup that this needs is in +[docs/backend/sycl.md](docs/backend/sycl.md). + If CMake cannot find CUDA, point the environment at it explicitly: ```bash @@ -67,8 +83,8 @@ export PATH=/usr/local/cuda/bin:$PATH export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH ``` -Check [docs/backend](docs/backend) for compiling `vla.cpp` with other platforms. -WLS2 and Apple Silicon has been tested. +Check [docs/backend](docs/backend) for compiling `vla.cpp` on other platforms. +WSL2 and Apple Silicon are both tested. --- @@ -77,18 +93,25 @@ WLS2 and Apple Silicon has been tested. Once the binaries are built, run one CPU prediction without a server or simulator: ```bash -pip install -U "huggingface_hub[cli]" gguf -hf download vrfai/smolvla-libero-gguf --local-dir models/smolvla +pip install -U "huggingface_hub[cli]" transformers + +# -hf fetches and caches the checkpoint (under $VLA_CACHE, default ~/.cache/vla) +./build/vla-cli -hf vrfai/smolvla-libero-gguf \ + --image assets/front.jpg --text "pick up the black bowl" --pretty -# One-shot CLI +# or point at a file you already have ./build/vla-cli --ckpt models/smolvla/smolvla-libero.gguf \ - --image assets/front.jpg --tokens 1,100,200,2 --pretty + --image assets/front.jpg --text "pick up the black bowl" --pretty ``` `vla-cli` runs a single prediction without a server or simulator: give it a model, -an image, and the tokenized instruction, and it prints the action chunk. Handy for +an image, and an instruction, and it prints the action chunk. Handy for smoke-testing a GGUF or scripting a quick inference. -`--tokens` are language token ids from the client tokenizer. + +There is no tokenizer in the C++ core, so `--text` calls +`scripts/tokenize_prompt.py` with the tokenizer the architecture was trained on +(`VLA_PYTHON` picks the interpreter, `VLA_TOKENIZE_SCRIPT` the script). Pass +`--tokens 1,100,200,2` instead if you already have ids. `--pretty` prints one action row per line; `--state` sets proprioception (defaults to zeros). @@ -143,6 +166,14 @@ vla-server: bound to tcp://*:5555. ready. Use `--bind` to change the address and port. Stop the server with `Ctrl-C`. +`vla-server` also takes `-hf user/repo[:file.gguf]` in place of a checkpoint path. + +Environment knobs that apply to every arch: + +- `VLA_N_THREADS` - CPU backend thread count, default core count capped at 16. +- `VLA_DEVICE` - GPU ordinal for CUDA and SYCL builds, default 0. +- `VLA_CACHE` - where `-hf` stores checkpoints, default `~/.cache/vla`. + --- ## Running the client @@ -228,18 +259,55 @@ pack the vision tower too (smaller, but more accuracy loss). ## Benchmarks -Latency in ms (inference plus transport), measured client-side on four targets: an -**RTX 3090**, an **NVIDIA Jetson AGX Orin**, an **NVIDIA Jetson Orin Nano (8 GB)**, -and an **Apple M4**. +`vla-bench` times `predict()` in-process on synthetic inputs: engine only, no +transport, no simulator, no claim about task success. + +```bash +./build/vla-bench -hf vrfai/smolvla-libero-gguf --images 2 --size 512 --markdown +``` -| Model | 3090 call (ms) | AGX Orin call (ms) | Orin Nano call (ms) | M4 call (ms) | -|---|---:|---:|---:|---:| -| `smolvla` | 86 | 262 | 567 | 888 | -| `pi0` | 264 | 893 | 1955 | 1135 | -| `gr00t_n1_5` | 109 | 461 | 1356 | - | -| `gr00t_n1_7` | 102 | 429 | - | 755 | -| `bitvla` | 145 | 809 | 2845 | - | -| `evo1` | 238 | 1048 | 3671 | - | +RTX 5090, driver 595.84, CUDA 13.2, 24-core host, weights as shipped, 20 reps +after 3 warmups, best of three sweeps, each model at its native input size and +view count. + +| Model | Views | Input | min ms | p50 ms | p90 ms | vision ms | +|---|--:|--:|--:|--:|--:|--:| +| VLA-Adapter | 1 | 224 | 18.2 | 19.8 | 21.1 | 9.4 | +| VLA-JEPA | 1 | 256 | 19.9 | 21.5 | 22.9 | 6.3 | +| BitVLA | 1 | 224 | 23.6 | 25.3 | 26.4 | 5.4 | +| GR00T N1.5 | 1 | 224 | 28.2 | 29.4 | 30.5 | 5.9 | +| GR00T N1.7 | 1 | 256 | 31.0 | 33.4 | 34.6 | 6.2 | +| GR00T N1.6 | 1 | 224 | 33.4 | 35.7 | 37.3 | 6.3 | +| OpenVLA-OFT | 1 | 224 | 47.4 | 49.2 | 50.2 | 10.3 | +| SmolVLA | 2 | 512 | 47.8 | 49.6 | 54.0 | 16.1 | +| pi0 | 2 | 224 | 48.9 | 52.1 | 55.0 | 11.6 | +| Evo-1 | 1 | 448 | 52.2 | 55.2 | 57.3 | 17.8 | +| pi0.5 | 2 | 224 | 53.4 | 56.1 | 59.3 | 11.4 | + +Jetson and Apple targets are absent: they have not been re-measured with +`vla-bench`. + +### Task success + +Latency says nothing about whether a policy works. LIBERO-Object, 10 tasks and 20 +episodes per model, terminated episodes counted as failures: + +| Model | Chunk replay | Success rate | +|---|--:|--:| +| BitVLA | 8 | 100.0% | +| GR00T N1.7 | 16 | 98.0% | +| GR00T N1.5 | 16 | 96.0% | +| Evo-1 | 8 | 94.5% | +| SmolVLA | 4 | 90.5% | +| π0 | 32 | 87.5% | +| GR00T N1.6 | 16 | 86.5% | + +From [eval/reports/report-rtx-3060.md](eval/reports/report-rtx-3060.md), swept on +an RTX 3060 at commit `dcc29a3` (2026-05-24). It predates π0.5, VLA-Adapter, +OpenVLA-OFT and VLA-JEPA, which have not been swept. Jetson AGX Orin and Orin +Nano runs are in the same directory. Success rate belongs to the checkpoint, not +the engine; `vla_predict_check` in [CONTRIBUTING.md](CONTRIBUTING.md) is how a +change is shown to leave it alone. --- @@ -248,19 +316,26 @@ and an **Apple M4**. Support matrix of models (rows) against platforms (columns). Legend: `Y` = supported (released and benchmarked), `~` = in progress, `-` = planned. -| Model | CPU (x86-64 / ARM) | CUDA | Metal | OpenVINO | Hexagon | -|---|:--:|:--:|:--:|:--:|:--:| -| [SmolVLA](https://hf.co/vrfai/smolvla-libero-gguf) | Y | Y | Y | - | - | -| [π0](https://hf.co/vrfai/pi0-libero-finetuned-v044-gguf) | Y | Y | Y | - | - | -| [π0.5](https://hf.co/vrfai/pi05-libero-gguf) | Y | Y | ~ | - | - | -| [GR00T N1.5](https://hf.co/vrfai/gr00tn1d5-libero-object-gguf) | Y | Y | ~ | - | - | -| [GR00T N1.6](https://hf.co/vrfai/gr00tn1d6-libero-gguf) | Y | Y | ~ | - | - | -| [GR00T N1.7](https://hf.co/vrfai/gr00tn1d7-libero-gguf) | Y | Y | Y | - | - | -| [BitVLA](https://hf.co/vrfai/bitvla-libero-gguf) | Y | Y | ~ | - | - | -| [Evo-1](https://hf.co/vrfai/evo1-libero-gguf) | Y | Y | ~ | - | - | -| [VLA-Adapter](https://hf.co/vrfai/vla-adapter-libero-gguf) | Y | Y | ~ | - | - | -| [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | ~ | - | - | -| [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | ~ | - | - | +| Model | CPU (x86-64 / ARM) | CUDA | SYCL (Intel) | Metal | OpenVINO | Hexagon | +|---|:--:|:--:|:--:|:--:|:--:|:--:| +| [SmolVLA](https://hf.co/vrfai/smolvla-libero-gguf) | Y | Y | Y | Y | - | - | +| [π0](https://hf.co/vrfai/pi0-libero-finetuned-v044-gguf) | Y | Y | - | Y | - | - | +| [π0.5](https://hf.co/vrfai/pi05-libero-gguf) | Y | Y | - | ~ | - | - | +| [GR00T N1.5](https://hf.co/vrfai/gr00tn1d5-libero-object-gguf) | Y | Y | - | ~ | - | - | +| [GR00T N1.6](https://hf.co/vrfai/gr00tn1d6-libero-gguf) | Y | Y | - | ~ | - | - | +| [GR00T N1.7](https://hf.co/vrfai/gr00tn1d7-libero-gguf) | Y | Y | - | Y | - | - | +| [BitVLA](https://hf.co/vrfai/bitvla-libero-gguf) | Y | Y | - | ~ | - | - | +| [Evo-1](https://hf.co/vrfai/evo1-libero-gguf) | Y | Y | Y | ~ | - | - | +| [VLA-Adapter](https://hf.co/vrfai/vla-adapter-libero-gguf) | Y | Y | ~ | ~ | - | - | +| [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | - | ~ | - | - | +| [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | - | ~ | - | - | + +--- + +## Contributing + +See [CONTRIBUTING.md](CONTRIBUTING.md) for how to prove a change is numerically +neutral, and the six sites you touch to add an architecture. --- diff --git a/bindings/python/README.md b/bindings/python/README.md new file mode 100644 index 0000000..ebf406c --- /dev/null +++ b/bindings/python/README.md @@ -0,0 +1,40 @@ +# vla-cpp + +Python bindings for [vla.cpp](https://github.com/VinRobotics/vla.cpp) over its +C ABI (`include/vla.h`). No PyTorch at inference time. + +```python +import vla_cpp + +model = vla_cpp.load("smolvla-libero.gguf") +actions = model.predict(frame_hwc_uint8, tokens=[1, 100, 200, 2]) +``` + +`predict` returns `[rows, max_action_dim]`; only the first +`model.config.real_action_dim` columns carry values. `model.config.denormalized` +says whether they are already in world units. + +## Finding the library + +Wheels bundle `libvla.so`. From a source checkout, build it and point at it: + +```bash +cmake -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build -j"$(nproc)" --target vla +VLA_LIBRARY=build/libvla.so LD_LIBRARY_PATH=build:build/bin python your_script.py +``` + +`LD_LIBRARY_PATH` is needed because `libvla.so` links `libvla_core.so` and the +ggml libraries from the same build tree. + +## API + +| | | +|---|---| +| `load(ckpt, mmproj=None, config=None)` | `mmproj` only for SmolVLA, pi0, pi0.5 | +| `Model.predict(images, tokens, state=None, noise=None, ...)` | `images` is one HWC array or a sequence | +| `Model.config` | resolved hyper-parameters | +| `Model.last_stats()` | per-phase timings, needs `timing=TIMING_PHASE` | +| `Model.close()` | or use as a context manager | + +Output is bit-identical to `vla-cli` on the same inputs. diff --git a/bindings/python/pyproject.toml b/bindings/python/pyproject.toml new file mode 100644 index 0000000..af666ad --- /dev/null +++ b/bindings/python/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "vla-cpp" +version = "0.2.0" +description = "Python bindings for vla.cpp, a C++ inference engine for Vision-Language-Action models." +readme = "README.md" +requires-python = ">=3.9" +license = { text = "Apache-2.0" } +dependencies = [] + +[project.optional-dependencies] +# Only for the array return type; the bindings work on plain lists without it. +numpy = ["numpy>=1.24"] + +[project.urls] +Homepage = "https://github.com/VinRobotics/vla.cpp" + +[tool.setuptools] +packages = ["vla_cpp"] + +# libvla is built by cmake, not by setuptools. A wheel bundles it next to the +# package; a source checkout finds it through VLA_LIBRARY or the loader path. +[tool.setuptools.package-data] +vla_cpp = ["*.so", "*.dylib", "*.dll", "lib/*"] diff --git a/bindings/python/vla_cpp/__init__.py b/bindings/python/vla_cpp/__init__.py new file mode 100644 index 0000000..5893866 --- /dev/null +++ b/bindings/python/vla_cpp/__init__.py @@ -0,0 +1,170 @@ +"""Python bindings for vla.cpp. + + import vla_cpp + model = vla_cpp.load("smolvla-libero.gguf") + actions = model.predict(image_hwc_uint8, tokens=[1, 100, 200, 2]) + +Actions are [rows, max_action_dim] float32; only the first +``model.config.real_action_dim`` columns carry values. +""" + +from __future__ import annotations + +import ctypes +from ctypes import POINTER, c_float, c_int32, c_int64 +from typing import Sequence + +from . import _ffi +from ._ffi import PIXEL_F32_RGB_01, PIXEL_U8, TIMING_NONE, TIMING_PHASE + +__all__ = ["Model", "load", "PIXEL_U8", "PIXEL_F32_RGB_01", "TIMING_NONE", "TIMING_PHASE"] + +_lib = None + + +def _lib_handle(): + global _lib + if _lib is None: + _lib = _ffi.load_library() + return _lib + + +def _as_f32_array(values, length: int, name: str): + """Accept a numpy array, a list, or None. Returns (ptr, keepalive).""" + if values is None: + return None, None + buf = (c_float * length)() + try: # numpy fast path without importing numpy as a hard dependency + mv = memoryview(values) + if mv.format == "f" and mv.nbytes == length * 4 and mv.c_contiguous: + ctypes.memmove(buf, (ctypes.c_char * mv.nbytes).from_buffer_copy(mv), mv.nbytes) + return buf, buf + except TypeError: + pass + seq = list(values) + if len(seq) != length: + raise ValueError(f"{name} has {len(seq)} values, model expects {length}") + for i, v in enumerate(seq): + buf[i] = float(v) + return buf, buf + + +class Model: + """A loaded checkpoint. Free it with ``close()`` or a ``with`` block.""" + + def __init__(self, handle, lib): + self._h = handle + self._lib = lib + cfg = _ffi.Config() + rc = lib.vla_model_config(handle, ctypes.byref(cfg)) + if rc != _ffi.OK: + raise RuntimeError(f"vla_model_config failed ({rc})") + self.config = cfg + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + return False + + def close(self): + if getattr(self, "_h", None): + self._lib.vla_model_free(self._h) + self._h = None + + def __del__(self): + self.close() + + def predict(self, images, tokens: Sequence[int], state=None, noise=None, + pixel_format: int = PIXEL_U8, timing: int = TIMING_NONE): + """Run one forward pass. + + images: one HWC array, or a sequence of them for multi-view. uint8 RGB by + default; pass pixel_format=PIXEL_F32_RGB_01 for float RGB in [0, 1]. + """ + if self._h is None: + raise RuntimeError("model is closed") + + views = images if isinstance(images, (list, tuple)) else [images] + if not views: + raise ValueError("at least one image is required") + + img_array = (_ffi.Image * len(views))() + keep = [] + for i, im in enumerate(views): + mv = memoryview(im) + if not mv.c_contiguous: + raise ValueError("image must be C-contiguous") + shape = mv.shape + if len(shape) != 3 or shape[2] != 3: + raise ValueError(f"image must be HxWx3, got {shape}") + raw = (ctypes.c_char * mv.nbytes).from_buffer_copy(mv) + keep.append(raw) + img_array[i].data = ctypes.cast(raw, ctypes.c_void_p) + img_array[i].h = int(shape[0]) + img_array[i].w = int(shape[1]) + img_array[i].format = int(pixel_format) + + tok = list(tokens) + if not tok: + raise ValueError("tokens must not be empty") + tok_buf = (c_int32 * len(tok))(*[int(t) for t in tok]) + + state_ptr, state_keep = _as_f32_array( + state if state is not None else [0.0] * int(self.config.max_state_dim), + int(self.config.max_state_dim), "state") + keep.append(state_keep) + + noise_len = int(self.config.max_action_dim) * int(self.config.n_suffix) + noise_ptr, noise_keep = _as_f32_array(noise, noise_len, "noise") + keep.append(noise_keep) + + cin = _ffi.Inputs() + cin.images = img_array + cin.n_images = len(views) + cin.lang_tokens = tok_buf + cin.n_lang = len(tok) + if state_ptr is not None: + cin.state = ctypes.cast(state_ptr, POINTER(c_float)) + if noise_ptr is not None: + cin.noise = ctypes.cast(noise_ptr, POINTER(c_float)) + cin.timing_detail = int(timing) + + out = POINTER(c_float)() + n = c_int64() + rc = self._lib.vla_predict(self._h, ctypes.byref(cin), ctypes.byref(out), ctypes.byref(n)) + if rc != _ffi.OK: + raise RuntimeError(f"vla_predict failed ({rc})") + try: + flat = [out[i] for i in range(n.value)] + finally: + self._lib.vla_free_actions(out) + + cols = int(self.config.max_action_dim) or 1 + rows = len(flat) // cols if cols else len(flat) + try: + import numpy as np + return np.asarray(flat, dtype="float32").reshape(rows, cols) + except ImportError: + return [flat[r * cols:(r + 1) * cols] for r in range(rows)] + + def last_stats(self) -> _ffi.Stats: + st = _ffi.Stats() + rc = self._lib.vla_last_stats(self._h, ctypes.byref(st)) + if rc != _ffi.OK: + raise RuntimeError(f"vla_last_stats failed ({rc})") + return st + + +def load(ckpt_path: str, mmproj_path: str | None = None, config_path: str | None = None) -> Model: + """Load a checkpoint. mmproj_path is only needed for SmolVLA, pi0 and pi0.5.""" + lib = _lib_handle() + handle = lib.vla_model_load( + mmproj_path.encode() if mmproj_path else None, + ckpt_path.encode(), + config_path.encode() if config_path else None, + ) + if not handle: + raise RuntimeError(f"could not load {ckpt_path}") + return Model(handle, lib) diff --git a/bindings/python/vla_cpp/_ffi.py b/bindings/python/vla_cpp/_ffi.py new file mode 100644 index 0000000..41ae61d --- /dev/null +++ b/bindings/python/vla_cpp/_ffi.py @@ -0,0 +1,163 @@ +"""ctypes declarations for libvla. Mirrors include/vla.h field for field.""" + +from __future__ import annotations + +import ctypes +import os +import sys +from ctypes import ( + POINTER, + c_char_p, + c_double, + c_float, + c_int32, + c_int64, + c_void_p, +) + +ABI_VERSION = 1 + +OK = 0 +ERR_ARG = -1 +ERR_PREDICT = -2 +ERR_EXCEPTION = -3 + +PIXEL_U8 = 0 +PIXEL_F32_RGB_01 = 1 + +TIMING_NONE = 0 +TIMING_PHASE = 1 + + +class Config(ctypes.Structure): + _fields_ = [ + ("n_img", c_int64), + ("n_lang", c_int64), + ("n_state", c_int64), + ("n_prefix", c_int64), + ("n_suffix", c_int64), + ("n_full", c_int64), + ("hidden", c_int64), + ("expert_h", c_int64), + ("intermediate", c_int64), + ("expert_inter", c_int64), + ("n_q_heads", c_int64), + ("n_kv_heads", c_int64), + ("head_dim", c_int64), + ("q_full_dim", c_int64), + ("kv_full_dim", c_int64), + ("n_layers", c_int64), + ("self_attn_every_n", c_int32), + ("max_state_dim", c_int64), + ("max_action_dim", c_int64), + ("real_state_dim", c_int64), + ("real_action_dim", c_int64), + ("norm_eps", c_float), + ("min_period", c_double), + ("max_period", c_double), + ("num_steps", c_int32), + ("rms_eps", c_float), + ("rope_n_dims", c_int32), + ("rope_mode", c_int32), + ("rope_freq_base", c_float), + ("denormalized", c_int32), + ] + + +class Image(ctypes.Structure): + _fields_ = [ + ("data", c_void_p), + ("w", c_int32), + ("h", c_int32), + ("format", c_int32), + ] + + +class Inputs(ctypes.Structure): + _fields_ = [ + ("images", POINTER(Image)), + ("n_images", c_int32), + ("precomputed_img_emb", POINTER(c_float)), + ("n_img_views", c_int32), + ("lang_tokens", POINTER(c_int32)), + ("n_lang", c_int32), + ("state", POINTER(c_float)), + ("noise", POINTER(c_float)), + ("attention_mask", POINTER(c_int32)), + ("attention_mask_n", c_int32), + ("timing_detail", c_int32), + ] + + +class Stats(ctypes.Structure): + _fields_ = [ + ("ms_total", c_float), + ("ms_vision", c_float), + ("ms_inference", c_float), + ("ms_prefill", c_float), + ("ms_denoise", c_float), + ] + + +def _library_name() -> str: + if sys.platform == "darwin": + return "libvla.dylib" + if sys.platform == "win32": + return "vla.dll" + return "libvla.so" + + +def _candidates() -> list[str]: + name = _library_name() + here = os.path.dirname(os.path.abspath(__file__)) + found = [] + env = os.environ.get("VLA_LIBRARY") + if env: + found.append(env) + # Bundled next to the package (what a wheel ships), then a local build tree. + found.append(os.path.join(here, name)) + found.append(os.path.join(here, "lib", name)) + found.append(name) # fall through to the loader search path + return found + + +def load_library() -> ctypes.CDLL: + errors = [] + lib = None + for path in _candidates(): + try: + lib = ctypes.CDLL(path) + break + except OSError as exc: + errors.append(f"{path}: {exc}") + if lib is None: + raise OSError( + "could not load libvla. Set VLA_LIBRARY to its path, or build it with " + "`cmake --build --target vla`.\nTried:\n " + "\n ".join(errors) + ) + + lib.vla_abi_version.restype = c_int32 + lib.vla_abi_version.argtypes = [] + + lib.vla_model_load.restype = c_void_p + lib.vla_model_load.argtypes = [c_char_p, c_char_p, c_char_p] + + lib.vla_model_free.restype = None + lib.vla_model_free.argtypes = [c_void_p] + + lib.vla_model_config.restype = c_int32 + lib.vla_model_config.argtypes = [c_void_p, POINTER(Config)] + + lib.vla_predict.restype = c_int32 + lib.vla_predict.argtypes = [c_void_p, POINTER(Inputs), POINTER(POINTER(c_float)), POINTER(c_int64)] + + lib.vla_free_actions.restype = None + lib.vla_free_actions.argtypes = [POINTER(c_float)] + + lib.vla_last_stats.restype = c_int32 + lib.vla_last_stats.argtypes = [c_void_p, POINTER(Stats)] + + got = lib.vla_abi_version() + if got != ABI_VERSION: + raise RuntimeError(f"libvla ABI {got}, this package expects {ABI_VERSION}") + return lib diff --git a/ci/config/matrix.env b/ci/config/matrix.env index 94da008..463ec76 100644 --- a/ci/config/matrix.env +++ b/ci/config/matrix.env @@ -43,7 +43,7 @@ models_for() { local v="MODELS_$1"; echo "${!v}"; } MULTISUITE_MODELS_rtx3090="bitvla gr00t_n1_7" MULTISUITE_SUITES="libero_spatial libero_object libero_goal libero_10" -multisuite_models_for() { local v=KNOWN_ISSUES"MULTISUITE_MODELS_$1"; echo "${!v:-}"; } +multisuite_models_for() { local v="MULTISUITE_MODELS_$1"; echo "${!v:-}"; } # Suites a given (platform, model) should run = DEFAULT_SUITE, plus the # multisuite set if the model is listed for that platform. diff --git a/docs/ADOPTION.md b/docs/ADOPTION.md index b6db733..1fd0ba0 100644 --- a/docs/ADOPTION.md +++ b/docs/ADOPTION.md @@ -1,29 +1,31 @@ # Adoption notes -vla.cpp is technically solid (7 architectures, self-contained GGUFs, CUDA + Jetson, -real benchmarks). The gap to llama.cpp-style reach is mostly distribution, not code. -Ordered by leverage: +The engine works; the gap is distribution. -1. **Publish on GitHub.** The repo lives on Bitbucket (`bitbucket.org/vinrobotics/vla.cpp`) - while the README already links a `github.com/VinRobotics/vla.cpp` URL. llama.cpp's reach - came from GitHub visibility, issues, and PRs. A public GitHub mirror is the single biggest - lever; nothing else here matters as much. +Done: -2. **Ship the models.** All seven GGUFs are already published under - [`vrfai`](https://huggingface.co/vrfai) on the Hub - the README's "coming soon" rows are - stale (now fixed). Keep the model table pointing at the real repos so the policies are - one `hf download` away. +1. **C ABI.** `include/vla.h` and `libvla`. `src/model.h` is C++ only, so + without it nothing outside C++ can link the engine. +2. **Python bindings.** `bindings/python`, ctypes over the ABI. +3. **Prebuilt binaries.** `.github/workflows/release.yml` publishes + linux-x86_64 (CPU and CUDA), linux-aarch64 (CPU, for Jetson-class boards), + macos-arm64-metal and a Docker image on tag. +4. **One-command model fetch.** `-hf user/repo[:file.gguf]` on `vla-cli`, + `vla-server` and `vla-bench`, cached under `$VLA_CACHE`. +5. **Reproducible benchmarks.** `vla-bench` emits the README table rows. +6. **Contributor path.** `CONTRIBUTING.md` has the six-site walkthrough for + adding an architecture, plus issue and PR templates. +7. **Instruction in, action out.** `vla-cli --text` tokenizes with the + architecture's own tokenizer, so the quickstart no longer needs raw ids. -3. **Cut releases.** `v0.1.0` is the first tag (see `CHANGELOG.md`). Tagged releases + - changelog give users something to pin and cite. +Left: -4. **Rotate the committed credential.** The local `.git/config` remote URL embeds an - access token (`https://@bitbucket.org/...`). It is never pushed (git config is not - tracked), so this is hygiene, not a live leak - but rotate it and use a credential helper - or SSH remote instead of an inline token. +- **Jetson CUDA binaries.** The aarch64 job is CPU only: the hosted arm64 image + carries no CUDA, so a Jetson GPU build still happens on the device. +- **PyPI.** The wheel is built from `bindings/python` but nothing publishes it. +- **`ci/baselines/rtx3090.json`** still disagrees with the README table, which is + now RTX 5090 numbers from `vla-bench`. Re-record the baselines on one machine. +- **Success rates.** The README table comes from a May 2026 RTX 3060 sweep and + covers seven of the eleven archs. A fresh sweep would cover the rest. -5. **Lower the build bar (optional).** A CUDA `Dockerfile` now exists; publishing a prebuilt - image (and, later, macOS/Metal or ROCm backends) removes the from-source step that stops - most drive-by users. - -None of these change inference behaviour; they change who can find and run it. +None of these change inference behaviour. diff --git a/docs/backend/metal.md b/docs/backend/metal.md index db6674c..e9adc61 100644 --- a/docs/backend/metal.md +++ b/docs/backend/metal.md @@ -17,9 +17,7 @@ To disable the Metal build at compile time use the `-DGGML_METAL=OFF` cmake opti When built with Metal support, you can explicitly disable GPU inference with the `--n-gpu-layers 0` command-line argument. ```bash -# Fetch llama.cpp at pinned tag and apply local patch -bash patches/patch.sh - +# cmake fetches llama.cpp at the pinned tag; no patch step. # On MacOS, Metal is enabled by default cmake -B build -DCMAKE_BUILD_TYPE=Release cmake --build build -j$(sysctl -n hw.ncpu) diff --git a/docs/backend/sycl.md b/docs/backend/sycl.md new file mode 100644 index 0000000..9017f93 --- /dev/null +++ b/docs/backend/sycl.md @@ -0,0 +1,222 @@ +# `vla.cpp` on Intel GPUs (SYCL backend) + +Notes for building and running `vla.cpp` on Intel discrete and integrated GPUs +through oneAPI SYCL. Unlike Metal, SYCL is **not** auto-detected: it needs the +oneAPI DPC++ compiler and an explicit `-DGGML_SYCL=ON`. + +Verified on an **Intel Arc A380** (DG2 / Xe-HPG, `8086:56a5`, 6 GB) on Ubuntu +22.04, kernel 6.8, with an AMD Ryzen host CPU. The same path covers the rest of +the Arc A/B series, Flex, Data Center Max, and the Xe iGPUs. + +## Prerequisites + +### 1. GPU compute runtime + +The kernel driver (`i915`, in-tree since 6.2 for DG2) is not enough - you also +need the userspace compute stack: Level Zero plus the NEO OpenCL runtime. + +```bash +wget -qO- https://repositories.intel.com/gpu/intel-graphics.key \ + | sudo gpg --yes --dearmor -o /usr/share/keyrings/intel-graphics.gpg +echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/gpu/ubuntu jammy client" \ + | sudo tee /etc/apt/sources.list.d/intel-gpu-jammy.list +sudo apt-get update +sudo apt-get install -y intel-opencl-icd libze-intel-gpu1 libze1 libze-dev intel-ocloc clinfo +``` + +Substitute your distro codename for `jammy`. Install the userspace packages +only - do **not** add `intel-i915-dkms` on a 6.8+ kernel, whose in-tree `i915` +already drives DG2. + +Then give your user access to the render node and re-login: + +```bash +sudo usermod -aG render,video "$USER" +``` + +Check it before going further - `clinfo -l` must name your GPU: + +``` +Platform #0: Intel(R) OpenCL Graphics + `-- Device #0: Intel(R) Arc(TM) A380 Graphics +``` + +### 2. oneAPI + +The SYCL backend needs the DPC++ compiler, oneMKL and oneDNN. **Deep Learning +Essentials** carries exactly those and is much smaller than the full Base +Toolkit. + +```bash +wget -qO- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \ + | sudo gpg --yes --dearmor -o /usr/share/keyrings/oneapi-archive-keyring.gpg +echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \ + | sudo tee /etc/apt/sources.list.d/oneAPI.list +sudo apt-get update +sudo apt-get install -y intel-deep-learning-essentials-2025.3 +``` + +2025.3 is the newest release verified by llama.cpp's own SYCL docs that still +supports Ubuntu 22.04; the 2026.x series dropped jammy. Confirm the toolchain +sees the GPU over Level Zero: + +```bash +source /opt/intel/oneapi/setvars.sh +sycl-ls +``` + +``` +[level_zero:gpu][level_zero:0] Intel(R) oneAPI Unified Runtime over Level-Zero, Intel(R) Arc(TM) A380 Graphics 12.56.5 [1.6.31294+20] +[opencl:gpu][opencl:1] Intel(R) OpenCL Graphics, Intel(R) Arc(TM) A380 Graphics OpenCL 3.0 NEO [24.39.31294] +``` + +Plus the usual host dependencies. Ubuntu 22.04 has no `cppzmq` package, so drop +its two headers in by hand: + +```bash +sudo apt-get install -y cmake ninja-build pkg-config \ + protobuf-compiler libprotobuf-dev libzmq3-dev +wget -q https://github.com/zeromq/cppzmq/archive/refs/tags/v4.10.0.tar.gz -O - | tar xz +sudo install -m644 cppzmq-4.10.0/zmq.hpp cppzmq-4.10.0/zmq_addon.hpp /usr/local/include/ +``` + +## Configure & build + +ggml's SYCL sources only compile under the oneAPI DPC++ driver, and +`CMAKE_CXX_COMPILER` is global, so the whole project - `vla_core`, the servers, +the CLI - is built by `icpx`. Configure a **fresh** build directory; switching +compilers in an existing one does not work. + +```bash +source /opt/intel/oneapi/setvars.sh + +cmake -B build-sycl -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DGGML_SYCL=ON \ + -DCMAKE_C_COMPILER=icx \ + -DCMAKE_CXX_COMPILER=icpx +cmake --build build-sycl -j$(nproc) +``` + +`setvars.sh` must be sourced in every shell that builds *or runs* the binaries - +`libsycl.so`, `libdnnl.so` and the oneMKL libraries live under `/opt/intel`. + +## GPU offload + +The core picks its backend at load time. Confirm from the startup banner: + +``` +vla: backend = SYCL (device 0: Intel(R) Arc(TM) A380 Graphics) +``` + +If you see `vla: backend = CPU (8 threads)` instead, the build did not pick up +SYCL, or no SYCL device was visible - re-check `sycl-ls` and your `render` group +membership. + +On a multi-GPU box, `VLA_DEVICE=` selects the ordinal (the same variable +selects the CUDA device). The index is range-checked against the SYCL device +count; an out-of-range value logs and falls back to CPU rather than running off +the end of the device array. + +> Single-backend, no per-op CPU fallback: the core drives one backend through +> `gallocr`, not a scheduler. An arch that hits an op the SYCL backend does not +> implement asserts at predict time rather than silently falling back. + +BitVLA is the one exception: it pins its ggml graph to the CPU backend by design +and offloads its LM through separate hand-written CUDA kernels, so a SYCL build +leaves it on the CPU. There is no SYCL port of those kernels. + +## Known issue: the SYCL VMM pool and oneDNN + +ggml-sycl's VMM pool hands out virtual-memory-backed pointers that oneDNN cannot +wrap in a `dnnl::memory`. When it happens the GEMM aborts the process: + +``` +could not create a memory object +SYCL error: ... in function ggml_sycl_op_mul_mat at .../ggml-sycl.cpp:3055 +``` + +It fires whenever `src0` is not already F32 - BF16, F16 and every quantized type +are converted into that pool before the GEMM - which is most checkpoints, +including the default BF16 weights of SmolVLA, π0, π0.5, Evo-1, VLA-Adapter and +OpenVLA-OFT. + +`vla.cpp` defaults `GGML_SYCL_ENABLE_VMM=0` when it brings up SYCL, which avoids +it and is the faster of the two workarounds (disabling oneDNN with +`GGML_SYCL_ENABLE_DNN=0` also clears the crash, but costs ~8%). It is only a +default: set `GGML_SYCL_ENABLE_VMM=1` explicitly to keep the pool on hardware +where it pays off. + +## Fixed upstream: `bf16 -> f32` copies + +ggml-sycl's copy table used to have `f16 -> f32` but no `bf16 -> f32`, so an arch +whose graph contained that copy aborted at predict time. VLA-Adapter hit it with +its default BF16 weights, and the workaround was `VLA_ADAPTER_F32_WEIGHTS=1`. + +llama.cpp b10326 adds the missing kernel (`cpy_1_bf16_f32` in +`ggml/src/ggml-sycl/cpy.cpp`), so VLA-Adapter should run on stock BF16 weights +now. Not yet re-tested on the A380 - if you hit the old abort, fall back to +`VLA_ADAPTER_F32_WEIGHTS=1` and file an issue. + +## Performance note: F32 weights + +BF16 has no native DPAS path on Xe-HPG, so BF16 weights are slower there than +plain F32 despite the extra bandwidth. Each arch exposes a switch +(`VLA_WEIGHT_DTYPE=f32` for SmolVLA, `VLA_PI0_F32_WEIGHTS=1` for π0, and so on), +and on the A380 it is worth ~16%: + +| SmolVLA weights | vision | inference | total | +|---|---:|---:|---:| +| BF16 (default) | 158 ms | 474 ms | **630 ms** | +| F32 (`VLA_WEIGHT_DTYPE=f32`) | 173 ms | 355 ms | **528 ms** | + +The tradeoff is memory - F32 doubles the resident weights (1.07 GiB -> 2.09 GiB +for SmolVLA), which matters on a 6 GB A380 for the larger checkpoints. The +default stays BF16 for that reason. + +## Results + +Measured with `vla_predict_check`, which is a test target - add +`-DVLA_BUILD_TESTS=ON` to the configure line above to get it. Fixed noise, so +runs are comparable; best of 5-10 iterations after 3 warmups. Host is an AMD Ryzen 5 5500 (CPU backend uses 8 +threads); GPU is the Arc A380. + +| Model | input | CPU | Arc A380 | speedup | +|---|---|---:|---:|---:| +| SmolVLA | 512 | 1,920 ms | **630 ms** | 3.0x | +| Evo-1 | 448 | 7,695 ms | **1,176 ms** | 6.5x | +| VLA-Adapter | 224 | 2,994 ms | **517 ms** | 5.8x | + +VLA-Adapter is measured with `VLA_ADAPTER_F32_WEIGHTS=1` on both sides, which was +required at the time (see the `bf16 -> f32` section above); the others run their +stock defaults. + +Per-stage for SmolVLA: + +| Stage | CPU | Arc A380 (SYCL) | +|--------------|-------------:|------------------:| +| vision | 1,119 ms | 158 ms | +| inference | 804 ms | 474 ms | +| **total/req**| **1,920 ms**| **630 ms** | + +SmolVLA gains least because its flow-matching denoise loop is a long chain of +small GEMMs that cannot fill 128 EUs; its vision tower alone is 7.1x. With +`VLA_WEIGHT_DTYPE=f32` it reaches 528 ms (3.6x). + +Outputs were checked against the CPU backend on every model above: max absolute +deviation 2.9e-3 on actions peaking at 0.99 (2.9e-6 for the all-F32 +VLA-Adapter run), RMS 2.4e-4 - BF16/F32 kernel rounding, not a numerical +regression. + +### Memory ceiling + +The A380 has 6 GB, of which ~5.7 GB is addressable. GR00T N1.7 (6.3 GB of F32 +weights) does not fit and dies in the allocator: + +``` +level_zero backend failed with error: 38 (UR_RESULT_ERROR_OUT_OF_HOST_MEMORY) +``` + +`VLA_GR00T_BF16_WEIGHTS=1` halves the weights but its activations still overflow +the card. There is no host-memory spill path - the core is single-backend - so +the larger checkpoints need an A770/B580-class card or better. diff --git a/docs/backend/wsl.md b/docs/backend/wsl.md index 5bcce3f..c715489 100644 --- a/docs/backend/wsl.md +++ b/docs/backend/wsl.md @@ -1,10 +1,9 @@ # `vla.cpp` on Windows (WSL2 + CUDA) `vla.cpp` targets Linux and macOS. On Windows the supported path is **WSL2** -with an Ubuntu distribution: the toolchain (`libzmq`, `protobuf`, `pkg-config`, -the bash `patches/patch.sh` script) and the CUDA build all run natively inside -the Linux environment, while still using the host NVIDIA GPU through the -WSL CUDA driver. +with an Ubuntu distribution: the toolchain (`libzmq`, `protobuf`, `pkg-config`) +and the CUDA build all run natively inside the Linux environment, while still +using the host NVIDIA GPU through the WSL CUDA driver. ## Prerequisites diff --git a/examples/chat/README.md b/examples/chat/README.md index 569c2f4..40a60ca 100644 --- a/examples/chat/README.md +++ b/examples/chat/README.md @@ -2,9 +2,9 @@ A minimal **streaming image+text chat** client for `vlm-server`, the llama.cpp + libmtmd chat runtime (`src/vlm/engine.cpp`) behind a ZMQ daemon. Send text and -images, get a streamed reply. The design rationale lives in -[docs/VLM-SERVER.md](../../docs/VLM-SERVER.md); this README is how to **run** it, -plus the validation numbers for the SmolVLM2-500M-Instruct setup. +images, get a streamed reply. The layering is sketched in +[docs/ARCHITECTURE.md](../../docs/ARCHITECTURE.md); this README is how to **run** +it, plus the validation numbers for the SmolVLM2-500M-Instruct setup. ``` examples/chat/ diff --git a/include/vla.h b/include/vla.h new file mode 100644 index 0000000..f278f46 --- /dev/null +++ b/include/vla.h @@ -0,0 +1,167 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Stable C ABI over src/model.h, so bindings do not have to match a C++ compiler +// or standard library. +// +// Ownership: vla_model_load pairs with vla_model_free, vla_predict pairs with +// vla_free_actions, and every pointer in vla_inputs is borrowed for the call. + +#ifndef VLA_H +#define VLA_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#if defined(_WIN32) +# define VLA_API __declspec(dllexport) +#else +# define VLA_API __attribute__((visibility("default"))) +#endif + +// Bumped on any incompatible change to the structs or functions below. +#define VLA_ABI_VERSION 1 + +typedef struct vla_model vla_model; + +typedef enum { + VLA_OK = 0, + VLA_ERR_ARG = -1, ///< Null handle or malformed vla_inputs. + VLA_ERR_PREDICT = -2, ///< The engine returned no actions. + VLA_ERR_EXCEPTION = -3, ///< A C++ exception was caught at the boundary. +} vla_status; + +typedef enum { + VLA_PIXEL_U8 = 0, ///< 8-bit interleaved RGB. + VLA_PIXEL_F32_RGB_01 = 1, ///< Float32 interleaved RGB in [0, 1]. +} vla_pixel_format; + +typedef enum { + VLA_TIMING_NONE = 0, + VLA_TIMING_PHASE = 1, +} vla_timing_detail; + +/// Mirrors vla::Config. See src/model.h for the per-field meaning. +typedef struct { + int64_t n_img; + int64_t n_lang; + int64_t n_state; + int64_t n_prefix; + int64_t n_suffix; + int64_t n_full; + + int64_t hidden; + int64_t expert_h; + int64_t intermediate; + int64_t expert_inter; + int64_t n_q_heads; + int64_t n_kv_heads; + int64_t head_dim; + int64_t q_full_dim; + int64_t kv_full_dim; + int64_t n_layers; + int32_t self_attn_every_n; + + int64_t max_state_dim; + int64_t max_action_dim; + int64_t real_state_dim; + int64_t real_action_dim; + float norm_eps; + double min_period; + double max_period; + int32_t num_steps; + + float rms_eps; + int32_t rope_n_dims; + int32_t rope_mode; + float rope_freq_base; + + /// Non-zero if vla_predict already applied the dataset statistics. Zero for + /// the GR00T family and VLA-JEPA, whose callers un-normalise themselves. + int32_t denormalized; +} vla_config; + +typedef struct { + const void * data; ///< First pixel, caller-owned. + int32_t w; + int32_t h; + int32_t format; ///< A vla_pixel_format value. +} vla_image; + +typedef struct { + const vla_image * images; + int32_t n_images; + + /// Optional, replaces the vision tower. Layout [n_img_views * n_img, hidden], + /// already in the scale that arch's LM expects. Pass NULL to use images. + const float * precomputed_img_emb; + int32_t n_img_views; + + const int32_t * lang_tokens; + int32_t n_lang; + + /// Length max_state_dim; pad real_state_dim..max with zeros. + const float * state; + /// Length n_suffix * max_action_dim, or NULL to sample internally. + const float * noise; + + /// Only Evo-1 reads this; other archs derive their own mask. + const int32_t * attention_mask; + int32_t attention_mask_n; + + int32_t timing_detail; ///< A vla_timing_detail value. +} vla_inputs; + +/// Milliseconds. Phase fields are zero unless timing_detail was VLA_TIMING_PHASE. +typedef struct { + float ms_total; + float ms_vision; + float ms_inference; + float ms_prefill; + float ms_denoise; +} vla_stats; + +/// VLA_ABI_VERSION of the loaded library, for a runtime compatibility check. +VLA_API int32_t vla_abi_version(void); + +/// mmproj_path may be NULL or "" for archs that bake vision into the checkpoint. +/// config_path may be NULL. Returns NULL on failure. +VLA_API vla_model * vla_model_load(const char * mmproj_path, + const char * ckpt_path, + const char * config_path); + +VLA_API void vla_model_free(vla_model * m); + +/// Fills out with the resolved config. Returns a vla_status. +VLA_API int32_t vla_model_config(const vla_model * m, vla_config * out); + +/// Runs one forward pass. On VLA_OK, *out_actions points to *out_n floats that +/// the caller releases with vla_free_actions. Row-major [num_steps or n_suffix, +/// max_action_dim]; only the first real_action_dim columns carry values. +VLA_API int32_t vla_predict(vla_model * m, const vla_inputs * in, + float ** out_actions, int64_t * out_n); + +VLA_API void vla_free_actions(float * actions); + +/// Timings of the most recent vla_predict on this handle. +VLA_API int32_t vla_last_stats(const vla_model * m, vla_stats * out); + +#ifdef __cplusplus +} +#endif + +#endif // VLA_H diff --git a/pyproject.toml b/pyproject.toml index 85ec419..eb1700b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "vla-cpp-tooling" -version = "0.1.1" +version = "0.2.0" description = "Python tooling for vla.cpp: HuggingFace -> GGUF converters and the ZeroMQ eval client." readme = "README.md" requires-python = ">=3.10" @@ -23,7 +23,10 @@ client = [ "msgpack-numpy>=0.4.8", "pillow>=10", "torch>=2.5", - "transformers>=4.51,<4.52", + # Only used for AutoTokenizer/AutoProcessor.from_pretrained (the pi0 PaliGemma + # tokenizer). Capped below 5.0: the v5 line is a breaking rewrite we have not + # tested against. + "transformers>=4.51,<5", "numpy>=1.24", ] diff --git a/scripts/tokenize_prompt.py b/scripts/tokenize_prompt.py new file mode 100644 index 0000000..4968224 --- /dev/null +++ b/scripts/tokenize_prompt.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Print the token ids for an instruction, using the tokenizer an arch was trained with. + + tokenize_prompt.py --arch smolvla --text "pick up the black bowl" + -> 1,4842,731,254,2482,7681,2 + +vla-cli --text calls this so the quickstart does not need raw ids. The eval +client keeps its own richer prompt handling; this only covers the plain case. +""" + +import argparse +import sys + +# Same tokenizers the eval client uses (eval/client/vla_cpp_client.py). +TOKENIZERS = { + "smolvla": "HuggingFaceTB/SmolVLM2-500M-Instruct", + "pi0": "google/paligemma-3b-pt-224", + "pi05": "google/paligemma-3b-pt-224", + "evo1": "OpenGVLab/InternVL3-1B", + "bitvla": "hongyuw/ft-bitvla-bitsiglipL-224px-libero_object-bf16", + "vla_adapter": "VLA-Adapter/LIBERO-Object-Pro", + "openvla_oft": "moojink/openvla-7b-oft-finetuned-libero-spatial-object-goal-10", + "vla_jepa": "Qwen/Qwen3-VL-2B-Instruct", + "gr00t_n1_5": "lerobot/eagle2hg-processor-groot-n1p5", + "gr00t_n1_7": "nvidia/Cosmos-Reason2-2B", +} +TRUST_REMOTE_CODE = {"evo1", "gr00t_n1_5"} + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--arch", required=True, choices=sorted(TOKENIZERS)) + ap.add_argument("--text", required=True) + ap.add_argument("--tokenizer", help="override the HuggingFace tokenizer id") + args = ap.parse_args() + + try: + from transformers import AutoTokenizer + except ImportError: + print("transformers is not installed: pip install -e \".[client]\"", file=sys.stderr) + return 1 + + name = args.tokenizer or TOKENIZERS[args.arch] + tok = AutoTokenizer.from_pretrained( + name, trust_remote_code=args.arch in TRUST_REMOTE_CODE) + ids = tok(args.text)["input_ids"] + print(",".join(str(int(i)) for i in ids)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/arch.h b/src/arch.h index 633e8ff..e66fa33 100644 --- a/src/arch.h +++ b/src/arch.h @@ -28,17 +28,26 @@ #include "model.h" #include +#include +#include #include #include #include namespace vla { -// Default CPU thread count for the in-tree loaders: all cores, capped at 8, -// with a safe fallback when the hardware count is unknown. +// CPU threads for the in-tree loaders; VLA_N_THREADS overrides. Cap measured on +// a 24-core host: 8 to 16 is 38-41% faster on vla_adapter, evo1 and gr00tn1d5, +// and 24 is slower than 16. Output is bit-identical either way. inline int default_cpu_threads() { + if (const char * e = std::getenv("VLA_N_THREADS")) { + char * end = nullptr; + const long n = std::strtol(e, &end, 10); + if (*end == '\0' && n > 0 && n <= 1024) return (int) n; + std::fprintf(stderr, "vla: ignoring VLA_N_THREADS='%s'\n", e); + } const unsigned hw = std::thread::hardware_concurrency(); - return hw == 0 ? 4 : (int) std::min(hw, 8u); + return hw == 0 ? 4 : (int) std::min(hw, 16u); } /** diff --git a/src/backend.h b/src/backend.h new file mode 100644 index 0000000..a4d89bb --- /dev/null +++ b/src/backend.h @@ -0,0 +1,163 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * @file backend.h + * @brief Compute-backend selection shared by every in-tree arch. + * + * Each arch used to open-code the same accelerator-then-CPU ladder, so adding a + * backend meant editing a dozen files. They all call @ref vla::backend_init + * instead; the ladder lives here once. + * + * Exactly one accelerator is compiled in, picked by the CMake flag that was + * used (`GGML_CUDA` / `GGML_SYCL` / `GGML_METAL`). There is no per-op CPU + * fallback: the core drives a single backend through `gallocr` rather than a + * scheduler, so an arch that hits an op the backend does not implement asserts + * at predict time instead of silently limping. + */ + +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#ifdef GGML_USE_CUDA +#include "ggml-cuda.h" +#endif +#ifdef GGML_USE_SYCL +#include "ggml-sycl.h" +#endif +#ifdef GGML_USE_METAL +#include "ggml-metal.h" +#endif + +#include +#include +#ifdef GGML_USE_SYCL +#include // setenv / _putenv_s +#include +#endif + +namespace vla { + +#ifdef GGML_USE_SYCL +// setenv is POSIX. _putenv_s has no "do not overwrite" mode, so check first. +inline void setenv_default(const char * key, const char * val) { +#ifdef _WIN32 + size_t len = 0; + if (getenv_s(&len, nullptr, 0, key) == 0 && len > 0) return; + _putenv_s(key, val); +#else + setenv(key, val, /*overwrite=*/0); +#endif +} +#endif + +/// Outcome of @ref backend_init. @c handle is null only if even the CPU backend +/// failed to come up, which callers treat as a fatal load error. +struct Backend { + ggml_backend_t handle = nullptr; +}; + +/// GPU ordinal for CUDA and SYCL; `VLA_DEVICE` overrides. Junk is rejected, not +/// silently read as device 0. +inline int backend_device_index() { + const char * e = std::getenv("VLA_DEVICE"); + if (!e || !*e) return 0; + char * end = nullptr; + const long idx = std::strtol(e, &end, 10); + if (*end != '\0' || idx < 0 || idx > 1024) { + std::fprintf(stderr, "vla: ignoring VLA_DEVICE='%s' (not a device index); using 0\n", e); + return 0; + } + return (int) idx; +} + +/** + * @brief Bring up the best compute backend available to this build. + * + * @param tag Log prefix identifying the arch, e.g. @c "vla(pi0)". + * @param n_threads Thread count handed to the CPU backend if it is used. + * @return The backend plus the flags the arch records about it. + */ +inline Backend backend_init(const char * tag, int n_threads) { + Backend b; + +#ifdef GGML_USE_CUDA + { + const int dev = backend_device_index(); + b.handle = ggml_backend_cuda_init(dev); + if (b.handle) { + std::printf("%s: backend = CUDA (device %d)\n", tag, dev); + } else { + std::fprintf(stderr, "%s: ggml_backend_cuda_init failed; falling back to CPU\n", tag); + } + } +#elif defined(GGML_USE_SYCL) + { + // ggml-sycl's VMM pool hands out virtual-memory-backed pointers that + // oneDNN cannot wrap in a dnnl::memory: the GEMM aborts the process with + // "could not create a memory object". Any src0 that is not already F32 + // (BF16, F16, and every quantized type) is converted into that pool + // first, so the crash hits most checkpoints. Turning the pool off is + // also the faster of the two workarounds -- measurably better than + // disabling oneDNN outright. Only a default: an explicit setting wins, + // for Intel GPUs where the pool is worth keeping. + // ggml reads this on the first SYCL entry point, so it must be set here. + // call_once: concurrent model_load would race on the environment. + static std::once_flag vmm_once; + std::call_once(vmm_once, [] { setenv_default("GGML_SYCL_ENABLE_VMM", "0"); }); + + // ggml_backend_sycl_init() guards the device index with assert(), which + // a Release build compiles out and then indexes past the device array. + // Range-check here so a SYCL build on a box with no Intel GPU (or a bad + // VLA_DEVICE) lands on CPU instead of corrupting memory. + const int dev = backend_device_index(); + const int n_dev = ggml_backend_sycl_get_device_count(); + if (dev >= n_dev) { + std::fprintf(stderr, "%s: SYCL device %d out of range (%d visible); falling back to CPU\n", + tag, dev, n_dev); + } else if ((b.handle = ggml_backend_sycl_init(dev)) != nullptr) { + char desc[256] = { 0 }; + ggml_backend_sycl_get_device_description(dev, desc, sizeof(desc)); + std::printf("%s: backend = SYCL (device %d: %s)\n", tag, dev, desc); + } else { + std::fprintf(stderr, "%s: ggml_backend_sycl_init failed; falling back to CPU\n", tag); + } + } +#elif defined(GGML_USE_METAL) + { + b.handle = ggml_backend_metal_init(); + if (b.handle) { + std::printf("%s: backend = Metal\n", tag); + } else { + std::fprintf(stderr, "%s: ggml_backend_metal_init failed; falling back to CPU\n", tag); + } + } +#endif + + if (!b.handle) { + b.handle = ggml_backend_cpu_init(); + if (!b.handle) { + std::fprintf(stderr, "%s: ggml_backend_cpu_init failed\n", tag); + return b; + } + ggml_backend_cpu_set_n_threads(b.handle, n_threads); + std::printf("%s: backend = CPU (%d threads)\n", tag, n_threads); + } + return b; +} + +} // namespace vla diff --git a/src/model.cpp b/src/model.cpp index b99ba3c..3b09ceb 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -87,6 +87,30 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { return ok; } +} // namespace + +bool config_is_sane(const Config& c) { + struct { const char* name; int64_t real; int64_t max; } pairs[] = { + { "state", c.real_state_dim, c.max_state_dim }, + { "action", c.real_action_dim, c.max_action_dim }, + }; + for (const auto& p : pairs) { + if (p.real < 0 || p.max < 0) { + std::fprintf(stderr, "vla: negative %s dim (real=%lld max=%lld)\n", + p.name, (long long) p.real, (long long) p.max); + return false; + } + if (p.max > 0 && p.real > p.max) { + std::fprintf(stderr, "vla: real_%s_dim %lld exceeds max_%s_dim %lld\n", + p.name, (long long) p.real, p.name, (long long) p.max); + return false; + } + } + return true; +} + +namespace { + bool detect_arch_safetensors(const std::string& path, Arch* out) { std::ifstream f(path, std::ios::binary); if (!f) return false; @@ -183,6 +207,10 @@ Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, break; } if (!impl) return nullptr; + if (!config_is_sane(impl->cfg)) { + std::fprintf(stderr, "vla: refusing to load %s\n", ckpt_path.c_str()); + return nullptr; + } auto* m = new Model(); m->impl = std::move(impl); diff --git a/src/model.h b/src/model.h index 9440f86..6265f34 100644 --- a/src/model.h +++ b/src/model.h @@ -76,6 +76,10 @@ struct Config { int rope_n_dims; ///< RoPE rotation width (per head). int rope_mode; ///< RoPE variant (NeoX / GPT-J / etc). float rope_freq_base; ///< RoPE base frequency. + + /// True if @ref predict already applied the dataset statistics. False for the + /// GR00T family and VLA-JEPA, whose callers un-normalise from a stats JSON. + bool denormalized = true; }; /// Opaque engine handle; created by @ref model_load and released by @@ -87,7 +91,10 @@ struct Model; */ enum class TimingDetail { NONE, ///< Only @c ms_total is populated. - PHASE, ///< Per-phase timings (vision, prefill, denoise, ...). + /// Per-phase timings (vision, prefill, denoise, ...). SmolVLA uses a second + /// builder here that does not pad the prefix to @c n_lang; same positions and + /// masking, so it differs from @c NONE only by float reduction order. + PHASE, }; /** @@ -120,7 +127,9 @@ struct Inputs { const ImageView* images; ///< Camera views (host memory). int n_images; ///< Number of @ref images. - /// Pre-computed image embeddings; bypasses the vision tower. + /// Pre-computed image embeddings, [n_img_views * n_img, hidden]; bypasses the + /// vision tower. Passed to the LM as-is, so the scale is arch-specific: pi0 + /// expects the projector output times 1/sqrt(hidden), pi0.5 expects it raw. const float* precomputed_img_emb = nullptr; int n_img_views = 0; ///< Number of views in /// @ref precomputed_img_emb. @@ -173,9 +182,8 @@ const Config& model_config(const Model* m); /** * @brief Run one forward pass. * - * Returns the predicted action chunk, normalised to the model's training - * statistics. The caller is responsible for un-normalising into world - * units. NaN/Inf inputs cause the call to abort. + * See @ref Config::denormalized for whether the result is in world units. + * NaN/Inf inputs cause the call to abort. * * @param m A handle from @ref model_load. * @param in Filled-in @ref Inputs struct. @@ -205,4 +213,12 @@ struct Stats { */ const Stats& last_stats(const Model* m); +/** + * @brief Reject a config whose real_* dims exceed its max_* dims. + * + * predict() sizes buffers from the max_* dims and loops to the real_* dims, so + * real > max writes out of bounds. Checked once for all archs at load. + */ +bool config_is_sane(const Config& c); + } diff --git a/src/models/bitvla.cpp b/src/models/bitvla.cpp index 57f23fe..b8a99e2 100644 --- a/src/models/bitvla.cpp +++ b/src/models/bitvla.cpp @@ -23,6 +23,7 @@ #endif #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" #ifdef VLA_BITVLA_CUDA_KERNELS #include "kernels/bitvla/bitvla_lm_cuda.h" @@ -97,9 +98,12 @@ struct BitvlaModelArch : public ModelArchBase { gguf_reader emb_reader{"bitvla"}; // stays open for per-step token-embedding row fetches std::vector stop_embed; // cached constant stop-token embedding row ggml_backend_t backend = nullptr; - bool is_cuda = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + scratch_ctx proprio_scratch; + scratch_ctx lm_scratch; + scratch_ctx head_scratch; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; bool packed_int2 = false; @@ -338,6 +342,22 @@ bool load_config(const gguf_reader & g, BitvlaModelArch & m, Config & cfg) { I("bitvla.tokens.stop_id", m.stop_id); m.packed_int2 = g.has("bitvla.quant.int2_packed") && g.u32("bitvla.quant.int2_packed") != 0; + // predict() sizes the patch buffer from n_patches but fills it by walking the + // image grid, so a KV that disagrees with the geometry overruns the buffer. + if (m.patch_size <= 0 || m.image_size <= 0 || m.image_size % m.patch_size != 0 || + m.n_patches != (m.image_size / m.patch_size) * (m.image_size / m.patch_size)) { + std::fprintf(stderr, "vla(bitvla): n_patches %lld does not match image %lld / patch %lld\n", + (long long) m.n_patches, (long long) m.image_size, (long long) m.patch_size); + return false; + } + // The CUDA LM writes seq*q_heads*head_dim into buffers sized seq*hidden. + if (m.lm_kv <= 0 || m.lm_head_dim <= 0 || m.lm_q % m.lm_kv != 0 || + m.lm_q * m.lm_head_dim != m.lm_hidden) { + std::fprintf(stderr, "vla(bitvla): lm q_heads %lld x head_dim %lld does not match hidden %lld\n", + (long long) m.lm_q, (long long) m.lm_head_dim, (long long) m.lm_hidden); + return false; + } + const std::string js = g.str("bitvla.statistics_json"); if (js.empty()) { std::fprintf(stderr, "vla(bitvla): bitvla.statistics_json KV missing - un-normalization will pass-through\n"); @@ -543,6 +563,8 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, (double) m->lm_rope_base, (long long) m->num_actions_chunk, (long long) m->action_dim, (long long) m->vocab_size, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); + // Not backend_init: the ggml graph stays on CPU and the LM offloads through + // the ternary CUDA kernels below. m->backend = ggml_backend_cpu_init(); if (!m->backend) { std::fprintf(stderr, "vla(bitvla): ggml_backend_cpu_init failed\n"); return nullptr; } ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); @@ -659,13 +681,23 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, if (cudaGetDeviceCount(&dev_count) == cudaSuccess && dev_count > 0) { cudaSetDevice(0); + // The ladder kernels dereference the scale pointer unconditionally, so + // a missing sidecar is a device-side OOB read, not a soft failure. + bool scales_ok = true; + auto load_bit = [&](ggml_tensor * t, int64_t N, int64_t K) -> std::pair { if (m->packed_int2) { int8_t * dp = upload_int8((const uint8_t*) t->data, ggml_nbytes(t), m->cuda_devptrs); std::string nm = ggml_get_name(t); std::string sn = nm.substr(0, nm.size() - 7) + ".scale"; std::vector sc = g.read_f32(sn.c_str()); - float * dws = sc.empty() ? nullptr : upload_f32_scales(sc.data(), (int) sc.size(), m->cuda_devptrs); + if (sc.empty()) { + std::fprintf(stderr, "vla(bitvla): int2 tensor %s has no %s sidecar\n", + nm.c_str(), sn.c_str()); + scales_ok = false; + return { dp, nullptr }; + } + float * dws = upload_f32_scales(sc.data(), (int) sc.size(), m->cuda_devptrs); return { dp, dws }; } float ws; @@ -679,7 +711,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, (int) m->lm_inter, (int) m->lm_layers, m->lm_rope_base, m->lm_rms_eps, max_seq); if (m->lm_cuda_ctx) { bool pack_ok = true; - for (int64_t L = 0; L < m->lm_layers && pack_ok; ++L) { + for (int64_t L = 0; L < m->lm_layers && pack_ok && scales_ok; ++L) { bitvla_lm_layer_cuda lyr{}; lyr.attn_norm_w = upload_bf16_from_f32((const float*) m->lm[L].attn_norm->data, m->lm_hidden, m->cuda_devptrs); @@ -697,8 +729,14 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, if (m->packed_int2) { lyr.gate_up_packed = upload_int8((const uint8_t*) m->lm[L].Wgate_up->data, ggml_nbytes(m->lm[L].Wgate_up), m->cuda_devptrs); - std::vector sc = g.read_f32(("lm.blk." + std::to_string(L) + ".ffn_gate_up.scale").c_str()); - lyr.gate_up_ws = upload_f32_scales(sc.data(), (int) sc.size(), m->cuda_devptrs); + const std::string sn = "lm.blk." + std::to_string(L) + ".ffn_gate_up.scale"; + std::vector sc = g.read_f32(sn.c_str()); + if (sc.empty()) { + std::fprintf(stderr, "vla(bitvla): missing %s\n", sn.c_str()); + scales_ok = false; + } else { + lyr.gate_up_ws = upload_f32_scales(sc.data(), (int) sc.size(), m->cuda_devptrs); + } } else { std::vector ws2; lyr.gate_up_packed = pack_and_upload_fused( @@ -711,7 +749,10 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, { auto r = load_bit(m->lm[L].Wdown, m->lm_hidden, m->lm_inter); lyr.down_packed = r.first; lyr.down_ws = r.second; } bitvla_lm_cuda_set_layer(m->lm_cuda_ctx, (int) L, &lyr); } - if (pack_ok) { + if (!scales_ok) { + std::fprintf(stderr, "vla(bitvla): int2 scale sidecars incomplete; refusing the CUDA LM\n"); + } + if (pack_ok && scales_ok) { __nv_bfloat16* onorm = upload_bf16_from_f32((const float*) m->lm_output_norm->data, m->lm_hidden, m->cuda_devptrs); bitvla_lm_cuda_set_output_norm(m->lm_cuda_ctx, onorm); @@ -888,6 +929,10 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, if (!t) continue; const size_t nb = ggml_nbytes(t); void* copy = std::malloc(nb); + if (!copy) { + std::fprintf(stderr, "vla(bitvla): out of memory keeping %zu bytes of CPU weights\n", nb); + return nullptr; + } std::memcpy(copy, t->data, nb); t->data = copy; m->cpu_kept_ptrs.push_back(copy); @@ -1014,9 +1059,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { #endif { - std::vector meta_buf((size_t) 24 * 1024 * 1024); - ggml_init_params gp = { meta_buf.size(), meta_buf.data(), true }; - ggml_context * ctx = ggml_init(gp); + ggml_context * ctx = vision_scratch.reset((size_t) 24 * 1024 * 1024); ggml_tensor * x_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, patch_flat, N); ggml_set_name(x_in, "patches"); @@ -1032,18 +1075,16 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { ggml_tensor * mm2 = ggml_add(ctx, ggml_mul_mat(ctx, mm_l2_w, mmg), mm_l2_b); ggml_set_name(mm2, "img_embeds"); - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 4096, false); ggml_build_forward_expand(gf, mm2); - if (!ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr_alloc_graph failed (view %lld)\n", (long long) v); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } + if (!vision_scratch.alloc(backend, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr_alloc_graph failed (view %lld)\n", (long long) v); return {}; } ggml_backend_tensor_set(x_in, patches.data(), 0, ggml_nbytes(x_in)); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): vision graph compute failed (view %lld)\n", (long long) v); - ggml_gallocr_free(galloc); ggml_free(ctx); return {}; + return {}; } ggml_backend_tensor_get(mm2, img_embeds_host.data() + (size_t) v * N * hidden_l, 0, (size_t) N * hidden_l * sizeof(float)); - ggml_gallocr_free(galloc); - ggml_free(ctx); } } stats.ms_vision = std::chrono::duration(clk::now() - t_v0).count(); @@ -1055,17 +1096,18 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { _dump_manifest(std::string("mm_proj_out fp32 ") + std::to_string(n_views) + " " + std::to_string(N) + " " + std::to_string(hidden_l)); std::vector proprio_embed_host((size_t) hidden_l); + // Like the other archs: a caller may leave the proprio vector out. + std::vector state_host((size_t) proprio_dim, 0.0f); + if (in.state) std::memcpy(state_host.data(), in.state, (size_t) proprio_dim * sizeof(float)); #ifdef VLA_BITVLA_CUDA_KERNELS if (cuda_fp32head_ready) { - if (bitvla_fp32head_proprio_forward(fp32head_cuda_ctx, in.state, proprio_embed_host.data(), 0) != 0) { + if (bitvla_fp32head_proprio_forward(fp32head_cuda_ctx, state_host.data(), proprio_embed_host.data(), 0) != 0) { std::fprintf(stderr, "vla(bitvla): CUDA proprio forward failed\n"); return {}; } } else #endif { - std::vector meta_buf((size_t) 4 * 1024 * 1024); - ggml_init_params gp = { meta_buf.size(), meta_buf.data(), true }; - ggml_context * ctx = ggml_init(gp); + ggml_context * ctx = proprio_scratch.reset((size_t) 4 * 1024 * 1024); ggml_tensor * x_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, proprio_dim, 1); ggml_set_name(x_in, "state"); ggml_tensor * h1 = ggml_add(ctx, ggml_mul_mat(ctx, pp_fc1_w, x_in), pp_fc1_b); @@ -1073,14 +1115,12 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { ggml_tensor * out = ggml_add(ctx, ggml_mul_mat(ctx, pp_fc2_w, h1_gel), pp_fc2_b); ggml_set_name(out, "proprio_embed"); - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); ggml_cgraph * gf = ggml_new_graph(ctx); ggml_build_forward_expand(gf, out); - if (!ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr failed (proprio)\n"); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } - ggml_backend_tensor_set(x_in, in.state, 0, ggml_nbytes(x_in)); - if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): proprio compute failed\n"); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } + if (!proprio_scratch.alloc(backend, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr failed (proprio)\n"); return {}; } + ggml_backend_tensor_set(x_in, state_host.data(), 0, ggml_nbytes(x_in)); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): proprio compute failed\n"); return {}; } ggml_backend_tensor_get(out, proprio_embed_host.data(), 0, (size_t) hidden_l * sizeof(float)); - ggml_gallocr_free(galloc); ggml_free(ctx); } _dump_bin("proprio_features", proprio_embed_host.data(), proprio_embed_host.size()); @@ -1201,9 +1241,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } else #endif { - std::vector meta_buf((size_t) 64 * 1024 * 1024); - ggml_init_params gp = { meta_buf.size(), meta_buf.data(), true }; - ggml_context * ctx = ggml_init(gp); + ggml_context * ctx = lm_scratch.reset((size_t) 64 * 1024 * 1024); ggml_tensor * x_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_l, seq); ggml_set_name(x_in, "inputs_embeds"); ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, seq); @@ -1221,10 +1259,9 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { ggml_tensor * action_hidden = ggml_get_rows(ctx, h_norm, action_ids); ggml_set_name(action_hidden, "action_hidden"); - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); ggml_build_forward_expand(gf, action_hidden); - if (!ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr failed (lm)\n"); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } + if (!lm_scratch.alloc(backend, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr failed (lm)\n"); return {}; } ggml_backend_tensor_set(x_in, inputs_embeds.data(), 0, ggml_nbytes(x_in)); std::vector pos_v(seq); @@ -1234,9 +1271,8 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { for (int64_t i = 0; i < n_action; ++i) aids[i] = (int32_t) (seq - 2 - n_action + i); ggml_backend_tensor_set(action_ids, aids.data(), 0, ggml_nbytes(action_ids)); - if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): lm prefill compute failed\n"); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): lm prefill compute failed\n"); return {}; } ggml_backend_tensor_get(action_hidden, last_hidden_at_actions.data(), 0, (size_t) n_action * hidden_l * sizeof(float)); - ggml_gallocr_free(galloc); ggml_free(ctx); } if (timing_phase) stats.ms_prefill = std::chrono::duration(clk::now() - t_p0).count(); @@ -1258,9 +1294,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } else #endif { - std::vector meta_buf((size_t) 8 * 1024 * 1024); - ggml_init_params gp = { meta_buf.size(), meta_buf.data(), true }; - ggml_context * ctx = ggml_init(gp); + ggml_context * ctx = head_scratch.reset((size_t) 8 * 1024 * 1024); ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, in_dim, chunk); ggml_set_name(x, "x"); @@ -1281,14 +1315,12 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { ggml_tensor * y = ggml_add(ctx, ggml_mul_mat(ctx, ah_fc2_w, ln2), ah_fc2_b); ggml_set_name(y, "y"); - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); ggml_cgraph * gf = ggml_new_graph(ctx); ggml_build_forward_expand(gf, y); - if (!ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr failed (action_head)\n"); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } + if (!head_scratch.alloc(backend, gf)) { std::fprintf(stderr, "vla(bitvla): gallocr failed (action_head)\n"); return {}; } ggml_backend_tensor_set(x, last_hidden_at_actions.data(), 0, ggml_nbytes(x)); - if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): action_head compute failed\n"); ggml_gallocr_free(galloc); ggml_free(ctx); return {}; } + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla): action_head compute failed\n"); return {}; } ggml_backend_tensor_get(y, normalized_actions.data(), 0, (size_t) chunk * action_dim * sizeof(float)); - ggml_gallocr_free(galloc); ggml_free(ctx); } if (timing_phase) stats.ms_denoise = std::chrono::duration(clk::now() - t_d0).count(); diff --git a/src/models/dit_common.h b/src/models/dit_common.h new file mode 100644 index 0000000..83f6cc7 --- /dev/null +++ b/src/models/dit_common.h @@ -0,0 +1,84 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// DiT head pieces shared verbatim by GR00T N1.5/N1.6/N1.7 and VLA-JEPA. dit_kv +// and build_dit_block take a per-arch struct and stay in their own files. + +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include + +namespace vla { + +// Row id of a stacked [out, in, n_embodiment] weight, applied to x. +inline ggml_tensor * cat_linear(ggml_context * C, ggml_tensor * W3d, ggml_tensor * b2d, int64_t id, ggml_tensor * x) { + const int64_t out = W3d->ne[0], in = W3d->ne[1]; + ggml_tensor * W_id = ggml_view_2d(C, W3d, out, in, W3d->nb[1], (size_t) id * W3d->nb[2]); + ggml_tensor * y = ggml_mul_mat(C, ggml_cont(C, ggml_transpose(C, W_id)), x); + return ggml_add(C, y, ggml_view_1d(C, b2d, out, (size_t) id * b2d->nb[1])); +} + +// Per-block AdaLN. The conditioning vector is (scale, shift) in that order; the +// final projection layer in each arch uses (shift, scale) instead. +inline ggml_tensor * adaln(ggml_context * C, ggml_tensor * x, ggml_tensor * temb, ggml_tensor * lw, ggml_tensor * lb, int64_t dim, float eps) { + ggml_tensor * cond = ggml_add(C, ggml_mul_mat(C, lw, ggml_silu(C, temb)), lb); + ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0), * sh = ggml_view_1d(C, cond, dim, (size_t) dim * sizeof(float)); + ggml_tensor * xn = ggml_norm(C, x, eps); + return ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); +} + +// cos first, then sin. Opposite order to action_sinusoid; both match the +// reference and are pinned by tests/test_dit_common.cpp. +inline void timesteps_proj(int64_t bucket, std::vector & out) { + const int64_t half = 128; const float lm = std::log(10000.0f); const float t = (float) bucket; + out.assign(256, 0.0f); + for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-lm * (float) i / (float) (half - 1)); out[i] = std::cos(emb); out[half + i] = std::sin(emb); } +} + +// Broadcast across the horizon. sin first, then cos. +inline void action_sinusoid(int64_t bucket, int64_t dim, int64_t T, std::vector & out) { + const int64_t half = dim / 2; const float step = std::log(10000.0f) / (float) half; const float t = (float) bucket; + out.assign((size_t) T * dim, 0.0f); + for (int64_t tk = 0; tk < T; ++tk) for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-(float) i * step); out[tk * dim + i] = std::sin(emb); out[tk * dim + half + i] = std::cos(emb); } +} + +// Log-spaced periods rather than frequencies, the openpi convention shared by +// pi0, pi0.5 and SmolVLA. +inline std::vector sinusoidal_time_emb(double t, int64_t dim, double min_p, double max_p) { + const int64_t half = dim / 2; + std::vector out(dim); + for (int64_t i = 0; i < half; ++i) { + const double frac = (half == 1) ? 0.0 : double(i) / double(half - 1); + const double period = min_p * std::pow(max_p / min_p, frac); + const double s = (2.0 * M_PI / period) * t; + out[i] = (float) std::sin(s); + out[half + i] = (float) std::cos(s); + } + return out; +} + +// Additive causal mask, -inf above the diagonal. +inline void build_causal_mask(int64_t seq, std::vector & out) { + out.assign((size_t) seq * seq, 0.0f); + const float NEG = -std::numeric_limits::infinity(); + for (int64_t q = 0; q < seq; ++q) + for (int64_t kv = q + 1; kv < seq; ++kv) out[q * seq + kv] = NEG; +} + +} // namespace vla diff --git a/src/models/dual_tower.h b/src/models/dual_tower.h new file mode 100644 index 0000000..20bc40a --- /dev/null +++ b/src/models/dual_tower.h @@ -0,0 +1,83 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// DINOv2 + SigLIP dual vision tower, shared by OpenVLA-OFT and VLA-Adapter. +// DINOv2 passes prefix=true (CLS + 4 register tokens, dropped after the blocks) +// and uses LayerScale; SigLIP passes prefix=false. + +#pragma once + +#include "ggml.h" +#include "model.h" + +#include +#include +#include + +namespace vla { + +struct ViTLayerW { ggml_tensor *n1w,*n1b,*n2w,*n2b,*ls1,*ls2,*Wqkv,*bqkv,*Wproj,*bproj,*Wfc1,*bfc1,*Wfc2,*bfc2; }; + +inline ggml_tensor * LN(ggml_context*C, ggml_tensor*x, ggml_tensor*w, ggml_tensor*b, float eps){ return ggml_add(C,ggml_mul(C,ggml_norm(C,x,eps),w),b); } + +inline ggml_tensor* vit_block(ggml_context*C, const ViTLayerW&w, ggml_tensor*x, int64_t N, int64_t hidden, int64_t heads, int64_t hd, float eps, bool ls){ + const float sc=1.0f/std::sqrt((float)hd); + ggml_tensor*xn=LN(C,x,w.n1w,w.n1b,eps); + ggml_tensor*qkv=ggml_add(C,ggml_mul_mat(C,w.Wqkv,xn),w.bqkv); + ggml_tensor*q=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],0*hidden*sizeof(float))); + ggml_tensor*k=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],1*hidden*sizeof(float))); + ggml_tensor*v=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],2*hidden*sizeof(float))); + ggml_tensor*Q=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,q,hd,heads,N),0,2,1,3)); + ggml_tensor*K=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,k,hd,heads,N),0,2,1,3)); + ggml_tensor*V=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,v,hd,heads,N),1,2,0,3)); + ggml_tensor*kq=ggml_mul_mat(C,K,Q); ggml_mul_mat_set_prec(kq,GGML_PREC_F32); + ggml_tensor*aw=ggml_soft_max_ext(C,kq,nullptr,sc,0.0f); + ggml_tensor*kqv=ggml_mul_mat(C,V,aw); + ggml_tensor*att=ggml_reshape_2d(C,ggml_cont(C,ggml_permute(C,kqv,0,2,1,3)),hidden,N); + ggml_tensor*ao=ggml_add(C,ggml_mul_mat(C,w.Wproj,att),w.bproj); + x=ggml_add(C,x,ls?ggml_mul(C,ao,w.ls1):ao); + ggml_tensor*xn2=LN(C,x,w.n2w,w.n2b,eps); + ggml_tensor*h=ggml_add(C,ggml_mul_mat(C,w.Wfc1,xn2),w.bfc1); h=ggml_gelu_erf(C,h); + h=ggml_add(C,ggml_mul_mat(C,w.Wfc2,h),w.bfc2); + return ggml_add(C,x,ls?ggml_mul(C,h,w.ls2):h); +} + +inline ggml_tensor* tower(ggml_context*C, ggml_tensor*pix, ggml_tensor*pw, ggml_tensor*pb, ggml_tensor*pos, + ggml_tensor*cls, ggml_tensor*reg, const std::vector&blk, + int64_t hidden, int64_t heads, int64_t hd, int64_t inter, int64_t patch, float eps, bool prefix){ + (void)inter; + ggml_tensor*conv=ggml_conv_2d(C,pw,pix,patch,patch,0,0,1,1); + // Patch count from the conv, not a constant: both callers run 224/14 today, + // and a different input size would otherwise reshape into the wrong grid. + const int64_t NP=conv->ne[0]*conv->ne[1], nprefix=prefix?5:0, N=NP+nprefix; + ggml_tensor*pt=ggml_cont(C,ggml_transpose(C,ggml_reshape_2d(C,conv,NP,hidden))); + pt=ggml_add(C,pt,pb); pt=ggml_add(C,pt,pos); + ggml_tensor*x=pt; + if(prefix){ ggml_tensor*tok=ggml_concat(C,ggml_reshape_2d(C,cls,hidden,1),reg,1); x=ggml_concat(C,tok,pt,1); } + for(size_t i=0;inb[1],nprefix*x->nb[1])); + return x; +} + +// HWC to CHW planar with per-channel mean/std: ImageNet for DINOv2, 0.5 for SigLIP. +inline void normalize_tower(const ImageView& v, int64_t S, const float mean[3], const float std_[3], std::vector& out){ + out.assign((size_t)3*S*S,0.0f); + for(int64_t h=0;h #include @@ -59,11 +55,22 @@ struct Evo1ModelArch : public ModelArchBase { ~Evo1ModelArch() override; std::string gguf_path; + // Opened once at load: reopening per predict re-parses the whole GGUF header. + gguf_reader io{"evo1"}; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + + struct MainKey { + int64_t seq=-1, nsteps=-1; + bool operator==(const MainKey & o) const { return seq==o.seq && nsteps==o.nsteps; } + }; + struct MainIO { + ggml_tensor *t_embeds=nullptr,*t_pos=nullptr,*t_lmmask=nullptr,*t_qmask=nullptr; + ggml_tensor *t_state=nullptr,*t_x=nullptr,*t_amask=nullptr,*x_action=nullptr; + }; + graph_cache main_graph; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_BF16; @@ -275,8 +282,8 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_EVO1_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g("evo1"); - if (!g.open(ckpt_path)) return nullptr; + if (!m->io.open(ckpt_path)) return nullptr; + gguf_reader & g = m->io; if (!g.has("evo1.architecture")) { std::fprintf(stderr, "vla(evo1): %s is not an evo1 GGUF (no evo1.architecture KV)\n", ckpt_path.c_str()); return nullptr; } @@ -288,20 +295,10 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, (long long) m->dit_layers, (long long) m->dit_heads, (long long) m->horizon, (long long) m->per_a, (long long) m->num_steps, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(evo1): backend = CUDA (device 0)\n"); } - else std::fprintf(stderr, "vla(evo1): ggml_backend_cuda_init failed; falling back to CPU\n"); -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(evo1): backend = Metal\n"); } - else std::fprintf(stderr, "vla(evo1): ggml_backend_metal_init failed; falling back to CPU\n"); -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(evo1): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(evo1): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(evo1)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t) 32 * 1024 * 1024, @@ -438,36 +435,30 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { } n_views = in.n_images; - ggml_init_params vp = { (size_t) 32 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 32 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(evo1): ggml_init(vision ctx) failed\n"); return {}; } ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, image_size, image_size, 3); ggml_set_input(t_px); ggml_tensor * t_ie = build_internvit_view(VC, *this, t_px); ggml_set_output(t_ie); ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vg, t_ie); - ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { + if (!vision_scratch.alloc(backend, vg)) { std::fprintf(stderr, "vla(evo1): vision ggml_gallocr_alloc_graph failed\n"); - if (vga) ggml_gallocr_free(vga); - ggml_free(VC); return {}; } img_emb_host.assign((size_t) n_views * num_image_token * lm_hidden, 0.0f); std::vector chw; const auto tv0 = std::chrono::steady_clock::now(); for (int64_t v = 0; v < n_views; ++v) { - if (!preprocess_image_chw(in.images[v], image_size, chw)) { ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!preprocess_image_chw(in.images[v], image_size, chw)) { return {}; } ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(evo1): vision graph compute failed (view %lld)\n", (long long) v); - ggml_gallocr_free(vga); ggml_free(VC); return {}; + return {}; } ggml_backend_tensor_get(t_ie, img_emb_host.data() + v * num_image_token * lm_hidden, 0, ggml_nbytes(t_ie)); } stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now() - tv0).count(); - ggml_gallocr_free(vga); - ggml_free(VC); img_emb_ptr = img_emb_host.data(); } else { std::fprintf(stderr, "vla(evo1): no images and no precomputed_img_emb in the request\n"); @@ -501,10 +492,8 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { input_ids.resize(max_text_length, pad_id); const int64_t SEQ = max_text_length; - gguf_reader g("evo1"); - if (!g.open(gguf_path)) return {}; std::vector inputs_embeds((size_t) SEQ * lm_hidden); - if (!g.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), lm_hidden)) return {}; + if (!io.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), lm_hidden)) return {}; { int64_t img_idx = 0; for (int64_t p = 0; p < SEQ; ++p) { @@ -536,7 +525,11 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { std::vector state_norm(per_a, 0.0f); for (int64_t i = 0; i < per_a; ++i) { const float lo = state_min[i], hi = state_max[i]; - float xn = 2.0f * (in.state[i] - lo) / (hi - lo + norm_eps_denom) - 1.0f; + // The converter zero-pads stats past real_state_dim, so lo == hi == 0 there + // and the affine below would map anything to -1. + if (hi <= lo) { state_norm[i] = 0.0f; continue; } + const float sv = in.state ? in.state[i] : 0.0f; + float xn = 2.0f * (sv - lo) / (hi - lo + norm_eps_denom) - 1.0f; if (xn < -1.0f) xn = -1.0f; if (xn > 1.0f) xn = 1.0f; state_norm[i] = xn; @@ -554,10 +547,10 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { } } - ggml_init_params cp = { (size_t) 96 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(evo1): ggml_init(ctx_compute) failed\n"); return {}; } - + // LM + DiT graph depends only on the padded length and step count. + const MainKey mkey{ SEQ, num_steps }; + const bool built = main_graph.ensure(backend, mkey, (size_t) 96 * 1024 * 1024, + [&](ggml_context * C, MainIO & gio) -> ggml_cgraph * { const int64_t E = embed_dim, hd_dit = E / dit_heads; const float scale_dit = 1.0f / std::sqrt((float) hd_dit); const int64_t Nctx = SEQ + 1; @@ -640,16 +633,21 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ggml_set_name(x_action, "x_final"); ggml_set_output(x_action); + gio.t_embeds=t_embeds; gio.t_pos=t_pos; gio.t_lmmask=t_lmmask; gio.t_qmask=t_qmask; + gio.t_state=t_state; gio.t_x=t_x; gio.t_amask=t_amask; gio.x_action=x_action; + ggml_cgraph * gf = ggml_new_graph_custom(C, 32768, false); ggml_build_forward_expand(gf, x_action); + return gf; + }); + if (!built) { std::fprintf(stderr, "vla(evo1): main graph build failed\n"); return {}; } + + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor * t_embeds = gio.t_embeds, * t_pos = gio.t_pos, * t_lmmask = gio.t_lmmask; + ggml_tensor * t_qmask = gio.t_qmask, * t_state = gio.t_state, * t_x = gio.t_x; + ggml_tensor * t_amask = gio.t_amask, * x_action = gio.x_action; - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { - std::fprintf(stderr, "vla(evo1): ggml_gallocr_alloc_graph failed\n"); - if (galloc) ggml_gallocr_free(galloc); - ggml_free(C); - return {}; - } ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { std::vector pp(SEQ); for (int64_t i = 0; i < SEQ; ++i) pp[i] = (int32_t) i; ggml_backend_tensor_set(t_pos, pp.data(), 0, ggml_nbytes(t_pos)); } { std::vector mk((size_t) SEQ * SEQ); const float NEG = -std::numeric_limits::infinity(); @@ -667,14 +665,12 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { const auto tc1 = std::chrono::steady_clock::now(); if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(evo1): ggml_backend_graph_compute failed (%d)\n", (int) st); - ggml_gallocr_free(galloc); ggml_free(C); return {}; + return {}; } stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); std::vector x_final((size_t) action_dim); ggml_backend_tensor_get(x_action, x_final.data(), 0, x_final.size() * sizeof(float)); - ggml_gallocr_free(galloc); - ggml_free(C); std::vector out((size_t) horizon * per_a); for (int64_t hstep = 0; hstep < horizon; ++hstep) diff --git a/src/models/gguf_reader.h b/src/models/gguf_reader.h index 3b3f739..ebe8944 100644 --- a/src/models/gguf_reader.h +++ b/src/models/gguf_reader.h @@ -58,11 +58,26 @@ struct gguf_reader { return true; } - bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } - uint32_t u32(const char * k) const { return gguf_get_val_u32(gctx, gguf_find_key(gctx, k)); } - float f32(const char * k) const { return gguf_get_val_f32(gctx, gguf_find_key(gctx, k)); } - double f64(const char * k) const { return gguf_get_val_f64(gctx, gguf_find_key(gctx, k)); } - std::string str(const char * k) const { const int64_t id = gguf_find_key(gctx, k); return id < 0 ? std::string() : std::string(gguf_get_val_str(gctx, id)); } + bool has(const char * k) const { return gguf_find_key(gctx, k) >= 0; } + + // gguf_get_val_* asserts on a type mismatch, killing the process on a bad + // file. Check the declared type first. + bool typed_key(const char * k, gguf_type want, int64_t * id_out) const { + const int64_t id = gguf_find_key(gctx, k); + if (id < 0) return false; + if (gguf_get_kv_type(gctx, id) != want) { + std::fprintf(stderr, "vla(%s): key %s has unexpected type %d\n", + arch, k, (int) gguf_get_kv_type(gctx, id)); + return false; + } + *id_out = id; + return true; + } + + uint32_t u32(const char * k) const { int64_t id; return typed_key(k, GGUF_TYPE_UINT32, &id) ? gguf_get_val_u32(gctx, id) : 0u; } + float f32(const char * k) const { int64_t id; return typed_key(k, GGUF_TYPE_FLOAT32, &id) ? gguf_get_val_f32(gctx, id) : 0.f; } + double f64(const char * k) const { int64_t id; return typed_key(k, GGUF_TYPE_FLOAT64, &id) ? gguf_get_val_f64(gctx, id) : 0.0; } + std::string str(const char * k) const { int64_t id; return typed_key(k, GGUF_TYPE_STRING, &id) ? std::string(gguf_get_val_str(gctx, id)) : std::string(); } const ggml_tensor * meta(const char * name) const { return ggml_get_tensor(meta_ctx, name); } // Resident type for a weight: keep a quantized source type (Q8_0, Q4_0, ...) @@ -72,11 +87,19 @@ struct gguf_reader { return (src && ggml_is_quantized(src->type)) ? src->type : prefer; } - bool read_raw(const char * name, void * buf) { + // cap must equal the declared tensor size: a tensor with the expected ne[0] but + // an extra dimension would write past the caller's vector. On failure buf may be + // partially written. + bool read_raw(const char * name, void * buf, size_t cap) { const int64_t id = gguf_find_tensor(gctx, name); if (id < 0) { std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch, name); return false; } const size_t off = data_off + gguf_get_tensor_offset(gctx, id); const size_t nb = gguf_get_tensor_size(gctx, id); + if (nb != cap) { + std::fprintf(stderr, "vla(%s): tensor %s is %zu bytes, caller expects %zu\n", + arch, name, nb, cap); + return false; + } if (fseeko(fp, (off_t) off, SEEK_SET) != 0) return false; return std::fread(buf, 1, nb, fp) == nb; } @@ -86,8 +109,8 @@ struct gguf_reader { if (!t) { std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch, name); return {}; } const int64_t n = ggml_nelements(t); std::vector out(n); - if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data())) return {}; } - else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data())) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } + if (t->type == GGML_TYPE_F32) { if (!read_raw(name, out.data(), out.size() * sizeof(float))) return {}; } + else if (t->type == GGML_TYPE_BF16) { std::vector tmp(n); if (!read_raw(name, tmp.data(), tmp.size() * sizeof(ggml_bf16_t))) return {}; ggml_bf16_to_fp32_row(tmp.data(), out.data(), n); } else { std::fprintf(stderr, "vla(%s): tensor %s unsupported type %d\n", arch, name, (int) t->type); return {}; } return out; } @@ -97,10 +120,17 @@ struct gguf_reader { // gemma_norm adds 1.0 per weight. std::vector read_convert(const char * name, ggml_type target, bool gemma_norm = false) { if (target != GGML_TYPE_F32 && target != GGML_TYPE_BF16) { + if (gemma_norm) { + // The +1 needs unpacked floats; skipping it would silently give wrong + // norm weights. + std::fprintf(stderr, "vla(%s): %s needs the Gemma norm +1 but the target type is packed\n", + arch, name); + return {}; + } const ggml_tensor * t = meta(name); if (!t) { std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch, name); return {}; } std::vector o(ggml_nbytes(t)); - if (!read_raw(name, o.data())) return {}; + if (!read_raw(name, o.data(), o.size())) return {}; return o; } std::vector f = read_f32(name); diff --git a/src/models/gr00tn1d5.cpp b/src/models/gr00tn1d5.cpp index 6c20589..45c946f 100644 --- a/src/models/gr00tn1d5.cpp +++ b/src/models/gr00tn1d5.cpp @@ -18,14 +18,12 @@ #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/vision_common.h" +#include "models/dit_common.h" #include #include @@ -55,11 +53,22 @@ struct Gr00tN1d5ModelArch : public ModelArchBase { ~Gr00tN1d5ModelArch() override; std::string gguf_path; + // Opened once at load: reopening per predict re-parses the whole GGUF header. + gguf_reader io{"gr00tn1d5"}; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + + struct MainKey { + int64_t seq=-1, nsteps=-1; + bool operator==(const MainKey & o) const { return seq==o.seq && nsteps==o.nsteps; } + }; + struct MainIO { + ggml_tensor *t_embeds=nullptr,*t_pos=nullptr,*t_lmmask=nullptr,*t_state=nullptr,*t_x0=nullptr,*actions=nullptr; + std::vector t_tau, t_tproj; + }; + graph_cache main_graph; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -94,20 +103,6 @@ struct Gr00tN1d5ModelArch : public ModelArchBase { namespace { -ggml_tensor * cat_linear(ggml_context * C, ggml_tensor * W3d, ggml_tensor * b2d, int64_t id, ggml_tensor * x) { - const int64_t out = W3d->ne[0], in = W3d->ne[1]; - ggml_tensor * W_id = ggml_view_2d(C, W3d, out, in, W3d->nb[1], (size_t) id * W3d->nb[2]); - ggml_tensor * y = ggml_mul_mat(C, ggml_cont(C, ggml_transpose(C, W_id)), x); - return ggml_add(C, y, ggml_view_1d(C, b2d, out, (size_t) id * b2d->nb[1])); -} - -ggml_tensor * adaln(ggml_context * C, ggml_tensor * x, ggml_tensor * temb, ggml_tensor * lw, ggml_tensor * lb, int64_t dim, float eps) { - ggml_tensor * cond = ggml_add(C, ggml_mul_mat(C, lw, ggml_silu(C, temb)), lb); - ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0), * sh = ggml_view_1d(C, cond, dim, (size_t) dim * sizeof(float)); - ggml_tensor * xn = ggml_norm(C, x, eps); - return ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); -} - ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, int64_t seq, int64_t heads, int64_t head_dim, int64_t hidden, float ln_eps) { const float scale = 1.0f / std::sqrt((float) head_dim); @@ -203,33 +198,6 @@ ggml_tensor * build_dit_block(ggml_context * C, const Gr00tN1d5ModelArch & m, co return ggml_add(C, h1, ff); } -bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(gr00tn1d5): image view is %dx%d, expected %lldx%lld\n", v.w, v.h, (long long) side, (long long) side); return false; - } - out.assign((size_t) 3 * side * side, 0.0f); - for (int64_t h = 0; h < side; ++h) - for (int64_t w = 0; w < side; ++w) - for (int64_t c = 0; c < 3; ++c) { - float px; - if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; - else px = ((const float *) v.data)[(h * side + w) * 3 + c]; - out[c * side * side + h * side + w] = px * 2.0f - 1.0f; - } - return true; -} - -void timesteps_proj(int64_t bucket, std::vector & out) { - const int64_t half = 128; const float lm = std::log(10000.0f); const float t = (float) bucket; - out.assign(256, 0.0f); - for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-lm * (float) i / (float) (half - 1)); out[i] = std::cos(emb); out[half + i] = std::sin(emb); } -} - -void action_sinusoid(int64_t bucket, int64_t dim, int64_t T, std::vector & out) { - const int64_t half = dim / 2; const float step = std::log(10000.0f) / (float) half; const float t = (float) bucket; - out.assign((size_t) T * dim, 0.0f); - for (int64_t tk = 0; tk < T; ++tk) for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-(float) i * step); out[tk * dim + i] = std::sin(emb); out[tk * dim + half + i] = std::cos(emb); } -} bool load_config(const gguf_reader & g, Gr00tN1d5ModelArch & m, Config & cfg) { auto U = [&](const char * k, int64_t & dst) { if (g.has(k)) dst = (int64_t) g.u32(k); }; @@ -268,6 +236,9 @@ bool load_config(const gguf_reader & g, Gr00tN1d5ModelArch & m, Config & cfg) { cfg.hidden = m.lm_hidden; cfg.n_q_heads = m.n_q; cfg.n_kv_heads = m.n_kv; cfg.head_dim = m.lm_head_dim; cfg.n_layers = m.lm_layers; cfg.num_steps = (int) m.num_steps; cfg.rms_eps = m.lm_rms_eps; cfg.rope_n_dims = (int) m.lm_head_dim; cfg.rope_mode = GGML_ROPE_TYPE_NEOX; cfg.rope_freq_base = m.lm_rope_base; + // Raw output: this arch expects the client to apply the dataset statistics + // (see the --stats-json flag in eval/client). + cfg.denormalized = false; cfg.norm_eps = 1e-8f; return true; } @@ -290,8 +261,8 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g("gr00tn1d5"); - if (!g.open(ckpt_path)) return nullptr; + if (!m->io.open(ckpt_path)) return nullptr; + gguf_reader & g = m->io; if (!g.has("gr00t_n1_5.architecture")) { std::fprintf(stderr, "vla(gr00tn1d5): %s is not a gr00t_n1_5 GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; std::printf("vla(gr00tn1d5): vit=%lldd×%lldL×%lldh n_img_tok=%lld lm=Qwen3 %lldd×%lldL (%lldq/%lldkv×%lld) " @@ -303,20 +274,10 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, (long long) m->action_horizon, (long long) m->action_dim, (long long) m->num_steps, (long long) m->embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(gr00tn1d5): backend = CUDA (device 0)\n"); } - else std::fprintf(stderr, "vla(gr00tn1d5): ggml_backend_cuda_init failed; falling back to CPU\n"); -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(gr00tn1d5): backend = Metal\n"); } - else std::fprintf(stderr, "vla(gr00tn1d5): ggml_backend_metal_init failed; falling back to CPU\n"); -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(gr00tn1d5): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(gr00tn1d5): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(gr00tn1d5)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; @@ -425,8 +386,7 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { } else if (in.images && in.n_images > 0) { n_views = in.n_images; img_emb_host.assign((size_t) n_views * K * H, 0.0f); - ggml_init_params vp = { (size_t) 64 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 64 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(gr00tn1d5): ggml_init(vision ctx) failed\n"); return {}; } const int64_t grid = image_size / patch_size; ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, image_size, image_size, 3); ggml_set_input(t_px); @@ -439,18 +399,16 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { ggml_set_output(vit_emb); ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vg, vit_emb); - ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { std::fprintf(stderr, "vla(gr00tn1d5): vision gallocr alloc failed\n"); if (vga) ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!vision_scratch.alloc(backend, vg)) { std::fprintf(stderr, "vla(gr00tn1d5): vision gallocr alloc failed\n"); return {}; } const auto tv0 = std::chrono::steady_clock::now(); std::vector chw; for (int64_t v = 0; v < n_views; ++v) { - if (!preprocess_image_chw(in.images[v], image_size, chw)) { ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!preprocess_image_chw("gr00tn1d5", in.images[v], image_size, chw)) { return {}; } ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); - if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(gr00tn1d5): vision compute failed\n"); ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(gr00tn1d5): vision compute failed\n"); return {}; } ggml_backend_tensor_get(vit_emb, img_emb_host.data() + v * K * H, 0, ggml_nbytes(vit_emb)); } stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now() - tv0).count(); - ggml_gallocr_free(vga); ggml_free(VC); img_emb_ptr = img_emb_host.data(); } else { std::fprintf(stderr, "vla(gr00tn1d5): no images and no precomputed_img_emb in the request\n"); return {}; @@ -475,10 +433,8 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { const int64_t SEQ = (int64_t) input_ids.size(); if (SEQ > max_seq_len) { std::fprintf(stderr, "vla(gr00tn1d5): prompt too long (%lld > %lld)\n", (long long) SEQ, (long long) max_seq_len); return {}; } - gguf_reader g("gr00tn1d5"); - if (!g.open(gguf_path)) return {}; std::vector inputs_embeds((size_t) SEQ * H); - if (!g.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), H)) return {}; + if (!io.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), H)) return {}; { int64_t k = 0; for (int64_t p = 0; p < SEQ; ++p) if (input_ids[p] == (int32_t) image_token_index) { if (k >= n_img) { std::fprintf(stderr, "vla(gr00tn1d5): more tokens than ViT embeds\n"); return {}; } @@ -491,10 +447,10 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { if (in.noise) std::memcpy(x_init.data(), in.noise, x_init.size() * sizeof(float)); else { std::mt19937 rng((uint32_t) std::chrono::steady_clock::now().time_since_epoch().count()); std::normal_distribution nd(0.f, 1.f); for (auto & v : x_init) v = nd(rng); } - ggml_init_params cp = { (size_t) 128 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(gr00tn1d5): ggml_init(ctx_compute) failed\n"); return {}; } - + // LM + VLSA + DiT graph depends only on the padded length and step count. + const MainKey mkey{ SEQ, num_steps }; + const bool built = main_graph.ensure(backend, mkey, (size_t) 128 * 1024 * 1024, + [&](ggml_context * C, MainIO & gio) -> ggml_cgraph * { ggml_tensor * t_embeds = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); ggml_set_input(t_embeds); ggml_tensor * t_pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ); ggml_set_input(t_pos); ggml_tensor * t_lmmask = ggml_new_tensor_2d(C, GGML_TYPE_F32, SEQ, SEQ); ggml_set_input(t_lmmask); @@ -552,11 +508,20 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { } ggml_set_name(actions, "action_pred"); ggml_set_output(actions); + gio.t_embeds=t_embeds; gio.t_pos=t_pos; gio.t_lmmask=t_lmmask; gio.t_state=t_state; + gio.t_x0=t_x0; gio.t_tau=t_tau; gio.t_tproj=t_tproj; gio.actions=actions; + ggml_cgraph * gf = ggml_new_graph_custom(C, 32768, false); ggml_build_forward_expand(gf, actions); + return gf; + }); + if (!built) { std::fprintf(stderr, "vla(gr00tn1d5): main graph build failed\n"); return {}; } - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(gr00tn1d5): gallocr alloc failed\n"); if (galloc) ggml_gallocr_free(galloc); ggml_free(C); return {}; } + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor * t_embeds = gio.t_embeds, * t_pos = gio.t_pos, * t_lmmask = gio.t_lmmask; + ggml_tensor * t_state = gio.t_state, * t_x0 = gio.t_x0, * actions = gio.actions; + std::vector & t_tau = gio.t_tau; std::vector & t_tproj = gio.t_tproj; ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { std::vector pp(SEQ); for (int64_t i = 0; i < SEQ; ++i) pp[i] = (int32_t) i; ggml_backend_tensor_set(t_pos, pp.data(), 0, ggml_nbytes(t_pos)); } @@ -575,12 +540,11 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { const auto tc0 = std::chrono::steady_clock::now(); const ggml_status st = ggml_backend_graph_compute(backend, gf); const auto tc1 = std::chrono::steady_clock::now(); - if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(gr00tn1d5): graph compute failed (%d)\n", (int) st); ggml_gallocr_free(galloc); ggml_free(C); return {}; } + if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(gr00tn1d5): graph compute failed (%d)\n", (int) st); return {}; } stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); std::vector out((size_t) AH * AD); ggml_backend_tensor_get(actions, out.data(), 0, out.size() * sizeof(float)); - ggml_gallocr_free(galloc); ggml_free(C); stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); return out; } diff --git a/src/models/gr00tn1d6.cpp b/src/models/gr00tn1d6.cpp index 7017d9a..aff2f29 100644 --- a/src/models/gr00tn1d6.cpp +++ b/src/models/gr00tn1d6.cpp @@ -18,14 +18,11 @@ #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/dit_common.h" #include #include @@ -54,11 +51,26 @@ struct Gr00tN1d6ModelArch : public ModelArchBase { ~Gr00tN1d6ModelArch() override; std::string gguf_path; + // Opened once at load: reopening per predict re-parses the whole GGUF header. + gguf_reader io{"gr00tn1d6"}; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + scratch_ctx merge_scratch; + + struct MainKey { + int64_t seq=-1, n_img=-1, seq_txt=-1, nsteps=-1; + bool operator==(const MainKey & o) const { + return seq==o.seq && n_img==o.n_img && seq_txt==o.seq_txt && nsteps==o.nsteps; + } + }; + struct MainIO { + ggml_tensor *t_embeds=nullptr,*t_pos=nullptr,*t_lmmask=nullptr,*t_state=nullptr,*t_x0=nullptr; + ggml_tensor *t_img_idx=nullptr,*t_txt_idx=nullptr,*actions=nullptr; + std::vector t_tau, t_tproj; + }; + graph_cache main_graph; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -92,20 +104,6 @@ struct Gr00tN1d6ModelArch : public ModelArchBase { namespace { -ggml_tensor * cat_linear(ggml_context * C, ggml_tensor * W3d, ggml_tensor * b2d, int64_t id, ggml_tensor * x) { - const int64_t out = W3d->ne[0], in = W3d->ne[1]; - ggml_tensor * W_id = ggml_view_2d(C, W3d, out, in, W3d->nb[1], (size_t) id * W3d->nb[2]); - ggml_tensor * y = ggml_mul_mat(C, ggml_cont(C, ggml_transpose(C, W_id)), x); - return ggml_add(C, y, ggml_view_1d(C, b2d, out, (size_t) id * b2d->nb[1])); -} - -ggml_tensor * adaln(ggml_context * C, ggml_tensor * x, ggml_tensor * temb, ggml_tensor * lw, ggml_tensor * lb, int64_t dim, float eps) { - ggml_tensor * cond = ggml_add(C, ggml_mul_mat(C, lw, ggml_silu(C, temb)), lb); - ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0), * sh = ggml_view_1d(C, cond, dim, (size_t) dim * sizeof(float)); - ggml_tensor * xn = ggml_norm(C, x, eps); - return ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); -} - ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, int64_t seq, int64_t heads, int64_t head_dim, int64_t hidden, float ln_eps) { const int64_t nv = x->ne[2]; @@ -219,18 +217,6 @@ void pixel_shuffle_back(const float * src, int64_t grid, int64_t hidden, int64_t } } -void timesteps_proj(int64_t bucket, std::vector & out) { - const int64_t half = 128; const float lm = std::log(10000.0f); const float t = (float) bucket; - out.assign(256, 0.0f); - for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-lm * (float) i / (float) (half - 1)); out[i] = std::cos(emb); out[half + i] = std::sin(emb); } -} - -void action_sinusoid(int64_t bucket, int64_t dim, int64_t T, std::vector & out) { - const int64_t half = dim / 2; const float step = std::log(10000.0f) / (float) half; const float t = (float) bucket; - out.assign((size_t) T * dim, 0.0f); - for (int64_t tk = 0; tk < T; ++tk) for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-(float) i * step); out[tk * dim + i] = std::sin(emb); out[tk * dim + half + i] = std::cos(emb); } -} - bool load_config(const gguf_reader & g, Gr00tN1d6ModelArch & m, Config & cfg) { auto U = [&](const char * k, int64_t & dst) { if (g.has(k)) dst = (int64_t) g.u32(k); }; auto F = [&](const char * k, float & dst) { if (g.has(k)) dst = g.f32(k); }; @@ -266,6 +252,23 @@ bool load_config(const gguf_reader & g, Gr00tN1d6ModelArch & m, Config & cfg) { } if (m.embodiment_id < 0 || m.embodiment_id >= m.max_embodiments) { std::fprintf(stderr, "vla(gr00tn1d6): embodiment id %lld out of range [0,%lld)\n", (long long) m.embodiment_id, (long long) m.max_embodiments); return false; } + // pixel_shuffle_back writes (grid/shuffle)^2 tokens into a buffer sized from + // n_img_tokens, so the KV has to agree with the grid it is derived from. + if (m.patch_size <= 0 || m.vit_pixel_shuffle <= 0 || m.image_size % m.patch_size != 0 || + (m.image_size / m.patch_size) % m.vit_pixel_shuffle != 0) { + std::fprintf(stderr, "vla(gr00tn1d6): image %lld / patch %lld / shuffle %lld do not divide evenly\n", + (long long) m.image_size, (long long) m.patch_size, (long long) m.vit_pixel_shuffle); + return false; + } + { + const int64_t g2 = (m.image_size / m.patch_size) / m.vit_pixel_shuffle; + if (m.n_img_tokens != g2 * g2) { + std::fprintf(stderr, "vla(gr00tn1d6): n_img_tokens %lld does not match the %lldx%lld shuffled grid\n", + (long long) m.n_img_tokens, (long long) g2, (long long) g2); + return false; + } + } + cfg = Config{}; cfg.n_img = m.n_img_tokens; cfg.n_lang = m.max_seq_len; cfg.n_state = 1; cfg.n_suffix = m.action_horizon; cfg.max_state_dim = m.max_state_dim; cfg.max_action_dim = m.action_dim; @@ -273,6 +276,9 @@ bool load_config(const gguf_reader & g, Gr00tN1d6ModelArch & m, Config & cfg) { cfg.hidden = m.lm_hidden; cfg.n_q_heads = m.n_q; cfg.n_kv_heads = m.n_kv; cfg.head_dim = m.lm_head_dim; cfg.n_layers = m.lm_layers; cfg.num_steps = (int) m.num_steps; cfg.rms_eps = m.lm_rms_eps; cfg.rope_n_dims = (int) m.lm_head_dim; cfg.rope_mode = GGML_ROPE_TYPE_NEOX; cfg.rope_freq_base = m.lm_rope_base; + // Raw output: this arch expects the client to apply the dataset statistics + // (see the --stats-json flag in eval/client). + cfg.denormalized = false; cfg.norm_eps = 1e-8f; return true; } @@ -295,8 +301,8 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, m->gguf_path = ckpt_path; m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; - gguf_reader g("gr00tn1d6"); - if (!g.open(ckpt_path)) return nullptr; + if (!m->io.open(ckpt_path)) return nullptr; + gguf_reader & g = m->io; if (!g.has("gr00t_n1_6.architecture")) { std::fprintf(stderr, "vla(gr00tn1d6): %s is not a gr00t_n1_6 GGUF\n", ckpt_path.c_str()); return nullptr; } if (!load_config(g, *m, m->cfg)) return nullptr; std::printf("vla(gr00tn1d6): vit=%lldd×%lldL×%lldh (Linear patch embed) pixel_shuffle÷%lld ⇒ n_img_tok=%lld mlp1=LN(%lld)→Linear→GELU→Linear " @@ -308,20 +314,10 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, (long long) m->action_horizon, (long long) m->action_dim, (long long) m->max_state_dim, (long long) m->num_steps, (long long) m->embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(gr00tn1d6): backend = CUDA (device 0)\n"); } - else std::fprintf(stderr, "vla(gr00tn1d6): ggml_backend_cuda_init failed; falling back to CPU\n"); -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(gr00tn1d6): backend = Metal\n"); } - else std::fprintf(stderr, "vla(gr00tn1d6): ggml_backend_metal_init failed; falling back to CPU\n"); -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(gr00tn1d6): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(gr00tn1d6)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; @@ -429,8 +425,7 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { n_views = in.n_images; img_emb_host.assign((size_t) n_views * K * H, 0.0f); - ggml_init_params vp = { (size_t) 64 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 64 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_init(vision ctx A) failed\n"); return {}; } ggml_tensor * t_patches = ggml_new_tensor_3d(VC, GGML_TYPE_F32, patch_dim, n_patches, n_views); ggml_set_input(t_patches); ggml_tensor * h = ggml_add(VC, ggml_add(VC, ggml_mul_mat(VC, vit_patch_w, t_patches), vit_patch_b), vit_pos); @@ -439,12 +434,10 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { ggml_set_output(post_ln); ggml_cgraph * vgA = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vgA, post_ln); - ggml_gallocr_t vgaA = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vgaA || !ggml_gallocr_alloc_graph(vgaA, vgA)) { std::fprintf(stderr, "vla(gr00tn1d6): vision gallocr A alloc failed\n"); if (vgaA) ggml_gallocr_free(vgaA); ggml_free(VC); return {}; } + if (!vision_scratch.alloc(backend, vgA)) { std::fprintf(stderr, "vla(gr00tn1d6): vision gallocr A alloc failed\n"); return {}; } - ggml_init_params mp = { (size_t) 16 * 1024 * 1024, nullptr, true }; - ggml_context * MC = ggml_init(mp); - if (!MC) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_init(vision ctx B) failed\n"); ggml_gallocr_free(vgaA); ggml_free(VC); return {}; } + ggml_context * MC = merge_scratch.reset((size_t) 16 * 1024 * 1024); + if (!MC) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_init(vision ctx B) failed\n"); return {}; } ggml_tensor * t_shuf = ggml_new_tensor_3d(MC, GGML_TYPE_F32, c4, K, n_views); ggml_set_input(t_shuf); ggml_tensor * mln = ggml_add(MC, ggml_mul(MC, ggml_norm(MC, t_shuf, connector_ln_eps), mm_ln_w), mm_ln_b); ggml_tensor * mz1 = ggml_add(MC, ggml_mul_mat(MC, mm_fc1_w, mln), mm_fc1_b); @@ -452,8 +445,7 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { ggml_set_output(vit_embeds); ggml_cgraph * vgB = ggml_new_graph(MC); ggml_build_forward_expand(vgB, vit_embeds); - ggml_gallocr_t vgaB = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vgaB || !ggml_gallocr_alloc_graph(vgaB, vgB)) { std::fprintf(stderr, "vla(gr00tn1d6): vision gallocr B alloc failed\n"); if (vgaB) ggml_gallocr_free(vgaB); ggml_gallocr_free(vgaA); ggml_free(MC); ggml_free(VC); return {}; } + if (!merge_scratch.alloc(backend, vgB)) { std::fprintf(stderr, "vla(gr00tn1d6): vision gallocr B alloc failed\n"); return {}; } const auto tv0 = std::chrono::steady_clock::now(); std::vector patches, @@ -478,7 +470,6 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { } if (vok) ggml_backend_tensor_get(vit_embeds, img_emb_host.data(), 0, ggml_nbytes(vit_embeds)); stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now() - tv0).count(); - ggml_gallocr_free(vgaB); ggml_gallocr_free(vgaA); ggml_free(MC); ggml_free(VC); if (!vok) return {}; img_emb_ptr = img_emb_host.data(); } else { @@ -503,10 +494,8 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { const int64_t SEQ = (int64_t) input_ids.size(); if (SEQ > max_seq_len) { std::fprintf(stderr, "vla(gr00tn1d6): prompt too long (%lld > %lld)\n", (long long) SEQ, (long long) max_seq_len); return {}; } - gguf_reader g("gr00tn1d6"); - if (!g.open(gguf_path)) return {}; std::vector inputs_embeds((size_t) SEQ * H); - if (!g.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), H)) return {}; + if (!io.fetch_rows_f32("token_embd.weight", input_ids, inputs_embeds.data(), H)) return {}; { int64_t k = 0; for (int64_t p = 0; p < SEQ; ++p) if (input_ids[p] == (int32_t) image_token_index) { if (k >= n_img) { std::fprintf(stderr, "vla(gr00tn1d6): more tokens than ViT embeds\n"); return {}; } @@ -530,10 +519,10 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { if (in.noise) std::memcpy(x_init.data(), in.noise, x_init.size() * sizeof(float)); else { std::mt19937 rng((uint32_t) std::chrono::steady_clock::now().time_since_epoch().count()); std::normal_distribution nd(0.f, 1.f); for (auto & v : x_init) v = nd(rng); } - ggml_init_params cp = { (size_t) 256 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_init(ctx_compute) failed\n"); return {}; } - + // LM + DiT graph depends only on the sequence split and step count. + const MainKey mkey{ SEQ, n_img, SEQ_TXT, num_steps }; + const bool built = main_graph.ensure(backend, mkey, (size_t) 256 * 1024 * 1024, + [&](ggml_context * C, MainIO & gio) -> ggml_cgraph * { ggml_tensor * t_embeds = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); ggml_set_input(t_embeds); ggml_tensor * t_pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ); ggml_set_input(t_pos); ggml_tensor * t_lmmask = ggml_new_tensor_2d(C, GGML_TYPE_F32, SEQ, SEQ); ggml_set_input(t_lmmask); @@ -601,11 +590,21 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { } ggml_set_name(actions, "action_pred"); ggml_set_output(actions); + gio.t_embeds=t_embeds; gio.t_pos=t_pos; gio.t_lmmask=t_lmmask; gio.t_state=t_state; gio.t_x0=t_x0; + gio.t_img_idx=t_img_idx; gio.t_txt_idx=t_txt_idx; gio.t_tau=t_tau; gio.t_tproj=t_tproj; gio.actions=actions; + ggml_cgraph * gf = ggml_new_graph_custom(C, 65536, false); ggml_build_forward_expand(gf, actions); + return gf; + }); + if (!built) { std::fprintf(stderr, "vla(gr00tn1d6): main graph build failed\n"); return {}; } - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(gr00tn1d6): gallocr alloc failed\n"); if (galloc) ggml_gallocr_free(galloc); ggml_free(C); return {}; } + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor * t_embeds = gio.t_embeds, * t_pos = gio.t_pos, * t_lmmask = gio.t_lmmask; + ggml_tensor * t_state = gio.t_state, * t_x0 = gio.t_x0; + ggml_tensor * t_img_idx = gio.t_img_idx, * t_txt_idx = gio.t_txt_idx, * actions = gio.actions; + std::vector & t_tau = gio.t_tau; std::vector & t_tproj = gio.t_tproj; ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { std::vector pp(SEQ); for (int64_t i = 0; i < SEQ; ++i) pp[i] = (int32_t) i; ggml_backend_tensor_set(t_pos, pp.data(), 0, ggml_nbytes(t_pos)); } @@ -626,12 +625,11 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { const auto tc0 = std::chrono::steady_clock::now(); const ggml_status st = ggml_backend_graph_compute(backend, gf); const auto tc1 = std::chrono::steady_clock::now(); - if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(gr00tn1d6): graph compute failed (%d)\n", (int) st); ggml_gallocr_free(galloc); ggml_free(C); return {}; } + if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(gr00tn1d6): graph compute failed (%d)\n", (int) st); return {}; } stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); std::vector out((size_t) AH * AD); ggml_backend_tensor_get(actions, out.data(), 0, out.size() * sizeof(float)); - ggml_gallocr_free(galloc); ggml_free(C); stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); return out; } diff --git a/src/models/gr00tn1d7.cpp b/src/models/gr00tn1d7.cpp index ca107db..d9742a7 100644 --- a/src/models/gr00tn1d7.cpp +++ b/src/models/gr00tn1d7.cpp @@ -18,14 +18,12 @@ #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/dit_common.h" +#include "models/qwen3vl_vit.h" #include #include @@ -44,11 +42,6 @@ namespace vla { namespace { -constexpr float CLIP_MEAN[3] = {0.5f, 0.5f, 0.5f}; -constexpr float CLIP_STD [3] = {0.5f, 0.5f, 0.5f}; - -struct VitLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wqkv,*bqkv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; -struct MergerW { ggml_tensor *nw,*nb,*fc1w,*fc1b,*fc2w,*fc2b; }; struct VlsaLayerW { ggml_tensor *n1w,*n1b,*n3w,*n3b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wff0,*bff0,*Wff2,*bff2; }; struct Qwen3LayerW { ggml_tensor *attn_norm,*Wq,*Wk,*Wv,*Wo,*q_norm,*k_norm,*ffn_norm,*Wgate,*Wup,*Wdown; }; struct DitLayerW { ggml_tensor *adaln_w,*adaln_b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wff0,*bff0,*Wff2,*bff2; @@ -62,10 +55,9 @@ struct Gr00tN1d7ModelArch : public ModelArchBase { std::string gguf_path; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -107,104 +99,32 @@ struct Gr00tN1d7ModelArch : public ModelArchBase { gguf_reader io; bool build_caches(); - struct MainGraph { - ggml_context * C = nullptr; ggml_gallocr_t galloc = nullptr; ggml_cgraph * gf = nullptr; + struct MainKey { + int64_t seq=-1, n_img=-1, seq_txt=-1, nsteps=-1; bool deepstack=false; + bool operator==(const MainKey & o) const { + return seq==o.seq && n_img==o.n_img && seq_txt==o.seq_txt && + nsteps==o.nsteps && deepstack==o.deepstack; + } + }; + struct MainIO { ggml_tensor *t_embeds=nullptr,*t_pos=nullptr,*t_lmmask=nullptr,*t_state=nullptr,*t_x0=nullptr; ggml_tensor *t_ds[3]={nullptr,nullptr,nullptr}; ggml_tensor *t_img_idx=nullptr,*t_txt_idx=nullptr,*actions=nullptr; std::vector t_tau, t_tproj; - int64_t seq=-1,n_img=-1,seq_txt=-1,nsteps=-1; bool deepstack=false; bool valid=false; - void release() { if (galloc) ggml_gallocr_free(galloc); if (C) ggml_free(C); - galloc=nullptr; C=nullptr; gf=nullptr; valid=false; t_tau.clear(); t_tproj.clear(); } - } mg; + }; + graph_cache mg; std::vector predict(const Inputs& in) override; }; namespace { -ggml_tensor * cat_linear(ggml_context * C, ggml_tensor * W3d, ggml_tensor * b2d, int64_t id, ggml_tensor * x) { - const int64_t out = W3d->ne[0], in = W3d->ne[1]; - ggml_tensor * W_id = ggml_view_2d(C, W3d, out, in, W3d->nb[1], (size_t) id * W3d->nb[2]); - ggml_tensor * y = ggml_mul_mat(C, ggml_cont(C, ggml_transpose(C, W_id)), x); - return ggml_add(C, y, ggml_view_1d(C, b2d, out, (size_t) id * b2d->nb[1])); -} - -ggml_tensor * adaln(ggml_context * C, ggml_tensor * x, ggml_tensor * temb, ggml_tensor * lw, ggml_tensor * lb, int64_t dim, float eps) { - ggml_tensor * cond = ggml_add(C, ggml_mul_mat(C, lw, ggml_silu(C, temb)), lb); - ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0), * sh = ggml_view_1d(C, cond, dim, (size_t) dim * sizeof(float)); - ggml_tensor * xn = ggml_norm(C, x, eps); - return ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); -} - -ggml_tensor * rope2d(ggml_context * C, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_t) { - const int64_t hd = x->ne[0], S = x->ne[1], Hh = x->ne[2]; const int64_t half = hd / 2; - ggml_tensor * x1 = ggml_cont(C, ggml_view_3d(C, x, half, S, Hh, x->nb[1], x->nb[2], 0)); - ggml_tensor * x2 = ggml_cont(C, ggml_view_3d(C, x, half, S, Hh, x->nb[1], x->nb[2], (size_t) half * x->nb[0])); - ggml_tensor * rot = ggml_concat(C, ggml_neg(C, x2), x1, 0); - return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, rot, sin_t)); -} - -inline bool fa_enabled() { static const bool e = (std::getenv("VLA_GR00T_FA") != nullptr); return e; } - ggml_tensor * head_view(ggml_context * C, ggml_tensor * proj, int64_t hd, int64_t heads, int64_t T, int64_t E, int nblk, int blk) { const size_t es = ggml_element_size(proj); return ggml_view_3d(C, proj, hd, heads, T, (size_t) hd * es, (size_t) nblk * E * es, (size_t) blk * E * es); } -ggml_tensor * flash_attn(ggml_context * C, ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, - ggml_tensor * mask, float scale, int64_t hidden) { - (void) hidden; - ggml_tensor * kf = (k->type == GGML_TYPE_F16) ? k : ggml_cast(C, k, GGML_TYPE_F16); - ggml_tensor * vf = (v->type == GGML_TYPE_F16) ? v : ggml_cast(C, v, GGML_TYPE_F16); - ggml_tensor * o = ggml_flash_attn_ext(C, q, kf, vf, mask, scale, 0.0f, 0.0f); - ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32); - return ggml_reshape_2d(C, o, o->ne[0] * o->ne[1], o->ne[2] * o->ne[3]); -} - -ggml_tensor * build_vit_layer(ggml_context * C, const VitLayerW & w, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_t, - int64_t seq, int64_t heads, int64_t hd, int64_t hidden, float ln_eps) { - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); - ggml_tensor * qkv = ggml_add(C, ggml_mul_mat(C, w.Wqkv, n1), w.bqkv); - ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], 0)); - ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) hidden * qkv->nb[0])); - ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) 2 * hidden * qkv->nb[0])); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, seq), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, seq), 0, 2, 1, 3)); - Q = rope2d(C, Q, cos_t, sin_t); K = rope2d(C, K, cos_t, sin_t); - ggml_tensor * att; - if (fa_enabled()) { - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 0, 2, 1, 3)); - att = flash_attn(C, Q, K, V, nullptr, scale, hidden); - } else { - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 1, 2, 0, 3)); - ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale, 0.0f); - att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq); - } - ggml_tensor * h1 = ggml_add(C, x, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); - ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.ln2w), w.ln2b); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wfc2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wfc1, n2), w.bfc1))), w.bfc2); - return ggml_add(C, h1, ff); -} - -ggml_tensor * build_merger(ggml_context * C, const MergerW & w, ggml_tensor * x, int64_t hidden, int64_t merge2, float ln_eps, bool pre_merge) { - ggml_tensor * m; - if (pre_merge) { - ggml_tensor * xn = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.nw), w.nb); - const int64_t n_patches = x->ne[1], c_merged = hidden * merge2 * merge2, n_merged = n_patches / (merge2 * merge2); - m = ggml_reshape_2d(C, ggml_cont(C, xn), c_merged, n_merged); - } else { - const int64_t n_patches = x->ne[1], c_merged = hidden * merge2 * merge2, n_merged = n_patches / (merge2 * merge2); - ggml_tensor * mr = ggml_reshape_2d(C, ggml_cont(C, x), c_merged, n_merged); - m = ggml_add(C, ggml_mul(C, ggml_norm(C, mr, ln_eps), w.nw), w.nb); - } - ggml_tensor * z1 = ggml_add(C, ggml_mul_mat(C, w.fc1w, m), w.fc1b); - return ggml_add(C, ggml_mul_mat(C, w.fc2w, ggml_gelu(C, z1)), w.fc2b); -} - ggml_tensor * build_vlsa_layer(ggml_context * C, const VlsaLayerW & w, ggml_tensor * x, int64_t seq, int64_t heads, int64_t hd, int64_t hidden, float ln_eps) { const float scale = 1.0f / std::sqrt((float) hd); @@ -217,7 +137,7 @@ ggml_tensor * build_vlsa_layer(ggml_context * C, const VlsaLayerW & w, ggml_tens ggml_tensor * att; if (fa_enabled()) { ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 0, 2, 1, 3)); - att = flash_attn(C, Q, K, V, nullptr, scale, hidden); + att = flash_attn(C, Q, K, V, nullptr, scale); } else { ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 1, 2, 0, 3)); ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); @@ -251,7 +171,7 @@ ggml_tensor * build_qwen3_layer(ggml_context * C, const Gr00tN1d7ModelArch & m, ggml_tensor * att; if (fa_enabled()) { ggml_tensor * V = ggml_cont(C, ggml_permute(C, vh, 0, 2, 1, 3)); - att = flash_attn(C, Q, K, V, ggml_cast(C, mask, GGML_TYPE_F16), scale, hq); + att = flash_attn(C, Q, K, V, ggml_cast(C, mask, GGML_TYPE_F16), scale); } else { ggml_tensor * V = ggml_cont(C, ggml_permute(C, vh, 1, 2, 0, 3)); ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); @@ -307,80 +227,6 @@ ggml_tensor * build_dit_block(ggml_context * C, const Gr00tN1d7ModelArch & m, co return ggml_add(C, h1, ff); } -void merge_block_coords(int64_t gh, int64_t gw, int64_t m, std::vector & row, std::vector & col) { - const int64_t S = gh * gw; row.assign(S, 0); col.assign(S, 0); - for (int64_t s = 0; s < S; ++s) { - int64_t t = s; const int64_t wj = t % m; t /= m; const int64_t wi = t % m; t /= m; - const int64_t bc = t % (gw / m); t /= (gw / m); const int64_t br = t; - row[s] = br * m + wi; col[s] = bc * m + wj; - } -} - -void vit_rope_tables(const std::vector & row, const std::vector & col, int64_t hd, double theta, - std::vector & cos_t, std::vector & sin_t) { - const int64_t S = (int64_t) row.size(), nf = hd / 4; - std::vector invf(nf); - for (int64_t i = 0; i < nf; ++i) invf[i] = 1.0 / std::pow(theta, (double)(2 * i) / (double)(hd / 2)); - cos_t.assign((size_t) S * hd, 0.0f); sin_t.assign((size_t) S * hd, 0.0f); - for (int64_t s = 0; s < S; ++s) { - std::vector emb(hd); - for (int64_t i = 0; i < nf; ++i) { emb[i] = (double) row[s] * invf[i]; emb[nf + i] = (double) col[s] * invf[i]; } - for (int64_t i = 0; i < hd / 2; ++i) emb[hd / 2 + i] = emb[i]; - for (int64_t i = 0; i < hd; ++i) { cos_t[s * hd + i] = (float) std::cos(emb[i]); sin_t[s * hd + i] = (float) std::sin(emb[i]); } - } -} - -void interp_pos_embed(const std::vector & table, int64_t num_side, int64_t hidden, - const std::vector & row, const std::vector & col, int64_t gh, int64_t gw, - std::vector & out) { - const int64_t S = (int64_t) row.size(); - out.assign((size_t) S * hidden, 0.0f); - auto src_coord = [&](int64_t k, int64_t g) -> double { return (g <= 1) ? 0.0 : (double) k * (double)(num_side - 1) / (double)(g - 1); }; - for (int64_t s = 0; s < S; ++s) { - const double hy = src_coord(row[s], gh), wx = src_coord(col[s], gw); - const int64_t h0 = (int64_t) std::floor(hy), w0 = (int64_t) std::floor(wx); - const int64_t h1 = std::min(h0 + 1, num_side - 1), w1 = std::min(w0 + 1, num_side - 1); - const double dh = hy - h0, dw = wx - w0; - const double c00 = (1 - dh) * (1 - dw), c01 = (1 - dh) * dw, c10 = dh * (1 - dw), c11 = dh * dw; - const float * T00 = &table[(h0 * num_side + w0) * hidden]; const float * T01 = &table[(h0 * num_side + w1) * hidden]; - const float * T10 = &table[(h1 * num_side + w0) * hidden]; const float * T11 = &table[(h1 * num_side + w1) * hidden]; - for (int64_t c = 0; c < hidden; ++c) out[s * hidden + c] = (float)(c00 * T00[c] + c01 * T01[c] + c10 * T10[c] + c11 * T11[c]); - } -} - -bool preprocess_image_patches(const ImageView & v, int64_t side, int64_t ps, int64_t tps, - const std::vector & row, const std::vector & col, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(gr00tn1d7): image view is %dx%d, expected %lldx%lld\n", v.w, v.h, (long long) side, (long long) side); return false; - } - const int64_t S = (int64_t) row.size(), pf = 3 * tps * ps * ps; - out.assign((size_t) pf * S, 0.0f); - auto px = [&](int64_t r, int64_t c, int64_t ch) -> float { - if (v.format == PixelFormat::U8) return ((const uint8_t *) v.data)[(r * side + c) * 3 + ch] / 255.0f; - return ((const float *) v.data)[(r * side + c) * 3 + ch]; - }; - for (int64_t s = 0; s < S; ++s) - for (int64_t ch = 0; ch < 3; ++ch) - for (int64_t ph = 0; ph < ps; ++ph) - for (int64_t pw = 0; pw < ps; ++pw) { - const float val = (px(row[s] * ps + ph, col[s] * ps + pw, ch) - CLIP_MEAN[ch]) / CLIP_STD[ch]; - for (int64_t t = 0; t < tps; ++t) out[s * pf + ch * tps * ps * ps + t * ps * ps + ph * ps + pw] = val; - } - return true; -} - -void timesteps_proj(int64_t bucket, std::vector & out) { - const int64_t half = 128; const float lm = std::log(10000.0f); const float t = (float) bucket; - out.assign(256, 0.0f); - for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-lm * (float) i / (float) (half - 1)); out[i] = std::cos(emb); out[half + i] = std::sin(emb); } -} - -void action_sinusoid(int64_t bucket, int64_t dim, int64_t T, std::vector & out) { - const int64_t half = dim / 2; const float step = std::log(10000.0f) / (float) half; const float t = (float) bucket; - out.assign((size_t) T * dim, 0.0f); - for (int64_t tk = 0; tk < T; ++tk) for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-(float) i * step); out[tk * dim + i] = std::sin(emb); out[tk * dim + half + i] = std::cos(emb); } -} - bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { auto U = [&](const char * k, int64_t & dst) { if (g.has(k)) dst = (int64_t) g.u32(k); }; auto F = [&](const char * k, float & dst) { if (g.has(k)) dst = g.f32(k); }; @@ -399,6 +245,15 @@ bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { U(fk("num_inference_timesteps"), m.num_steps); U(fk("num_timestep_buckets"), m.num_buckets); U(fk("max_num_embodiments"), m.max_embodiments); U(fk("max_seq_len"), m.max_seq_len); U(fk("image_target_size"), m.image_target_size); + // merge_block_coords only enumerates the patch grid exactly when the spatial + // merge divides it; otherwise it emits rows past the position table. + if (m.patch_size <= 0 || m.spatial_merge <= 0 || m.image_target_size % m.patch_size != 0 || + (m.image_target_size / m.patch_size) % m.spatial_merge != 0) { + std::fprintf(stderr, "vla(gr00tn1d7): image %lld / patch %lld / merge %lld do not divide evenly\n", + (long long) m.image_target_size, (long long) m.patch_size, (long long) m.spatial_merge); + return false; + } + if (const char * ns = std::getenv("VLA_NUM_STEPS")) { char * end = nullptr; long v = std::strtol(ns, &end, 10); if (end && *end == '\0' && v >= 1) { m.num_steps = (int64_t) v; std::fprintf(stderr, "vla(gr00tn1d7): VLA_NUM_STEPS override → num_steps=%lld\n", (long long) v); } @@ -432,6 +287,9 @@ bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { cfg.hidden = m.lm_hidden; cfg.n_q_heads = m.n_q; cfg.n_kv_heads = m.n_kv; cfg.head_dim = m.lm_head_dim; cfg.n_layers = m.lm_layers; cfg.num_steps = (int) m.num_steps; cfg.rms_eps = m.lm_rms_eps; cfg.rope_n_dims = (int) m.lm_head_dim; cfg.rope_mode = GGML_ROPE_TYPE_NEOX; cfg.rope_freq_base = m.lm_rope_base; + // Raw output: this arch expects the client to apply the dataset statistics + // (see the --stats-json flag in eval/client). + cfg.denormalized = false; cfg.norm_eps = 1e-8f; return true; } @@ -470,20 +328,10 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, (long long) m->action_horizon, (long long) m->action_dim, (long long) m->max_state_dim, (long long) m->num_steps, (long long) m->embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(gr00tn1d7): backend = CUDA (device 0)\n"); } - else std::fprintf(stderr, "vla(gr00tn1d7): ggml_backend_cuda_init failed; falling back to CPU\n"); -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(gr00tn1d7): backend = Metal\n"); } - else std::fprintf(stderr, "vla(gr00tn1d7): ggml_backend_metal_init failed; falling back to CPU\n"); -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(gr00tn1d7): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(gr00tn1d7): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(gr00tn1d7)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; @@ -688,8 +536,7 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { img_emb_host.assign((size_t) n_views * K * H, 0.0f); for (int j = 0; j < 3; ++j) ds_host[j].assign((size_t) n_views * K * H, 0.0f); - ggml_init_params vp = { (size_t) 512 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 512 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(gr00tn1d7): ggml_init(vision ctx) failed\n"); return {}; } ggml_tensor * t_patches = ggml_new_tensor_2d(VC, GGML_TYPE_F32, vit_patch_flat, n_patches); ggml_set_input(t_patches); ggml_tensor * t_pos = ggml_new_tensor_2d(VC, GGML_TYPE_F32, vit_hidden, n_patches); ggml_set_input(t_pos); @@ -711,13 +558,12 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { ggml_cgraph * vg = ggml_new_graph_custom(VC, 16384, false); ggml_build_forward_expand(vg, vit_embeds); for (int j = 0; j < 3; ++j) ggml_build_forward_expand(vg, ds_out[j]); - ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { std::fprintf(stderr, "vla(gr00tn1d7): vision gallocr alloc failed\n"); if (vga) ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!vision_scratch.alloc(backend, vg)) { std::fprintf(stderr, "vla(gr00tn1d7): vision gallocr alloc failed\n"); return {}; } const auto tv0 = std::chrono::steady_clock::now(); std::vector patches; bool vok = true; for (int64_t v = 0; v < n_views && vok; ++v) { - if (!preprocess_image_patches(in.images[v], side, ps, temporal_patch, grow, gcol, patches)) { vok = false; break; } + if (!preprocess_image_patches("gr00tn1d7", in.images[v], side, ps, temporal_patch, grow, gcol, patches)) { vok = false; break; } ggml_backend_tensor_set(t_pos, pos_interp.data(), 0, ggml_nbytes(t_pos)); ggml_backend_tensor_set(t_cos, rope_cos.data(), 0, ggml_nbytes(t_cos)); @@ -728,7 +574,6 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { for (int j = 0; j < 3; ++j) ggml_backend_tensor_get(ds_out[j], ds_host[j].data() + v * K * H, 0, ggml_nbytes(ds_out[j])); } stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now() - tv0).count(); - ggml_gallocr_free(vga); ggml_free(VC); if (!vok) return {}; img_emb_ptr = img_emb_host.data(); } else { @@ -787,18 +632,17 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { if (in.noise) std::memcpy(x_init.data(), in.noise, x_init.size() * sizeof(float)); else { std::mt19937 rng((uint32_t) std::chrono::steady_clock::now().time_since_epoch().count()); std::normal_distribution nd(0.f, 1.f); for (auto & v : x_init) v = nd(rng); } - const bool use_cache = (std::getenv("VLA_GR00T_GRAPH_CACHE") != nullptr) && !do_dump; - const bool reuse = use_cache && mg.valid && mg.seq == SEQ && mg.n_img == n_img && - mg.seq_txt == SEQ_TXT && mg.nsteps == num_steps && mg.deepstack == inject_deepstack; + // On by default: 16% faster, bit-identical. Set VLA_GR00T_GRAPH_CACHE=0 to opt out. + // Dumping adds graph outputs, so it always rebuilds. + const char * gc = std::getenv("VLA_GR00T_GRAPH_CACHE"); + const bool use_cache = (!gc || std::strcmp(gc, "0") != 0) && !do_dump; + if (!use_cache) mg.release(); ggml_tensor * eagle = nullptr, * vl_embs = nullptr; std::vector lm_h_dump, vlsa_dump; - if (!reuse) { - if (mg.C) mg.release(); - ggml_init_params cp = { (size_t) 256 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(gr00tn1d7): ggml_init(ctx_compute) failed\n"); return {}; } - + const MainKey mkey{ SEQ, n_img, SEQ_TXT, num_steps, inject_deepstack }; + const bool built = mg.ensure(backend, mkey, (size_t) 256 * 1024 * 1024, + [&](ggml_context * C, MainIO & gio) -> ggml_cgraph * { ggml_tensor * t_embeds = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); ggml_set_input(t_embeds); ggml_tensor * t_pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, 4 * SEQ); ggml_set_input(t_pos); ggml_tensor * t_lmmask = ggml_new_tensor_2d(C, GGML_TYPE_F32, SEQ, SEQ); ggml_set_input(t_lmmask); @@ -876,25 +720,22 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { } ggml_set_name(actions, "action_pred"); ggml_set_output(actions); + gio.t_embeds=t_embeds; gio.t_pos=t_pos; gio.t_lmmask=t_lmmask; gio.t_state=t_state; gio.t_x0=t_x0; + gio.t_ds[0]=t_ds[0]; gio.t_ds[1]=t_ds[1]; gio.t_ds[2]=t_ds[2]; + gio.t_img_idx=t_img_idx; gio.t_txt_idx=t_txt_idx; gio.t_tau=t_tau; gio.t_tproj=t_tproj; gio.actions=actions; + ggml_cgraph * gf = ggml_new_graph_custom(C, 65536, false); ggml_build_forward_expand(gf, actions); + return gf; + }); + if (!built) { std::fprintf(stderr, "vla(gr00tn1d7): main graph build failed\n"); return {}; } - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { std::fprintf(stderr, "vla(gr00tn1d7): gallocr alloc failed\n"); if (galloc) ggml_gallocr_free(galloc); ggml_free(C); return {}; } - - mg.C=C; mg.galloc=galloc; mg.gf=gf; - mg.t_embeds=t_embeds; mg.t_pos=t_pos; mg.t_lmmask=t_lmmask; mg.t_state=t_state; mg.t_x0=t_x0; - mg.t_ds[0]=t_ds[0]; mg.t_ds[1]=t_ds[1]; mg.t_ds[2]=t_ds[2]; - mg.t_img_idx=t_img_idx; mg.t_txt_idx=t_txt_idx; mg.t_tau=t_tau; mg.t_tproj=t_tproj; mg.actions=actions; - mg.seq=SEQ; mg.n_img=n_img; mg.seq_txt=SEQ_TXT; mg.nsteps=num_steps; mg.deepstack=inject_deepstack; - mg.valid = use_cache; - } - - ggml_context * C = mg.C; ggml_cgraph * gf = mg.gf; ggml_gallocr_t galloc = mg.galloc; (void) C; (void) galloc; - ggml_tensor * t_embeds = mg.t_embeds, * t_pos = mg.t_pos, * t_lmmask = mg.t_lmmask, * t_state = mg.t_state, * t_x0 = mg.t_x0; - ggml_tensor * t_ds[3] = { mg.t_ds[0], mg.t_ds[1], mg.t_ds[2] }; - ggml_tensor * t_img_idx = mg.t_img_idx, * t_txt_idx = mg.t_txt_idx, * actions = mg.actions; - std::vector & t_tau = mg.t_tau; std::vector & t_tproj = mg.t_tproj; + MainIO & gio = mg.io(); + ggml_cgraph * gf = mg.graph(); + ggml_tensor * t_embeds = gio.t_embeds, * t_pos = gio.t_pos, * t_lmmask = gio.t_lmmask, * t_state = gio.t_state, * t_x0 = gio.t_x0; + ggml_tensor * t_ds[3] = { gio.t_ds[0], gio.t_ds[1], gio.t_ds[2] }; + ggml_tensor * t_img_idx = gio.t_img_idx, * t_txt_idx = gio.t_txt_idx, * actions = gio.actions; + std::vector & t_tau = gio.t_tau; std::vector & t_tproj = gio.t_tproj; ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { @@ -947,11 +788,7 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { std::memcpy(pp.data() + (size_t) 3 * SEQ, pp.data() + (size_t) 0 * SEQ, (size_t) SEQ * sizeof(int32_t)); ggml_backend_tensor_set(t_pos, pp.data(), 0, ggml_nbytes(t_pos)); } - if (c_mask_seq != SEQ) { - c_mask.assign((size_t) SEQ * SEQ, 0.0f); const float NEG = -std::numeric_limits::infinity(); - for (int64_t q = 0; q < SEQ; ++q) for (int64_t kv = 0; kv < SEQ; ++kv) c_mask[q * SEQ + kv] = (kv <= q) ? 0.0f : NEG; - c_mask_seq = SEQ; - } + if (c_mask_seq != SEQ) { build_causal_mask(SEQ, c_mask); c_mask_seq = SEQ; } ggml_backend_tensor_set(t_lmmask, c_mask.data(), 0, ggml_nbytes(t_lmmask)); { std::vector st(max_state_dim, 0.0f); for (int64_t i = 0; i < max_state_dim; ++i) st[i] = in.state ? in.state[i] : 0.0f; ggml_backend_tensor_set(t_state, st.data(), 0, ggml_nbytes(t_state)); } ggml_backend_tensor_set(t_x0, x_init.data(), 0, ggml_nbytes(t_x0)); diff --git a/src/models/openvla_oft.cpp b/src/models/openvla_oft.cpp index c91d6ca..4650272 100644 --- a/src/models/openvla_oft.cpp +++ b/src/models/openvla_oft.cpp @@ -15,18 +15,15 @@ #include "arch.h" #include "model.h" #include "vision_common.h" +#include "models/dual_tower.h" #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" #include #include @@ -42,7 +39,6 @@ namespace vla { namespace { - bool parse_stats(const std::string & js, int64_t want, std::vector & q01, std::vector & q99, std::vector & mask, std::string & suite) { auto find_key = [&](size_t from, const std::string & key) -> size_t { @@ -81,7 +77,6 @@ bool parse_stats(const std::string & js, int64_t want, std::vector & q01, return (int64_t) q01.size() == want && (int64_t) q99.size() == want; } -struct ViTLayerW { ggml_tensor *n1w,*n1b,*n2w,*n2b,*ls1,*ls2,*Wqkv,*bqkv,*Wproj,*bproj,*Wfc1,*bfc1,*Wfc2,*bfc2; }; struct LMLayerW { ggml_tensor *attn_norm,*Wq,*Wk,*Wv,*Wo,*ffn_norm,*Wg,*Wu,*Wd; }; struct HeadBlkW { ggml_tensor *lnw,*lnb,*linw,*linb; }; @@ -96,8 +91,18 @@ struct OpenVlaOftModelArch : public ModelArchBase { } ggml_backend_t backend = nullptr; - bool is_gpu = false; int n_threads = default_cpu_threads(); + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + + struct MainKey { + int64_t seq=-1, n_views=-1, n_lang=-1; + bool operator==(const MainKey & o) const { return seq==o.seq && n_views==o.n_views && n_lang==o.n_lang; } + }; + struct MainIO { + ggml_tensor *t_ids=nullptr,*t_state=nullptr,*t_proj=nullptr,*act0=nullptr,*t_pos=nullptr,*norm_actions=nullptr; + }; + graph_cache main_graph; ggml_backend_buffer_t weight_buf = nullptr; ggml_type mt = GGML_TYPE_BF16; @@ -109,7 +114,7 @@ struct OpenVlaOftModelArch : public ModelArchBase { float lm_rope_base=1e4f, lm_rms_eps=1e-6f; int64_t chunk=8,action_dim=7,proprio_dim=8,head_hidden=4096,head_blocks=2; float head_ln_eps=1e-5f; - int64_t stop_id=2,empty_id=29871; + int64_t stop_id=2; ggml_tensor *d_patch_w,*d_patch_b,*d_cls,*d_reg,*d_pos; std::vector dvit; ggml_tensor *s_patch_w,*s_patch_b,*s_pos; std::vector svit; @@ -123,49 +128,6 @@ struct OpenVlaOftModelArch : public ModelArchBase { std::vector predict(const Inputs& in) override; }; -namespace { - -static ggml_tensor * LN(ggml_context*C, ggml_tensor*x, ggml_tensor*w, ggml_tensor*b, float eps){ return ggml_add(C,ggml_mul(C,ggml_norm(C,x,eps),w),b); } - -static ggml_tensor* vit_block(ggml_context*C, const ViTLayerW&w, ggml_tensor*x, int64_t N, int64_t hidden, int64_t heads, int64_t hd, float eps, bool ls){ - const float sc=1.0f/std::sqrt((float)hd); - ggml_tensor*xn=LN(C,x,w.n1w,w.n1b,eps); - ggml_tensor*qkv=ggml_add(C,ggml_mul_mat(C,w.Wqkv,xn),w.bqkv); - ggml_tensor*q=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],0*hidden*sizeof(float))); - ggml_tensor*k=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],1*hidden*sizeof(float))); - ggml_tensor*v=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],2*hidden*sizeof(float))); - ggml_tensor*Q=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,q,hd,heads,N),0,2,1,3)); - ggml_tensor*K=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,k,hd,heads,N),0,2,1,3)); - ggml_tensor*V=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,v,hd,heads,N),1,2,0,3)); - ggml_tensor*kq=ggml_mul_mat(C,K,Q); ggml_mul_mat_set_prec(kq,GGML_PREC_F32); - ggml_tensor*aw=ggml_soft_max_ext(C,kq,nullptr,sc,0.0f); - ggml_tensor*kqv=ggml_mul_mat(C,V,aw); - ggml_tensor*att=ggml_reshape_2d(C,ggml_cont(C,ggml_permute(C,kqv,0,2,1,3)),hidden,N); - ggml_tensor*ao=ggml_add(C,ggml_mul_mat(C,w.Wproj,att),w.bproj); - x=ggml_add(C,x,ls?ggml_mul(C,ao,w.ls1):ao); - ggml_tensor*xn2=LN(C,x,w.n2w,w.n2b,eps); - ggml_tensor*h=ggml_add(C,ggml_mul_mat(C,w.Wfc1,xn2),w.bfc1); h=ggml_gelu_erf(C,h); - h=ggml_add(C,ggml_mul_mat(C,w.Wfc2,h),w.bfc2); - return ggml_add(C,x,ls?ggml_mul(C,h,w.ls2):h); -} - -static ggml_tensor* tower(ggml_context*C, ggml_tensor*pix, ggml_tensor*pw, ggml_tensor*pb, ggml_tensor*pos, - ggml_tensor*cls, ggml_tensor*reg, const std::vector&blk, - int64_t hidden, int64_t heads, int64_t hd, int64_t inter, int64_t patch, float eps, bool prefix){ - (void)inter; - const int64_t NP=256, nprefix=prefix?5:0, N=NP+nprefix; - ggml_tensor*conv=ggml_conv_2d(C,pw,pix,patch,patch,0,0,1,1); - ggml_tensor*pt=ggml_cont(C,ggml_transpose(C,ggml_reshape_2d(C,conv,NP,hidden))); - pt=ggml_add(C,pt,pb); pt=ggml_add(C,pt,pos); - ggml_tensor*x=pt; - if(prefix){ ggml_tensor*tok=ggml_concat(C,ggml_reshape_2d(C,cls,hidden,1),reg,1); x=ggml_concat(C,tok,pt,1); } - for(size_t i=0;inb[1],nprefix*x->nb[1])); - return x; -} - -} - std::unique_ptr openvla_oft_create(const std::string& mmproj_path, const std::string& ckpt_path, const std::string& ) { @@ -194,7 +156,9 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path U("openvla_oft.action.chunk",m->chunk); U("openvla_oft.action.action_dim",m->action_dim); U("openvla_oft.action.proprio_dim",m->proprio_dim); U("openvla_oft.action.head_hidden",m->head_hidden); U("openvla_oft.action.head_blocks",m->head_blocks); F("openvla_oft.action.head_ln_eps",m->head_ln_eps); - U("openvla_oft.tokens.stop_id",m->stop_id); U("openvla_oft.tokens.empty_id",m->empty_id); + // No empty_id: the reference zeroes the action-slot embeddings instead + // (modeling_prismatic.py:891), which is what act0 below does. + U("openvla_oft.tokens.stop_id",m->stop_id); if (m->lm_head_dim==0) m->lm_head_dim = m->lm_hidden / m->n_q; if (g.has("openvla_oft.statistics_json")) { @@ -203,18 +167,10 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path std::printf("vla(openvla_oft): unnorm suite = %s (q99 dim %zu)\n", m->suite.c_str(), m->q99.size()); } -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_gpu=true; std::printf("vla(openvla_oft): backend = CUDA (device 0)\n"); } -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu=true; std::printf("vla(openvla_oft): backend = Metal\n"); } -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(openvla_oft): cpu backend init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(openvla_oft): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(openvla_oft)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t)64*1024*1024, nullptr, true }; @@ -288,18 +244,6 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path return m; } -namespace { - -void normalize_tower(const ImageView& v, int64_t S, const float mean[3], const float std_[3], std::vector& out){ - out.assign((size_t)3*S*S,0.0f); - for(int64_t h=0;h OpenVlaOftModelArch::predict(const Inputs& in) { using clock = std::chrono::steady_clock; const auto t0 = clock::now(); @@ -325,7 +269,7 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { std::vector proj_host((size_t)HC*NPATCH); { const auto tv=clock::now(); - ggml_init_params vp={(size_t)64*1024*1024,nullptr,true}; ggml_context*C=ggml_init(vp); + ggml_context*C=vision_scratch.reset((size_t)64*1024*1024); std::vector px_d(n_views), px_s(n_views), cmb(n_views); for(int v=0; v OpenVlaOftModelArch::predict(const Inputs& in) { ph=ggml_add(C,ggml_mul_mat(C,pj_fc2w,ph),pj_fc2b); ph=ggml_gelu_erf(C,ph); ggml_tensor*proj=ggml_add(C,ggml_mul_mat(C,pj_fc3w,ph),pj_fc3b); ggml_set_output(proj); ggml_cgraph*vg=ggml_new_graph_custom(C,16384,false); ggml_build_forward_expand(vg,proj); - ggml_gallocr_t ga=ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if(!ga||!ggml_gallocr_alloc_graph(ga,vg)){ std::fprintf(stderr,"vla(openvla_oft): vision gallocr failed\n"); if(ga)ggml_gallocr_free(ga); ggml_free(C); return {}; } + if(!vision_scratch.alloc(backend,vg)){ std::fprintf(stderr,"vla(openvla_oft): vision gallocr failed\n"); return {}; } std::vector dbuf, sbuf; for(int v=0;v(clock::now()-tv).count(); } @@ -369,8 +311,10 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { const int64_t ACT_START = NUM_PATCHES + NUM_PROMPT_TOKENS; const int64_t SEQ = 1 + NUM_PATCHES + (L-1) + n_act + 1; const auto ti=clock::now(); - ggml_init_params mp={(size_t)256*1024*1024,nullptr,true}; ggml_context*C=ggml_init(mp); - + // LM + action head graph depends only on the sequence layout. + const MainKey mkey{ SEQ, n_views, L }; + const bool built = main_graph.ensure(backend, mkey, (size_t)256*1024*1024, + [&](ggml_context*C, MainIO & gio)->ggml_cgraph*{ ggml_tensor*t_ids=ggml_new_tensor_1d(C,GGML_TYPE_I32,L+1); ggml_set_input(t_ids); ggml_tensor*emb=ggml_get_rows(C,token_embd,t_ids); if(emb->type!=GGML_TYPE_F32) emb=ggml_cast(C,emb,GGML_TYPE_F32); @@ -403,6 +347,8 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { ggml_tensor*kr=ggml_rope_ext(C,kh,t_pos,nullptr,(int)lm_head_dim,GGML_ROPE_TYPE_NEOX,0,lm_rope_base,1.0f,0.0f,1.0f,32.0f,1.0f); ggml_tensor*Q=ggml_cont(C,ggml_permute(C,qr,0,2,1,3)),*K=ggml_cont(C,ggml_permute(C,kr,0,2,1,3)),*V=ggml_cont(C,ggml_permute(C,vh,1,2,0,3)); ggml_tensor*kq=ggml_mul_mat(C,K,Q); ggml_mul_mat_set_prec(kq,GGML_PREC_F32); + // Unmasked on purpose: OpenVLA-OFT patches transformers to replace the + // causal mask across the whole sequence (modeling_llama.py:719-723). ggml_tensor*aw=ggml_soft_max_ext(C,kq,nullptr,lsc,0.0f); ggml_tensor*kqv=ggml_mul_mat(C,V,aw); ggml_tensor*att=ggml_reshape_2d(C,ggml_cont(C,ggml_permute(C,kqv,0,2,1,3)),HC,SEQ); @@ -426,9 +372,18 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { hh=LN(C,hh,h_ln2w,h_ln2b,head_ln_eps); ggml_tensor*norm_actions=ggml_add(C,ggml_mul_mat(C,h_fc2w,hh),h_fc2b); ggml_set_output(norm_actions); + gio.t_ids=t_ids; gio.t_state=t_state; gio.t_proj=t_proj; gio.act0=act0; + gio.t_pos=t_pos; gio.norm_actions=norm_actions; + ggml_cgraph*gf=ggml_new_graph_custom(C,16384,false); ggml_build_forward_expand(gf,norm_actions); - ggml_gallocr_t ga=ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if(!ga||!ggml_gallocr_alloc_graph(ga,gf)){ std::fprintf(stderr,"vla(openvla_oft): main gallocr failed\n"); if(ga)ggml_gallocr_free(ga); ggml_free(C); return {}; } + return gf; + }); + if(!built){ std::fprintf(stderr,"vla(openvla_oft): main graph build failed\n"); return {}; } + + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor*t_ids=gio.t_ids,*t_state=gio.t_state,*t_proj=gio.t_proj; + ggml_tensor*act0=gio.act0,*t_pos=gio.t_pos,*norm_actions=gio.norm_actions; { std::vector ids(L+1); for(int64_t i=0;i OpenVlaOftModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(t_state,sv.data(),0,ggml_nbytes(t_state)); } { std::vector z((size_t)HC*n_act,0.0f); ggml_backend_tensor_set(act0,z.data(),0,ggml_nbytes(act0)); } - if(ggml_backend_graph_compute(backend,gf)!=GGML_STATUS_SUCCESS){ std::fprintf(stderr,"vla(openvla_oft): main compute failed\n"); ggml_gallocr_free(ga); ggml_free(C); return {}; } + if(ggml_backend_graph_compute(backend,gf)!=GGML_STATUS_SUCCESS){ std::fprintf(stderr,"vla(openvla_oft): main compute failed\n"); return {}; } std::vector na((size_t)action_dim*chunk); ggml_backend_tensor_get(norm_actions,na.data(),0,na.size()*sizeof(float)); - ggml_gallocr_free(ga); ggml_free(C); stats.ms_inference = std::chrono::duration(clock::now()-ti).count(); const int64_t Wd = cfg.max_action_dim>0 ? cfg.max_action_dim : action_dim; diff --git a/src/models/pi0.cpp b/src/models/pi0.cpp index 0c38116..b905b1c 100644 --- a/src/models/pi0.cpp +++ b/src/models/pi0.cpp @@ -19,14 +19,12 @@ #include "ggml-cpu.h" #include "ggml-backend.h" #include "ggml-alloc.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/dit_common.h" +#include "models/vision_common.h" #include #include @@ -68,19 +66,6 @@ bool is_gemma_norm(const std::string & name) { return lm && name.find("norm.weight") != std::string::npos; } -std::vector sinusoidal_time_emb(double t, int64_t dim, double min_p, double max_p) { - const int64_t half = dim / 2; - std::vector out(dim); - for (int64_t i = 0; i < half; ++i) { - const double frac = (half == 1) ? 0.0 : double(i) / double(half - 1); - const double period = min_p * std::pow(max_p / min_p, frac); - const double s = (2.0 * M_PI / period) * t; - out[i] = (float) std::sin(s); - out[half + i] = (float) std::cos(s); - } - return out; -} - bool ends_with(const std::string & s, const char * sfx) { const size_t n = std::strlen(sfx); return s.size() >= n && s.compare(s.size() - n, n, sfx) == 0; @@ -95,11 +80,23 @@ struct Pi0ModelArch : public ModelArchBase { std::vector predict(const Inputs& in) override; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; ggml_backend_buffer_t weight_buf = nullptr; ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + + struct MainKey { + int64_t n_img=-1, n_lang=-1, nsteps=-1; + bool operator==(const MainKey & o) const { return n_img==o.n_img && n_lang==o.n_lang && nsteps==o.nsteps; } + }; + struct MainIO { + ggml_tensor *t_image_emb=nullptr,*t_lang_emb=nullptr,*t_prefix_pos=nullptr,*t_state=nullptr; + ggml_tensor *t_x0=nullptr,*t_suffix_pos=nullptr,*t_full_mask=nullptr,*x_final=nullptr; + std::vector t_time; + }; + graph_cache main_graph; std::string ckpt_path_; + // Opened once at load: reopening per predict re-parses the whole GGUF header. + gguf_reader io{"pi0"}; ggml_type matmul_type = GGML_TYPE_BF16; // In-tree SigLIP-So400m/14 vision tower (was llama.cpp clip.cpp mmproj). @@ -125,7 +122,7 @@ struct Pi0ModelArch : public ModelArchBase { std::vector state_mean, state_std, action_mean, action_std; std::mt19937 rng{std::random_device{}()}; - int n_threads = 4; + int n_threads = default_cpu_threads(); }; namespace { @@ -153,23 +150,6 @@ ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_ } // CHW-planar float image in [-1,1] for ggml_conv_2d (SigLIP mean/std 0.5). -bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(pi0): image view is %dx%d, expected %lldx%lld\n", - v.w, v.h, (long long) side, (long long) side); - return false; - } - out.assign((size_t) 3 * side * side, 0.0f); - for (int64_t h = 0; h < side; ++h) - for (int64_t w = 0; w < side; ++w) - for (int64_t c = 0; c < 3; ++c) { - float px; - if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; - else px = ((const float *) v.data)[(h * side + w) * 3 + c]; - out[c * side * side + h * side + w] = px * 2.0f - 1.0f; - } - return true; -} ggml_tensor * build_gemma_layer( ggml_context * ctx, const GemmaLayerW & w, @@ -302,7 +282,12 @@ bool load_stats(gguf_reader & g, Pi0ModelArch & m) { const ggml_tensor * t = g.meta(name); if (!t) { std::printf("vla(pi0): %s missing - identity\n", name); return; } if (t->ne[0] != (int64_t) dst.size()) { std::printf("vla(pi0): %s dim mismatch - identity\n", name); return; } - if (!g.read_raw(name, dst.data())) std::printf("vla(pi0): %s read failed - identity\n", name); + const std::vector identity = dst; + if (!g.read_raw(name, dst.data(), dst.size() * sizeof(float))) { + // A short read leaves dst half-overwritten. + dst = identity; + std::printf("vla(pi0): %s read failed - identity\n", name); + } }; read1d("state_mean", m.state_mean); read1d("state_std", m.state_std); @@ -336,8 +321,8 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, m->ckpt_path_ = ckpt_path; m->matmul_type = std::getenv("VLA_PI0_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g("pi0"); - if (!g.open(ckpt_path)) return nullptr; + if (!m->io.open(ckpt_path)) return nullptr; + gguf_reader & g = m->io; if (!g.has("pi0.architecture") || g.str("pi0.architecture") != "pi0") { std::fprintf(stderr, "vla(pi0): '%s' is not a π₀ GGUF (pi0.architecture missing/wrong)\n", ckpt_path.c_str()); @@ -354,24 +339,11 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, cfg.num_steps, (long long) cfg.real_state_dim, (long long) cfg.real_action_dim, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init( 0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(pi0): backend = CUDA (device 0)\n"); } - else { std::fprintf(stderr, "vla(pi0): ggml_backend_cuda_init failed; falling back to CPU\n"); } -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(pi0): backend = Metal\n"); } - else { std::fprintf(stderr, "vla(pi0): ggml_backend_metal_init failed; falling back to CPU\n"); } -#endif + m->n_threads = default_cpu_threads(); { - const unsigned hw = std::thread::hardware_concurrency(); - m->n_threads = (hw == 0) ? 4 : (int) std::min(hw, 8u); - } - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(pi0): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(pi0): backend = CPU (%d threads)\n", m->n_threads); + const Backend b = backend_init("vla(pi0)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } // The SigLIP tower is now bundled in the ckpt GGUF; mmproj_path is ignored. @@ -397,10 +369,12 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, } ggml_context * W = m->ctx_weights; std::vector weights; + // A miss returns before pushing, so the null scan below cannot see it. + bool missing = false; auto mk = [&](const char * name, ggml_type type, int n_dims, const int64_t * ne) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return nullptr; } + if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); missing = true; return nullptr; } ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), n_dims, ne); ggml_set_name(t, name); weights.push_back(t); @@ -409,12 +383,12 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, auto mk_mm = [&](const char * name) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return nullptr; } + if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); missing = true; return nullptr; } return mk(name, m->matmul_type, GGML_MAX_DIMS, gt->ne); }; auto mk_f32 = [&](const char * name) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); return nullptr; } + if (!gt) { std::fprintf(stderr, "vla(pi0): missing tensor %s\n", name); missing = true; return nullptr; } return mk(name, GGML_TYPE_F32, GGML_MAX_DIMS, gt->ne); }; @@ -464,6 +438,7 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, m->W_at1 = mk_f32("action_time_mlp_in.weight"); m->b_at1 = mk_f32("action_time_mlp_in.bias"); m->W_at2 = mk_f32("action_time_mlp_out.weight"); m->b_at2 = mk_f32("action_time_mlp_out.bias"); m->W_aout = mk_f32("action_out_proj.weight"); m->b_aout = mk_f32("action_out_proj.bias"); + if (missing) { std::fprintf(stderr, "vla(pi0): checkpoint is missing weights\n"); return nullptr; } for (ggml_tensor * t : weights) if (!t) { std::fprintf(stderr, "vla(pi0): weight tensor creation failed\n"); return nullptr; } if (!m->ex_final_norm || !m->W_sp || !m->b_sp || !m->W_ain || !m->b_ain || !m->W_at1 || !m->b_at1 || !m->W_at2 || !m->b_at2 || !m->W_aout || !m->b_aout) { @@ -521,8 +496,7 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { n_img_tokens = (int64_t) in.n_images * K; img_emb_host.assign((size_t) in.n_images * K * H, 0.0f); - ggml_init_params vp = { (size_t) 128 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 128 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(pi0): ggml_init(vision ctx) failed\n"); return {}; } ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, vit_image_size, vit_image_size, 3); ggml_set_input(t_px); ggml_tensor * conv = ggml_conv_2d(VC, vit_patch_w, t_px, (int) vit_patch_size, (int) vit_patch_size, 0, 0, 1, 1); @@ -539,26 +513,23 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vg, vit_emb); - ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { + + if (!vision_scratch.alloc(backend, vg)) { std::fprintf(stderr, "vla(pi0): vision gallocr alloc failed\n"); - if (vga) ggml_gallocr_free(vga); - ggml_free(VC); return {}; } const auto tv0 = clk::now(); std::vector chw; for (int v = 0; v < in.n_images; ++v) { - if (!preprocess_image_chw(in.images[v], vit_image_size, chw)) { ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!preprocess_image_chw("pi0", in.images[v], vit_image_size, chw)) { return {}; } ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(pi0): vision compute failed (view %d)\n", v); - ggml_gallocr_free(vga); ggml_free(VC); return {}; + return {}; } ggml_backend_tensor_get(vit_emb, img_emb_host.data() + (size_t) v * K * H, 0, ggml_nbytes(vit_emb)); } stats.ms_vision = std::chrono::duration(clk::now() - tv0).count(); - ggml_gallocr_free(vga); ggml_free(VC); } if (in.n_lang < 1 || !in.lang_tokens) { @@ -572,15 +543,13 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { std::vector lang_ids(in.lang_tokens, in.lang_tokens + n_lang); std::vector lang_rows((size_t) n_lang * hidden_pl); { - gguf_reader g("pi0"); - if (!g.open(ckpt_path_)) return {}; - if (!g.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; + if (!io.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; } - ggml_init_params cp = { (size_t) 64 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(pi0): ggml_init(ctx_compute) failed\n"); return {}; } - + // Prefix + expert graph depends only on the token counts and step count. + const MainKey mkey{ n_img_tokens, n_lang, num_steps }; + const bool built = main_graph.ensure(backend, mkey, (size_t) 64 * 1024 * 1024, + [&](ggml_context * C, MainIO & gio) -> ggml_cgraph * { ggml_tensor * t_image_emb = ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden_pl, n_img_tokens); ggml_set_input(t_image_emb); ggml_tensor * t_lang_emb = ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden_pl, n_lang); ggml_set_input(t_lang_emb); ggml_tensor * t_prefix_pos= ggml_new_tensor_1d(C, GGML_TYPE_I32, n_prefix); ggml_set_input(t_prefix_pos); @@ -627,16 +596,23 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { ggml_tensor * x_final = x_t; ggml_set_output(x_final); + gio.t_image_emb=t_image_emb; gio.t_lang_emb=t_lang_emb; gio.t_prefix_pos=t_prefix_pos; + gio.t_state=t_state; gio.t_x0=t_x0; gio.t_suffix_pos=t_suffix_pos; + gio.t_full_mask=t_full_mask; gio.t_time=t_time; gio.x_final=x_final; + ggml_cgraph * gf = ggml_new_graph_custom(C, 16384, false); ggml_build_forward_expand(gf, x_final); - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { - std::fprintf(stderr, "vla(pi0): ggml_gallocr_alloc_graph failed (out of memory?)\n"); - if (galloc) ggml_gallocr_free(galloc); - ggml_free(C); - return {}; - } + return gf; + }); + if (!built) { std::fprintf(stderr, "vla(pi0): main graph build failed\n"); return {}; } + + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor * t_image_emb = gio.t_image_emb, * t_lang_emb = gio.t_lang_emb; + ggml_tensor * t_prefix_pos = gio.t_prefix_pos, * t_state = gio.t_state, * t_x0 = gio.t_x0; + ggml_tensor * t_suffix_pos = gio.t_suffix_pos, * t_full_mask = gio.t_full_mask; + ggml_tensor * x_final = gio.x_final; + std::vector & t_time = gio.t_time; ggml_backend_tensor_set(t_image_emb, img_emb_host.data(), 0, ggml_nbytes(t_image_emb)); ggml_backend_tensor_set(t_lang_emb, lang_rows.data(), 0, ggml_nbytes(t_lang_emb)); @@ -685,15 +661,11 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { stats.ms_inference = std::chrono::duration(clk::now() - ti0).count(); if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(pi0): ggml_backend_graph_compute failed (%d)\n", (int) st); - ggml_gallocr_free(galloc); - ggml_free(C); return {}; } std::vector out((size_t) chunk * max_ad); ggml_backend_tensor_get(x_final, out.data(), 0, out.size() * sizeof(float)); - ggml_gallocr_free(galloc); - ggml_free(C); for (int64_t t = 0; t < chunk; ++t) { float * row = out.data() + (size_t) t * max_ad; for (int64_t j = 0; j < max_ad; ++j) diff --git a/src/models/pi05.cpp b/src/models/pi05.cpp index 9a30087..7e8c877 100644 --- a/src/models/pi05.cpp +++ b/src/models/pi05.cpp @@ -19,14 +19,12 @@ #include "ggml-cpu.h" #include "ggml-backend.h" #include "ggml-alloc.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/dit_common.h" +#include "models/vision_common.h" #include #include @@ -75,19 +73,6 @@ struct ExpertLayerW { // SigLIP-So400m vision block weights (PaliGemma tower, built in-tree like gr00tn1d5). struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; -std::vector sinusoidal_time_emb(double t, int64_t dim, double min_p, double max_p) { - const int64_t half = dim / 2; - std::vector out(dim); - for (int64_t i = 0; i < half; ++i) { - const double frac = (half == 1) ? 0.0 : double(i) / double(half - 1); - const double period = min_p * std::pow(max_p / min_p, frac); - const double s = (2.0 * M_PI / period) * t; - out[i] = (float) std::sin(s); - out[half + i] = (float) std::cos(s); - } - return out; -} - bool ends_with(const std::string & s, const char * sfx) { const size_t n = std::strlen(sfx); return s.size() >= n && s.compare(s.size() - n, n, sfx) == 0; @@ -110,11 +95,23 @@ struct Pi05ModelArch : public ModelArchBase { std::vector predict(const Inputs& in) override; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; ggml_backend_buffer_t weight_buf = nullptr; ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + + struct MainKey { + int64_t n_img=-1, n_lang=-1, nsteps=-1; + bool operator==(const MainKey & o) const { return n_img==o.n_img && n_lang==o.n_lang && nsteps==o.nsteps; } + }; + struct MainIO { + ggml_tensor *t_image_emb=nullptr,*t_lang_emb=nullptr,*t_prefix_pos=nullptr; + ggml_tensor *t_x0=nullptr,*t_suffix_pos=nullptr,*x_final=nullptr; + std::vector t_time; + }; + graph_cache main_graph; std::string ckpt_path_; + // Opened once at load: reopening per predict re-parses the whole GGUF header. + gguf_reader io{"pi05"}; ggml_type matmul_type = GGML_TYPE_BF16; int64_t adarms_cond_dim = 0; @@ -143,7 +140,7 @@ struct Pi05ModelArch : public ModelArchBase { bool quantile_norm = false; std::mt19937 rng{std::random_device{}()}; - int n_threads = 4; + int n_threads = default_cpu_threads(); }; namespace { @@ -171,23 +168,6 @@ ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_ } // CHW-planar float image in [-1,1] for ggml_conv_2d (SigLIP mean/std 0.5). -bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(pi05): image view is %dx%d, expected %lldx%lld\n", - v.w, v.h, (long long) side, (long long) side); - return false; - } - out.assign((size_t) 3 * side * side, 0.0f); - for (int64_t h = 0; h < side; ++h) - for (int64_t w = 0; w < side; ++w) - for (int64_t c = 0; c < 3; ++c) { - float px; - if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; - else px = ((const float *) v.data)[(h * side + w) * 3 + c]; - out[c * side * side + h * side + w] = px * 2.0f - 1.0f; - } - return true; -} ggml_tensor * build_vlm_layer( ggml_context * ctx, const VlmLayerW & w, @@ -374,7 +354,12 @@ bool load_stats(gguf_reader & g, Pi05ModelArch & m) { const ggml_tensor * t = g.meta(name); if (!t) { std::printf("vla(pi05): %s missing - identity\n", name); return; } if (t->ne[0] != (int64_t) dst.size()) { std::printf("vla(pi05): %s dim mismatch - identity\n", name); return; } - if (!g.read_raw(name, dst.data())) std::printf("vla(pi05): %s read failed - identity\n", name); + const std::vector identity = dst; + if (!g.read_raw(name, dst.data(), dst.size() * sizeof(float))) { + // A short read leaves dst half-overwritten. + dst = identity; + std::printf("vla(pi05): %s read failed - identity\n", name); + } }; read1d("state_mean", m.state_mean); read1d("state_std", m.state_std); @@ -414,8 +399,8 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, m->ckpt_path_ = ckpt_path; m->matmul_type = std::getenv("VLA_PI05_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; - gguf_reader g("pi05"); - if (!g.open(ckpt_path)) return nullptr; + if (!m->io.open(ckpt_path)) return nullptr; + gguf_reader & g = m->io; if (!g.has("pi05.architecture") || g.str("pi05.architecture") != "pi05") { std::fprintf(stderr, "vla(pi05): '%s' is not a π0.5 GGUF (pi05.architecture missing/wrong)\n", ckpt_path.c_str()); @@ -435,24 +420,11 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, (long long) cfg.n_lang, (long long) m->adarms_cond_dim, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init( 0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(pi05): backend = CUDA (device 0)\n"); } - else { std::fprintf(stderr, "vla(pi05): ggml_backend_cuda_init failed; falling back to CPU\n"); } -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(pi05): backend = Metal\n"); } - else { std::fprintf(stderr, "vla(pi05): ggml_backend_metal_init failed; falling back to CPU\n"); } -#endif + m->n_threads = default_cpu_threads(); { - const unsigned hw = std::thread::hardware_concurrency(); - m->n_threads = (hw == 0) ? 4 : (int) std::min(hw, 8u); - } - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(pi05): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(pi05): backend = CPU (%d threads)\n", m->n_threads); + const Backend b = backend_init("vla(pi05)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } // The SigLIP tower is now bundled in the ckpt GGUF; mmproj_path is ignored. @@ -478,10 +450,12 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, } ggml_context * W = m->ctx_weights; std::vector weights; + // A miss returns before pushing, so the null scan below cannot see it. + bool missing = false; auto mk = [&](const char * name, ggml_type type, int n_dims, const int64_t * ne) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return nullptr; } + if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); missing = true; return nullptr; } ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), n_dims, ne); ggml_set_name(t, name); weights.push_back(t); @@ -489,12 +463,12 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, }; auto mk_mm = [&](const char * name) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return nullptr; } + if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); missing = true; return nullptr; } return mk(name, m->matmul_type, GGML_MAX_DIMS, gt->ne); }; auto mk_f32 = [&](const char * name) -> ggml_tensor * { const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); return nullptr; } + if (!gt) { std::fprintf(stderr, "vla(pi05): missing tensor %s\n", name); missing = true; return nullptr; } return mk(name, GGML_TYPE_F32, GGML_MAX_DIMS, gt->ne); }; @@ -561,6 +535,7 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, m->W_tin = mk_f32("time_mlp_in.weight"); m->b_tin = mk_f32("time_mlp_in.bias"); m->W_tout = mk_f32("time_mlp_out.weight"); m->b_tout = mk_f32("time_mlp_out.bias"); m->W_aout = mk_f32("action_out_proj.weight"); m->b_aout = mk_f32("action_out_proj.bias"); + if (missing) { std::fprintf(stderr, "vla(pi05): checkpoint is missing weights\n"); return nullptr; } for (ggml_tensor * t : weights) if (!t) { std::fprintf(stderr, "vla(pi05): weight tensor creation failed\n"); return nullptr; } if (!m->ex_final_w || !m->ex_final_b || !m->W_ain || !m->b_ain || !m->W_tin || !m->b_tin || !m->W_tout || !m->b_tout || !m->W_aout || !m->b_aout) { @@ -617,8 +592,7 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { n_img_tokens = (int64_t) in.n_images * K; img_emb_host.assign((size_t) in.n_images * K * H, 0.0f); - ggml_init_params vp = { (size_t) 128 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 128 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(pi05): ggml_init(vision ctx) failed\n"); return {}; } ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, vit_image_size, vit_image_size, 3); ggml_set_input(t_px); ggml_tensor * conv = ggml_conv_2d(VC, vit_patch_w, t_px, (int) vit_patch_size, (int) vit_patch_size, 0, 0, 1, 1); @@ -635,29 +609,27 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vg, vit_emb); - ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { + + if (!vision_scratch.alloc(backend, vg)) { std::fprintf(stderr, "vla(pi05): vision gallocr alloc failed\n"); - if (vga) ggml_gallocr_free(vga); - ggml_free(VC); return {}; } const auto tv0 = clk::now(); std::vector chw; for (int v = 0; v < in.n_images; ++v) { - if (!preprocess_image_chw(in.images[v], vit_image_size, chw)) { ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!preprocess_image_chw("pi05", in.images[v], vit_image_size, chw)) { return {}; } ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(pi05): vision compute failed (view %d)\n", v); - ggml_gallocr_free(vga); ggml_free(VC); return {}; + return {}; } ggml_backend_tensor_get(vit_emb, img_emb_host.data() + (size_t) v * K * H, 0, ggml_nbytes(vit_emb)); } stats.ms_vision = std::chrono::duration(clk::now() - tv0).count(); - ggml_gallocr_free(vga); ggml_free(VC); - // π0.5's image tokens are the raw PaliGemma projector features: this undoes - // the 1/sqrt(hidden) the shared vision graph applies (π0 keeps them scaled). + // Undo the 1/sqrt(hidden) the shared vision graph applies; pi05 wants raw + // projector features. Inside this branch on purpose: precomputed_img_emb + // replaces the tower and is already LM-ready. const float img_scale = (float) std::sqrt((double) hidden_pl); for (float & x : img_emb_host) x *= img_scale; } @@ -672,15 +644,13 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { std::vector lang_ids(in.lang_tokens, in.lang_tokens + n_lang); std::vector lang_rows((size_t) n_lang * hidden_pl); { - gguf_reader g("pi05"); - if (!g.open(ckpt_path_)) return {}; - if (!g.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; + if (!io.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; } - ggml_init_params cp = { (size_t) 64 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(pi05): ggml_init(ctx_compute) failed\n"); return {}; } - + // Prefix + expert graph depends only on the token counts and step count. + const MainKey mkey{ n_img_tokens, n_lang, num_steps }; + const bool built = main_graph.ensure(backend, mkey, (size_t) 64 * 1024 * 1024, + [&](ggml_context * C, MainIO & gio) -> ggml_cgraph * { ggml_tensor * t_image_emb = ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden_pl, n_img_tokens); ggml_set_input(t_image_emb); ggml_tensor * t_lang_emb = ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden_pl, n_lang); ggml_set_input(t_lang_emb); ggml_tensor * t_prefix_pos= ggml_new_tensor_1d(C, GGML_TYPE_I32, n_prefix); ggml_set_input(t_prefix_pos); @@ -725,16 +695,21 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { ggml_tensor * x_final = x_t; ggml_set_output(x_final); + gio.t_image_emb=t_image_emb; gio.t_lang_emb=t_lang_emb; gio.t_prefix_pos=t_prefix_pos; + gio.t_x0=t_x0; gio.t_suffix_pos=t_suffix_pos; gio.t_time=t_time; gio.x_final=x_final; + ggml_cgraph * gf = ggml_new_graph_custom(C, 16384, false); ggml_build_forward_expand(gf, x_final); + return gf; + }); + if (!built) { std::fprintf(stderr, "vla(pi05): main graph build failed\n"); return {}; } - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!galloc || !ggml_gallocr_alloc_graph(galloc, gf)) { - std::fprintf(stderr, "vla(pi05): ggml_gallocr_alloc_graph failed (out of memory?)\n"); - if (galloc) ggml_gallocr_free(galloc); - ggml_free(C); - return {}; - } + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor * t_image_emb = gio.t_image_emb, * t_lang_emb = gio.t_lang_emb; + ggml_tensor * t_prefix_pos = gio.t_prefix_pos, * t_x0 = gio.t_x0; + ggml_tensor * t_suffix_pos = gio.t_suffix_pos, * x_final = gio.x_final; + std::vector & t_time = gio.t_time; ggml_backend_tensor_set(t_image_emb, img_emb_host.data(), 0, ggml_nbytes(t_image_emb)); ggml_backend_tensor_set(t_lang_emb, lang_rows.data(), 0, ggml_nbytes(t_lang_emb)); @@ -761,15 +736,11 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { stats.ms_inference = std::chrono::duration(clk::now() - ti0).count(); if (st != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(pi05): ggml_backend_graph_compute failed (%d)\n", (int) st); - ggml_gallocr_free(galloc); - ggml_free(C); return {}; } std::vector out((size_t) chunk * max_ad); ggml_backend_tensor_get(x_final, out.data(), 0, out.size() * sizeof(float)); - ggml_gallocr_free(galloc); - ggml_free(C); if (!std::getenv("VLA_PI05_SKIP_UNNORM")) { for (int64_t t = 0; t < chunk; ++t) { diff --git a/src/models/qwen3vl_vit.h b/src/models/qwen3vl_vit.h new file mode 100644 index 0000000..9abc09a --- /dev/null +++ b/src/models/qwen3vl_vit.h @@ -0,0 +1,172 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Qwen3-VL vision tower, shared by GR00T N1.7 and VLA-JEPA. + +#pragma once + +#include "model.h" + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +namespace vla { + +constexpr float QWEN3VL_MEAN[3] = {0.5f, 0.5f, 0.5f}; +constexpr float QWEN3VL_STD [3] = {0.5f, 0.5f, 0.5f}; + +struct VitLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wqkv,*bqkv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; +struct MergerW { ggml_tensor *nw,*nb,*fc1w,*fc1b,*fc2w,*fc2b; }; + +inline ggml_tensor * rope2d(ggml_context * C, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_t) { + const int64_t hd = x->ne[0], S = x->ne[1], Hh = x->ne[2]; const int64_t half = hd / 2; + ggml_tensor * x1 = ggml_cont(C, ggml_view_3d(C, x, half, S, Hh, x->nb[1], x->nb[2], 0)); + ggml_tensor * x2 = ggml_cont(C, ggml_view_3d(C, x, half, S, Hh, x->nb[1], x->nb[2], (size_t) half * x->nb[0])); + ggml_tensor * rot = ggml_concat(C, ggml_neg(C, x2), x1, 0); + return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, rot, sin_t)); +} + +inline bool fa_enabled() { static const bool e = (std::getenv("VLA_FLASH_ATTN") != nullptr); return e; } + +inline ggml_tensor * flash_attn(ggml_context * C, ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, + ggml_tensor * mask, float scale) { + ggml_tensor * kf = (k->type == GGML_TYPE_F16) ? k : ggml_cast(C, k, GGML_TYPE_F16); + ggml_tensor * vf = (v->type == GGML_TYPE_F16) ? v : ggml_cast(C, v, GGML_TYPE_F16); + ggml_tensor * o = ggml_flash_attn_ext(C, q, kf, vf, mask, scale, 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32); + return ggml_reshape_2d(C, o, o->ne[0] * o->ne[1], o->ne[2] * o->ne[3]); +} + +inline ggml_tensor * build_vit_layer(ggml_context * C, const VitLayerW & w, ggml_tensor * x, + ggml_tensor * cos_t, ggml_tensor * sin_t, + int64_t seq, int64_t heads, int64_t hd, int64_t hidden, float ln_eps) { + const float scale = 1.0f / std::sqrt((float) hd); + ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); + ggml_tensor * qkv = ggml_add(C, ggml_mul_mat(C, w.Wqkv, n1), w.bqkv); + ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], 0)); + ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) hidden * qkv->nb[0])); + ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) 2 * hidden * qkv->nb[0])); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, seq), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, seq), 0, 2, 1, 3)); + Q = rope2d(C, Q, cos_t, sin_t); K = rope2d(C, K, cos_t, sin_t); + ggml_tensor * att; + if (fa_enabled()) { + ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 0, 2, 1, 3)); + att = flash_attn(C, Q, K, V, nullptr, scale); + } else { + ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 1, 2, 0, 3)); + ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale, 0.0f); + att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq); + } + ggml_tensor * h1 = ggml_add(C, x, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); + ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.ln2w), w.ln2b); + ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wfc2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wfc1, n2), w.bfc1))), w.bfc2); + return ggml_add(C, h1, ff); +} + +// pre_merge normalizes before the reshape, the deepstack taps after. +inline ggml_tensor * build_merger(ggml_context * C, const MergerW & w, ggml_tensor * x, + int64_t hidden, int64_t merge2, float ln_eps, bool pre_merge) { + const int64_t n_patches = x->ne[1], c_merged = hidden * merge2 * merge2, n_merged = n_patches / (merge2 * merge2); + ggml_tensor * m; + if (pre_merge) { + ggml_tensor * xn = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.nw), w.nb); + m = ggml_reshape_2d(C, ggml_cont(C, xn), c_merged, n_merged); + } else { + ggml_tensor * mr = ggml_reshape_2d(C, ggml_cont(C, x), c_merged, n_merged); + m = ggml_add(C, ggml_mul(C, ggml_norm(C, mr, ln_eps), w.nw), w.nb); + } + ggml_tensor * z1 = ggml_add(C, ggml_mul_mat(C, w.fc1w, m), w.fc1b); + return ggml_add(C, ggml_mul_mat(C, w.fc2w, ggml_gelu(C, z1)), w.fc2b); +} + +// Patch row/col after the spatial merge. +inline void merge_block_coords(int64_t gh, int64_t gw, int64_t m, std::vector & row, std::vector & col) { + const int64_t S = gh * gw; row.assign(S, 0); col.assign(S, 0); + for (int64_t s = 0; s < S; ++s) { + int64_t t = s; const int64_t wj = t % m; t /= m; const int64_t wi = t % m; t /= m; + const int64_t bc = t % (gw / m); t /= (gw / m); const int64_t br = t; + row[s] = br * m + wi; col[s] = bc * m + wj; + } +} + +inline void vit_rope_tables(const std::vector & row, const std::vector & col, int64_t hd, double theta, + std::vector & cos_t, std::vector & sin_t) { + const int64_t S = (int64_t) row.size(), nf = hd / 4; + std::vector invf(nf); + for (int64_t i = 0; i < nf; ++i) invf[i] = 1.0 / std::pow(theta, (double)(2 * i) / (double)(hd / 2)); + cos_t.assign((size_t) S * hd, 0.0f); sin_t.assign((size_t) S * hd, 0.0f); + for (int64_t s = 0; s < S; ++s) { + std::vector emb(hd); + for (int64_t i = 0; i < nf; ++i) { emb[i] = (double) row[s] * invf[i]; emb[nf + i] = (double) col[s] * invf[i]; } + for (int64_t i = 0; i < hd / 2; ++i) emb[hd / 2 + i] = emb[i]; + for (int64_t i = 0; i < hd; ++i) { cos_t[s * hd + i] = (float) std::cos(emb[i]); sin_t[s * hd + i] = (float) std::sin(emb[i]); } + } +} + +// Bilinear resample of the pretrained position table onto gh x gw. +inline void interp_pos_embed(const std::vector & table, int64_t num_side, int64_t hidden, + const std::vector & row, const std::vector & col, int64_t gh, int64_t gw, + std::vector & out) { + const int64_t S = (int64_t) row.size(); + out.assign((size_t) S * hidden, 0.0f); + auto src_coord = [&](int64_t k, int64_t g) -> double { return (g <= 1) ? 0.0 : (double) k * (double)(num_side - 1) / (double)(g - 1); }; + for (int64_t s = 0; s < S; ++s) { + // Clamped, not just h1/w1: a grid that the spatial merge does not divide + // pushes row/col past gh-1 and would index off the end of the table. + const double lim = (double) (num_side - 1); + const double hy = std::min(src_coord(row[s], gh), lim), wx = std::min(src_coord(col[s], gw), lim); + const int64_t h0 = (int64_t) std::floor(hy), w0 = (int64_t) std::floor(wx); + const int64_t h1 = std::min(h0 + 1, num_side - 1), w1 = std::min(w0 + 1, num_side - 1); + const double dh = hy - h0, dw = wx - w0; + const double c00 = (1 - dh) * (1 - dw), c01 = (1 - dh) * dw, c10 = dh * (1 - dw), c11 = dh * dw; + const float * T00 = &table[(h0 * num_side + w0) * hidden]; const float * T01 = &table[(h0 * num_side + w1) * hidden]; + const float * T10 = &table[(h1 * num_side + w0) * hidden]; const float * T11 = &table[(h1 * num_side + w1) * hidden]; + for (int64_t c = 0; c < hidden; ++c) out[s * hidden + c] = (float)(c00 * T00[c] + c01 * T01[c] + c10 * T10[c] + c11 * T11[c]); + } +} + +// HWC to flat patches. No resize: the view must already be side x side. +inline bool preprocess_image_patches(const char * arch, const ImageView & v, int64_t side, int64_t ps, int64_t tps, + const std::vector & row, const std::vector & col, + std::vector & out) { + if (v.w != (int) side || v.h != (int) side || !v.data) { + std::fprintf(stderr, "vla(%s): image view is %dx%d, expected %lldx%lld\n", + arch, v.w, v.h, (long long) side, (long long) side); + return false; + } + const int64_t S = (int64_t) row.size(), pf = 3 * tps * ps * ps; + out.assign((size_t) pf * S, 0.0f); + auto px = [&](int64_t r, int64_t c, int64_t ch) -> float { + if (v.format == PixelFormat::U8) return ((const uint8_t *) v.data)[(r * side + c) * 3 + ch] / 255.0f; + return ((const float *) v.data)[(r * side + c) * 3 + ch]; + }; + for (int64_t s = 0; s < S; ++s) + for (int64_t ch = 0; ch < 3; ++ch) + for (int64_t ph = 0; ph < ps; ++ph) + for (int64_t pw = 0; pw < ps; ++pw) { + const float val = (px(row[s] * ps + ph, col[s] * ps + pw, ch) - QWEN3VL_MEAN[ch]) / QWEN3VL_STD[ch]; + for (int64_t t = 0; t < tps; ++t) out[s * pf + ch * tps * ps * ps + t * ps * ps + ph * ps + pw] = val; + } + return true; +} + +} // namespace vla diff --git a/src/models/scratch_ctx.h b/src/models/scratch_ctx.h new file mode 100644 index 0000000..9c11650 --- /dev/null +++ b/src/models/scratch_ctx.h @@ -0,0 +1,111 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Compute context and graph allocator reused across predict calls; rebuilding +// them costs 2-4 ms on the larger graphs. One scratch per graph role, and +// tensors die at the next reset. +// +// graph_cache keeps the built graph too, for the archs whose shape depends on a +// small key. ggml-cuda can then capture and replay it, which needs the node list +// to stay put. + +#pragma once + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include + +namespace vla { + +class scratch_ctx { +public: + scratch_ctx() = default; + scratch_ctx(const scratch_ctx &) = delete; + scratch_ctx & operator=(const scratch_ctx &) = delete; + ~scratch_ctx() { release(); } + + ggml_context * reset(size_t arena) { + if (ctx_) { ggml_reset(ctx_); return ctx_; } + ggml_init_params p = { arena, nullptr, true }; + ctx_ = ggml_init(p); + return ctx_; + } + + bool alloc(ggml_backend_t backend, ggml_cgraph * gf) { + if (!galloc_) galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + return galloc_ && ggml_gallocr_alloc_graph(galloc_, gf); + } + + void release() { + if (galloc_) { ggml_gallocr_free(galloc_); galloc_ = nullptr; } + if (ctx_) { ggml_free(ctx_); ctx_ = nullptr; } + } + +private: + ggml_context * ctx_ = nullptr; + ggml_gallocr_t galloc_ = nullptr; +}; + +// Key is whatever shape the graph depends on (it needs operator==); IO holds the +// input/output tensor handles the arch fills in each call. build(ctx, io) emits +// the graph and returns it, or null to fail the call. +template +class graph_cache { +public: + graph_cache() = default; + graph_cache(const graph_cache &) = delete; + graph_cache & operator=(const graph_cache &) = delete; + ~graph_cache() { release(); } + + template + bool ensure(ggml_backend_t backend, const Key & key, size_t arena, Build && build) { + if (valid_ && key_ == key) return true; + release(); + ggml_init_params p = { arena, nullptr, true }; + ctx_ = ggml_init(p); + if (!ctx_) return false; + io_ = IO{}; + gf_ = build(ctx_, io_); + if (!gf_) { release(); return false; } + galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!galloc_ || !ggml_gallocr_alloc_graph(galloc_, gf_)) { release(); return false; } + key_ = key; + valid_ = true; + return true; + } + + IO & io() { return io_; } + ggml_cgraph * graph() { return gf_; } + + void release() { + if (galloc_) { ggml_gallocr_free(galloc_); galloc_ = nullptr; } + if (ctx_) { ggml_free(ctx_); ctx_ = nullptr; } + gf_ = nullptr; + io_ = IO{}; + valid_ = false; + } + +private: + ggml_context * ctx_ = nullptr; + ggml_gallocr_t galloc_ = nullptr; + ggml_cgraph * gf_ = nullptr; + Key key_{}; + IO io_{}; + bool valid_ = false; +}; + +} // namespace vla diff --git a/src/models/smolvla.cpp b/src/models/smolvla.cpp index 357e00d..4cd478f 100644 --- a/src/models/smolvla.cpp +++ b/src/models/smolvla.cpp @@ -18,17 +18,14 @@ #include "arch.h" #include "model.h" #include "vision_common.h" +#include "scratch_ctx.h" +#include "dit_common.h" #include "ggml.h" #include "ggml-backend.h" #include "ggml-cpu.h" #include "gguf.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "nlohmann/json.hpp" @@ -95,20 +92,33 @@ struct safetensors { std::fprintf(stderr, "vla: shape mismatch for %s\n", name.c_str()); return false; } + if (info.dtype != "BF16" && info.dtype != "F32") { + std::fprintf(stderr, "vla: unsupported dtype for %s: %s\n", + name.c_str(), info.dtype.c_str()); + return false; + } + // dst is sized from expected_shape, so the declared span has to match it. + // Without this a file can name the right shape and a longer span. + const size_t elsz = (info.dtype == "BF16") ? sizeof(ggml_bf16_t) : sizeof(float); + size_t want = elsz; + for (const int64_t d : info.shape) { + if (d < 0) { std::fprintf(stderr, "vla: negative dim for %s\n", name.c_str()); return false; } + want *= (size_t) d; + } + if (info.off_end < info.off_begin || info.off_end - info.off_begin != want) { + std::fprintf(stderr, "vla: bad data_offsets for %s\n", name.c_str()); + return false; + } const size_t bytes = info.off_end - info.off_begin; file.seekg(data_blob_start + info.off_begin, std::ios::beg); if (info.dtype == "BF16") { std::vector tmp(bytes / sizeof(ggml_bf16_t)); file.read(reinterpret_cast(tmp.data()), bytes); ggml_bf16_to_fp32_row(tmp.data(), dst, tmp.size()); - } else if (info.dtype == "F32") { - file.read(reinterpret_cast(dst), bytes); } else { - std::fprintf(stderr, "vla: unsupported dtype for %s: %s\n", - name.c_str(), info.dtype.c_str()); - return false; + file.read(reinterpret_cast(dst), bytes); } - return true; + return !file.fail(); } bool read_raw(const std::string & name, void * dst, size_t expected_bytes, @@ -189,7 +199,7 @@ struct gguf_source { const int64_t id = gguf_find_tensor(gctx, name.c_str()); const size_t offset = data_off + gguf_get_tensor_offset(gctx, id); const size_t bytes = gguf_get_tensor_size(gctx, id); - if (std::fseek(fp, (long) offset, SEEK_SET) != 0) { + if (fseeko(fp, (off_t) offset, SEEK_SET) != 0) { std::fprintf(stderr, "vla: fseek failed for %s\n", name.c_str()); return false; } @@ -223,7 +233,7 @@ struct gguf_source { return false; } const size_t offset = data_off + gguf_get_tensor_offset(gctx, id); - if (std::fseek(fp, (long) offset, SEEK_SET) != 0) return false; + if (fseeko(fp, (off_t) offset, SEEK_SET) != 0) return false; return std::fread(dst, 1, bytes, fp) == bytes; } @@ -290,12 +300,12 @@ struct SmolVLAModelArch : public ModelArchBase { ggml_backend_t backend = nullptr; ggml_backend_buffer_t weight_buf = nullptr; - bool is_cuda = false; - bool is_gpu = false; ggml_type weight_dtype = GGML_TYPE_BF16; ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + scratch_ctx connector_scratch; ggml_tensor * E_lang = nullptr; ggml_tensor * Wstate = nullptr; @@ -367,23 +377,6 @@ ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_ } // CHW-planar float image in [-1,1] for ggml_conv_2d (SigLIP mean/std 0.5). -bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(smolvla): image view is %dx%d, expected %lldx%lld\n", - v.w, v.h, (long long) side, (long long) side); - return false; - } - out.assign((size_t) 3 * side * side, 0.0f); - for (int64_t h = 0; h < side; ++h) - for (int64_t w = 0; w < side; ++w) - for (int64_t c = 0; c < 3; ++c) { - float px; - if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; - else px = ((const float *) v.data)[(h * side + w) * 3 + c]; - out[c * side * side + h * side + w] = px * 2.0f - 1.0f; - } - return true; -} bool load_config_from_json(const std::string & path, Config & cfg) { std::ifstream f(path); @@ -693,20 +686,6 @@ std::string hf_to_gguf(const std::string & n) { return n; } -std::vector sinusoidal_time_emb(double timestep, int64_t dim, - double min_period, double max_period) { - const int64_t half = dim / 2; - std::vector out(dim); - for (int64_t i = 0; i < half; ++i) { - const double frac = (half == 1) ? 0.0 : double(i) / double(half - 1); - const double period = min_period * std::pow(max_period / min_period, frac); - const double scale = 2.0 * M_PI / period; - const double s = scale * timestep; - out[i] = static_cast(std::sin(s)); - out[half + i] = static_cast(std::cos(s)); - } - return out; -} ggml_tensor * rope_q_or_k(ggml_context * ctx, ggml_tensor * x, ggml_tensor * positions, const Config & cfg) { @@ -927,33 +906,10 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, std::printf("vla: config = %s\n", cfg_path.c_str()); } -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init( 0); - if (m->backend) { - m->is_cuda = true; - m->is_gpu = true; - std::printf("vla: backend = CUDA (device 0)\n"); - } else { - std::fprintf(stderr, "vla: ggml_backend_cuda_init failed; falling back to CPU\n"); - } -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { - m->is_gpu = true; - std::printf("vla: backend = Metal\n"); - } else { - std::fprintf(stderr, "vla: ggml_backend_metal_init failed; falling back to CPU\n"); - } -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { - std::fprintf(stderr, "vla: ggml_backend_cpu_init failed\n"); - delete m; - return nullptr; - } - ggml_backend_cpu_set_n_threads(m->backend, default_cpu_threads()); - std::printf("vla: backend = CPU (%d threads)\n", default_cpu_threads()); + { + const Backend b = backend_init("vla", default_cpu_threads()); + if (!b.handle) { delete m; return nullptr; } + m->backend = b.handle; } vram_probe(m->backend, "after backend init"); @@ -976,7 +932,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, if (k * k != m->vit_n_tokens) { std::fprintf(stderr, "vla: smolvla vit geometry mismatch (grid=%lld scale=%lld -> %lld tokens, KV says %lld)\n", (long long) grid, (long long) m->vit_scale, (long long) (k * k), (long long) m->vit_n_tokens); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -988,7 +943,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, if (!use_gguf) { if (!st.open(ckpt_path)) { std::fprintf(stderr, "vla: failed to open %s\n", ckpt_path.c_str()); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -1003,7 +957,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, } if (max_layer < 0) { std::fprintf(stderr, "vla: cannot infer n_layers from %s\n", ckpt_path.c_str()); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -1013,7 +966,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, const auto it = st.tensors.find("model.vlm_with_expert.lm_expert.layers.0.mlp.gate_proj.weight"); if (it == st.tensors.end() || it->second.shape.size() != 2) { std::fprintf(stderr, "vla: missing/malformed expert gate_proj for shape derivation\n"); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -1023,7 +975,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, std::fprintf(stderr, "vla: expert_h mismatch - config implies %lld, " "checkpoint gate_proj has %lld\n", (long long) m->cfg.expert_h, (long long) it->second.shape[1]); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -1043,7 +994,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, if (use_gguf) { if (!load_normalizer_stats_from_gguf(gst, *m)) { std::fprintf(stderr, "vla: failed to load normalizer stats from gguf\n"); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -1064,7 +1014,6 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, m->ctx_weights = ggml_init(gparams); if (!m->ctx_weights) { std::fprintf(stderr, "vla: ggml_init (weights) failed\n"); - ggml_backend_free(m->backend); delete m; return nullptr; } @@ -1486,8 +1435,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { const auto t_vision_begin = clock::now(); // Graph A: SigLIP ViT (conv patch-embed -> +pos -> layers -> post_ln), plain sequential positions. - ggml_init_params vpA = { size_t(256) * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vpA); + ggml_context * VC = m->vision_scratch.reset(size_t(256) * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(smolvla): ggml_init(vision ctx) failed\n"); return {}; } ggml_tensor * t_px = ggml_new_tensor_3d(VC, GGML_TYPE_F32, m->vit_image, m->vit_image, 3); ggml_set_input(t_px); ggml_tensor * conv = ggml_conv_2d(VC, m->vit_patch_w, t_px, (int) m->vit_patch, (int) m->vit_patch, 0, 0, 1, 1); @@ -1499,35 +1447,28 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_set_output(post_ln); ggml_cgraph * gA = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(gA, post_ln); - ggml_gallocr_t vgA = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m->backend)); - if (!vgA || !ggml_gallocr_alloc_graph(vgA, gA)) { + if (!m->vision_scratch.alloc(m->backend, gA)) { std::fprintf(stderr, "vla(smolvla): vision gallocr A alloc failed\n"); - if (vgA) ggml_gallocr_free(vgA); - ggml_free(VC); return {}; } // Graph B: pixel-shuffle connector, a single bias-free matmul (c4 -> hidden). - ggml_init_params vpB = { size_t(64) * 1024 * 1024, nullptr, true }; - ggml_context * MC = ggml_init(vpB); - if (!MC) { std::fprintf(stderr, "vla(smolvla): ggml_init(connector ctx) failed\n"); ggml_gallocr_free(vgA); ggml_free(VC); return {}; } + ggml_context * MC = m->connector_scratch.reset(size_t(64) * 1024 * 1024); + if (!MC) { std::fprintf(stderr, "vla(smolvla): ggml_init(connector ctx) failed\n"); return {}; } ggml_tensor * t_shuf = ggml_new_tensor_2d(MC, GGML_TYPE_F32, c4, K); ggml_set_input(t_shuf); ggml_tensor * img_embeds = ggml_mul_mat(MC, m->mm_fc, t_shuf); ggml_set_output(img_embeds); ggml_cgraph * gB = ggml_new_graph(MC); ggml_build_forward_expand(gB, img_embeds); - ggml_gallocr_t vgB = ggml_gallocr_new(ggml_backend_get_default_buffer_type(m->backend)); - if (!vgB || !ggml_gallocr_alloc_graph(vgB, gB)) { + if (!m->connector_scratch.alloc(m->backend, gB)) { std::fprintf(stderr, "vla(smolvla): vision gallocr B alloc failed\n"); - if (vgB) ggml_gallocr_free(vgB); - ggml_gallocr_free(vgA); ggml_free(MC); ggml_free(VC); return {}; } std::vector chw, post_host((size_t) H * n_patches), shuf_host((size_t) c4 * K); bool vok = true; for (int v = 0; v < n_views && vok; ++v) { - if (!preprocess_image_chw(in.images[v], m->vit_image, chw)) { vok = false; break; } + if (!preprocess_image_chw("smolvla", in.images[v], m->vit_image, chw)) { vok = false; break; } ggml_backend_tensor_set(t_px, chw.data(), 0, ggml_nbytes(t_px)); if (ggml_backend_graph_compute(m->backend, gA) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(smolvla): vision compute A failed (view %d)\n", v); vok = false; break; @@ -1540,7 +1481,6 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { } ggml_backend_tensor_get(img_embeds, img_emb_pre.data() + size_t(v) * per_view_n, 0, ggml_nbytes(img_embeds)); } - ggml_gallocr_free(vgB); ggml_free(MC); ggml_gallocr_free(vgA); ggml_free(VC); if (!vok) return {}; m->stats.ms_vision = std::chrono::duration(clock::now() - t_vision_begin).count(); } @@ -1588,9 +1528,9 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { const int64_t pad_start = cfg.n_img + in.n_lang; const int64_t pad_end = cfg.n_img + n_lang_max; - std::vector state_host(cfg.max_state_dim); - std::memcpy(state_host.data(), in.state, cfg.max_state_dim * sizeof(float)); - for (int64_t i = 0; i < cfg.real_state_dim; ++i) { + std::vector state_host(cfg.max_state_dim, 0.0f); + if (in.state) std::memcpy(state_host.data(), in.state, cfg.max_state_dim * sizeof(float)); + for (int64_t i = 0; i < cfg.real_state_dim && i < cfg.max_state_dim; ++i) { state_host[i] = (state_host[i] - m->state_mean[i]) / (m->state_std[i] + cfg.norm_eps); } @@ -1715,9 +1655,9 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_tensor * pos_full = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, cfg.n_suffix); ggml_tensor * pos_rebased = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, cfg.n_suffix); - std::vector state_host(cfg.max_state_dim); - std::memcpy(state_host.data(), in.state, cfg.max_state_dim * sizeof(float)); - for (int64_t i = 0; i < cfg.real_state_dim; ++i) { + std::vector state_host(cfg.max_state_dim, 0.0f); + if (in.state) std::memcpy(state_host.data(), in.state, cfg.max_state_dim * sizeof(float)); + for (int64_t i = 0; i < cfg.real_state_dim && i < cfg.max_state_dim; ++i) { state_host[i] = (state_host[i] - m->state_mean[i]) / (m->state_std[i] + cfg.norm_eps); } diff --git a/src/models/vision_common.h b/src/models/vision_common.h index 35ad863..9e33f7b 100644 --- a/src/models/vision_common.h +++ b/src/models/vision_common.h @@ -17,8 +17,12 @@ #pragma once +#include "model.h" + #include +#include #include +#include namespace vla { @@ -46,4 +50,26 @@ inline void pixel_shuffle_hf(const float * src, float * dst, } } +// HWC to CHW planar in [-1, 1], the SigLIP convention used by SmolVLA, pi0, pi0.5 +// and GR00T N1.5. No resize: the view must already be side x side. arch only +// labels the error. +inline bool preprocess_image_chw(const char * arch, const ImageView & v, int64_t side, + std::vector & out) { + if (v.w != (int) side || v.h != (int) side || !v.data) { + std::fprintf(stderr, "vla(%s): image view is %dx%d, expected %lldx%lld\n", + arch, v.w, v.h, (long long) side, (long long) side); + return false; + } + out.assign((size_t) 3 * side * side, 0.0f); + for (int64_t h = 0; h < side; ++h) + for (int64_t w = 0; w < side; ++w) + for (int64_t c = 0; c < 3; ++c) { + float px; + if (v.format == PixelFormat::U8) px = ((const uint8_t *) v.data)[(h * side + w) * 3 + c] / 255.0f; + else px = ((const float *) v.data)[(h * side + w) * 3 + c]; + out[c * side * side + h * side + w] = px * 2.0f - 1.0f; + } + return true; +} + } // namespace vla diff --git a/src/models/vla_adapter.cpp b/src/models/vla_adapter.cpp index fca8d0d..9351784 100644 --- a/src/models/vla_adapter.cpp +++ b/src/models/vla_adapter.cpp @@ -15,18 +15,16 @@ #include "arch.h" #include "model.h" #include "vision_common.h" +#include "models/dual_tower.h" #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/dit_common.h" #include #include @@ -42,7 +40,6 @@ namespace vla { namespace { - bool parse_stats(const std::string & js, int64_t want, std::vector & q01, std::vector & q99, std::vector & mask, std::string & suite) { auto find_key = [&](size_t from, const std::string & key) -> size_t { @@ -82,7 +79,6 @@ bool parse_stats(const std::string & js, int64_t want, std::vector & q01, return (int64_t) q01.size() == want && (int64_t) q99.size() == want; } -struct ViTLayerW { ggml_tensor *n1w,*n1b,*n2w,*n2b,*ls1,*ls2,*Wqkv,*bqkv,*Wproj,*bproj,*Wfc1,*bfc1,*Wfc2,*bfc2; }; struct LMLayerW { ggml_tensor *attn_norm,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*ffn_norm,*Wg,*Wu,*Wd; }; struct HeadBlkW { ggml_tensor *Wq,*bq,*Wks,*bks,*Wvs,*bvs,*Wka,*bka,*Wva,*bva,*Wkt,*bkt,*Wvt,*bvt,*Wo,*bo,*flnw,*flnb,*flw,*flb; float rg; }; @@ -97,8 +93,20 @@ struct VlaAdapterModelArch : public ModelArchBase { } ggml_backend_t backend = nullptr; - bool is_gpu = false; int n_threads = default_cpu_threads(); + int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + + struct MainKey { + int64_t seq=-1, n_views=-1, nprompt=-1; + bool operator==(const MainKey & o) const { return seq==o.seq && n_views==o.n_views && nprompt==o.nprompt; } + }; + struct MainIO { + ggml_tensor *t_ids=nullptr,*t_proj=nullptr,*t_pos=nullptr,*t_mask=nullptr; + ggml_tensor *t_state=nullptr,*t_x0=nullptr,*norm_actions=nullptr; + ggml_tensor *cT=nullptr,*sT=nullptr,*cA=nullptr,*sA=nullptr,*cK=nullptr,*sK=nullptr; + }; + graph_cache main_graph; ggml_backend_buffer_t weight_buf = nullptr; ggml_type mt = GGML_TYPE_BF16; @@ -126,45 +134,9 @@ struct VlaAdapterModelArch : public ModelArchBase { namespace { -static ggml_tensor * LN(ggml_context*C, ggml_tensor*x, ggml_tensor*w, ggml_tensor*b, float eps){ return ggml_add(C,ggml_mul(C,ggml_norm(C,x,eps),w),b); } - -static ggml_tensor* vit_block(ggml_context*C, const ViTLayerW&w, ggml_tensor*x, int64_t N, int64_t hidden, int64_t heads, int64_t hd, float eps, bool ls){ - const float sc=1.0f/std::sqrt((float)hd); - ggml_tensor*xn=LN(C,x,w.n1w,w.n1b,eps); - ggml_tensor*qkv=ggml_add(C,ggml_mul_mat(C,w.Wqkv,xn),w.bqkv); - ggml_tensor*q=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],0*hidden*sizeof(float))); - ggml_tensor*k=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],1*hidden*sizeof(float))); - ggml_tensor*v=ggml_cont(C,ggml_view_2d(C,qkv,hidden,N,qkv->nb[1],2*hidden*sizeof(float))); - ggml_tensor*Q=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,q,hd,heads,N),0,2,1,3)); - ggml_tensor*K=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,k,hd,heads,N),0,2,1,3)); - ggml_tensor*V=ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,v,hd,heads,N),1,2,0,3)); - ggml_tensor*kq=ggml_mul_mat(C,K,Q); ggml_mul_mat_set_prec(kq,GGML_PREC_F32); - ggml_tensor*aw=ggml_soft_max_ext(C,kq,nullptr,sc,0.0f); - ggml_tensor*kqv=ggml_mul_mat(C,V,aw); - ggml_tensor*att=ggml_reshape_2d(C,ggml_cont(C,ggml_permute(C,kqv,0,2,1,3)),hidden,N); - ggml_tensor*ao=ggml_add(C,ggml_mul_mat(C,w.Wproj,att),w.bproj); - x=ggml_add(C,x,ls?ggml_mul(C,ao,w.ls1):ao); - ggml_tensor*xn2=LN(C,x,w.n2w,w.n2b,eps); - ggml_tensor*h=ggml_add(C,ggml_mul_mat(C,w.Wfc1,xn2),w.bfc1); h=ggml_gelu_erf(C,h); - h=ggml_add(C,ggml_mul_mat(C,w.Wfc2,h),w.bfc2); - return ggml_add(C,x,ls?ggml_mul(C,h,w.ls2):h); -} - -static ggml_tensor* tower(ggml_context*C, ggml_tensor*pix, ggml_tensor*pw, ggml_tensor*pb, ggml_tensor*pos, - ggml_tensor*cls, ggml_tensor*reg, const std::vector&blk, - int64_t hidden, int64_t heads, int64_t hd, int64_t inter, int64_t patch, float eps, bool prefix){ - (void)inter; - const int64_t NP=256, nprefix=prefix?5:0, N=NP+nprefix; - ggml_tensor*conv=ggml_conv_2d(C,pw,pix,patch,patch,0,0,1,1); - ggml_tensor*pt=ggml_cont(C,ggml_transpose(C,ggml_reshape_2d(C,conv,NP,hidden))); - pt=ggml_add(C,pt,pb); pt=ggml_add(C,pt,pos); - ggml_tensor*x=pt; - if(prefix){ ggml_tensor*tok=ggml_concat(C,ggml_reshape_2d(C,cls,hidden,1),reg,1); x=ggml_concat(C,tok,pt,1); } - for(size_t i=0;inb[1],nprefix*x->nb[1])); - return x; -} - +// Interleaved rotation, paired with the half-split frequency table in fill_cs, so +// a rotation pair gets two different angles. The reference does the same +// (action_heads.py:163 vs :137-140) and the weights were trained on it. Leave it. static ggml_tensor* hrot(ggml_context*C, ggml_tensor*x, int64_t HD){ int64_t L=x->ne[1],H=x->ne[2]; ggml_tensor*xp=ggml_reshape_4d(C,x,2,HD/2,L,H); @@ -211,24 +183,24 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path U("vla_adapter.action.head_dim",m->head_dim); F("vla_adapter.action.head_rope_base",m->head_rope_base); F("vla_adapter.action.ln_eps",m->head_ln_eps); U("vla_adapter.tokens.stop_id",m->stop_id); + // predict() taps one LM layer per head block, so head_blocks past lm_layers + // would read off the end of the layer-output vector. + if(m->lm_layers<1 || m->head_blocks<1 || m->head_blocks>m->lm_layers){ + std::fprintf(stderr,"vla(vla_adapter): head_blocks %lld outside [1, lm_layers %lld]\n", + (long long)m->head_blocks,(long long)m->lm_layers); + return nullptr; + } + if (g.has("vla_adapter.statistics_json")) { if (!parse_stats(g.str("vla_adapter.statistics_json"), m->action_dim, m->q01, m->q99, m->unnorm_mask, m->suite)) { std::fprintf(stderr, "vla(vla_adapter): failed to parse statistics_json\n"); return nullptr; } std::printf("vla(vla_adapter): unnorm suite = %s (q99 dim %zu)\n", m->suite.c_str(), m->q99.size()); } -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_gpu=true; std::printf("vla(vla_adapter): backend = CUDA (device 0)\n"); } -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu=true; std::printf("vla(vla_adapter): backend = Metal\n"); } -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(vla_adapter): cpu backend init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(vla_adapter): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(vla_adapter)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t)64*1024*1024, nullptr, true }; @@ -312,14 +284,6 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path namespace { -void normalize_tower(const ImageView& v, int64_t S, const float mean[3], const float std_[3], std::vector& out){ - out.assign((size_t)3*S*S,0.0f); - for(int64_t h=0;h VlaAdapterModelArch::predict(const Inputs& in) { @@ -346,7 +310,7 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { std::vector proj_host((size_t)HC*NP*n_views); { const auto tv=clock::now(); - ggml_init_params vp={(size_t)64*1024*1024,nullptr,true}; ggml_context*C=ggml_init(vp); + ggml_context*C=vision_scratch.reset((size_t)64*1024*1024); std::vector px_d(n_views), px_s(n_views); std::vector cmb(n_views); for(int v=0; v VlaAdapterModelArch::predict(const Inputs& in) { ph=ggml_add(C,ggml_mul_mat(C,pj_fc2w,ph),pj_fc2b); ph=ggml_gelu_erf(C,ph); ggml_tensor*proj=ggml_add(C,ggml_mul_mat(C,pj_fc3w,ph),pj_fc3b); ggml_set_output(proj); ggml_cgraph*vg=ggml_new_graph_custom(C,16384,false); ggml_build_forward_expand(vg,proj); - ggml_gallocr_t ga=ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if(!ga||!ggml_gallocr_alloc_graph(ga,vg)){ std::fprintf(stderr,"vla(vla_adapter): vision gallocr failed\n"); if(ga)ggml_gallocr_free(ga); ggml_free(C); return {}; } + if(!vision_scratch.alloc(backend,vg)){ std::fprintf(stderr,"vla(vla_adapter): vision gallocr failed\n"); return {}; } std::vector dbuf, sbuf; for(int v=0;v(clock::now()-tv).count(); } @@ -389,8 +351,10 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { const int64_t NPATCH = NP * n_views; const int64_t SEQ = 1 + NPATCH + (NPROMPT-1) + num_tokens + 1; const auto ti=clock::now(); - ggml_init_params mp={(size_t)128*1024*1024,nullptr,true}; ggml_context*C=ggml_init(mp); - + // LM + action head graph depends only on the sequence layout. + const MainKey mkey{ SEQ, n_views, NPROMPT }; + const bool built = main_graph.ensure(backend, mkey, (size_t)128*1024*1024, + [&](ggml_context*C, MainIO & gio)->ggml_cgraph*{ ggml_tensor*t_ids=ggml_new_tensor_1d(C,GGML_TYPE_I32,NPROMPT+num_tokens+1); ggml_set_input(t_ids); ggml_tensor*emb=ggml_get_rows(C,token_embd,t_ids); if(emb->type!=GGML_TYPE_F32) emb=ggml_cast(C,emb,GGML_TYPE_F32); @@ -474,9 +438,20 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { ggml_tensor*xn=LN(C,hx,h_ln2w,h_ln2b,head_ln_eps); ggml_tensor*norm_actions=ggml_add(C,ggml_mul_mat(C,h_fc2w,xn),h_fc2b); ggml_set_output(norm_actions); + gio.t_ids=t_ids; gio.t_proj=t_proj; gio.t_pos=t_pos; gio.t_mask=t_mask; + gio.t_state=t_state; gio.t_x0=t_x0; gio.norm_actions=norm_actions; + gio.cT=cT; gio.sT=sT; gio.cA=cA; gio.sA=sA; gio.cK=cK; gio.sK=sK; + ggml_cgraph*gf=ggml_new_graph_custom(C,65536,false); ggml_build_forward_expand(gf,norm_actions); - ggml_gallocr_t ga=ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if(!ga||!ggml_gallocr_alloc_graph(ga,gf)){ std::fprintf(stderr,"vla(vla_adapter): main gallocr failed\n"); if(ga)ggml_gallocr_free(ga); ggml_free(C); return {}; } + return gf; + }); + if(!built){ std::fprintf(stderr,"vla(vla_adapter): main graph build failed\n"); return {}; } + + MainIO & gio = main_graph.io(); + ggml_cgraph * gf = main_graph.graph(); + ggml_tensor*t_ids=gio.t_ids,*t_proj=gio.t_proj,*t_pos=gio.t_pos,*t_mask=gio.t_mask; + ggml_tensor*t_state=gio.t_state,*t_x0=gio.t_x0,*norm_actions=gio.norm_actions; + ggml_tensor*cT=gio.cT,*sT=gio.sT,*cA=gio.cA,*sA=gio.sA,*cK=gio.cK,*sK=gio.sK; { std::vector ids(NPROMPT+num_tokens+1); for(int64_t i=0;i VlaAdapterModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(t_ids,ids.data(),0,ggml_nbytes(t_ids)); } ggml_backend_tensor_set(t_proj,proj_host.data(),0,ggml_nbytes(t_proj)); { std::vector pp(SEQ); for(int64_t i=0;i mk((size_t)SEQ*SEQ); const float NI=-std::numeric_limits::infinity(); - for(int64_t q=0;q mk; build_causal_mask(SEQ, mk); ggml_backend_tensor_set(t_mask,mk.data(),0,ggml_nbytes(t_mask)); } { std::vector sv(proprio_dim,0.0f); for(int64_t i=0;i VlaAdapterModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(cc,cb.data(),0,ggml_nbytes(cc)); ggml_backend_tensor_set(ss,sb.data(),0,ggml_nbytes(ss)); }; fill_cs(cT,sT,chunk); fill_cs(cA,sA,num_tokens+1); fill_cs(cK,sK,NPATCH); - if(ggml_backend_graph_compute(backend,gf)!=GGML_STATUS_SUCCESS){ std::fprintf(stderr,"vla(vla_adapter): main compute failed\n"); ggml_gallocr_free(ga); ggml_free(C); return {}; } + if(ggml_backend_graph_compute(backend,gf)!=GGML_STATUS_SUCCESS){ std::fprintf(stderr,"vla(vla_adapter): main compute failed\n"); return {}; } std::vector na((size_t)action_dim*chunk); ggml_backend_tensor_get(norm_actions,na.data(),0,na.size()*sizeof(float)); - ggml_gallocr_free(ga); ggml_free(C); stats.ms_inference = std::chrono::duration(clock::now()-ti).count(); const int64_t W = cfg.max_action_dim>0 ? cfg.max_action_dim : action_dim; diff --git a/src/models/vla_jepa.cpp b/src/models/vla_jepa.cpp index 6622c55..91549cf 100644 --- a/src/models/vla_jepa.cpp +++ b/src/models/vla_jepa.cpp @@ -18,14 +18,12 @@ #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" -#ifdef GGML_USE_CUDA -#include "ggml-cuda.h" -#endif -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif +#include "backend.h" #include "gguf.h" #include "models/gguf_reader.h" +#include "models/scratch_ctx.h" +#include "models/dit_common.h" +#include "models/qwen3vl_vit.h" #include #include @@ -44,11 +42,6 @@ namespace vla { namespace { -constexpr float CLIP_MEAN[3] = {0.5f, 0.5f, 0.5f}; -constexpr float CLIP_STD [3] = {0.5f, 0.5f, 0.5f}; - -struct VitLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wqkv,*bqkv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; -struct MergerW { ggml_tensor *nw,*nb,*fc1w,*fc1b,*fc2w,*fc2b; }; struct Qwen3LayerW { ggml_tensor *attn_norm,*Wq,*Wk,*Wv,*Wo,*q_norm,*k_norm,*ffn_norm,*Wgate,*Wup,*Wdown; }; struct DitLayerW { ggml_tensor *adaln_w,*adaln_b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wff0,*bff0,*Wff2,*bff2; }; @@ -60,10 +53,28 @@ struct VlaJepaModelArch : public ModelArchBase { std::string gguf_path; ggml_backend_t backend = nullptr; - bool is_cuda = false; - bool is_gpu = false; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + scratch_ctx vision_scratch; + struct LmKey { + int64_t seq=-1, nfuture=-1; + bool operator==(const LmKey & o) const { return seq==o.seq && nfuture==o.nfuture; } + }; + struct LmIO { + ggml_tensor *t_embeds=nullptr,*t_pos2=nullptr,*t_lmmask=nullptr,*t_emb_idx=nullptr; + ggml_tensor *t_ds[3]={nullptr,nullptr,nullptr}; + ggml_tensor *eagle=nullptr,*conditioning=nullptr; + }; + struct HeadKey { + int64_t nsteps=-1; + bool operator==(const HeadKey & o) const { return nsteps==o.nsteps; } + }; + struct HeadIO { + ggml_tensor *t_cond=nullptr,*t_state=nullptr,*t_x0=nullptr,*actions=nullptr; + std::vector t_tau, t_tproj; + }; + graph_cache lm_graph; + graph_cache head_graph; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; @@ -107,57 +118,6 @@ struct VlaJepaModelArch : public ModelArchBase { namespace { -ggml_tensor * adaln(ggml_context * C, ggml_tensor * x, ggml_tensor * temb, ggml_tensor * lw, ggml_tensor * lb, int64_t dim, float eps) { - ggml_tensor * cond = ggml_add(C, ggml_mul_mat(C, lw, ggml_silu(C, temb)), lb); - ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0), * sh = ggml_view_1d(C, cond, dim, (size_t) dim * sizeof(float)); - ggml_tensor * xn = ggml_norm(C, x, eps); - return ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); -} - -ggml_tensor * rope2d(ggml_context * C, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_t) { - const int64_t hd = x->ne[0], S = x->ne[1], Hh = x->ne[2]; const int64_t half = hd / 2; - ggml_tensor * x1 = ggml_cont(C, ggml_view_3d(C, x, half, S, Hh, x->nb[1], x->nb[2], 0)); - ggml_tensor * x2 = ggml_cont(C, ggml_view_3d(C, x, half, S, Hh, x->nb[1], x->nb[2], (size_t) half * x->nb[0])); - ggml_tensor * rot = ggml_concat(C, ggml_neg(C, x2), x1, 0); - return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, rot, sin_t)); -} - -ggml_tensor * build_vit_layer(ggml_context * C, const VitLayerW & w, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_t, - int64_t seq, int64_t heads, int64_t hd, int64_t hidden, float ln_eps) { - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); - ggml_tensor * qkv = ggml_add(C, ggml_mul_mat(C, w.Wqkv, n1), w.bqkv); - ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], 0)); - ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) hidden * qkv->nb[0])); - ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) 2 * hidden * qkv->nb[0])); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, seq), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, seq), 0, 2, 1, 3)); - Q = rope2d(C, Q, cos_t, sin_t); K = rope2d(C, K, cos_t, sin_t); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, seq), 1, 2, 0, 3)); - ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale, 0.0f); - ggml_tensor * att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq); - ggml_tensor * h1 = ggml_add(C, x, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); - ggml_tensor * n2 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.ln2w), w.ln2b); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wfc2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wfc1, n2), w.bfc1))), w.bfc2); - return ggml_add(C, h1, ff); -} - -ggml_tensor * build_merger(ggml_context * C, const MergerW & w, ggml_tensor * x, int64_t hidden, int64_t merge2, float ln_eps, bool pre_merge) { - ggml_tensor * m; - if (pre_merge) { - ggml_tensor * xn = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.nw), w.nb); - const int64_t n_patches = x->ne[1], c_merged = hidden * merge2 * merge2, n_merged = n_patches / (merge2 * merge2); - m = ggml_reshape_2d(C, ggml_cont(C, xn), c_merged, n_merged); - } else { - const int64_t n_patches = x->ne[1], c_merged = hidden * merge2 * merge2, n_merged = n_patches / (merge2 * merge2); - ggml_tensor * mr = ggml_reshape_2d(C, ggml_cont(C, x), c_merged, n_merged); - m = ggml_add(C, ggml_mul(C, ggml_norm(C, mr, ln_eps), w.nw), w.nb); - } - ggml_tensor * z1 = ggml_add(C, ggml_mul_mat(C, w.fc1w, m), w.fc1b); - return ggml_add(C, ggml_mul_mat(C, w.fc2w, ggml_gelu(C, z1)), w.fc2b); -} - ggml_tensor * build_qwen3_layer(ggml_context * C, const VlaJepaModelArch & m, const Qwen3LayerW & w, ggml_tensor * h, ggml_tensor * positions, ggml_tensor * mask, int64_t seq) { const int64_t hd = m.lm_head_dim, n_q = m.n_q, n_kv = m.n_kv, hq = n_q * hd; @@ -209,80 +169,6 @@ ggml_tensor * build_dit_block(ggml_context * C, const VlaJepaModelArch & m, cons return ggml_add(C, h1, ff); } -void merge_block_coords(int64_t gh, int64_t gw, int64_t m, std::vector & row, std::vector & col) { - const int64_t S = gh * gw; row.assign(S, 0); col.assign(S, 0); - for (int64_t s = 0; s < S; ++s) { - int64_t t = s; const int64_t wj = t % m; t /= m; const int64_t wi = t % m; t /= m; - const int64_t bc = t % (gw / m); t /= (gw / m); const int64_t br = t; - row[s] = br * m + wi; col[s] = bc * m + wj; - } -} - -void vit_rope_tables(const std::vector & row, const std::vector & col, int64_t hd, double theta, - std::vector & cos_t, std::vector & sin_t) { - const int64_t S = (int64_t) row.size(), nf = hd / 4; - std::vector invf(nf); - for (int64_t i = 0; i < nf; ++i) invf[i] = 1.0 / std::pow(theta, (double)(2 * i) / (double)(hd / 2)); - cos_t.assign((size_t) S * hd, 0.0f); sin_t.assign((size_t) S * hd, 0.0f); - for (int64_t s = 0; s < S; ++s) { - std::vector emb(hd); - for (int64_t i = 0; i < nf; ++i) { emb[i] = (double) row[s] * invf[i]; emb[nf + i] = (double) col[s] * invf[i]; } - for (int64_t i = 0; i < hd / 2; ++i) emb[hd / 2 + i] = emb[i]; - for (int64_t i = 0; i < hd; ++i) { cos_t[s * hd + i] = (float) std::cos(emb[i]); sin_t[s * hd + i] = (float) std::sin(emb[i]); } - } -} - -void interp_pos_embed(const std::vector & table, int64_t num_side, int64_t hidden, - const std::vector & row, const std::vector & col, int64_t gh, int64_t gw, - std::vector & out) { - const int64_t S = (int64_t) row.size(); - out.assign((size_t) S * hidden, 0.0f); - auto src_coord = [&](int64_t k, int64_t g) -> double { return (g <= 1) ? 0.0 : (double) k * (double)(num_side - 1) / (double)(g - 1); }; - for (int64_t s = 0; s < S; ++s) { - const double hy = src_coord(row[s], gh), wx = src_coord(col[s], gw); - const int64_t h0 = (int64_t) std::floor(hy), w0 = (int64_t) std::floor(wx); - const int64_t h1 = std::min(h0 + 1, num_side - 1), w1 = std::min(w0 + 1, num_side - 1); - const double dh = hy - h0, dw = wx - w0; - const double c00 = (1 - dh) * (1 - dw), c01 = (1 - dh) * dw, c10 = dh * (1 - dw), c11 = dh * dw; - const float * T00 = &table[(h0 * num_side + w0) * hidden]; const float * T01 = &table[(h0 * num_side + w1) * hidden]; - const float * T10 = &table[(h1 * num_side + w0) * hidden]; const float * T11 = &table[(h1 * num_side + w1) * hidden]; - for (int64_t c = 0; c < hidden; ++c) out[s * hidden + c] = (float)(c00 * T00[c] + c01 * T01[c] + c10 * T10[c] + c11 * T11[c]); - } -} - -bool preprocess_image_patches(const ImageView & v, int64_t side, int64_t ps, int64_t tps, - const std::vector & row, const std::vector & col, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(vla_jepa): image view is %dx%d, expected %lldx%lld\n", v.w, v.h, (long long) side, (long long) side); return false; - } - const int64_t S = (int64_t) row.size(), pf = 3 * tps * ps * ps; - out.assign((size_t) pf * S, 0.0f); - auto px = [&](int64_t r, int64_t c, int64_t ch) -> float { - if (v.format == PixelFormat::U8) return ((const uint8_t *) v.data)[(r * side + c) * 3 + ch] / 255.0f; - return ((const float *) v.data)[(r * side + c) * 3 + ch]; - }; - for (int64_t s = 0; s < S; ++s) - for (int64_t ch = 0; ch < 3; ++ch) - for (int64_t ph = 0; ph < ps; ++ph) - for (int64_t pw = 0; pw < ps; ++pw) { - const float val = (px(row[s] * ps + ph, col[s] * ps + pw, ch) - CLIP_MEAN[ch]) / CLIP_STD[ch]; - for (int64_t t = 0; t < tps; ++t) out[s * pf + ch * tps * ps * ps + t * ps * ps + ph * ps + pw] = val; - } - return true; -} - -void timesteps_proj(int64_t bucket, std::vector & out) { - const int64_t half = 128; const float lm = std::log(10000.0f); const float t = (float) bucket; - out.assign(256, 0.0f); - for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-lm * (float) i / (float) (half - 1)); out[i] = std::cos(emb); out[half + i] = std::sin(emb); } -} - -void action_sinusoid(int64_t bucket, int64_t dim, int64_t T, std::vector & out) { - const int64_t half = dim / 2; const float step = std::log(10000.0f) / (float) half; const float t = (float) bucket; - out.assign((size_t) T * dim, 0.0f); - for (int64_t tk = 0; tk < T; ++tk) for (int64_t i = 0; i < half; ++i) { const float emb = t * std::exp(-(float) i * step); out[tk * dim + i] = std::sin(emb); out[tk * dim + half + i] = std::cos(emb); } -} - bool load_config(const gguf_reader & g, VlaJepaModelArch & m, Config & cfg) { auto U = [&](const char * k, int64_t & dst) { if (g.has(k)) dst = (int64_t) g.u32(k); }; auto F = [&](const char * k, float & dst) { if (g.has(k)) dst = g.f32(k); }; @@ -307,6 +193,20 @@ bool load_config(const gguf_reader & g, VlaJepaModelArch & m, Config & cfg) { F(fk("vit_rope_theta"), m.vit_rope_base); F(fk("dit_ln_eps"), m.dit_ln_eps); F(fk("dit_norm_out_eps"), m.dit_norm_out_eps); if (g.has(fk("lm_rope_theta"))) m.lm_rope_base = (float) g.f64(fk("lm_rope_theta")); + // merge_block_coords only enumerates the patch grid exactly when the spatial + // merge divides it; otherwise it emits rows past the position table. + if (m.patch_size <= 0 || m.spatial_merge <= 0 || m.image_target_size % m.patch_size != 0 || + (m.image_target_size / m.patch_size) % m.spatial_merge != 0) { + std::fprintf(stderr, "vla(vla_jepa): image %lld / patch %lld / merge %lld do not divide evenly\n", + (long long) m.image_target_size, (long long) m.patch_size, (long long) m.spatial_merge); + return false; + } + // timesteps_proj always emits 256 floats into the time-projection input. + if (m.time_proj_dim != 256) { + std::fprintf(stderr, "vla(vla_jepa): time_proj_dim %lld, expected 256\n", (long long) m.time_proj_dim); + return false; + } + cfg = Config{}; cfg.n_img = (m.image_target_size / m.patch_size / m.spatial_merge) * (m.image_target_size / m.patch_size / m.spatial_merge); cfg.n_lang = 1024; cfg.n_state = 1; @@ -315,6 +215,9 @@ bool load_config(const gguf_reader & g, VlaJepaModelArch & m, Config & cfg) { cfg.hidden = m.lm_hidden; cfg.n_q_heads = m.n_q; cfg.n_kv_heads = m.n_kv; cfg.head_dim = m.lm_head_dim; cfg.n_layers = m.lm_layers; cfg.num_steps = (int) m.num_steps; cfg.rms_eps = m.lm_rms_eps; cfg.rope_n_dims = (int) m.lm_head_dim; cfg.rope_mode = GGML_ROPE_TYPE_IMROPE; cfg.rope_freq_base = m.lm_rope_base; + // Raw output: this arch expects the client to apply the dataset statistics + // (see the --stats-json flag in eval/client). + cfg.denormalized = false; cfg.norm_eps = 1e-8f; return true; } @@ -349,19 +252,10 @@ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, (long long) m->action_horizon, (long long) m->action_dim, (long long) m->state_dim, (long long) m->num_future, (long long) m->num_steps, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); -#ifdef GGML_USE_CUDA - m->backend = ggml_backend_cuda_init(0); - if (m->backend) { m->is_cuda = true; m->is_gpu = true; std::printf("vla(vla_jepa): backend = CUDA (device 0)\n"); } - else std::fprintf(stderr, "vla(vla_jepa): ggml_backend_cuda_init failed; falling back to CPU\n"); -#elif defined(GGML_USE_METAL) - m->backend = ggml_backend_metal_init(); - if (m->backend) { m->is_gpu = true; std::printf("vla(vla_jepa): backend = Metal\n"); } -#endif - if (!m->backend) { - m->backend = ggml_backend_cpu_init(); - if (!m->backend) { std::fprintf(stderr, "vla(vla_jepa): ggml_backend_cpu_init failed\n"); return nullptr; } - ggml_backend_cpu_set_n_threads(m->backend, m->n_threads); - std::printf("vla(vla_jepa): backend = CPU (%d threads)\n", m->n_threads); + { + const Backend b = backend_init("vla(vla_jepa)", m->n_threads); + if (!b.handle) { return nullptr; } + m->backend = b.handle; } ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; @@ -534,8 +428,7 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { } if (inj_patches.empty() && !in.images) { std::fprintf(stderr, "vla(vla_jepa): n_images=%d but the images pointer is null\n", in.n_images); return {}; } - ggml_init_params vp = { (size_t) 512 * 1024 * 1024, nullptr, true }; - ggml_context * VC = ggml_init(vp); + ggml_context * VC = vision_scratch.reset((size_t) 512 * 1024 * 1024); if (!VC) { std::fprintf(stderr, "vla(vla_jepa): ggml_init(vision ctx) failed\n"); return {}; } ggml_tensor * t_patches = ggml_new_tensor_2d(VC, GGML_TYPE_F32, vit_patch_flat, n_patches); ggml_set_input(t_patches); ggml_tensor * t_pos = ggml_new_tensor_2d(VC, GGML_TYPE_F32, vit_hidden, n_patches); ggml_set_input(t_pos); @@ -556,8 +449,7 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { ggml_cgraph * vg = ggml_new_graph_custom(VC, 16384, false); ggml_build_forward_expand(vg, vit_embeds); for (int j = 0; j < 3; ++j) ggml_build_forward_expand(vg, ds_out[j]); - ggml_gallocr_t vga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!vga || !ggml_gallocr_alloc_graph(vga, vg)) { std::fprintf(stderr, "vla(vla_jepa): vision gallocr alloc failed\n"); if (vga) ggml_gallocr_free(vga); ggml_free(VC); return {}; } + if (!vision_scratch.alloc(backend, vg)) { std::fprintf(stderr, "vla(vla_jepa): vision gallocr alloc failed\n"); return {}; } const auto tv0 = std::chrono::steady_clock::now(); std::vector patches; bool vok = true; @@ -565,7 +457,7 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { if (!inj_patches.empty()) { ggml_backend_tensor_set(t_patches, inj_patches.data() + v * n_patches * vit_patch_flat, 0, ggml_nbytes(t_patches)); } else { - if (!preprocess_image_patches(in.images[v], side, ps, temporal_patch, c_grow, c_gcol, patches)) { vok = false; break; } + if (!preprocess_image_patches("vla_jepa", in.images[v], side, ps, temporal_patch, c_grow, c_gcol, patches)) { vok = false; break; } ggml_backend_tensor_set(t_patches, patches.data(), 0, ggml_nbytes(t_patches)); } ggml_backend_tensor_set(t_pos, c_pos_interp.data(), 0, ggml_nbytes(t_pos)); @@ -577,7 +469,6 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { if (dump_prefix) { char nm[32]; std::snprintf(nm, sizeof(nm), "vit_view%lld", (long long) v); char path[1024]; std::snprintf(path, sizeof(path), "%s_%s_%lldx%lld.f32", dump_prefix, nm, (long long) H, (long long) K); FILE * fp = std::fopen(path, "wb"); if (fp) { std::fwrite(img_emb_host.data() + v * K * H, sizeof(float), (size_t) K * H, fp); std::fclose(fp); } } } stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now() - tv0).count(); - ggml_gallocr_free(vga); ggml_free(VC); if (!vok) return {}; const int64_t n_img = n_views * K; @@ -609,9 +500,9 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { std::vector> ds_pad(3); for (int j = 0; j < 3; ++j) { ds_pad[j].assign((size_t) SEQ * H, 0.0f); for (int64_t k = 0; k < n_img; ++k) std::memcpy(ds_pad[j].data() + (size_t) image_pos_idx[k] * H, ds_host[j].data() + (size_t) k * H, H * sizeof(float)); } - ggml_init_params cp = { (size_t) 512 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(cp); - if (!C) { std::fprintf(stderr, "vla(vla_jepa): ggml_init(LM ctx) failed\n"); return {}; } + const LmKey lkey{ SEQ, num_future }; + const bool lm_built = lm_graph.ensure(backend, lkey, (size_t) 512 * 1024 * 1024, + [&](ggml_context * C, LmIO & gio) -> ggml_cgraph * { ggml_tensor * t_embeds = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); ggml_set_input(t_embeds); ggml_tensor * t_pos2 = ggml_new_tensor_1d(C, GGML_TYPE_I32, 4 * SEQ); ggml_set_input(t_pos2); ggml_tensor * t_lmmask = ggml_new_tensor_2d(C, GGML_TYPE_F32, SEQ, SEQ); ggml_set_input(t_lmmask); @@ -627,10 +518,22 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { ggml_set_output(eagle); ggml_tensor * conditioning = ggml_get_rows(C, eagle, t_emb_idx); ggml_set_output(conditioning); + gio.t_embeds=t_embeds; gio.t_pos2=t_pos2; gio.t_lmmask=t_lmmask; gio.t_emb_idx=t_emb_idx; + gio.t_ds[0]=t_ds[0]; gio.t_ds[1]=t_ds[1]; gio.t_ds[2]=t_ds[2]; + gio.eagle=eagle; gio.conditioning=conditioning; + ggml_cgraph * lg = ggml_new_graph_custom(C, 32768, false); ggml_build_forward_expand(lg, conditioning); - ggml_gallocr_t lga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!lga || !ggml_gallocr_alloc_graph(lga, lg)) { std::fprintf(stderr, "vla(vla_jepa): LM gallocr alloc failed\n"); if (lga) ggml_gallocr_free(lga); ggml_free(C); return {}; } + return lg; + }); + if (!lm_built) { std::fprintf(stderr, "vla(vla_jepa): LM graph build failed\n"); return {}; } + + LmIO & gio = lm_graph.io(); + ggml_cgraph * lg = lm_graph.graph(); + ggml_tensor * t_embeds = gio.t_embeds, * t_pos2 = gio.t_pos2, * t_lmmask = gio.t_lmmask; + ggml_tensor * t_emb_idx = gio.t_emb_idx; + ggml_tensor * t_ds[3] = { gio.t_ds[0], gio.t_ds[1], gio.t_ds[2] }; + ggml_tensor * eagle = gio.eagle, * conditioning = gio.conditioning; ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); @@ -659,22 +562,24 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { std::memcpy(pp.data() + (size_t) 3 * SEQ, pp.data(), (size_t) SEQ * sizeof(int32_t)); ggml_backend_tensor_set(t_pos2, pp.data(), 0, ggml_nbytes(t_pos2)); } - if (c_mask_seq != SEQ) { c_mask.assign((size_t) SEQ * SEQ, 0.0f); const float NEG = -std::numeric_limits::infinity(); for (int64_t q = 0; q < SEQ; ++q) for (int64_t kv = 0; kv < SEQ; ++kv) c_mask[q * SEQ + kv] = (kv <= q) ? 0.0f : NEG; c_mask_seq = SEQ; } + if (c_mask_seq != SEQ) { build_causal_mask(SEQ, c_mask); c_mask_seq = SEQ; } ggml_backend_tensor_set(t_lmmask, c_mask.data(), 0, ggml_nbytes(t_lmmask)); ggml_backend_tensor_set(t_emb_idx, emb_pos_idx.data(), 0, ggml_nbytes(t_emb_idx)); for (int j = 0; j < 3; ++j) ggml_backend_tensor_set(t_ds[j], ds_pad[j].data(), 0, ggml_nbytes(t_ds[j])); const auto tp0 = std::chrono::steady_clock::now(); - if (ggml_backend_graph_compute(backend, lg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(vla_jepa): LM compute failed\n"); ggml_gallocr_free(lga); ggml_free(C); return {}; } + if (ggml_backend_graph_compute(backend, lg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(vla_jepa): LM compute failed\n"); return {}; } stats.ms_prefill = std::chrono::duration(std::chrono::steady_clock::now() - tp0).count(); if (dump_prefix) { dump_t("eagle", eagle); dump_t("conditioning", conditioning); } ggml_backend_tensor_get(conditioning, cond_host.data(), 0, cond_host.size() * sizeof(float)); - ggml_gallocr_free(lga); ggml_free(C); } - ggml_init_params hp = { (size_t) 256 * 1024 * 1024, nullptr, true }; - ggml_context * C = ggml_init(hp); - if (!C) { std::fprintf(stderr, "vla(vla_jepa): ggml_init(head ctx) failed\n"); return {}; } + // Dumping adds graph outputs, so it always rebuilds. + std::vector step_seq, step_pred, step_vel, step_act; + if (dump_prefix) head_graph.release(); + const HeadKey hkey{ num_steps }; + const bool head_built = head_graph.ensure(backend, hkey, (size_t) 256 * 1024 * 1024, + [&](ggml_context * C, HeadIO & gio) -> ggml_cgraph * { ggml_tensor * t_cond = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, num_future); ggml_set_input(t_cond); ggml_tensor * t_state = ggml_new_tensor_2d(C, GGML_TYPE_F32, state_dim, 1); ggml_set_input(t_state); ggml_tensor * t_x0 = ggml_new_tensor_2d(C, GGML_TYPE_F32, AD, AH); ggml_set_input(t_x0); @@ -684,7 +589,8 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { ggml_tensor * state_features = ggml_add(C, ggml_mul_mat(C, se_l2W, ggml_relu(C, ggml_add(C, ggml_mul_mat(C, se_l1W, t_state), se_l1b))), se_l2b); ggml_tensor * future = future_tokens; const float dt = 1.0f / (float) num_steps; - std::vector step_seq(num_steps), step_pred(num_steps), step_vel(num_steps), step_act(num_steps); + step_seq.assign(num_steps, nullptr); step_pred.assign(num_steps, nullptr); + step_vel.assign(num_steps, nullptr); step_act.assign(num_steps, nullptr); ggml_tensor * actions = t_x0; for (int64_t s = 0; s < num_steps; ++s) { @@ -719,11 +625,20 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { if (dump_prefix) { ggml_set_output(step_seq[s]); ggml_set_output(step_pred[s]); ggml_set_output(step_vel[s]); ggml_set_output(step_act[s]); } } ggml_set_output(actions); + gio.t_cond=t_cond; gio.t_state=t_state; gio.t_x0=t_x0; gio.actions=actions; + gio.t_tau=t_tau; gio.t_tproj=t_tproj; + ggml_cgraph * hg = ggml_new_graph_custom(C, 65536, false); ggml_build_forward_expand(hg, actions); if (dump_prefix) for (int64_t s = 0; s < num_steps; ++s) { ggml_build_forward_expand(hg, step_seq[s]); ggml_build_forward_expand(hg, step_pred[s]); ggml_build_forward_expand(hg, step_vel[s]); ggml_build_forward_expand(hg, step_act[s]); } - ggml_gallocr_t hga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!hga || !ggml_gallocr_alloc_graph(hga, hg)) { std::fprintf(stderr, "vla(vla_jepa): head gallocr alloc failed\n"); if (hga) ggml_gallocr_free(hga); ggml_free(C); return {}; } + return hg; + }); + if (!head_built) { std::fprintf(stderr, "vla(vla_jepa): head graph build failed\n"); return {}; } + + HeadIO & hio = head_graph.io(); + ggml_cgraph * hg = head_graph.graph(); + ggml_tensor * t_cond = hio.t_cond, * t_state = hio.t_state, * t_x0 = hio.t_x0, * actions = hio.actions; + std::vector & t_tau = hio.t_tau; std::vector & t_tproj = hio.t_tproj; ggml_backend_tensor_set(t_cond, cond_host.data(), 0, ggml_nbytes(t_cond)); { std::vector st(state_dim, 0.0f); for (int64_t i = 0; i < state_dim; ++i) st[i] = in.state ? in.state[i] : 0.0f; ggml_backend_tensor_set(t_state, st.data(), 0, ggml_nbytes(t_state)); } @@ -731,7 +646,7 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { for (int64_t s = 0; s < num_steps; ++s) { ggml_backend_tensor_set(t_tau[s], c_tau[(size_t) s].data(), 0, ggml_nbytes(t_tau[s])); ggml_backend_tensor_set(t_tproj[s], c_tproj[(size_t) s].data(), 0, ggml_nbytes(t_tproj[s])); } const auto td0 = std::chrono::steady_clock::now(); - if (ggml_backend_graph_compute(backend, hg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(vla_jepa): head compute failed\n"); ggml_gallocr_free(hga); ggml_free(C); return {}; } + if (ggml_backend_graph_compute(backend, hg) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(vla_jepa): head compute failed\n"); return {}; } stats.ms_denoise = std::chrono::duration(std::chrono::steady_clock::now() - td0).count(); stats.ms_inference = stats.ms_prefill + stats.ms_denoise; @@ -745,7 +660,6 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { std::vector out((size_t) AH * AD); ggml_backend_tensor_get(actions, out.data(), 0, out.size() * sizeof(float)); - ggml_gallocr_free(hga); ggml_free(C); stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); return out; } diff --git a/src/serving/hf_fetch.h b/src/serving/hf_fetch.h new file mode 100644 index 0000000..cfe6c33 --- /dev/null +++ b/src/serving/hf_fetch.h @@ -0,0 +1,122 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Resolves -hf user/repo[:file.gguf] to a local path, shelling out to the hf +// CLI on a miss. + +#pragma once + +#include +#include +#include +#include + +namespace vla { + +// Repo ids reach a shell command, so reject anything outside this set. +inline bool hf_token_ok(const std::string & s, bool allow_slash) { + if (s.empty() || s.size() > 200) return false; + // A leading '/' would make fs::path join replace the cache root instead of + // extending it, putting the download anywhere on disk. + if (s.front() == '-' || s.front() == '/' || s.find("..") != std::string::npos) return false; + for (const char c : s) { + const bool ok = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c == '.' || c == '_' || c == '-' || + (allow_slash && c == '/'); + if (!ok) return false; + } + return true; +} + +inline std::string hf_cache_root() { + if (const char * e = std::getenv("VLA_CACHE"); e && *e) return e; + if (const char * h = std::getenv("HOME"); h && *h) return std::string(h) + "/.cache/vla"; + return ".vla-cache"; +} + +// Largest non-mmproj .gguf in dir, or "". +inline std::string hf_pick_gguf(const std::filesystem::path & dir, const std::string & want) { + namespace fs = std::filesystem; + std::error_code ec; + std::string best; + uintmax_t best_size = 0; + for (fs::recursive_directory_iterator it(dir, ec), end; it != end && !ec; it.increment(ec)) { + if (!it->is_regular_file(ec) || it->path().extension() != ".gguf") continue; + const std::string name = it->path().filename().string(); + if (!want.empty()) { + if (name == want) return it->path().string(); + continue; + } + if (name.rfind("mmproj", 0) == 0) continue; + const uintmax_t sz = it->file_size(ec); + if (ec) continue; + if (sz > best_size) { best_size = sz; best = it->path().string(); } + } + return best; +} + +// Returns "" and explains on stderr. +inline std::string hf_resolve(const std::string & spec) { + namespace fs = std::filesystem; + + const size_t colon = spec.find(':'); + const std::string repo = spec.substr(0, colon); + const std::string file = (colon == std::string::npos) ? "" : spec.substr(colon + 1); + + if (repo.find('/') == std::string::npos || !hf_token_ok(repo, true) || + (!file.empty() && !hf_token_ok(file, false))) { + std::fprintf(stderr, "vla: -hf expects user/repo[:file.gguf], got '%s'\n", spec.c_str()); + return ""; + } + + const fs::path dir = fs::path(hf_cache_root()) / repo; + + std::error_code ec; + if (fs::is_directory(dir, ec)) { + const std::string hit = hf_pick_gguf(dir, file); + if (!hit.empty()) return hit; + } + + fs::create_directories(dir, ec); + if (ec) { + std::fprintf(stderr, "vla: cannot create %s: %s\n", dir.string().c_str(), ec.message().c_str()); + return ""; + } + + // The cache root reaches the shell too, and it comes from VLA_CACHE or HOME. + // A single quote in either would close the quoting and run the rest. + if (dir.string().find('\'') != std::string::npos) { + std::fprintf(stderr, "vla: refusing a cache path containing a quote: %s\n", dir.string().c_str()); + return ""; + } + + std::string cmd = "hf download " + repo; + if (!file.empty()) cmd += " " + file; + cmd += " --local-dir '" + dir.string() + "'"; + std::fprintf(stderr, "vla: %s\n", cmd.c_str()); + + const int rc = std::system(cmd.c_str()); + if (rc != 0) { + std::fprintf(stderr, + "vla: download failed (exit %d). Install the CLI with\n" + " pip install -U \"huggingface_hub[cli]\"\n", rc); + return ""; + } + + const std::string hit = hf_pick_gguf(dir, file); + if (hit.empty()) std::fprintf(stderr, "vla: no .gguf under %s after download\n", dir.string().c_str()); + return hit; +} + +} // namespace vla diff --git a/src/serving/server.cpp b/src/serving/server.cpp index 1cecaad..63bcb23 100644 --- a/src/serving/server.cpp +++ b/src/serving/server.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include "model.h" +#include "serving/hf_fetch.h" #include "serving/vla.pb.h" #define STB_IMAGE_IMPLEMENTATION @@ -57,6 +58,20 @@ bool decode_image(const vla::Image & img, data.size()); return false; } + // Header first: stbi_load allocates 3*w*h before returning, so a small JPEG + // declaring huge dimensions would allocate gigabytes before any check. + if (!stbi_info_from_memory( + reinterpret_cast(data.data()), + static_cast(data.size()), &w, &h, &ch)) { + std::fprintf(stderr, "vla-server: stbi_info_from_memory failed: %s\n", + stbi_failure_reason()); + return false; + } + if (w <= 0 || h <= 0 || w > int(kMaxImageDim) || h > int(kMaxImageDim)) { + std::fprintf(stderr, "vla-server: JPEG dims %dx%d out of range (max %u)\n", + w, h, kMaxImageDim); + return false; + } unsigned char * px = stbi_load_from_memory( reinterpret_cast(data.data()), static_cast(data.size()), @@ -110,6 +125,14 @@ bool decode_image(const vla::Image & img, f32.resize(pixels); std::memcpy(f32.data(), img.data().data(), expected); + // State and noise are swept for NaN/Inf; pixels were not, so a bad pixel + // came back out as a robot action. + for (size_t i = 0; i < pixels; ++i) { + if (!std::isfinite(f32[i])) { + std::fprintf(stderr, "vla-server: F32_RGB_01 pixel %zu is not finite\n", i); + return false; + } + } view = { f32.data(), int(img.width()), int(img.height()), vla::PixelFormat::F32_RGB_01 }; return true; @@ -127,6 +150,21 @@ std::string make_error_response(uint64_t request_id, const std::string & msg) { return resp.SerializeAsString(); } +// Discard frames after the first. Must run to completion: a queued frame keeps +// REP in receive state and send throws EFSM. Stalled means the peer announced a +// frame it never sent, so no reply is possible until the rest arrives. +enum class Drain { Clean, Extra, Stalled }; + +Drain drain_extra_frames(zmq::socket_t & sock) { + Drain d = Drain::Clean; + while (sock.get(zmq::sockopt::rcvmore)) { + zmq::message_t junk; + if (!sock.recv(junk, zmq::recv_flags::none)) return Drain::Stalled; + d = Drain::Extra; + } + return d; +} + int find_non_finite(const float * data, int n) { for (int i = 0; i < n; ++i) { if (!std::isfinite(data[i])) return i; @@ -137,11 +175,13 @@ int find_non_finite(const float * data, int n) { void usage(const char * prog) { std::fprintf(stderr, "usage: %s [--bind ADDR] [--timing-detail none|phase] [--config PATH] " - "[] \n" + "[] ( | -hf user/repo[:file.gguf])\n" " vision-tower mmproj GGUF (SigLIP / PaliGemma /\n" " connector). Required for SmolVLA, π0, Evo-1, GR00T.\n" " Omit for BitVLA - its vision tower is baked into\n" " the combined ckpt GGUF.\n" + " -hf HuggingFace repo, user/repo[:file.gguf]; downloaded\n" + " on a miss and cached under $VLA_CACHE.\n" " SmolVLA .safetensors or .gguf, or any of the other\n" " supported architectures' .gguf; the architecture is\n" " auto-detected from the checkpoint.\n" @@ -166,6 +206,7 @@ int main(int argc, char ** argv) { std::string bind_addr = "tcp://*:5555"; std::string mmproj_path; std::string ckpt_path; + std::string hf_spec; std::string config_path; vla::TimingDetail timing_detail = vla::TimingDetail::NONE; @@ -174,6 +215,8 @@ int main(int argc, char ** argv) { std::string a = argv[i]; if (a == "--bind" && i + 1 < argc) { bind_addr = argv[++i]; + } else if (a == "-hf" && i + 1 < argc) { + hf_spec = argv[++i]; } else if (a == "--config" && i + 1 < argc) { config_path = argv[++i]; } else if (a == "--timing-detail" && i + 1 < argc) { @@ -192,14 +235,17 @@ int main(int argc, char ** argv) { positionals.push_back(std::move(a)); } } - if (positionals.size() == 1) { + if (!hf_spec.empty() && positionals.empty()) { + ckpt_path = vla::hf_resolve(hf_spec); + if (ckpt_path.empty()) return 1; + } else if (positionals.size() == 1) { ckpt_path = positionals[0]; } else if (positionals.size() == 2) { mmproj_path = positionals[0]; ckpt_path = positionals[1]; } else { std::fprintf(stderr, - "vla-server: expected 1 or 2 positional args " + "vla-server: expected -hf, or 1 or 2 positional args " "( for SmolVLA/π0/Evo-1/GR00T, " "or just for BitVLA), got %zu\n", positionals.size()); @@ -230,8 +276,12 @@ int main(int argc, char ** argv) { zmq::context_t zctx( 1); zmq::socket_t sock(zctx, zmq::socket_type::rep); sock.set(zmq::sockopt::linger, 0); - // cap inbound messages so one oversized request cannot exhaust memory. - sock.set(zmq::sockopt::maxmsgsize, int64_t(256) * 1024 * 1024); + // 64 MiB is above any real request (16 views of 512x512 F32 RGB is ~50 MiB) and + // low enough to bound protobuf's expansion during ParseFromArray. + sock.set(zmq::sockopt::maxmsgsize, int64_t(64) * 1024 * 1024); + // A peer that sends a frame with SNDMORE and then stalls would otherwise park + // this single-threaded loop in recv for good, starving every other client. + sock.set(zmq::sockopt::rcvtimeo, 5000); sock.bind(bind_addr); std::printf("vla-server: bound to %s. ready.\n", bind_addr.c_str()); @@ -285,6 +335,19 @@ int main(int argc, char ** argv) { continue; } + // Without this an unauthenticated client shuts the server down with one + // two-frame request: the reply fails and send_reply sets g_shutdown. + const Drain drained = drain_extra_frames(sock); + if (drained == Drain::Stalled) { + // Back to the poll rather than blocking here, so shutdown still works. + std::fprintf(stderr, "vla-server: peer stalled mid-request\n"); + continue; + } + if (drained == Drain::Extra) { + send_reply(make_error_response(0, "expected a single-frame request")); + continue; + } + vla::PredictRequest req; if (!req.ParseFromArray(req_msg.data(), static_cast(req_msg.size()))) { std::fprintf(stderr, "vla-server: PredictRequest parse failed (size=%zu)\n", diff --git a/src/serving/vla-bench.cpp b/src/serving/vla-bench.cpp new file mode 100644 index 0000000..1a070eb --- /dev/null +++ b/src/serving/vla-bench.cpp @@ -0,0 +1,165 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Latency for one checkpoint. Synthetic inputs: engine only, no task success. + +#include "model.h" +#include "serving/hf_fetch.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +void usage(const char * prog) { + std::fprintf(stderr, + "usage: %s (--ckpt c.gguf | -hf user/repo) [--mmproj m.gguf]\n" + " [--label name] [--images N] [--size N] [--tokens N]\n" + " [--extra-token ID] [--extra-count N] [--warmup N] [--reps N] [--markdown]\n" + " --label row label (default: the checkpoint filename)\n" + " --images camera views (default 1)\n" + " --size square input side in pixels (default 224)\n" + " --tokens language token count (default 16)\n" + " --extra-token token id appended --extra-count times (VLA-JEPA needs its\n" + " tokens)\n" + " --warmup untimed calls before measuring (default 3)\n" + " --reps timed calls (default 20)\n" + " --markdown print a markdown table row instead of a plain summary\n", + prog); +} + +// v must be sorted. +double percentile(const std::vector & v, double p) { + if (v.empty()) return 0.0; + const double idx = p * (double) (v.size() - 1); + const size_t lo = (size_t) std::floor(idx), hi = (size_t) std::ceil(idx); + return v[lo] + (v[hi] - v[lo]) * (idx - (double) lo); +} + +} // namespace + +int main(int argc, char ** argv) { + std::string ckpt, mmproj, hf, label; + int n_images = 1, side = 224, n_tokens = 16, warmup = 3, reps = 20; + int extra_token = -1, extra_count = 0; + bool markdown = false; + + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto need = [&](const char * name) -> const char * { + if (i + 1 >= argc) { std::fprintf(stderr, "vla-bench: %s needs a value\n", name); std::exit(1); } + return argv[++i]; + }; + if (a == "--ckpt") ckpt = need("--ckpt"); + else if (a == "-hf") hf = need("-hf"); + else if (a == "--mmproj") mmproj = need("--mmproj"); + else if (a == "--label") label = need("--label"); + else if (a == "--images") n_images = std::atoi(need("--images")); + else if (a == "--size") side = std::atoi(need("--size")); + else if (a == "--tokens") n_tokens = std::atoi(need("--tokens")); + else if (a == "--extra-token") extra_token = std::atoi(need("--extra-token")); + else if (a == "--extra-count") extra_count = std::atoi(need("--extra-count")); + else if (a == "--warmup") warmup = std::atoi(need("--warmup")); + else if (a == "--reps") reps = std::atoi(need("--reps")); + else if (a == "--markdown") markdown = true; + else if (a == "-h" || a == "--help") { usage(argv[0]); return 0; } + else { std::fprintf(stderr, "vla-bench: unknown argument %s\n", a.c_str()); usage(argv[0]); return 1; } + } + + if (!hf.empty()) { + if (!ckpt.empty()) { std::fprintf(stderr, "vla-bench: pass --ckpt or -hf, not both\n"); return 1; } + ckpt = vla::hf_resolve(hf); + if (ckpt.empty()) return 1; + } + if (ckpt.empty()) { usage(argv[0]); return 1; } + if (n_images < 1 || side < 16 || n_tokens < 1 || warmup < 0 || reps < 1) { + std::fprintf(stderr, "vla-bench: --images/--size/--tokens/--reps must be positive\n"); + return 1; + } + if (label.empty()) { + const size_t slash = ckpt.find_last_of('/'); + label = (slash == std::string::npos) ? ckpt : ckpt.substr(slash + 1); + } + + vla::Model * m = vla::model_load(mmproj, ckpt, ""); + if (!m) { std::fprintf(stderr, "vla-bench: model_load failed\n"); return 1; } + const vla::Config & cfg = vla::model_config(m); + + std::vector> pixels(n_images, std::vector((size_t) 3 * side * side)); + std::vector views(n_images); + for (int v = 0; v < n_images; ++v) { + for (int y = 0; y < side; ++y) + for (int x = 0; x < side; ++x) + for (int c = 0; c < 3; ++c) + pixels[v][((size_t) y * side + x) * 3 + c] = (uint8_t) ((x + 2 * y + 40 * c + 17 * v) & 0xFF); + views[v] = vla::ImageView{ pixels[v].data(), side, side, vla::PixelFormat::U8 }; + } + + std::vector lang((size_t) n_tokens); + for (int i = 0; i < n_tokens; ++i) lang[i] = 1 + (i % 100); + if (extra_token >= 0 && extra_count > 0) lang.insert(lang.end(), (size_t) extra_count, extra_token); + + std::vector state((size_t) cfg.max_state_dim, 0.0f); + for (int64_t i = 0; i < cfg.real_state_dim && i < cfg.max_state_dim; ++i) state[i] = 0.01f * (float) (i + 1); + + std::vector noise((size_t) cfg.max_action_dim * (size_t) cfg.n_suffix); + for (size_t i = 0; i < noise.size(); ++i) noise[i] = 0.001f * (float) ((i * 2654435761u) % 1000) - 0.5f; + + vla::Inputs in{}; + in.images = views.data(); + in.n_images = n_images; + in.lang_tokens = lang.data(); + in.n_lang = (int) lang.size(); + in.state = state.data(); + in.noise = noise.data(); + + for (int i = 0; i < warmup; ++i) { + if (vla::predict(m, in).empty()) { std::fprintf(stderr, "vla-bench: predict failed\n"); vla::model_free(m); return 1; } + } + + std::vector ms; + ms.reserve((size_t) reps); + double vision_sum = 0.0; + for (int i = 0; i < reps; ++i) { + const auto t0 = std::chrono::steady_clock::now(); + const std::vector out = vla::predict(m, in); + const auto t1 = std::chrono::steady_clock::now(); + if (out.empty()) { std::fprintf(stderr, "vla-bench: predict failed at rep %d\n", i); vla::model_free(m); return 1; } + ms.push_back(std::chrono::duration(t1 - t0).count()); + vision_sum += vla::last_stats(m).ms_vision; + } + + std::sort(ms.begin(), ms.end()); + const double lo = ms.front(); + const double p50 = percentile(ms, 0.50); + const double p90 = percentile(ms, 0.90); + const double vision = vision_sum / (double) reps; + + if (markdown) { + std::printf("| %s | %d | %d | %d | %.1f | %.1f | %.1f | %.1f |\n", + label.c_str(), n_images, side, n_tokens, lo, p50, p90, vision); + } else { + std::printf("%s: min %.1f ms p50 %.1f ms p90 %.1f ms vision %.1f ms (%d views, %dx%d, %d tokens, %d reps)\n", + label.c_str(), lo, p50, p90, vision, n_images, side, side, n_tokens, reps); + } + + vla::model_free(m); + return 0; +} diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp index 597a010..c2fa374 100644 --- a/src/serving/vla-cli.cpp +++ b/src/serving/vla-cli.cpp @@ -13,14 +13,16 @@ // limitations under the License. // One-shot action prediction from the command line. Loads a model, decodes an -// image plus an already-tokenized instruction, runs one predict(), and prints -// the action chunk. No server, no simulator. Tokenization stays in the Python -// client, so language is passed as token ids here. +// image plus an instruction, runs one predict(), and prints the action chunk. +// No server, no simulator. There is no tokenizer in the C++ core, so --text +// shells out to scripts/tokenize_prompt.py; --tokens takes ids directly. // // vla-cli [--mmproj m.gguf] --ckpt c.gguf --image img.jpg [--image img2.jpg] -// --tokens id,id,... [--state f,f,...] [--pretty] +// (--text "pick up the bowl" | --tokens id,id,...) [--state f,f,...] [--pretty] +#include "arch.h" #include "model.h" +#include "serving/hf_fetch.h" #define STB_IMAGE_IMPLEMENTATION #define STB_IMAGE_STATIC @@ -92,14 +94,84 @@ bool load_image(const char * path, std::vector & buf, int & w, int & h) return true; } +const char * arch_slug(Arch a) { + switch (a) { + case Arch::SMOLVLA: return "smolvla"; + case Arch::PI0: return "pi0"; + case Arch::PI05: return "pi05"; + case Arch::EVO1: return "evo1"; + case Arch::GR00T_N1_5: return "gr00t_n1_5"; + case Arch::GR00T_N1_6: return "gr00t_n1_6"; + case Arch::GR00T_N1_7: return "gr00t_n1_7"; + case Arch::BITVLA: return "bitvla"; + case Arch::VLA_ADAPTER: return "vla_adapter"; + case Arch::OPENVLA_OFT: return "openvla_oft"; + case Arch::VLA_JEPA: return "vla_jepa"; + } + return ""; +} + +// The instruction reaches a shell command, so keep it to plain prose. +bool text_ok(const std::string & s) { + if (s.empty() || s.size() > 512) return false; + for (const char c : s) { + const bool ok = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c == ' ' || c == '.' || c == ',' || + c == '-' || c == '_' || c == '\''; + if (!ok) return false; + } + return true; +} + +// Ask scripts/tokenize_prompt.py for the ids, using the tokenizer the arch was +// trained with. Returns "" and explains on stderr. +std::string tokenize_text(const std::string & ckpt, const std::string & text) { + Arch arch; + if (!detect_arch_from_ckpt(ckpt, &arch)) { + std::fprintf(stderr, "vla-cli: cannot detect the arch of %s for --text\n", ckpt.c_str()); + return ""; + } + if (!text_ok(text)) { + std::fprintf(stderr, "vla-cli: --text takes plain prose (letters, digits, space . , - _ ')\n"); + return ""; + } + std::string esc; + for (const char c : text) { if (c == '\'') esc += "'\\''"; else esc += c; } + // Env first so a packaged binary can point at its own copy of the script. + const char * env = std::getenv("VLA_TOKENIZE_SCRIPT"); + const std::string script = (env && *env) ? std::string(env) + : std::string(VLA_SOURCE_DIR) + "/scripts/tokenize_prompt.py"; + const char * py = std::getenv("VLA_PYTHON"); + const std::string interp = (py && *py) ? std::string(py) : std::string("python3"); + const std::string cmd = "'" + interp + "' '" + script + "' --arch " + arch_slug(arch) + + " --text '" + esc + "'"; + + FILE * fp = popen(cmd.c_str(), "r"); + if (!fp) { std::fprintf(stderr, "vla-cli: cannot run %s\n", cmd.c_str()); return ""; } + std::string out; + char buf[4096]; + while (std::fgets(buf, sizeof(buf), fp)) out += buf; + if (pclose(fp) != 0) { + std::fprintf(stderr, + "vla-cli: tokenizing failed. Install the client extras with\n" + " pip install -e \".[client]\"\n" + " (VLA_PYTHON selects a different interpreter)\n"); + return ""; + } + while (!out.empty() && (out.back() == '\n' || out.back() == '\r')) out.pop_back(); + return out; +} + void usage(const char * prog) { std::fprintf(stderr, - "usage: %s [--mmproj m.gguf] --ckpt c.gguf --image img.jpg [--image ...]\n" - " --tokens id,id,... [--state f,f,...] [--pretty]\n" + "usage: %s [--mmproj m.gguf] (--ckpt c.gguf | -hf user/repo) --image img.jpg [--image ...]\n" + " (--text \"...\" | --tokens id,id,...) [--state f,f,...] [--pretty]\n" " --mmproj vision-tower GGUF (SmolVLA/pi0/pi0.5); omit for baked-vision archs\n" " --ckpt model checkpoint GGUF\n" + " -hf HuggingFace repo, user/repo[:file.gguf], cached under $VLA_CACHE\n" " --image image file, repeat for multi-view (decoded via stb_image)\n" - " --tokens language token ids, comma-separated (tokenize in the client)\n" + " --text instruction, tokenized by scripts/tokenize_prompt.py (needs transformers)\n" + " --tokens language token ids, comma-separated, if you tokenized already\n" " --state proprioception floats, comma-separated (default zeros)\n" " --pretty print one action row (max_action_dim values) per line\n", prog); @@ -108,7 +180,7 @@ void usage(const char * prog) { } // namespace int main(int argc, char ** argv) { - std::string mmproj, ckpt, tokens_s, state_s; + std::string mmproj, ckpt, hf, tokens_s, state_s, text_s; std::vector image_paths; bool pretty = false; @@ -120,14 +192,27 @@ int main(int argc, char ** argv) { }; if (a == "--mmproj") mmproj = need("--mmproj"); else if (a == "--ckpt") ckpt = need("--ckpt"); + else if (a == "-hf") hf = need("-hf"); else if (a == "--image") image_paths.push_back(need("--image")); else if (a == "--tokens") tokens_s = need("--tokens"); + else if (a == "--text") text_s = need("--text"); else if (a == "--state") state_s = need("--state"); else if (a == "--pretty") pretty = true; else if (a == "-h" || a == "--help") { usage(argv[0]); return 0; } else { std::fprintf(stderr, "vla-cli: unknown argument %s\n", a.c_str()); usage(argv[0]); return 1; } } - if (ckpt.empty() || image_paths.empty() || tokens_s.empty()) { usage(argv[0]); return 1; } + if (!hf.empty()) { + if (!ckpt.empty()) { std::fprintf(stderr, "vla-cli: pass --ckpt or -hf, not both\n"); return 1; } + ckpt = vla::hf_resolve(hf); + if (ckpt.empty()) return 1; + } + if (ckpt.empty() || image_paths.empty() || (tokens_s.empty() && text_s.empty())) { usage(argv[0]); return 1; } + if (!tokens_s.empty() && !text_s.empty()) { std::fprintf(stderr, "vla-cli: pass --text or --tokens, not both\n"); return 1; } + if (!text_s.empty()) { + tokens_s = tokenize_text(ckpt, text_s); + if (tokens_s.empty()) return 1; + std::fprintf(stderr, "vla-cli: --text tokenized to %s\n", tokens_s.c_str()); + } // Validate the cheap args before loading the model. std::vector lang; diff --git a/src/serving/vlm-server.cpp b/src/serving/vlm-server.cpp index 3fa3e5d..2f8f75f 100644 --- a/src/serving/vlm-server.cpp +++ b/src/serving/vlm-server.cpp @@ -15,11 +15,20 @@ #include "vlm/engine.h" #include "serving/vlm.pb.h" +// stbi_info_from_memory only, to preflight JPEG dimensions before mtmd decodes. +#define STB_IMAGE_IMPLEMENTATION +#define STB_IMAGE_STATIC +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-function" +#include "stb_image.h" +#pragma GCC diagnostic pop + #include #include #include #include +#include #include #include #include @@ -101,8 +110,11 @@ int main(int argc, char ** argv) { zmq::context_t zctx( 1); zmq::socket_t sock(zctx, zmq::socket_type::router); sock.set(zmq::sockopt::linger, 0); - // cap inbound messages so one oversized request cannot exhaust memory. - sock.set(zmq::sockopt::maxmsgsize, int64_t(256) * 1024 * 1024); + // Per frame only; the recv loop caps the multipart total. + sock.set(zmq::sockopt::maxmsgsize, int64_t(64) * 1024 * 1024); + // A peer that sends a frame with SNDMORE and then stalls would otherwise park + // this single-threaded loop in recv for good, starving every other client. + sock.set(zmq::sockopt::rcvtimeo, 5000); sock.bind(bind_addr); std::printf("vlm-server: bound to %s. ready.\n", bind_addr.c_str()); @@ -132,9 +144,15 @@ int main(int argc, char ** argv) { } if (!(poll[0].revents & ZMQ_POLLIN)) continue; + // maxmsgsize bounds each frame but not how many, so a peer could stream + // sub-limit frames until memory runs out. + constexpr size_t kMaxEnvFrames = 8; + constexpr size_t kMaxEnvBytes = 64 * 1024; + std::vector env; std::string payload; - bool recv_ok = true, have_payload = false; + size_t env_bytes = 0; + bool recv_ok = true, have_payload = false, env_overflow = false; for (;;) { zmq::message_t part; try { @@ -146,13 +164,23 @@ int main(int argc, char ** argv) { recv_ok = false; break; } if (sock.get(zmq::sockopt::rcvmore)) { - env.emplace_back(static_cast(part.data()), part.size()); + env_bytes += part.size(); + if (env.size() >= kMaxEnvFrames || env_bytes > kMaxEnvBytes) { + // Keep draining so the socket stays sane, but stop accumulating. + env_overflow = true; + } else { + env.emplace_back(static_cast(part.data()), part.size()); + } } else { payload.assign(static_cast(part.data()), part.size()); have_payload = true; break; } } + if (env_overflow) { + std::fprintf(stderr, "vlm-server: oversized ROUTER envelope; request dropped\n"); + continue; + } if (!recv_ok || !have_payload || env.empty()) continue; auto send_reply = [&](const std::string & body) { @@ -181,6 +209,20 @@ int main(int argc, char ** argv) { send_reply(make_error_stream(rid, "ChatRequest has no messages")); continue; } + // One 60 MiB payload of tiny messages would cost template formatting and + // tokenization far beyond anything n_ctx could consume. + constexpr int kMaxMessages = 512; + constexpr size_t kMaxTextBytes = 4u * 1024 * 1024; + if (req.messages_size() > kMaxMessages) { + send_reply(make_error_stream(rid, "too many messages (max 512)")); + continue; + } + size_t text_bytes = 0; + for (const auto & m : req.messages()) text_bytes += m.content().size(); + if (text_bytes > kMaxTextBytes) { + send_reply(make_error_stream(rid, "message text too large (max 4 MiB)")); + continue; + } if (req.images_size() > 16) { send_reply(make_error_stream(rid, "too many image views (max 16)")); continue; @@ -194,6 +236,18 @@ int main(int argc, char ** argv) { vlm::Image out; if (im.encoding() == vlm_chat::Image::JPEG) { const auto & d = im.data(); + // Header first: mtmd decodes with no dimension guard. + int jw = 0, jh = 0, jc = 0; + if (d.size() > size_t(INT_MAX) || + !stbi_info_from_memory(reinterpret_cast(d.data()), + static_cast(d.size()), &jw, &jh, &jc) || + jw <= 0 || jh <= 0 || + jw > int(kMaxImageDim) || jh > int(kMaxImageDim)) { + char buf[96]; std::snprintf(buf, sizeof(buf), + "image[%d] JPEG dims %dx%d rejected (max %u)", v, jw, jh, kMaxImageDim); + send_reply(make_error_stream(rid, buf)); + decode_ok = false; break; + } if (!engine.decode_image_buf( reinterpret_cast(d.data()), d.size(), out)) { char buf[64]; std::snprintf(buf, sizeof(buf), "image[%d] JPEG decode failed", v); diff --git a/src/vla_c_api.cpp b/src/vla_c_api.cpp new file mode 100644 index 0000000..74c10bc --- /dev/null +++ b/src/vla_c_api.cpp @@ -0,0 +1,177 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Type conversion at the C boundary, and a barrier for C++ exceptions. + +#include "vla.h" + +#include "model.h" + +#include +#include +#include +#include + +namespace { + +// vla_model is opaque to C, so it can just wrap the C++ handle. +struct vla_model_impl { + vla::Model * m = nullptr; +}; + +vla::PixelFormat to_pixel_format(int32_t f) { + return f == VLA_PIXEL_F32_RGB_01 ? vla::PixelFormat::F32_RGB_01 + : vla::PixelFormat::U8; +} + +} // namespace + +struct vla_model : vla_model_impl {}; + +extern "C" { + +int32_t vla_abi_version(void) { return VLA_ABI_VERSION; } + +vla_model * vla_model_load(const char * mmproj_path, const char * ckpt_path, + const char * config_path) { + if (!ckpt_path) return nullptr; + try { + vla::Model * m = vla::model_load(mmproj_path ? mmproj_path : "", + ckpt_path, + config_path ? config_path : ""); + if (!m) return nullptr; + // Owns the engine until the handle exists, so a throwing new does not + // strand the whole model. + std::unique_ptr guard(m, vla::model_free); + auto * h = new vla_model(); + h->m = guard.release(); + return h; + } catch (...) { + return nullptr; + } +} + +void vla_model_free(vla_model * h) { + if (!h) return; + vla::model_free(h->m); + delete h; +} + +int32_t vla_model_config(const vla_model * h, vla_config * out) { + if (!h || !h->m || !out) return VLA_ERR_ARG; + try { + const vla::Config & c = vla::model_config(h->m); + *out = vla_config{}; + out->n_img = c.n_img; + out->n_lang = c.n_lang; + out->n_state = c.n_state; + out->n_prefix = c.n_prefix; + out->n_suffix = c.n_suffix; + out->n_full = c.n_full; + out->hidden = c.hidden; + out->expert_h = c.expert_h; + out->intermediate = c.intermediate; + out->expert_inter = c.expert_inter; + out->n_q_heads = c.n_q_heads; + out->n_kv_heads = c.n_kv_heads; + out->head_dim = c.head_dim; + out->q_full_dim = c.q_full_dim; + out->kv_full_dim = c.kv_full_dim; + out->n_layers = c.n_layers; + out->self_attn_every_n = c.self_attn_every_n; + out->max_state_dim = c.max_state_dim; + out->max_action_dim = c.max_action_dim; + out->real_state_dim = c.real_state_dim; + out->real_action_dim = c.real_action_dim; + out->norm_eps = c.norm_eps; + out->min_period = c.min_period; + out->max_period = c.max_period; + out->num_steps = c.num_steps; + out->rms_eps = c.rms_eps; + out->rope_n_dims = c.rope_n_dims; + out->rope_mode = c.rope_mode; + out->rope_freq_base = c.rope_freq_base; + out->denormalized = c.denormalized ? 1 : 0; + return VLA_OK; + } catch (...) { + return VLA_ERR_EXCEPTION; + } +} + +int32_t vla_predict(vla_model * h, const vla_inputs * in, + float ** out_actions, int64_t * out_n) { + if (!h || !h->m || !in || !out_actions || !out_n) return VLA_ERR_ARG; + *out_actions = nullptr; + *out_n = 0; + if (in->n_images < 0 || in->n_lang < 0 || in->n_img_views < 0) return VLA_ERR_ARG; + if (in->n_images > 0 && !in->images) return VLA_ERR_ARG; + if (in->n_lang > 0 && !in->lang_tokens) return VLA_ERR_ARG; + + try { + std::vector views((size_t) (in->n_images > 0 ? in->n_images : 0)); + for (size_t i = 0; i < views.size(); ++i) { + views[i] = vla::ImageView{ in->images[i].data, + in->images[i].w, + in->images[i].h, + to_pixel_format(in->images[i].format) }; + } + + vla::Inputs ci{}; + ci.images = views.empty() ? nullptr : views.data(); + ci.n_images = (int) views.size(); + ci.precomputed_img_emb = in->precomputed_img_emb; + ci.n_img_views = in->n_img_views; + ci.lang_tokens = in->lang_tokens; + ci.n_lang = in->n_lang; + ci.state = in->state; + ci.noise = in->noise; + ci.attention_mask = in->attention_mask; + ci.attention_mask_n = in->attention_mask_n; + ci.timing_detail = in->timing_detail == VLA_TIMING_PHASE + ? vla::TimingDetail::PHASE + : vla::TimingDetail::NONE; + + const std::vector act = vla::predict(h->m, ci); + if (act.empty()) return VLA_ERR_PREDICT; + + // malloc pairs with vla_free_actions, which callers may replace. + float * buf = (float *) std::malloc(act.size() * sizeof(float)); + if (!buf) return VLA_ERR_EXCEPTION; + std::memcpy(buf, act.data(), act.size() * sizeof(float)); + *out_actions = buf; + *out_n = (int64_t) act.size(); + return VLA_OK; + } catch (...) { + return VLA_ERR_EXCEPTION; + } +} + +void vla_free_actions(float * actions) { std::free(actions); } + +int32_t vla_last_stats(const vla_model * h, vla_stats * out) { + if (!h || !h->m || !out) return VLA_ERR_ARG; + try { + const vla::Stats & s = vla::last_stats(h->m); + out->ms_total = s.ms_total; + out->ms_vision = s.ms_vision; + out->ms_inference = s.ms_inference; + out->ms_prefill = s.ms_prefill; + out->ms_denoise = s.ms_denoise; + return VLA_OK; + } catch (...) { + return VLA_ERR_EXCEPTION; + } +} + +} // extern "C" diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3e20804..7a227ab 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -9,3 +9,33 @@ add_executable(test_vision_common test_vision_common.cpp) target_include_directories(test_vision_common PRIVATE ${CMAKE_SOURCE_DIR}/src) target_compile_options(test_vision_common PRIVATE -Wall -Wextra) add_test(NAME vision_common COMMAND test_vision_common) + +add_executable(test_rope_conventions test_rope_conventions.cpp) +target_include_directories(test_rope_conventions PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_compile_options(test_rope_conventions PRIVATE -Wall -Wextra) +add_test(NAME rope_conventions COMMAND test_rope_conventions) + +# Built as C so the public header stays C-clean. +add_executable(test_c_api test_c_api.c) +target_link_libraries(test_c_api PRIVATE vla) +target_compile_options(test_c_api PRIVATE -Wall -Wextra) +add_test(NAME c_api COMMAND test_c_api) + +add_executable(test_dit_common test_dit_common.cpp) +target_include_directories(test_dit_common PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(test_dit_common PRIVATE ggml) +target_compile_options(test_dit_common PRIVATE -Wall -Wextra) +add_test(NAME dit_common COMMAND test_dit_common) + +add_executable(test_qwen3vl_vit test_qwen3vl_vit.cpp) +target_include_directories(test_qwen3vl_vit PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(test_qwen3vl_vit PRIVATE ggml) +target_compile_options(test_qwen3vl_vit PRIVATE -Wall -Wextra) +add_test(NAME qwen3vl_vit COMMAND test_qwen3vl_vit) + +# Links vla_core: it calls the real config_is_sane rather than a copy. +add_executable(test_config_guard test_config_guard.cpp) +target_include_directories(test_config_guard PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(test_config_guard PRIVATE vla_core) +target_compile_options(test_config_guard PRIVATE -Wall -Wextra) +add_test(NAME config_guard COMMAND test_config_guard) diff --git a/tests/predict_check.cpp b/tests/predict_check.cpp index 1d71659..f27f7f0 100644 --- a/tests/predict_check.cpp +++ b/tests/predict_check.cpp @@ -18,13 +18,15 @@ // others read it, so fixed noise is reproducible for all archs). // // predict_check [mmproj.gguf] [n_images] -// env: VLA_IMG_SIZE (square input, default 224), VLA_BENCH_ITERS (>0 = time it) +// env: VLA_IMG_SIZE (square input, default 224), VLA_BENCH_ITERS (>0 = time it), +// VLA_TIMING=phase, VLA_EXTRA_TOKEN / VLA_EXTRA_COUNT #include "model.h" #include #include #include +#include #include #include @@ -64,7 +66,12 @@ int main(int argc, char** argv) { views[v] = ImageView{ imgbuf[v].data(), W, H, PixelFormat::U8 }; } + // VLA-JEPA needs its tokens; the others ignore the extras. std::vector lang = {1, 100, 200, 300, 400, 2}; + if (const char* tok = std::getenv("VLA_EXTRA_TOKEN")) { + const char* cnt = std::getenv("VLA_EXTRA_COUNT"); + lang.insert(lang.end(), (size_t)(cnt ? std::atoi(cnt) : 1), (int32_t)std::atoi(tok)); + } std::vector state((size_t)cfg.max_state_dim, 0.0f); for (int i = 0; i < (int)cfg.real_state_dim; ++i) state[i] = 0.01f * (float)(i + 1); @@ -80,7 +87,8 @@ int main(int argc, char** argv) { in.n_lang = (int)lang.size(); in.state = state.data(); in.noise = noise.data(); - in.timing_detail = TimingDetail::NONE; + const char* td = std::getenv("VLA_TIMING"); + in.timing_detail = (td && std::string(td) == "phase") ? TimingDetail::PHASE : TimingDetail::NONE; std::vector act = predict(m, in); std::printf("action_len=%zu\n", act.size()); diff --git a/tests/py/test_bindings.py b/tests/py/test_bindings.py new file mode 100644 index 0000000..c44fa62 --- /dev/null +++ b/tests/py/test_bindings.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Checks the ctypes structs against include/vla.h. + +A field added on one side and not the other silently misreads every value after +it. Set VLA_LIBRARY to also load the library and check its ABI version. +""" + +import ctypes +import os +import re +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) +sys.path.insert(0, os.path.join(ROOT, "bindings", "python")) + +from vla_cpp import _ffi # noqa: E402 + +HEADER = os.path.join(ROOT, "include", "vla.h") + +C_TO_CTYPES = { + "int32_t": ctypes.c_int32, + "int64_t": ctypes.c_int64, + "float": ctypes.c_float, + "double": ctypes.c_double, +} + + +def header_struct_fields(text, name): + """Field (name, c_type) pairs of `typedef struct { ... } name;`.""" + # [^{}] so the body cannot swallow the structs in between. + m = re.search(r"typedef struct \{([^{}]*)\}\s*" + name + r"\s*;", text, re.S) + assert m, f"{name} not found in vla.h" + fields = [] + for line in m.group(1).splitlines(): + line = re.sub(r"/\*.*?\*/", "", line) + line = line.split("///")[0].split("//")[0].strip() + if not line or not line.endswith(";"): + continue + decl = line[:-1].strip() + parts = decl.split() + if len(parts) < 2: + continue + ctype, fname = parts[0], parts[-1].lstrip("*") + fields.append((fname, ctype)) + return fields + + +def check_struct(text, c_name, py_struct, skip_types=()): + hdr = header_struct_fields(text, c_name) + py = [(n, t) for n, t in py_struct._fields_] + assert len(hdr) == len(py), ( + f"{c_name}: vla.h has {len(hdr)} fields, python has {len(py)}\n" + f" header: {[n for n, _ in hdr]}\n python: {[n for n, _ in py]}" + ) + for (hn, ht), (pn, pt) in zip(hdr, py): + assert hn == pn, f"{c_name}: field order differs, {hn!r} vs {pn!r}" + if ht in skip_types: + continue + want = C_TO_CTYPES.get(ht) + if want is not None: + assert pt is want, f"{c_name}.{hn}: header {ht}, python {pt}" + print(f" {c_name}: {len(hdr)} fields match") + + +def main(): + text = open(HEADER).read() + + m = re.search(r"#define VLA_ABI_VERSION\s+(\d+)", text) + assert m, "VLA_ABI_VERSION not found" + assert int(m.group(1)) == _ffi.ABI_VERSION, ( + f"vla.h ABI {m.group(1)}, _ffi.ABI_VERSION {_ffi.ABI_VERSION}") + print(f" ABI version {_ffi.ABI_VERSION} matches") + + check_struct(text, "vla_config", _ffi.Config) + check_struct(text, "vla_stats", _ffi.Stats) + # Pointer fields differ in spelling between the two; only order is checked. + check_struct(text, "vla_image", _ffi.Image, skip_types=("void",)) + check_struct(text, "vla_inputs", _ffi.Inputs, + skip_types=("vla_image", "float", "int32_t")) + + for name in ("VLA_OK", "VLA_ERR_ARG", "VLA_ERR_PREDICT", "VLA_ERR_EXCEPTION"): + assert name in text, f"{name} missing from vla.h" + assert _ffi.OK == 0 and _ffi.ERR_ARG == -1 + assert _ffi.ERR_PREDICT == -2 and _ffi.ERR_EXCEPTION == -3 + print(" status codes match") + + if os.environ.get("VLA_LIBRARY"): + lib = _ffi.load_library() + assert lib.vla_abi_version() == _ffi.ABI_VERSION + print(" libvla loads and reports a matching ABI") + + print("bindings: ok") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_c_api.c b/tests/test_c_api.c new file mode 100644 index 0000000..7371e4a --- /dev/null +++ b/tests/test_c_api.c @@ -0,0 +1,113 @@ +/* Copyright 2026 VinRobotics + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* Compiled as C, not C++, so it fails if the header ever stops being C-clean. + * Argument checks run with no model. Set VLA_TEST_GGUF to also run a real + * predict and compare against vla-cli on the same inputs. */ + +#include "vla.h" + +#include +#include +#include + +static int failures = 0; + +static void check(int cond, const char * what) { + if (!cond) { printf("FAIL: %s\n", what); failures++; } +} + +int main(void) { + check(vla_abi_version() == VLA_ABI_VERSION, "abi version matches header"); + + /* Null handling must not crash. */ + vla_model_free(NULL); + vla_free_actions(NULL); + check(vla_model_load(NULL, NULL, NULL) == NULL, "load(NULL) returns NULL"); + + vla_config cfg; + check(vla_model_config(NULL, &cfg) == VLA_ERR_ARG, "config(NULL) is ERR_ARG"); + vla_stats st; + check(vla_last_stats(NULL, &st) == VLA_ERR_ARG, "stats(NULL) is ERR_ARG"); + + float * act = NULL; + int64_t n = 0; + vla_inputs in; + memset(&in, 0, sizeof(in)); + check(vla_predict(NULL, &in, &act, &n) == VLA_ERR_ARG, "predict(NULL) is ERR_ARG"); + + const char * gguf = getenv("VLA_TEST_GGUF"); + if (!gguf) { + printf("c api: ok (argument checks only; set VLA_TEST_GGUF for a real run)\n"); + return failures ? 1 : 0; + } + + vla_model * m = vla_model_load(getenv("VLA_TEST_MMPROJ"), gguf, NULL); + check(m != NULL, "load real checkpoint"); + if (!m) return 1; + + check(vla_model_config(m, &cfg) == VLA_OK, "config on a loaded model"); + printf("c api: max_action_dim=%lld n_suffix=%lld denormalized=%d\n", + (long long) cfg.max_action_dim, (long long) cfg.n_suffix, cfg.denormalized); + + /* Same synthetic inputs predict_check uses, so the numbers are comparable. */ + int side = getenv("VLA_IMG_SIZE") ? atoi(getenv("VLA_IMG_SIZE")) : 224; + unsigned char * px = (unsigned char *) malloc((size_t) 3 * side * side); + check(px != NULL, "pixel buffer"); + if (!px) return 1; + for (int y = 0; y < side; ++y) + for (int x = 0; x < side; ++x) + for (int c = 0; c < 3; ++c) + px[((size_t) y * side + x) * 3 + c] = (unsigned char) ((x + 2 * y + 40 * c) & 0xFF); + + vla_image img; + img.data = px; img.w = side; img.h = side; img.format = VLA_PIXEL_U8; + + int32_t lang[6] = { 1, 100, 200, 300, 400, 2 }; + float * state = (float *) calloc((size_t) (cfg.max_state_dim > 0 ? cfg.max_state_dim : 1), sizeof(float)); + for (int i = 0; i < (int) cfg.real_state_dim; ++i) state[i] = 0.01f * (float) (i + 1); + + size_t noise_n = (size_t) cfg.max_action_dim * (size_t) cfg.n_suffix; + float * noise = noise_n ? (float *) malloc(noise_n * sizeof(float)) : NULL; + for (size_t i = 0; i < noise_n; ++i) + noise[i] = 0.001f * (float) ((i * 2654435761u) % 1000) - 0.5f; + + memset(&in, 0, sizeof(in)); + in.images = &img; + in.n_images = 1; + in.lang_tokens = lang; + in.n_lang = 6; + in.state = state; + in.noise = noise; + + int32_t rc = vla_predict(m, &in, &act, &n); + check(rc == VLA_OK, "predict returns OK"); + check(act != NULL && n > 0, "predict returns a buffer"); + if (rc == VLA_OK) { + printf("action_len=%lld\n", (long long) n); + for (int64_t i = 0; i < n; ++i) printf("%.9g\n", (double) act[i]); + vla_free_actions(act); + } + + check(vla_last_stats(m, &st) == VLA_OK, "stats after predict"); + check(st.ms_total > 0.0f, "total time is positive"); + + free(noise); free(state); free(px); + vla_model_free(m); + + if (failures) { printf("c api: %d FAILURES\n", failures); return 1; } + printf("c api: ok\n"); + return 0; +} diff --git a/tests/test_config_guard.cpp b/tests/test_config_guard.cpp new file mode 100644 index 0000000..7b4072d --- /dev/null +++ b/tests/test_config_guard.cpp @@ -0,0 +1,64 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// predict() sizes host buffers from the max_* dims and loops to the real_* dims, +// so real > max writes out of bounds. SmolVLA shipped that loop unbounded. + +#include "model.h" + +#undef NDEBUG // keep assert() live even in Release builds +#include +#include + +// The real guard, not a copy: a reimplementation here could not fail on a +// regression in src/model.cpp. +static bool sane(const vla::Config & c) { return vla::config_is_sane(c); } + +int main() { + vla::Config c{}; + + // A realistic SmolVLA config passes. + c.max_state_dim = 32; c.real_state_dim = 8; + c.max_action_dim = 32; c.real_action_dim = 7; + assert(sane(c)); + + // Equal is fine: the loops are half-open. + c.real_state_dim = 32; c.real_action_dim = 32; + assert(sane(c)); + + // real > max is the heap-overflow shape. + c.real_state_dim = 33; + assert(!sane(c)); + c.real_state_dim = 8; + c.real_action_dim = 33; + assert(!sane(c)); + c.real_action_dim = 7; + assert(sane(c)); + + // Negative dims come from a garbage or hostile GGUF. + c.real_state_dim = -1; + assert(!sane(c)); + c.real_state_dim = 8; + c.max_action_dim = -1; + assert(!sane(c)); + c.max_action_dim = 32; + assert(sane(c)); + + // max == 0 means "arch does not use this dim"; do not reject those. + c.max_state_dim = 0; c.real_state_dim = 0; + assert(sane(c)); + + std::printf("config guard: ok\n"); + return 0; +} diff --git a/tests/test_dit_common.cpp b/tests/test_dit_common.cpp new file mode 100644 index 0000000..6387f37 --- /dev/null +++ b/tests/test_dit_common.cpp @@ -0,0 +1,89 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Pins the two DiT time embeddings shared by GR00T N1.5/N1.6/N1.7 and VLA-JEPA. +// Pure host functions, so they can be tested without a checkpoint, which the +// end-to-end sha oracle cannot do for three of the four archs. The trig orders +// are opposite and both match the reference. + +#include "models/dit_common.h" + +#undef NDEBUG // keep assert() live even in Release builds +#include +#include +#include +#include + +static bool close(float a, float b) { return std::fabs(a - b) <= 1e-6f * (1.0f + std::fabs(b)); } + +int main() { + // --- timesteps_proj: 256 wide, cos in [0,128), sin in [128,256). --- + { + std::vector out; + vla::timesteps_proj(0, out); + assert(out.size() == 256); + // bucket 0 -> emb == 0 for every i -> cos(0)=1, sin(0)=0. + for (int i = 0; i < 128; ++i) { assert(close(out[i], 1.0f)); assert(close(out[128 + i], 0.0f)); } + + vla::timesteps_proj(250, out); + const float lm = std::log(10000.0f); + for (int i : { 0, 1, 63, 127 }) { + const float emb = 250.0f * std::exp(-lm * (float) i / 127.0f); + assert(close(out[i], std::cos(emb))); // cos first + assert(close(out[128 + i], std::sin(emb))); // sin second + } + // The i=0 term has exp(0)=1, so it is the raw timestep. + assert(close(out[0], std::cos(250.0f))); + } + + // --- action_sinusoid: [T, dim], sin in the low half, cos in the high half. --- + { + const int64_t dim = 8, T = 3; + std::vector out; + vla::action_sinusoid(0, dim, T, out); + assert(out.size() == (size_t) (T * dim)); + for (int64_t tk = 0; tk < T; ++tk) + for (int64_t i = 0; i < dim / 2; ++i) { + assert(close(out[tk * dim + i], 0.0f)); // sin(0) + assert(close(out[tk * dim + dim / 2 + i], 1.0f)); // cos(0) + } + + vla::action_sinusoid(500, dim, T, out); + const float step = std::log(10000.0f) / (float) (dim / 2); + for (int64_t i = 0; i < dim / 2; ++i) { + const float emb = 500.0f * std::exp(-(float) i * step); + assert(close(out[i], std::sin(emb))); // sin first + assert(close(out[dim / 2 + i], std::cos(emb))); // cos second + } + // Every horizon step carries the same embedding. + for (int64_t tk = 1; tk < T; ++tk) + for (int64_t i = 0; i < dim; ++i) + assert(close(out[tk * dim + i], out[i])); + } + + // --- The orders are opposite; a consistency cleanup must fail here. --- + { + std::vector tp, as; + vla::timesteps_proj(7, tp); + vla::action_sinusoid(7, 256, 1, as); + // Same width and same bucket, but tp[0] is a cosine and as[0] is a sine. + assert(as.size() == tp.size()); + assert(!close(tp[0], as[0])); + assert(close(tp[0], std::cos(7.0f))); + assert(close(as[0], std::sin(7.0f))); + } + + std::printf("dit common: ok\n"); + return 0; +} diff --git a/tests/test_qwen3vl_vit.cpp b/tests/test_qwen3vl_vit.cpp new file mode 100644 index 0000000..1605480 --- /dev/null +++ b/tests/test_qwen3vl_vit.cpp @@ -0,0 +1,118 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Pins the Qwen3-VL patch geometry. predict_check covers the graph builders. + +#include "models/qwen3vl_vit.h" + +#include +#include +#include + +using namespace vla; + +static int fails = 0; +#define CHECK(c) do { if (!(c)) { std::printf("FAIL %s:%d %s\n", __FILE__, __LINE__, #c); ++fails; } } while (0) + +// 4x4 grid, 2x2 merge: patches group by block, not raster order. +static void test_merge_block_coords() { + std::vector row, col; + merge_block_coords(4, 4, 2, row, col); + CHECK(row.size() == 16 && col.size() == 16); + const int64_t want_r[16] = {0,0,1,1, 0,0,1,1, 2,2,3,3, 2,2,3,3}; + const int64_t want_c[16] = {0,1,0,1, 2,3,2,3, 0,1,0,1, 2,3,2,3}; + for (int i = 0; i < 16; ++i) { CHECK(row[i] == want_r[i]); CHECK(col[i] == want_c[i]); } +} + +// Second half of the table repeats the first. +static void test_vit_rope_tables() { + std::vector row = {0, 1}, col = {0, 2}; + std::vector c, s; + const int64_t hd = 8; + vit_rope_tables(row, col, hd, 10000.0, c, s); + CHECK((int64_t) c.size() == 2 * hd && (int64_t) s.size() == 2 * hd); + for (int64_t p = 0; p < 2; ++p) + for (int64_t i = 0; i < hd / 2; ++i) { + CHECK(c[p * hd + i] == c[p * hd + hd / 2 + i]); + CHECK(s[p * hd + i] == s[p * hd + hd / 2 + i]); + } + for (int64_t i = 0; i < hd; ++i) { CHECK(c[i] == 1.0f); CHECK(s[i] == 0.0f); } + CHECK(std::fabs(c[hd + 0] - std::cos(1.0f)) < 1e-6f); + CHECK(std::fabs(c[hd + hd / 4] - std::cos(2.0f)) < 1e-6f); +} + +// Same grid in and out is an identity resample. +static void test_interp_pos_embed_identity() { + const int64_t side = 2, hidden = 2; + std::vector table = {0,10, 1,11, 2,12, 3,13}; + std::vector row, col; + merge_block_coords(side, side, 1, row, col); + std::vector out; + interp_pos_embed(table, side, hidden, row, col, side, side, out); + CHECK((int64_t) out.size() == side * side * hidden); + for (size_t s = 0; s < row.size(); ++s) { + const int64_t src = row[s] * side + col[s]; + CHECK(std::fabs(out[s * hidden + 0] - table[src * hidden + 0]) < 1e-6f); + CHECK(std::fabs(out[s * hidden + 1] - table[src * hidden + 1]) < 1e-6f); + } +} + +// 2x2 to 3x3 puts the midpoint at the mean of the corners. +static void test_interp_pos_embed_bilinear() { + const int64_t side = 2, hidden = 1; + std::vector table = {0, 2, 4, 6}; + std::vector row, col; + merge_block_coords(3, 3, 1, row, col); + std::vector out; + interp_pos_embed(table, side, hidden, row, col, 3, 3, out); + for (size_t s = 0; s < row.size(); ++s) + if (row[s] == 1 && col[s] == 1) CHECK(std::fabs(out[s] - 3.0f) < 1e-6f); +} + +// Wrong-sized view must be rejected before any pixel is read. +static void test_preprocess_rejects_bad_view() { + std::vector px(3 * 4 * 4, 0); + std::vector row, col; + merge_block_coords(2, 2, 1, row, col); + std::vector out; + ImageView bad{px.data(), 3, 4, PixelFormat::U8}; + CHECK(!preprocess_image_patches("test", bad, 4, 2, 2, row, col, out)); + ImageView null{nullptr, 4, 4, PixelFormat::U8}; + CHECK(!preprocess_image_patches("test", null, 4, 2, 2, row, col, out)); +} + +// Every temporal slice repeats the same value. +static void test_preprocess_values() { + const int64_t side = 2, ps = 1, tps = 2; + std::vector px((size_t) 3 * side * side, 128); + std::vector row, col; + merge_block_coords(side, side, 1, row, col); + std::vector out; + ImageView v{px.data(), (int) side, (int) side, PixelFormat::U8}; + CHECK(preprocess_image_patches("test", v, side, ps, tps, row, col, out)); + const int64_t pf = 3 * tps * ps * ps; + CHECK((int64_t) out.size() == pf * (int64_t) row.size()); + for (float f : out) CHECK(std::fabs(f - (128.0f / 255.0f * 2.0f - 1.0f)) < 1e-6f); +} + +int main() { + test_merge_block_coords(); + test_vit_rope_tables(); + test_interp_pos_embed_identity(); + test_interp_pos_embed_bilinear(); + test_preprocess_rejects_bad_view(); + test_preprocess_values(); + if (fails == 0) std::printf("qwen3vl_vit: all checks passed\n"); + return fails == 0 ? 0 : 1; +} diff --git a/tests/test_rope_conventions.cpp b/tests/test_rope_conventions.cpp new file mode 100644 index 0000000..8d6ab02 --- /dev/null +++ b/tests/test_rope_conventions.cpp @@ -0,0 +1,101 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Pins the two rotary conventions so a cleanup cannot swap one for the other. +// +// VLA-Adapter's action head pairs an interleaved rotation with a half-split +// frequency table, so a rotation pair gets two angles. The reference does the same +// (action_heads.py:163 vs :137-140) and the weights were trained on it, so making +// it self-consistent would break the shipped checkpoints. + +#undef NDEBUG // keep assert() live even in Release builds +#include +#include +#include +#include + +namespace { + +// src/models/vla_adapter.cpp hrot(): out[2k] = -x[2k+1], out[2k+1] = x[2k]. +std::vector rotate_interleaved(const std::vector & x) { + const size_t hd = x.size(); + std::vector out(hd); + for (size_t k = 0; k < hd / 2; ++k) { + out[2 * k] = -x[2 * k + 1]; + out[2 * k + 1] = x[2 * k]; + } + return out; +} + +// HuggingFace rotate_half, the NeoX convention ggml_rope_ext implements. +std::vector rotate_half(const std::vector & x) { + const size_t hd = x.size(), half = hd / 2; + std::vector out(hd); + for (size_t i = 0; i < half; ++i) { + out[i] = -x[half + i]; + out[half + i] = x[i]; + } + return out; +} + +// src/models/vla_adapter.cpp fill_cs(): j = mi % half, inv = base^(-2j/HD). +double adapter_angle(size_t mi, size_t hd, double base, double t) { + const size_t j = mi % (hd / 2); + return t * (1.0 / std::pow(base, (2.0 * (double) j) / (double) hd)); +} + +} // namespace + +int main() { + const size_t hd = 8; + const double base = 10000.0; + const double t = 3.0; + + const std::vector x = { 1, 2, 3, 4, 5, 6, 7, 8 }; + + // 1. The two rotations differ, so a swap is observable. + const std::vector ri = rotate_interleaved(x); + const std::vector rh = rotate_half(x); + assert(ri != rh); + assert(ri[0] == -2.0f && ri[1] == 1.0f); // interleaved pairs (x0,x1) + assert(rh[0] == -5.0f && rh[4] == 1.0f); // half-split pairs (x0,x4) + + // 2. The adapter frequency table is half-split: index i and i+half share an + // angle. That is what makes it mismatch the interleaved rotation above. + for (size_t i = 0; i < hd / 2; ++i) { + assert(adapter_angle(i, hd, base, t) == adapter_angle(i + hd / 2, hd, base, t)); + } + + // 3. The mismatch: an interleaved pair (2k, 2k+1) does not share an angle + // under this table. Pinned on purpose, see the header. + bool any_pair_differs = false; + for (size_t k = 0; k < hd / 2; ++k) { + if (adapter_angle(2 * k, hd, base, t) != adapter_angle(2 * k + 1, hd, base, t)) { + any_pair_differs = true; + } + } + assert(any_pair_differs); + + // 4. An interleaved table (j = mi / 2) would make every pair agree. That is + // the fix that must not be applied. + for (size_t k = 0; k < hd / 2; ++k) { + const auto interleaved_angle = [&](size_t mi) { + return t * (1.0 / std::pow(base, (2.0 * (double) (mi / 2)) / (double) hd)); + }; + assert(interleaved_angle(2 * k) == interleaved_angle(2 * k + 1)); + } + + std::printf("rope conventions: ok\n"); + return 0; +}