diff --git a/CMakeLists.txt b/CMakeLists.txt index 54fcca6..89da1ee 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -36,10 +36,15 @@ if(GGML_CUDA) set(_vla_llama_patch PATCH_COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/scripts/patch_ggml_cuda_ext_hook.py ) endif() +# Overridable so a regression can be bisected against another tag in a separate +# build dir (-DVLA_LLAMA_TAG=b10326) without editing this file. The patch +# anchors in scripts/patch_ggml_cuda_ext_hook.py are checked against the default. +set(VLA_LLAMA_TAG "b10331" CACHE STRING "llama.cpp tag to fetch") + include(FetchContent) FetchContent_Declare(llama GIT_REPOSITORY https://github.com/ggml-org/llama.cpp - GIT_TAG b10331 + GIT_TAG ${VLA_LLAMA_TAG} GIT_SHALLOW TRUE ${_vla_llama_patch} ) @@ -64,6 +69,14 @@ vla_exclude_fetched_targets(${llama_SOURCE_DIR}) add_library(vla_core src/model.cpp + src/loader.cpp + src/options.cpp + src/modules/action_expert.cpp + src/modules/dit_head.cpp + src/modules/encoder.cpp + src/modules/prompt.cpp + src/modules/qwen3_lm.cpp + src/modules/siglip_vit.cpp src/models/smolvla.cpp src/models/pi0.cpp src/models/pi05.cpp @@ -114,6 +127,9 @@ if(GGML_CUDA) CUDA_SEPARABLE_COMPILATION ON POSITION_INDEPENDENT_CODE ON ) + target_include_directories(bitvla_cuda_kernels PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ) target_compile_features(bitvla_cuda_kernels PRIVATE cxx_std_17) target_compile_options(bitvla_cuda_kernels PRIVATE $<$:-O3 --use_fast_math -Xptxas=-O3> diff --git a/ci/lib/common.sh b/ci/lib/common.sh index f2bff2c..7d0aa46 100755 --- a/ci/lib/common.sh +++ b/ci/lib/common.sh @@ -139,7 +139,7 @@ apply_gr00t_env() { local arch="$1" case "$arch" in gr00t_n1_5|gr00t_n1_6|gr00t_n1_7) - export VLA_GR00T_BF16_WEIGHTS="${VLA_GR00T_BF16_WEIGHTS:-1}" ;; + : ;; esac case "$arch" in gr00t_n1_5) export VLA_GR00T_EMBODIMENT="${VLA_GR00T_EMBODIMENT:-new_embodiment}" ;; diff --git a/docs/backend/sycl.md b/docs/backend/sycl.md index 9017f93..08dc291 100644 --- a/docs/backend/sycl.md +++ b/docs/backend/sycl.md @@ -217,6 +217,6 @@ 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 +`--weight-dtype bf16` (now the default) 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/eval/README.md b/eval/README.md index 0aba3d4..6d98282 100644 --- a/eval/README.md +++ b/eval/README.md @@ -56,7 +56,7 @@ Notes: accept the licence, or point `--tokenizer` at a local copy. - **GR00T** arches need `--stats-json /dataset_statistics.json` and an embodiment selected server-side via `VLA_GR00T_EMBODIMENT` (`new_embodiment` for N1.5, `libero_panda` for N1.6, - `libero_sim` for N1.7), plus `VLA_GR00T_BF16_WEIGHTS=1` to fit an 8 GB card. + `libero_sim` for N1.7). BF16 weights are the default, which is also what fits an 8 GB card. To sweep every model over `libero_object` tasks 0–9, use `eval/run_libero.sh -i `. @@ -66,7 +66,7 @@ So far only **GR00T-N1.6** is wired (the `gr00t-n1d6-bridge` checkpoint with the embodiment). Serve it, then drive from the SimplerEnv venv: ```bash -VLA_GR00T_BF16_WEIGHTS=1 VLA_GR00T_EMBODIMENT=oxe_widowx \ +VLA_GR00T_EMBODIMENT=oxe_widowx \ ./build/vla-server "$GR00T_N1D6_GGUF" eval/sim/simpler/simpler_uv/.venv/bin/python eval/client/run_simpler_client_direct.py \ diff --git a/eval/refactor_verify.sh b/eval/refactor_verify.sh new file mode 100755 index 0000000..0ec5c33 --- /dev/null +++ b/eval/refactor_verify.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +# Copyright 2026 VinRobotics - Apache-2.0 +# +# Bit-exactness and latency harness for the src/ layer/module/model refactor. +# +# eval/refactor_verify.sh actions only +# BENCH=20 eval/refactor_verify.sh actions + predict() timing +# +# Each arch runs twice: at its shipping defaults, and under the alternate +# precision. Both must stay byte-identical across a refactor, and neither may +# regress in latency. +# +# eval/refactor_verify.sh outputs/refactor/before +# ...change... +# cmake --build build -j"$(nproc)" --target vla_predict_check +# eval/refactor_verify.sh outputs/refactor/after +# diff -r outputs/refactor/before outputs/refactor/after +# +# Never rebuild while a sweep is running: relinking libvla_core.so under it +# makes every remaining arch fail to load. +# +# ARCHS=... restricts the sweep. The square input side is probed rather than +# hardcoded, because a tower fed the wrong side returns action_len=0 instead of +# failing, and a wrong side would silently "pass" a diff. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +BIN="${BIN:-${REPO_ROOT}/build/tests/vla_predict_check}" +HF="${HF:-/mnt/data/hf_data/vrfai}" +OUT="${1:-${REPO_ROOT}/outputs/refactor/baseline}" +SIDES="${SIDES:-224 256 448 512}" +BENCH="${BENCH:-0}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" + +# arch|ckpt|mmproj|n_images|env|alternate-config CLI flags. +# openvla_oft has no alternate: at f32 its weights need 30 GB. +MODELS=( + "smolvla|${HF}/smolvla-libero-gguf/smolvla-libero.gguf|${HF}/backup/mmproj-smolvla-libero.gguf|2||--flash-attn --mm-prec default" + "pi0|${HF}/pi0-libero-finetuned-v044-gguf/pi0-libero-finetuned-v044.gguf|${HF}/backup/mmproj-pi0-libero-finetuned-v044.gguf|2||--act-dtype bf16 --flash-attn" + "pi05|${HF}/pi05-libero-gguf/pi05-libero.gguf|${HF}/backup/mmproj-pi05-libero.gguf|2||--weight-dtype f32" + "evo1|${HF}/evo1-libero-gguf/evo1-libero.gguf||2||--act-dtype bf16 --flash-attn" + "gr00t_n1_5|${HF}/gr00tn1d5-libero-object-gguf/gr00tn1d5-libero-object.gguf||2||--weight-dtype f32" + "gr00t_n1_6|${HF}/gr00tn1d6-libero-gguf/gr00tn1d6-libero.gguf||2||--weight-dtype f32" + "gr00t_n1_7|${HF}/gr00tn1d7-libero-gguf/libero_object/gr00tn1d7-libero-object.gguf||2||--weight-dtype f32" + "bitvla|${HF}/bitvla-libero-gguf/libero_object/bitvla-libero-object.gguf||2||--weight-dtype bf16" + "vla_adapter|${HF}/vla-adapter-libero-object-gguf/libero_object/vla-adapter-libero-object.gguf||2||--weight-dtype f32" + "openvla_oft|${HF}/openvla-oft-libero-gguf/openvla-oft-libero.gguf||2||" + "vla_jepa|${HF}/vla-jepa-libero/vla-jepa.gguf||2|VLA_EXTRA_TOKEN=151697 VLA_EXTRA_COUNT=32|--weight-dtype f32" +) + +[[ -x "${BIN}" ]] || { echo "ERROR: missing ${BIN} (cmake -DVLA_BUILD_TESTS=ON)" >&2; exit 1; } +mkdir -p "${OUT}" + +run_one() { + local arch="$1" ckpt="$2" mmproj="$3" nimg="$4" env_str="$5" tag="$6" side="$7" cli="$8" + # shellcheck disable=SC2086 + env ${env_str} VLA_IMG_SIZE="${side}" VLA_BENCH_ITERS="${BENCH}" \ + "${BIN}" "${ckpt}" "${mmproj}" "${nimg}" ${cli} \ + > "${OUT}/${arch}${tag}.actions.txt" 2> "${OUT}/${arch}${tag}.log" +} + +fail=0 +for row in "${MODELS[@]}"; do + IFS='|' read -r arch ckpt mmproj nimg always fastest <<< "${row}" + + if [[ -n "${ARCHS:-}" && " ${ARCHS} " != *" ${arch} "* ]]; then + continue + fi + if [[ ! -e "${ckpt}" ]]; then + echo "[skip] ${arch}: no checkpoint at ${ckpt}" + continue + fi + + side="" + for s in ${SIDES}; do + if run_one "${arch}" "${ckpt}" "${mmproj}" "${nimg}" "${always}" "" "${s}" "" \ + && ! grep -q '^action_len=0$' "${OUT}/${arch}.actions.txt"; then + side="${s}" + echo "${s}" > "${OUT}/${arch}.side" + break + fi + done + if [[ -z "${side}" ]]; then + echo "[FAIL] ${arch}: no input side in '${SIDES}' produced a chunk; see ${OUT}/${arch}.log" >&2 + fail=1 + continue + fi + + line="[ok ] ${arch} side=${side}" + [[ "${BENCH}" -gt 0 ]] && line+=" default=$(grep -oP 'min=\K[0-9.]+' "${OUT}/${arch}.log" | head -1)ms" + + if [[ -n "${fastest}" ]]; then + if run_one "${arch}" "${ckpt}" "${mmproj}" "${nimg}" "${always}" ".alt" "${side}" "${fastest}" \ + && ! grep -q '^action_len=0$' "${OUT}/${arch}.alt.actions.txt"; then + line+=" alt=ok" + [[ "${BENCH}" -gt 0 ]] && line+=" $(grep -oP 'min=\K[0-9.]+' "${OUT}/${arch}.alt.log" | head -1)ms" + else + echo "[FAIL] ${arch}: alternate config produced no chunk; see ${OUT}/${arch}.fast.log" >&2 + fail=1 + fi + fi + echo "${line}" +done + +echo +echo "written to ${OUT}" +exit "${fail}" diff --git a/eval/run_libero.sh b/eval/run_libero.sh index 658fe9c..5c589ed 100644 --- a/eval/run_libero.sh +++ b/eval/run_libero.sh @@ -298,9 +298,10 @@ run_model() { client_extra+=(--stats-json "${stats_json}") fi + # BF16 weights are the shipping default now, and are what every published + # GR00T success rate was measured under; passed explicitly so the log records it. if [[ "${arch}" == gr00t_n1_5 || "${arch}" == gr00t_n1_6 || "${arch}" == gr00t_n1_7 ]]; then - export VLA_GR00T_BF16_WEIGHTS="${VLA_GR00T_BF16_WEIGHTS:-1}" - echo "[${arch}] VLA_GR00T_BF16_WEIGHTS=${VLA_GR00T_BF16_WEIGHTS}" + server_args+=(--weight-dtype "${WEIGHT_DTYPE:-bf16}") fi if [[ -n "${_USER_VLA_GR00T_EMBODIMENT}" ]]; then export VLA_GR00T_EMBODIMENT="${_USER_VLA_GR00T_EMBODIMENT}" diff --git a/eval/run_libero_client.sh b/eval/run_libero_client.sh index 140bda9..19cc7f8 100644 --- a/eval/run_libero_client.sh +++ b/eval/run_libero_client.sh @@ -142,21 +142,20 @@ EOF ;; gr00t_n1_5) cat <&2 return 1 fi - - export VLA_GR00T_BF16_WEIGHTS="${VLA_GR00T_BF16_WEIGHTS:-1}" + # bf16 weights are the shipping default; the env switch was retired. export VLA_GR00T_EMBODIMENT="${VLA_GR00T_EMBODIMENT:-${EMBODIMENT}}" - echo "[${arch}] VLA_GR00T_BF16_WEIGHTS=${VLA_GR00T_BF16_WEIGHTS} VLA_GR00T_EMBODIMENT=${VLA_GR00T_EMBODIMENT}" local log="${LOG_DIR}/${arch}.log" echo "====================" diff --git a/scripts/patch_ggml_cuda_ext_hook.py b/scripts/patch_ggml_cuda_ext_hook.py index dff4bd5..d02bcb4 100644 --- a/scripts/patch_ggml_cuda_ext_hook.py +++ b/scripts/patch_ggml_cuda_ext_hook.py @@ -30,13 +30,19 @@ What it changes (ggml/src/ggml-cuda/ggml-cuda.cu only) ------------------------------------------------------ - 1. An exported function pointer, null by default. - 2. One call to it at the top of ggml_cuda_compute_forward. Returning false - means "not mine", and ggml runs the op exactly as before. + 1. Two exported function pointers, null by default. + 2. One call to the first at the top of ggml_cuda_compute_forward. Returning + false means "not mine", and ggml runs the op exactly as before. 3. The RMS_NORM+MUL fusion check GGML_ASSERTs F32 rather than declining, so a BF16 rms_norm aborts the process before dispatch is ever reached. Those two asserts become a return, which is what the surrounding checks already do for every other unsupported type. + 4. One call to the second in the ADD/MUL fusion branch of ggml_cuda_try_fuse. + Fusion happens in ggml_backend_cuda_graph_compute, upstream of + ggml_cuda_compute_forward, so the hook in (2) never sees a fused node -- + and ggml_cuda_op_fused_binbcast_impl handles F32/F16 only and GGML_ABORTs + on BF16. Without this the choice is a crash or no fusion at all for BF16 + activations, and the unfused path costs ~18 ms/call on evo1. With the pointer left null this is a no-op, so an unpatched-but-hooked ggml behaves identically to a stock one. @@ -57,6 +63,11 @@ extern "C" { typedef bool (*ggml_cuda_ext_forward_t)(struct ggml_tensor * dst, void * stream); __attribute__((visibility("default"))) ggml_cuda_ext_forward_t ggml_cuda_ext_forward = nullptr; + +// Same contract for a fused ADD/MUL run: dst carries src[0] plus n_fuse addends +// in src[1..n_fuse], all sharing one layout, and dst->data is the final output. +typedef bool (*ggml_cuda_ext_fused_binbcast_t)(struct ggml_tensor * dst, int n_fuse, void * stream); +__attribute__((visibility("default"))) ggml_cuda_ext_fused_binbcast_t ggml_cuda_ext_fused_binbcast = nullptr; } static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { @@ -80,6 +91,25 @@ }""", ) +# The fused ADD/MUL run is assembled here and handed to a kernel that supports +# F32/F16 only. Offer it to the extension first; declining costs one null check. +FUSED_BINBCAST_GUARD = ( + """ if (node->op == GGML_OP_ADD) { + ggml_cuda_op_fused_add(*cuda_ctx, &fused_node, n_fuse); + } else { + ggml_cuda_op_fused_mul(*cuda_ctx, &fused_node, n_fuse); + }""", + """ // vla.cpp: CUDA extension hook - first refusal on the fused node. + if (!(ggml_cuda_ext_fused_binbcast && + ggml_cuda_ext_fused_binbcast(&fused_node, n_fuse, (void *) cuda_ctx->stream()))) { + if (node->op == GGML_OP_ADD) { + ggml_cuda_op_fused_add(*cuda_ctx, &fused_node, n_fuse); + } else { + ggml_cuda_op_fused_mul(*cuda_ctx, &fused_node, n_fuse); + } + }""", +) + def main(): src = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else ".").resolve() @@ -91,7 +121,7 @@ def main(): if MARKER in text: return # idempotent: re-configure over an already-patched tree - for old, new in (HOOK_DECL, FUSION_GUARD): + for old, new in (HOOK_DECL, FUSION_GUARD, FUSED_BINBCAST_GUARD): n = text.count(old) if n != 1: raise SystemExit( diff --git a/scripts/restyle.py b/scripts/restyle.py new file mode 100644 index 0000000..06b8eb6 --- /dev/null +++ b/scripts/restyle.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics - Apache-2.0 +# +# Applies the structural half of the house style to C/C++ sources: one +# statement per line, no brace-wrapped one-liner blocks, no statement stapled +# to the end of an if/else. +# +# Whitespace only. Strings, character literals, comments and preprocessor lines +# are copied through untouched, and the splitter never runs inside parentheses, +# so a for-header's semicolons stay put. +# +# scripts/restyle.py src/models/vla_adapter.cpp ... + +import re +import sys + +def spans_to_skip(line): + """Character ranges of string/char literals and comments in `line`.""" + out, i, n = [], 0, len(line) + while i < n: + c = line[i] + if c in '"\'': + j = i + 1 + while j < n: + if line[j] == '\\': + j += 2 + continue + if line[j] == c: + break + j += 1 + out.append((i, min(j + 1, n))) + i = j + 1 + elif line.startswith('//', i): + out.append((i, n)) + break + elif line.startswith('/*', i): + j = line.find('*/', i + 2) + j = n if j < 0 else j + 2 + out.append((i, j)) + i = j + else: + i += 1 + return out + +def code_positions(line): + """Indices of `line` that are real code.""" + skip = spans_to_skip(line) + return [i for i in range(len(line)) if not any(a <= i < b for a, b in skip)] + +def split_statements(body, indent): + """`a; b; c` -> one statement per line, ignoring semicolons inside parens.""" + parts, depth, cur = [], 0, '' + live = set(code_positions(body)) + for i, ch in enumerate(body): + if i in live: + if ch in '([': + depth += 1 + elif ch in ')]': + depth -= 1 + elif ch == ';' and depth == 0: + parts.append(cur.strip() + ';') + cur = '' + continue + cur += ch + if cur.strip(): + parts.append(cur.strip()) + return [indent + p for p in parts if p.strip(';').strip()] + +ONE_LINE_BLOCK = re.compile(r'^(?P\s*)(?P.*?\{)\s*(?P[^{}]*?;)\s*\}\s*$') +CONTROL_HEAD = re.compile(r'^(?P\s*)(?:\}\s*else\s+)?(?:if|for|while)\s*\(') + +def control_body(line): + """Split `if/for/while (...) stmt;` into (head, body); None if it is not one.""" + m = CONTROL_HEAD.match(line) + if not m: + return None + live = set(code_positions(line)) + i = line.index('(', m.end() - 1) + depth = 0 + for j in range(i, len(line)): + if j not in live: + continue + if line[j] == '(': + depth += 1 + elif line[j] == ')': + depth -= 1 + if depth == 0: + head, body = line[:j + 1], line[j + 1:].strip() + if body and body.endswith(';') and '{' not in body and ';' not in body[:-1]: + return head, body + return None + return None + +def restyle_line(line): + if line.lstrip().startswith('#'): + return [line] + if '"' in line or "'" in line: + # A literal containing a brace or semicolon would confuse the splitters. + lit = ''.join(line[a:b] for a, b in spans_to_skip(line)) + if any(c in lit for c in '{};'): + return [line] + + m = ONE_LINE_BLOCK.match(line) + if m and 'namespace' not in m.group('head'): + inner = m.group('indent') + ' ' + stmts = split_statements(m.group('body'), inner) + if len(stmts) >= 1: + return [m.group('indent') + m.group('head')] + stmts + [m.group('indent') + '}'] + + m = re.match(r'^(?P\s*)(?:\}\s*)?else\s+(?P[^{};]*?;)\s*$', line) + if m and not re.match(r'^\s*(?:\}\s*)?else\s+if\b', line): + head = line[:line.index('else') + 4] + return [head.rstrip(), m.group('indent') + ' ' + m.group('body').strip()] + + cb = control_body(line) + if cb: + head, body = cb + indent = re.match(r'\s*', line).group(0) + return [head, indent + ' ' + body] + + return [line] + +def restyle(text): + total = 0 + for _ in range(6): + out, changed = [], 0 + for line in text.split('\n'): + new = restyle_line(line) + if new != [line]: + changed += 1 + out.extend(new) + text = '\n'.join(out) + total += changed + if not changed: + break + return text, total + +if __name__ == '__main__': + total = 0 + for path in sys.argv[1:]: + src = open(path).read() + dst, n = restyle(src) + if n: + open(path, 'w').write(dst) + print(f'{path}: {n} lines expanded') + total += n + print(f'total {total}') diff --git a/scripts/tighten_ops.py b/scripts/tighten_ops.py new file mode 100644 index 0000000..2cb3cb7 --- /dev/null +++ b/scripts/tighten_ops.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics - Apache-2.0 +# +# Removes the spaces around binary operators, per the house style. Whitespace +# only, and deliberately conservative: it edits an operator only when both +# sides are unambiguously values. +# +# Left untouched on purpose: +# * & unless one side is a literal, a ')' / ']', or a '.'/'->' member, +# because `ggml_tensor * t` is a declaration, not a product +# < > template brackets and includes are indistinguishable here +# = `a = b` reads better spaced outside a for-header +# any operator whose tightening would merge two tokens (`a - -b`) +# +# Strings, character literals, comments and preprocessor lines are copied +# through untouched. + +import re +import sys + +sys.path.insert(0, __file__.rsplit('/', 1)[0]) +from restyle import spans_to_skip + +VALUE_END = re.compile(r'(?:[A-Za-z_]\w*|\d[\w.]*|[)\]])$') +VALUE_START = re.compile(r'^(?:[A-Za-z_]\w*|\d|[(\[])') +MEMBER = re.compile(r'(?:\.\s*[A-Za-z_]\w*|->\s*[A-Za-z_]\w*|\)|\]|\d)$') + +OPS = ['+', '-', '/', '%'] + +def tighten(line): + if line.lstrip().startswith('#'): + return line + + skip = spans_to_skip(line) + live = lambda i: not any(a <= i < b for a, b in skip) + + out = line + for _ in range(40): + changed = False + for m in re.finditer(r'(?<=\S) (' + '|'.join(re.escape(o) for o in OPS + ['*']) + r') (?=\S)', out): + i, op = m.start(), m.group(1) + if not live(i): + continue + left, right = out[:i], out[i + 3:] + if not VALUE_END.search(left) or not VALUE_START.match(right): + continue + # `a * b` is only a product when one side is plainly a value. + if op == '*' and not (MEMBER.search(left) or right[0].isdigit()): + continue + # never let the operator glue onto a neighbour + if left[-1] == op or right[0] in '=&|<>+-' or right[0] == op: + continue + out = left + op + right + skip = spans_to_skip(out) + changed = True + break + if not changed: + break + return out + +if __name__ == '__main__': + total = 0 + for path in sys.argv[1:]: + src = open(path).read() + dst = '\n'.join(tighten(l) for l in src.split('\n')) + n = sum(1 for a, b in zip(src.split('\n'), dst.split('\n')) if a != b) + if n: + open(path, 'w').write(dst) + total += n + print(f'tightened {total} lines') diff --git a/src/models/act_dtype.h b/src/act_dtype.h similarity index 98% rename from src/models/act_dtype.h rename to src/act_dtype.h index 7463676..d74cd84 100644 --- a/src/models/act_dtype.h +++ b/src/act_dtype.h @@ -62,7 +62,7 @@ inline ggml_tensor * as_type(ggml_context * C, ggml_tensor * t, ggml_type ty) { inline ggml_tensor * mul_mat_t(ggml_context * C, ggml_tensor * a, ggml_tensor * b, ggml_type type) { // ggml_can_mul_mat is internal to ggml; this is the same condition. GGML_ASSERT(a->ne[0] == b->ne[0]); - GGML_ASSERT(b->ne[2] % a->ne[2] == 0 && b->ne[3] % a->ne[3] == 0); + GGML_ASSERT(b->ne[2]%a->ne[2] == 0 && b->ne[3]%a->ne[3] == 0); GGML_ASSERT(!ggml_is_transposed(a)); const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; diff --git a/src/arch.h b/src/arch.h index e66fa33..45db734 100644 --- a/src/arch.h +++ b/src/arch.h @@ -26,6 +26,7 @@ #pragma once #include "model.h" +#include "options.h" #include #include @@ -43,7 +44,8 @@ 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; + 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(); @@ -93,9 +95,9 @@ class ModelArchBase { /** * @brief Run a full forward pass and return one chunk of normalised actions. - * @param in Vision + language + state inputs (see @ref Inputs). + * @param in Vision+language+state inputs (see @ref Inputs). * @return Flattened action chunk of length - * @c cfg.num_steps * cfg.real_action_dim. + * @c cfg.num_steps*cfg.real_action_dim. */ virtual std::vector predict(const Inputs& in) = 0; }; @@ -103,13 +105,14 @@ class ModelArchBase { /** * @brief Build a SmolVLA model from its mmproj and checkpoint GGUFs. * @param mmproj_path Path to the vision-tower GGUF. - * @param ckpt_path Path to the LM + action-expert GGUF. + * @param ckpt_path Path to the LM+action-expert GGUF. * @param config_path Optional JSON override; pass empty to use bundled config. * @return Owning pointer to the constructed model. */ std::unique_ptr smolvla_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a pi0 model from its mmproj and checkpoint GGUFs. @@ -117,7 +120,8 @@ std::unique_ptr smolvla_create(const std::string& mmproj_path, */ std::unique_ptr pi0_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a pi0.5 model from its mmproj and checkpoint GGUFs. @@ -125,7 +129,8 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, */ std::unique_ptr pi05_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build an Evo-1 model. Vision is baked into @p ckpt_path; pass @@ -134,7 +139,8 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, */ std::unique_ptr evo1_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a GR00T N1.5 model. Vision is baked into @p ckpt_path. @@ -142,7 +148,8 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, */ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a GR00T N1.6 model. Vision is baked into @p ckpt_path. @@ -150,7 +157,8 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, */ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a GR00T N1.7 model. Vision is baked into @p ckpt_path. @@ -158,7 +166,8 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, */ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a BitVLA model. Vision is baked into @p ckpt_path. @@ -166,7 +175,8 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, */ std::unique_ptr bitvla_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a VLA-Adapter model. Vision is baked into @p ckpt_path. @@ -174,7 +184,8 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, */ std::unique_ptr vla_adapter_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a OpenVLA-OFT model. Vision is baked into @p ckpt_path. @@ -182,7 +193,8 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path */ std::unique_ptr openvla_oft_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Build a VLA-JEPA model. Vision is baked into @p ckpt_path. @@ -190,7 +202,8 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path */ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path); + const std::string& config_path, + const Options& opts); /** * @brief Inspect a GGUF and identify the architecture tag. diff --git a/src/backend.h b/src/backend.h index 597275b..33377d4 100644 --- a/src/backend.h +++ b/src/backend.h @@ -57,7 +57,8 @@ namespace vla { 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; + if (getenv_s(&len, nullptr, 0, key) == 0 && len > 0) + return; _putenv_s(key, val); #else setenv(key, val, /*overwrite=*/0); @@ -79,7 +80,8 @@ struct Backend { /// silently read as device 0. inline int backend_device_index() { const char * e = std::getenv("VLA_DEVICE"); - if (!e || !*e) return 0; + if (!e || !*e) + return 0; char * end = nullptr; const long idx = std::strtol(e, &end, 10); if (*end != '\0' || idx < 0 || idx > 1024) { diff --git a/src/cuda/vla_cuda_bf16.cu b/src/cuda/vla_cuda_bf16.cu index 6ff6b2f..70a3d63 100644 --- a/src/cuda/vla_cuda_bf16.cu +++ b/src/cuda/vla_cuda_bf16.cu @@ -36,6 +36,7 @@ // never a reduction. #include "ggml.h" +#include "env_flag.h" #include #include @@ -48,14 +49,21 @@ extern "C" { typedef bool (*ggml_cuda_ext_forward_t)(struct ggml_tensor * dst, void * stream); extern ggml_cuda_ext_forward_t ggml_cuda_ext_forward; + +typedef bool (*ggml_cuda_ext_fused_binbcast_t)(struct ggml_tensor * dst, int n_fuse, void * stream); +extern ggml_cuda_ext_fused_binbcast_t ggml_cuda_ext_fused_binbcast; } namespace { constexpr int BLOCK = 256; -inline __device__ float bf2f(const __nv_bfloat16 v) { return __bfloat162float(v); } -inline __device__ __nv_bfloat16 f2bf(const float v) { return __float2bfloat16(v); } +inline __device__ float bf2f(const __nv_bfloat16 v) { + return __bfloat162float(v); +} +inline __device__ __nv_bfloat16 f2bf(const float v) { + return __float2bfloat16(v); +} // --------------------------------------------------------------------------- // elementwise binary: dst = op(src0, src1), src1 broadcast per ggml_can_repeat @@ -67,67 +75,382 @@ inline __device__ __nv_bfloat16 f2bf(const float v) { return __float2bfloat16(v) enum class BinOp { Add, Mul }; +inline __device__ float apply_bin(BinOp op, float a, float b) { + return op == BinOp::Add ? a+b : a * b; +} + +// Broadcast index along one dimension. ggml's rule is a modulo, but the only +// shapes that occur are "same extent" and "extent 1"; both operands are kernel +// arguments so the branch is uniform across the block and the 64-bit modulo is +// left for the general case that never fires in practice. +inline __device__ int64_t bcast_idx(int64_t i, int64_t ne_src, int64_t ne_dst) { + if (ne_src == ne_dst) + return i; + if (ne_src == 1) + return 0; + return i%ne_src; +} + +// Rows are addressed through blockIdx.y/z rather than recovered from a flat +// index. A flat grid-stride loop costs a 64-bit division and three modulos per +// element to rebuild (i0,i1,i2,i3), which is what made this kernel 12.8 us/launch +// against ggml's 5.7 for the same work -- and at ~3,700 elementwise nodes per +// evo1 call that difference was larger than everything BF16 activations saved. template -__global__ void k_bin_bcast_bf16( - const __nv_bfloat16 * __restrict__ src0, const S1 * __restrict__ src1, - __nv_bfloat16 * __restrict__ dst, +__global__ void k_bin_bcast_bf16_rows( + const __nv_bfloat16*__restrict__ src0, const S1*__restrict__ src1, + __nv_bfloat16*__restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13, + const int64_t s10, const int64_t s11, const int64_t s12, const int64_t s13, + const int64_t d0, const int64_t d1, const int64_t d2, const int64_t d3) { + const int64_t i1 = blockIdx.y; + const int64_t i23 = blockIdx.z; + const int64_t i2 = i23%ne2; // once per block, not per element + const int64_t i3 = i23/ne2; + + const int64_t j1 = bcast_idx(i1, ne11, ne1); + const int64_t j2 = bcast_idx(i2, ne12, ne2); + const int64_t j3 = bcast_idx(i3, ne13, ne3); + + const __nv_bfloat16*__restrict__ r0 = src0+i1*s01+i2*s02+i3*s03; + const S1 * __restrict__ r1 = src1+j1*s11+j2*s12+j3*s13; + __nv_bfloat16 * __restrict__ rd = dst + i1*d1 + i2*d2 + i3*d3; + + for (int64_t i0 = (int64_t) blockIdx.x*blockDim.x+threadIdx.x; i0 < ne0; + i0 += (int64_t) gridDim.x*blockDim.x) { + const float a = bf2f(r0[i0*s00]); + const float b = (float) r1[bcast_idx(i0, ne10, ne0)*s10]; + rd[i0*d0] = f2bf(apply_bin(op, a, b)); + } +} + +// Vectorized row kernel: 8 BF16 per thread (one uint4 load/store), for the +// shape that dominates -- dst/src0 contiguous along dim 0 and src1 a full-width +// row repeated over the other dims, i.e. every bias add and norm-weight mul. +// The arithmetic is per element in float exactly as the scalar path does it, so +// results are bit-identical; only the memory access widens. Worth doing because +// a graph-node-level profile puts these kernels at ~27 ms of evo1's 147, and +// their mean/median split (6.3 vs 2.2 us) says the big tensors carry the total. +template +__global__ void k_bin_bcast_bf16_vec8( + const __nv_bfloat16*__restrict__ src0, const S1*__restrict__ src1, + __nv_bfloat16*__restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne11, const int64_t ne12, const int64_t ne13, + const int64_t s11, const int64_t s12, const int64_t s13, + const int64_t d1, const int64_t d2, const int64_t d3) { + const int64_t i1 = blockIdx.y; + const int64_t i23 = blockIdx.z; + const int64_t i2 = i23%ne2; + const int64_t i3 = i23/ne2; + + const __nv_bfloat16*__restrict__ r0 = src0+i1*s01+i2*s02+i3*s03; + const S1 * __restrict__ r1 = src1+bcast_idx(i1, ne11, ne1)*s11 + + bcast_idx(i2, ne12, ne2)*s12 + + bcast_idx(i3, ne13, ne3)*s13; + __nv_bfloat16 * __restrict__ rd = dst + i1*d1 + i2*d2 + i3*d3; + + const int64_t nvec = ne0/8; + for (int64_t v = (int64_t) blockIdx.x*blockDim.x+threadIdx.x; v < nvec; + v += (int64_t) gridDim.x*blockDim.x) { + const int64_t i0 = v*8; + + uint4 a = *reinterpret_cast(r0+i0); + __nv_bfloat16*av = reinterpret_cast<__nv_bfloat16 *>(&a); + + #pragma unroll + for (int k=0; k<8; ++k) { + const float b = (float) r1[i0+k]; + av[k] = f2bf(apply_bin(op, bf2f(av[k]), b)); + } + *reinterpret_cast(rd+i0) = a; + } +} + +// Fallback for shapes the row grid cannot address (gridDim.y/z cap at 65535). +template +__global__ void k_bin_bcast_bf16_flat( + const __nv_bfloat16*__restrict__ src0, const S1*__restrict__ src1, + __nv_bfloat16*__restrict__ dst, const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, const int64_t s00, const int64_t s01, const int64_t s02, const int64_t s03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13, const int64_t s10, const int64_t s11, const int64_t s12, const int64_t s13, const int64_t d0, const int64_t d1, const int64_t d2, const int64_t d3) { const int64_t total = ne0*ne1*ne2*ne3; - for (int64_t idx = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; idx < total; + for (int64_t idx = (int64_t) blockIdx.x*blockDim.x+threadIdx.x; idx < total; idx += (int64_t) gridDim.x*blockDim.x) { const int64_t i0 = idx % ne0; - const int64_t i1 = (idx / ne0) % ne1; - const int64_t i2 = (idx / (ne0*ne1)) % ne2; - const int64_t i3 = idx / (ne0*ne1*ne2); + const int64_t i1 = (idx/ne0) % ne1; + const int64_t i2 = (idx/(ne0*ne1)) % ne2; + const int64_t i3 = idx/(ne0*ne1*ne2); - const float a = bf2f(src0[i0*s00 + i1*s01 + i2*s02 + i3*s03]); - const float b = (float) src1[(i0 % ne10)*s10 + (i1 % ne11)*s11 + - (i2 % ne12)*s12 + (i3 % ne13)*s13]; + const float a = bf2f(src0[i0*s00+i1*s01+i2*s02+i3*s03]); + const float b = (float) src1[(i0%ne10)*s10+(i1%ne11)*s11 + + (i2%ne12)*s12+(i3%ne13)*s13]; - dst[i0*d0 + i1*d1 + i2*d2 + i3*d3] = f2bf(op == BinOp::Add ? a + b : a * b); + dst[i0*d0+i1*d1+i2*d2+i3*d3] = f2bf(apply_bin(op, a, b)); } } // element strides (ggml stores byte strides) -inline int64_t es(const ggml_tensor * t, int i) { return t->nb[i] / ggml_type_size(t->type); } +inline int64_t es(const ggml_tensor * t, int i) { + return t->nb[i]/ggml_type_size(t->type); +} + +// Launch shape for the row-addressed kernels; ok=false means fall back to flat. +struct RowGrid { + dim3 grid; + dim3 block; + bool ok; +}; + +// VLA_BF16_FLAT=1 forces the flat kernel, for bisecting. The two paths do the +// same per-element arithmetic in the same order, so they are bit-identical -- +// which is what makes the A/B a correctness check rather than a smoke test. +inline bool force_flat() { + static const bool v = vla::env_flag("VLA_BF16_FLAT"); + return v; +} + +inline RowGrid row_grid(int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + RowGrid g{}; + const int64_t nz = ne2*ne3; + if (ne1 > 65535 || nz > 65535 || ne1 < 1 || nz < 1) { + g.ok = false; + return g; + } + + unsigned bx = 32; + while (bx < (unsigned) BLOCK && (int64_t) bx < ne0) + bx *= 2; + int64_t gx = (ne0+bx-1)/bx; + if (gx > 65535) + gx = 65535; + if (gx < 1) + gx = 1; + + g.block = dim3(bx, 1, 1); + g.grid = dim3((unsigned) gx, (unsigned) ne1, (unsigned) nz); + g.ok = true; + return g; +} template bool bin_bcast(ggml_tensor * dst, cudaStream_t stream) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - if (!src0 || !src1) return false; - if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) return false; - if (src1->type != GGML_TYPE_BF16 && src1->type != GGML_TYPE_F32) return false; - if (!ggml_are_same_shape(src0, dst)) return false; - if (!ggml_can_repeat(src1, src0)) return false; - - const int64_t total = ggml_nelements(dst); - const int64_t blocks = (total + BLOCK - 1) / BLOCK; - const int grid = (int) (blocks < 65535 ? blocks : 65535); + if (!src0 || !src1) + return false; + if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) + return false; + if (src1->type != GGML_TYPE_BF16 && src1->type != GGML_TYPE_F32) + return false; + if (!ggml_are_same_shape(src0, dst)) + return false; + if (!ggml_can_repeat(src1, src0)) + return false; + + // Only the unfused path honours VLA_BF16_FLAT: the fused path has no flat + // fallback, and declining there hands the run to ggml's aborting kernel. + const RowGrid g = force_flat() ? RowGrid{} + : row_grid(dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]); + + // 8-wide path: dim 0 contiguous on both sides and src1 a full-width row. + // Requires 16 B alignment for the uint4 accesses, which the ggml allocator + // gives at tensor start but not necessarily at a strided row offset. + const bool vec8_shape = + g.ok && dst->ne[0]%8 == 0 && + es(src0, 0) == 1 && es(dst, 0) == 1 && es(src1, 0) == 1 && + src1->ne[0] == dst->ne[0] && + es(src0, 1)%8 == 0 && es(dst, 1)%8 == 0 && es(src1, 1)%8 == 0 && + ((uintptr_t) src0->data%16) == 0 && ((uintptr_t) dst->data%16) == 0; + if (vec8_shape) { + const int64_t nvec = dst->ne[0]/8; + unsigned bx = 32; + while (bx < (unsigned) BLOCK && (int64_t) bx < nvec) + bx *= 2; + int64_t gx = (nvec+bx-1)/bx; + if (gx > 65535) + gx = 65535; + if (gx < 1) + gx = 1; + const dim3 vgrid((unsigned) gx, g.grid.y, g.grid.z); + const dim3 vblock(bx, 1, 1); + +#define VLA_LAUNCH_VEC8(TYPE) \ + k_bin_bcast_bf16_vec8<<>>( \ + (const __nv_bfloat16 *) src0->data, (const TYPE *) src1->data, \ + (__nv_bfloat16 *) dst->data, \ + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], \ + es(src0,1), es(src0,2), es(src0,3), \ + src1->ne[1], src1->ne[2], src1->ne[3], \ + es(src1,1), es(src1,2), es(src1,3), \ + es(dst,1), es(dst,2), es(dst,3)) + + if (src1->type == GGML_TYPE_BF16) { + VLA_LAUNCH_VEC8(__nv_bfloat16); + } + else { + VLA_LAUNCH_VEC8(float); + } +#undef VLA_LAUNCH_VEC8 + return true; + } + +#define VLA_LAUNCH_BIN(TYPE) \ + do { \ + if (g.ok) { \ + k_bin_bcast_bf16_rows<<>>( \ + (const __nv_bfloat16 *) src0->data, (const TYPE *) src1->data, \ + (__nv_bfloat16 *) dst->data, \ + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], \ + es(src0,0), es(src0,1), es(src0,2), es(src0,3), \ + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], \ + es(src1,0), es(src1,1), es(src1,2), es(src1,3), \ + es(dst,0), es(dst,1), es(dst,2), es(dst,3)); \ + } else { \ + const int64_t blocks = (ggml_nelements(dst)+BLOCK-1)/BLOCK; \ + const int flat = (int) (blocks < 65535 ? blocks : 65535); \ + k_bin_bcast_bf16_flat<<>>( \ + (const __nv_bfloat16 *) src0->data, (const TYPE *) src1->data, \ + (__nv_bfloat16 *) dst->data, \ + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], \ + es(src0,0), es(src0,1), es(src0,2), es(src0,3), \ + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], \ + es(src1,0), es(src1,1), es(src1,2), es(src1,3), \ + es(dst,0), es(dst,1), es(dst,2), es(dst,3)); \ + } \ + } while (0) + + if (src1->type == GGML_TYPE_BF16) { + VLA_LAUNCH_BIN(__nv_bfloat16); + } else { + VLA_LAUNCH_BIN(float); + } +#undef VLA_LAUNCH_BIN + return true; +} + +// --------------------------------------------------------------------------- +// fused elementwise binary: dst = (((src0 op src1) op src2) ... op src[n_fuse]) +// --------------------------------------------------------------------------- +// +// ggml fuses runs of up to 8 ADD or MUL nodes in ggml_cuda_try_fuse and hands +// the run over as one synthetic node: src[0] is the base, src[1..n_fuse] are the +// addends, dst->data is the last node's output. The fusion check upstream only +// admits a run when every addend has the same layout, which is why one stride +// set covers all of them. +// +// Accumulation is in float with a single conversion at the end, matching ggml's +// own k_bin_bcast: `float result = (float) src0[..]; result = bin_op(result, ..); +// dst[i0] = (dst_t) result;`. Doing it any other way would make a fused run +// disagree with the unfused one. + +constexpr int MAX_FUSE = 8; + +template +struct SrcPtrs { const S1*p[MAX_FUSE]; }; + +template +__global__ void k_fused_bin_bcast_bf16( + const __nv_bfloat16*__restrict__ src0, const SrcPtrs srcs, const int n_fuse, + __nv_bfloat16*__restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13, + const int64_t s10, const int64_t s11, const int64_t s12, const int64_t s13, + const int64_t d0, const int64_t d1, const int64_t d2, const int64_t d3) { + const int64_t i1 = blockIdx.y; + const int64_t i23 = blockIdx.z; + const int64_t i2 = i23%ne2; + const int64_t i3 = i23/ne2; + + const int64_t row1 = bcast_idx(i1, ne11, ne1)*s11 + + bcast_idx(i2, ne12, ne2)*s12 + + bcast_idx(i3, ne13, ne3)*s13; + + const __nv_bfloat16*__restrict__ r0 = src0+i1*s01+i2*s02+i3*s03; + __nv_bfloat16 * __restrict__ rd = dst + i1*d1 + i2*d2 + i3*d3; + + for (int64_t i0 = (int64_t) blockIdx.x*blockDim.x+threadIdx.x; i0 < ne0; + i0 += (int64_t) gridDim.x*blockDim.x) { + const int64_t j = row1+bcast_idx(i0, ne10, ne0)*s10; + + float acc = bf2f(r0[i0*s00]); + for (int k=0; k +bool fused_bin_bcast(ggml_tensor * dst, int n_fuse, cudaStream_t stream) { + if (n_fuse < 2 || n_fuse > MAX_FUSE) + return false; + + const ggml_tensor * src0 = dst->src[0]; + if (!src0) + return false; + if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) + return false; + if (!ggml_are_same_shape(src0, dst)) + return false; + + // src[1] fixes the layout and type every other addend must match; the + // upstream fusion check guarantees it, and this re-checks rather than + // trusting it, because getting it wrong reads out of bounds. + const ggml_tensor * src1 = dst->src[1]; + if (!src1) + return false; + if (src1->type != GGML_TYPE_BF16 && src1->type != GGML_TYPE_F32) + return false; + if (!ggml_can_repeat(src1, src0)) + return false; + + for (int k=1; ksrc[k+1]; + if (!s || s->type != src1->type) + return false; + if (!ggml_are_same_shape(s, src1)) + return false; + for (int d=0; dnb[d] != src1->nb[d]) + return false; + } + } + + // No flat fallback here: declining just sends the run to ggml's fused + // kernel, which aborts on BF16, so an unaddressable shape must decline + // before the hook claims it. + const RowGrid g = row_grid(dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]); + if (!g.ok) + return false; + +#define VLA_LAUNCH_FUSED(TYPE) \ + do { \ + SrcPtrs srcs{}; \ + for (int k=0; ksrc[k+1]->data; \ + k_fused_bin_bcast_bf16<<>>( \ + (const __nv_bfloat16 *) src0->data, srcs, n_fuse, \ + (__nv_bfloat16 *) dst->data, \ + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], \ + es(src0,0), es(src0,1), es(src0,2), es(src0,3), \ + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], \ + es(src1,0), es(src1,1), es(src1,2), es(src1,3), \ + es(dst,0), es(dst,1), es(dst,2), es(dst,3)); \ + } while (0) if (src1->type == GGML_TYPE_BF16) { - k_bin_bcast_bf16<<>>( - (const __nv_bfloat16 *) src0->data, (const __nv_bfloat16 *) src1->data, - (__nv_bfloat16 *) dst->data, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - es(src0,0), es(src0,1), es(src0,2), es(src0,3), - src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - es(src1,0), es(src1,1), es(src1,2), es(src1,3), - es(dst,0), es(dst,1), es(dst,2), es(dst,3)); + VLA_LAUNCH_FUSED(__nv_bfloat16); } else { - k_bin_bcast_bf16<<>>( - (const __nv_bfloat16 *) src0->data, (const float *) src1->data, - (__nv_bfloat16 *) dst->data, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - es(src0,0), es(src0,1), es(src0,2), es(src0,3), - src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - es(src1,0), es(src1,1), es(src1,2), es(src1,3), - es(dst,0), es(dst,1), es(dst,2), es(dst,3)); + VLA_LAUNCH_FUSED(float); } +#undef VLA_LAUNCH_FUSED return true; } @@ -139,18 +462,21 @@ enum class UnOp { Silu, Relu, Gelu, GeluErf }; template inline __device__ float apply_unary(const float x) { - if (op == UnOp::Silu) return x / (1.0f + expf(-x)); - if (op == UnOp::Relu) return x > 0.0f ? x : 0.0f; - if (op == UnOp::GeluErf) return 0.5f*x*(1.0f + erff(x*0.70710678118654752440f)); + if (op == UnOp::Silu) + return x/(1.0f+expf(-x)); + if (op == UnOp::Relu) + return x > 0.0f ? x : 0.0f; + if (op == UnOp::GeluErf) + return 0.5f*x*(1.0f+erff(x*0.70710678118654752440f)); // tanh approximation, matching ggml's GGML_UNARY_OP_GELU const float c = 0.79788456080286535588f; // sqrt(2/pi) - return 0.5f*x*(1.0f + tanhf(c*(x + 0.044715f*x*x*x))); + return 0.5f*x*(1.0f+tanhf(c*(x+0.044715f*x*x*x))); } template -__global__ void k_unary_bf16(const __nv_bfloat16 * __restrict__ x, - __nv_bfloat16 * __restrict__ dst, const int64_t n) { - for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < n; +__global__ void k_unary_bf16(const __nv_bfloat16*__restrict__ x, + __nv_bfloat16*__restrict__ dst, const int64_t n) { + for (int64_t i = (int64_t) blockIdx.x*blockDim.x+threadIdx.x; i < n; i += (int64_t) gridDim.x*blockDim.x) { dst[i] = f2bf(apply_unary(bf2f(x[i]))); } @@ -159,12 +485,14 @@ __global__ void k_unary_bf16(const __nv_bfloat16 * __restrict__ x, template bool unary(ggml_tensor * dst, cudaStream_t stream) { const ggml_tensor * src0 = dst->src[0]; - if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) return false; + if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) + return false; // The elementwise index math above assumes a dense buffer. - if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) return false; + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) + return false; const int64_t n = ggml_nelements(dst); - const int64_t blocks = (n + BLOCK - 1) / BLOCK; + const int64_t blocks = (n+BLOCK-1)/BLOCK; const int grid = (int) (blocks < 65535 ? blocks : 65535); k_unary_bf16<<>>( (const __nv_bfloat16 *) src0->data, (__nv_bfloat16 *) dst->data, n); @@ -175,25 +503,27 @@ bool unary(ggml_tensor * dst, cudaStream_t stream) { // scale: dst = x*scale + bias // --------------------------------------------------------------------------- -__global__ void k_scale_bf16(const __nv_bfloat16 * __restrict__ x, __nv_bfloat16 * __restrict__ dst, +__global__ void k_scale_bf16(const __nv_bfloat16*__restrict__ x, __nv_bfloat16*__restrict__ dst, const float scale, const float bias, const int64_t n) { - for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < n; + for (int64_t i = (int64_t) blockIdx.x*blockDim.x+threadIdx.x; i < n; i += (int64_t) gridDim.x*blockDim.x) { - dst[i] = f2bf(scale*bf2f(x[i]) + bias); + dst[i] = f2bf(scale*bf2f(x[i])+bias); } } bool scale(ggml_tensor * dst, cudaStream_t stream) { const ggml_tensor * src0 = dst->src[0]; - if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) return false; - if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) return false; + if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) + return false; + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) + return false; float s = 1.0f, b = 0.0f; - memcpy(&s, (const float *) dst->op_params + 0, sizeof(float)); - memcpy(&b, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&s, (const float *) dst->op_params+0, sizeof(float)); + memcpy(&b, (const float *) dst->op_params+1, sizeof(float)); const int64_t n = ggml_nelements(dst); - const int64_t blocks = (n + BLOCK - 1) / BLOCK; + const int64_t blocks = (n+BLOCK-1)/BLOCK; const int grid = (int) (blocks < 65535 ? blocks : 65535); k_scale_bf16<<>>( (const __nv_bfloat16 *) src0->data, (__nv_bfloat16 *) dst->data, s, b, n); @@ -208,57 +538,66 @@ __device__ inline float block_sum(float v, float * shared) { const int tid = threadIdx.x; shared[tid] = v; __syncthreads(); - for (int s = blockDim.x / 2; s > 0; s >>= 1) { - if (tid < s) shared[tid] += shared[tid + s]; + for (int s=blockDim.x/2; s>0; s >>= 1) { + if (tid < s) + shared[tid] += shared[tid+s]; __syncthreads(); } return shared[0]; } template -__global__ void k_norm_bf16(const __nv_bfloat16 * __restrict__ x, __nv_bfloat16 * __restrict__ dst, +__global__ void k_norm_bf16(const __nv_bfloat16*__restrict__ x, __nv_bfloat16*__restrict__ dst, const int64_t ncols, const int64_t sx1, const int64_t sd1, const float eps) { __shared__ float shared[BLOCK]; const int64_t row = blockIdx.x; - const __nv_bfloat16 * xr = x + row*sx1; - __nv_bfloat16 * dr = dst + row*sd1; + const __nv_bfloat16*xr = x + row*sx1; + __nv_bfloat16 * dr = dst+row*sd1; float sum = 0.0f, sumsq = 0.0f; - for (int64_t c = threadIdx.x; c < ncols; c += blockDim.x) { + for (int64_t c=threadIdx.x; c bool norm(ggml_tensor * dst, cudaStream_t stream) { const ggml_tensor * src0 = dst->src[0]; - if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) return false; + if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16) + return false; // Rows must be dense; higher dims are handled by flattening into the row index. - if (src0->nb[0] != ggml_type_size(src0->type)) return false; - if (dst->nb[0] != ggml_type_size(dst->type)) return false; - if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) return false; + if (src0->nb[0] != ggml_type_size(src0->type)) + return false; + if (dst->nb[0] != ggml_type_size(dst->type)) + return false; + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) + return false; float eps = 0.0f; memcpy(&eps, dst->op_params, sizeof(float)); const int64_t ncols = src0->ne[0]; - const int64_t nrows = ggml_nelements(src0) / ncols; - if (nrows > 2147483647) return false; + const int64_t nrows = ggml_nelements(src0)/ncols; + if (nrows > 2147483647) + return false; k_norm_bf16<<<(int) nrows, BLOCK, 0, stream>>>( (const __nv_bfloat16 *) src0->data, (__nv_bfloat16 *) dst->data, @@ -285,7 +624,8 @@ cublasHandle_t g_handle = nullptr; bool mul_mat(ggml_tensor * dst, cudaStream_t stream) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - if (!src0 || !src1) return false; + if (!src0 || !src1) + return false; if (dst->type != GGML_TYPE_BF16 || src0->type != GGML_TYPE_BF16 || src1->type != GGML_TYPE_BF16) { return false; } @@ -295,18 +635,21 @@ bool mul_mat(ggml_tensor * dst, cudaStream_t stream) { // src0 is either shared across the whole batch or batched 1:1 with src1 const bool batch_ok = (src0->ne[2] == 1 && src0->ne[3] == 1) || (src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3]); - if (!batch_ok) return false; + if (!batch_ok) + return false; - if (!g_handle && cublasCreate(&g_handle) != CUBLAS_STATUS_SUCCESS) return false; - if (cublasSetStream(g_handle, stream) != CUBLAS_STATUS_SUCCESS) return false; + if (!g_handle && cublasCreate(&g_handle) != CUBLAS_STATUS_SUCCESS) + return false; + if (cublasSetStream(g_handle, stream) != CUBLAS_STATUS_SUCCESS) + return false; const int64_t ne00 = src0->ne[0], ne01 = src0->ne[1]; const int64_t ne10 = src1->ne[0], ne11 = src1->ne[1]; const int64_t ne12 = src1->ne[2], ne13 = src1->ne[3]; const int64_t ne0 = dst->ne[0], ne1 = dst->ne[1]; - const __nv_bfloat16 * a = (const __nv_bfloat16 *) src0->data; - const __nv_bfloat16 * b = (const __nv_bfloat16 *) src1->data; + const __nv_bfloat16*a = (const __nv_bfloat16 *) src0->data; + const __nv_bfloat16*b = (const __nv_bfloat16 *) src1->data; __nv_bfloat16 * c = (__nv_bfloat16 *) dst->data; const float alpha = 1.0f, beta = 0.0f; @@ -340,8 +683,21 @@ bool mul_mat(ggml_tensor * dst, cudaStream_t stream) { // hook entry point // --------------------------------------------------------------------------- +extern "C" bool vla_cuda_bf16_fused_binbcast(ggml_tensor * dst, int n_fuse, void * stream_v) { + if (!dst) + return false; + cudaStream_t stream = (cudaStream_t) stream_v; + + switch (dst->op) { + case GGML_OP_ADD: return fused_bin_bcast(dst, n_fuse, stream); + case GGML_OP_MUL: return fused_bin_bcast(dst, n_fuse, stream); + default: return false; + } +} + extern "C" bool vla_cuda_bf16_forward(ggml_tensor * dst, void * stream_v) { - if (!dst) return false; + if (!dst) + return false; cudaStream_t stream = (cudaStream_t) stream_v; switch (dst->op) { @@ -368,6 +724,14 @@ namespace vla { // Called once, after the CUDA backend is up. Idempotent. void cuda_register_bf16_ops() { ggml_cuda_ext_forward = vla_cuda_bf16_forward; + + // Fusion runs in ggml_backend_cuda_graph_compute, upstream of + // ggml_cuda_compute_forward, so the pointer above never sees a fused node. + // Without this second one a fused BF16 add reaches a kernel that handles + // F32/F16 only and GGML_ABORTs ("unsupported types for fusion: dst: bf16, + // src0: bf16, src1: f32"), and declining fusion instead costs ~18 ms/call + // on evo1 -- enough to make BF16 activations a net loss. + ggml_cuda_ext_fused_binbcast = vla_cuda_bf16_fused_binbcast; } } // namespace vla diff --git a/src/env_flag.h b/src/env_flag.h new file mode 100644 index 0000000..718a009 --- /dev/null +++ b/src/env_flag.h @@ -0,0 +1,62 @@ +// 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 env_flag.h + * @brief Boolean environment switches with a value, not just presence. + */ + +#pragma once + +#include +#include +#include + +namespace vla { + +/** + * @brief Read a boolean environment switch. + * + * Presence alone used to mean "on", so @c VLA_EVO1_FA=0 turned flash attention + * *on* -- the opposite of what anyone typing it intends, and a quiet way to + * publish a benchmark of the wrong configuration. + * + * @c 0, @c false, @c off, @c no and the empty string are false (case + * insensitive); any other value is true; unset returns @p def. This matches + * @c VLA_GR00T_GRAPH_CACHE, which already parsed its value rather than its + * presence. + * + * @param name Variable to read. + * @param def Value to use when the variable is not set at all. + */ +inline bool env_flag(const char * name, bool def = false) { + const char * v = std::getenv(name); + if (!v) + return def; + if (!*v) return false; // FOO= reads as "unset it" + + char buf[8] = {}; + size_t n = 0; + for (; n= 0; } + 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 (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)); @@ -74,11 +86,25 @@ struct gguf_reader { 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); } + 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, ...) // so it stays packed and ggml_mul_mat dequantizes at compute; otherwise use @@ -92,15 +118,19 @@ struct gguf_reader { // 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); + 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; + if (fseeko(fp, (off_t) off, SEEK_SET) != 0) + return false; return std::fread(buf, 1, nb, fp) == nb; } @@ -109,8 +139,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(), 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); } + 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; } @@ -136,29 +166,50 @@ struct gguf_reader { std::vector f = read_f32(name); if (f.empty()) return {}; const int64_t n = (int64_t) f.size(); - if (gemma_norm) for (int64_t i = 0; i < n; ++i) f[i] += 1.0f; - if (target == GGML_TYPE_F32) { std::vector o(n * sizeof(float)); std::memcpy(o.data(), f.data(), o.size()); return o; } - if (target == GGML_TYPE_BF16) { std::vector o(n * sizeof(ggml_bf16_t)); ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); return o; } + if (gemma_norm) for (int64_t i=0; i o(n * sizeof(float)); + std::memcpy(o.data(), f.data(), o.size()); + return o; + } + if (target == GGML_TYPE_BF16) { + std::vector o(n * sizeof(ggml_bf16_t)); + ggml_fp32_to_bf16_row(f.data(), reinterpret_cast(o.data()), n); + return o; + } std::fprintf(stderr, "vla(%s): unsupported resident type %d for %s\n", arch, (int) target, name); return {}; } bool fetch_rows_f32(const char * name, const std::vector & row_ids, float * dst, int64_t cols) { const ggml_tensor * t = meta(name); - if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { std::fprintf(stderr, "vla(%s): %s shape unfit for row-fetch\n", arch, name); return false; } - if (t->type != GGML_TYPE_F32 && t->type != GGML_TYPE_BF16) { std::fprintf(stderr, "vla(%s): %s type %d not f32/bf16 for row-fetch\n", arch, name, (int) t->type); return false; } + if (!t || t->ne[0] != cols || t->ne[2] != 1 || t->ne[3] != 1) { + std::fprintf(stderr, "vla(%s): %s shape unfit for row-fetch\n", arch, name); + return false; + } + if (t->type != GGML_TYPE_F32 && t->type != GGML_TYPE_BF16) { + std::fprintf(stderr, "vla(%s): %s type %d not f32/bf16 for row-fetch\n", arch, name, (int) t->type); + return false; + } const int64_t rows = t->ne[1]; const int64_t id = gguf_find_tensor(gctx, name); - const size_t base = data_off + gguf_get_tensor_offset(gctx, id); + const size_t base = data_off+gguf_get_tensor_offset(gctx, id); const size_t elsz = (t->type == GGML_TYPE_F32) ? 4u : 2u; const size_t rb = (size_t) cols * elsz; std::vector row(rb); - for (size_t k = 0; k < row_ids.size(); ++k) { + for (size_t k=0; k= rows) { std::fprintf(stderr, "vla(%s): row %d out of range for %s\n", arch, r, name); return false; } - if (fseeko(fp, (off_t) (base + (size_t) r * rb), SEEK_SET) != 0) return false; - if (std::fread(row.data(), 1, rb, fp) != rb) return false; - if (elsz == 4) std::memcpy(dst + k * cols, row.data(), rb); - else ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst + k * cols, cols); + if (r < 0 || r >= rows) { + std::fprintf(stderr, "vla(%s): row %d out of range for %s\n", arch, r, name); + return false; + } + if (fseeko(fp, (off_t) (base+(size_t) r * rb), SEEK_SET) != 0) + return false; + if (std::fread(row.data(), 1, rb, fp) != rb) + return false; + if (elsz == 4) + std::memcpy(dst+k * cols, row.data(), rb); + else + ggml_bf16_to_fp32_row(reinterpret_cast(row.data()), dst+k * cols, cols); } return true; } diff --git a/src/kernels/bitvla/bitnet_kernels.cu b/src/kernels/bitvla/bitnet_kernels.cu index 5669cd5..2204474 100644 --- a/src/kernels/bitvla/bitnet_kernels.cu +++ b/src/kernels/bitvla/bitnet_kernels.cu @@ -14,6 +14,8 @@ #include "bitnet_kernels.h" +#include "env_flag.h" + #include #include @@ -67,7 +69,7 @@ extern "C" void bitlinear_int8xint2(int8_t* input0, int8_t* input1, __nv_bfloat1 // VLA_BITVLA_NARROW_GEMM=1 to fall back to the one-tile-per-CTA kernel, which // is what the A/B correctness harness and any regression bisect want. static bool bitlinear_use_wide() { - static const bool wide = (std::getenv("VLA_BITVLA_NARROW_GEMM") == nullptr); + static const bool wide = !vla::env_flag("VLA_BITVLA_NARROW_GEMM"); return wide; } @@ -81,7 +83,8 @@ extern "C" void bitlinear_int8xint2_m( #define WIDE(NN, KK, WS) \ launch_ladder_int8xint2_m_wide( \ input0, input1, output0, s, ws, M, stream) - if (N == 2560 && K == 2560) WIDE(2560, 2560, 1); + if (N == 2560 && K == 2560) + WIDE(2560, 2560, 1); else if (N == 640 && K == 2560) WIDE(640, 2560, 1); else if (N == 13824 && K == 2560) WIDE(13824, 2560, 2); else if (N == 2560 && K == 6912) WIDE(2560, 6912, 1); @@ -97,7 +100,8 @@ extern "C" void bitlinear_int8xint2_m( return; } - if (N == 2560 && K == 2560) launch_ladder_int8xint2_m<2560, 2560, 1, 128>(input0, input1, output0, s, ws, M, stream); + if (N == 2560 && K == 2560) + launch_ladder_int8xint2_m<2560, 2560, 1, 128>(input0, input1, output0, s, ws, M, stream); else if (N == 640 && K == 2560) launch_ladder_int8xint2_m<640, 2560, 1, 128>(input0, input1, output0, s, ws, M, stream); else if (N == 13824 && K == 2560) launch_ladder_int8xint2_m<13824, 2560, 2, 128>(input0, input1, output0, s, ws, M, stream); else if (N == 2560 && K == 6912) launch_ladder_int8xint2_m<2560, 6912, 1, 128>(input0, input1, output0, s, ws, M, stream); diff --git a/src/kernels/bitvla/bitnet_kernels.h b/src/kernels/bitvla/bitnet_kernels.h index 73e114c..bae049d 100644 --- a/src/kernels/bitvla/bitnet_kernels.h +++ b/src/kernels/bitvla/bitnet_kernels.h @@ -24,8 +24,8 @@ * fragment multiply per tile. * * Two GEMM entry points are provided: - * * @ref ladder_int8xint2_kernel - single-row (M=1) decode kernel - * used for next-token / single-query inference. + * * @ref ladder_int8xint2_kernel-single-row (M=1) decode kernel + * used for next-token/single-query inference. * * @ref ladder_int8xint2_kernel_m + @ref launch_ladder_int8xint2_m * - multi-row (M>1) variant for prefill and ViT batches. * @@ -80,11 +80,11 @@ __device__ void decode_i2s_to_i8s(T1 *_i2s, T2 *_i8s, const int N = 16) static constexpr uint I4s_TO_I8s_MAGIC_NUM = 0x00000000; #pragma unroll - for (int i = 0; i < (N / 4); i++) + for (int i=0; i<(N/4); i++) { asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(i8s[i]) - : "r"(i2s >> (2 * i)), "n"(BOTTOM_MASK), "n"(I4s_TO_I8s_MAGIC_NUM), "n"(immLut)); + : "r"(i2s >> (2*i)), "n"(BOTTOM_MASK), "n"(I4s_TO_I8s_MAGIC_NUM), "n"(immLut)); i8s[i] = __vsubss4(i8s[i], 0x02020202); } } @@ -92,7 +92,7 @@ __device__ void decode_i2s_to_i8s(T1 *_i2s, T2 *_i8s, const int N = 16) /** * @brief Single-row ternary GEMM kernel (M = 1). * - * Computes one row of @c dtype_transform[0,:] = (A * B^T) / s[0] * ws, + * Computes one row of @c dtype_transform[0,:] = (A * B^T)/s[0]*ws, * with @c A in int8, @c B packed as int2 (decoded on the fly), accumulated * in int32 via @c __dp4a, then scaled back to bf16. The output bias * @c ws is applied per @c ws_num column groups. @@ -117,38 +117,38 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel(int8_t* __restric int red_buf0[1]; in_thread_C_local[0] = 0; #pragma unroll - for (int k_0 = 0; k_0 < K/(K_per_loop * K_block_size); ++k_0) { - *(int4*)(A_local + 0) = *(int4*)(A + ((k_0 * K_per_loop * K_block_size) + (((int)threadIdx.x) * K_per_loop))); + for (int k_0=0; k_0> 1) * wmma_K * wmma_N / 4) + - ((((int)threadIdx.y) >> 3) * (wmma_K * wmma_N / 2) / 4) + - ((((int)threadIdx.x) & 1) * (wmma_K * wmma_N / 4) / 4) + - ((((int)threadIdx.y) & 7) * (wmma_K / 2) / 4) + (((int)blockIdx.x)*N_block_size * K/4) + + (k_0*K_block_size * K_per_loop * wmma_N/4) + + ((((int)threadIdx.x) >> 1)*wmma_K * wmma_N/4) + + ((((int)threadIdx.y) >> 3)*(wmma_K * wmma_N/2)/4) + + ((((int)threadIdx.x) & 1)*(wmma_K * wmma_N/4)/4) + + ((((int)threadIdx.y) & 7)*(wmma_K/2)/4) ); decode_i2s_to_i8s(B_reshape_local, B_decode_local, 16); #pragma unroll - for (int k_2_0 = 0; k_2_0 < 4; ++k_2_0) { - in_thread_C_local[0] = __dp4a(*(int *)&A_local[((k_2_0 * 4))],*(int *)&B_decode_local[((k_2_0 * 4))], in_thread_C_local[0]); + for (int k_2_0=0; k_2_0<4; ++k_2_0) { + in_thread_C_local[0] = __dp4a(*(int *)&A_local[((k_2_0*4))],*(int *)&B_decode_local[((k_2_0*4))], in_thread_C_local[0]); } } red_buf0[0] = in_thread_C_local[0]; #pragma unroll - for (int offset = K_block_size/2; offset > 0; offset /= 2) { + for (int offset=K_block_size/2; offset>0; offset /= 2) { red_buf0[0] += __shfl_down_sync(__activemask(), red_buf0[0], offset, K_block_size); } - int out_idx = ((((int)blockIdx.x) * N_block_size) + ((int)threadIdx.y)); - int ws_idx = out_idx / (N / ws_num); + int out_idx = ((((int)blockIdx.x)*N_block_size)+((int)threadIdx.y)); + int ws_idx = out_idx/(N/ws_num); if (threadIdx.x == 0) - dtype_transform[out_idx] = __float2bfloat16(((float)red_buf0[0]) / s[0] * ws[ws_idx]); + dtype_transform[out_idx] = __float2bfloat16(((float)red_buf0[0])/s[0]*ws[ws_idx]); } /** * @brief Multi-row ternary GEMM kernel (M > 1) using @c wmma fragments. * * Tiles M rows in chunks of @p M_ROWS per CTA. Each warp processes - * @c TILES_PER_WARP = @p M_ROWS / 64 row-tiles, accumulating int32 via + * @c TILES_PER_WARP = @p M_ROWS/64 row-tiles, accumulating int32 via * @c wmma::mma_sync against B fragments decoded from the i2 pack into * shared memory. * @@ -159,7 +159,7 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel(int8_t* __restric * @param A,B Activation (int8) and weight (int2-packed int8) buffers. * @param out bf16 output (M x N). * @param s,ws Per-row activation scale and per-column-group output scale. - * @param M Actual row count (may be < blockDim.y * M_ROWS - tail + * @param M Actual row count (may be < blockDim.y*M_ROWS-tail * threads predicate against this). */ template @@ -172,13 +172,13 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel_m( constexpr int N_block_size = 16; constexpr int K_per_loop = 16, wmma_K = 32, wmma_N = 16; constexpr int K_CHUNK = 128; - constexpr int TILES_PER_WARP = M_ROWS / (4 * 16); + constexpr int TILES_PER_WARP = M_ROWS/(4*16); const int tx = (int)threadIdx.x; const int ty = (int)threadIdx.y; - const int tid = ty * 8 + tx; + const int tid = ty*8+tx; const int warp = tid >> 5; - const int m_base = (int)blockIdx.y * M_ROWS; + const int m_base = (int)blockIdx.y*M_ROWS; __shared__ signed char A_smem[M_ROWS][K_CHUNK]; __shared__ signed char W_smem[16][K_CHUNK]; @@ -191,35 +191,36 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel_m( wmma::fragment b_frag; wmma::fragment acc[TILES_PER_WARP]; #pragma unroll - for (int t = 0; t < TILES_PER_WARP; ++t) wmma::fill_fragment(acc[t], 0); + for (int t=0; t> 3; - const int kk = (idx & 7) * 16; - const int mrow = m_base + m; - const int8_t* aptr = A + ((mrow < M ? mrow : 0) * K) + k_0 * K_CHUNK + kk; + const int kk = (idx & 7)*16; + const int mrow = m_base+m; + const int8_t* aptr = A+((mrow < M ? mrow : 0)*K)+k_0*K_CHUNK+kk; *(int4*)(&A_smem[m][kk]) = *(const int4*)aptr; } B_reshape_local[0] = *(int*)(B + - (((int)blockIdx.x) * N_block_size * K / 4) + - (k_0 * 8 * K_per_loop * wmma_N / 4) + - ((tx >> 1) * wmma_K * wmma_N / 4) + - ((ty >> 3) * (wmma_K * wmma_N / 2) / 4) + - ((tx & 1) * (wmma_K * wmma_N / 4) / 4) + - ((ty & 7) * (wmma_K / 2) / 4)); + (((int)blockIdx.x)*N_block_size * K/4) + + (k_0*8*K_per_loop * wmma_N/4) + + ((tx >> 1)*wmma_K * wmma_N/4) + + ((ty >> 3)*(wmma_K * wmma_N/2)/4) + + ((tx & 1)*(wmma_K * wmma_N/4)/4) + + ((ty & 7)*(wmma_K/2)/4)); decode_i2s_to_i8s(B_reshape_local, B_decode_local, 16); - *(int4*)(&W_smem[ty][tx * 16]) = *(int4*)(&B_decode_local[0]); + *(int4*)(&W_smem[ty][tx*16]) = *(int4*)(&B_decode_local[0]); __syncthreads(); #pragma unroll - for (int k16 = 0; k16 < K_CHUNK / 16; ++k16) { - wmma::load_matrix_sync(b_frag, &W_smem[0][k16 * 16], K_CHUNK); + for (int k16=0; k16> 4; const int col = lin & 15; - const int m = m_base + ml; + const int m = m_base+ml; if (m < M) - out[m * N + n_base + col] = __float2bfloat16(((float)C_smem[ml][col]) / s[m] * wsv); + out[m * N+n_base+col] = __float2bfloat16(((float)C_smem[ml][col])/s[m]*wsv); } } @@ -256,7 +257,7 @@ static inline void launch_ladder_int8xint2_m( int8_t* A, int8_t* B, __nv_bfloat16* out, float* s, float* ws, int M, cudaStream_t stream) { ladder_int8xint2_kernel_m - <<>>(A, B, out, s, ws, M); + <<>>(A, B, out, s, ws, M); } /** @@ -291,29 +292,29 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel_m_wide( constexpr int K_per_loop = 16, wmma_K = 32, wmma_N = 16; constexpr int K_CHUNK = 128; constexpr int WARPS = 4; - constexpr int M_TILES = M_ROWS / 16; - constexpr int N_BLOCKS = N / 16; // total 16-wide column tiles in the matrix + constexpr int M_TILES = M_ROWS/16; + constexpr int N_BLOCKS = N/16; // total 16-wide column tiles in the matrix // Warps split two ways. N_TILES of them take different column tiles (that is // the A-reuse win); the remaining WARPS/N_TILES take different row ranges // (that is parallelism, which matters when N is small enough that column // tiles alone cannot fill the GPU). N_TILES == 1 reproduces the original // kernel's mapping exactly. - constexpr int M_GROUPS = WARPS / N_TILES; - constexpr int M_PER_WARP = M_TILES / M_GROUPS; + constexpr int M_GROUPS = WARPS/N_TILES; + constexpr int M_PER_WARP = M_TILES/M_GROUPS; // Row stride padded to break shared-memory bank conflicts: at a stride of // 128 B every row of a 16-row fragment starts on bank 0, so each // load_matrix_sync serialises 16 ways. 144 B (still a multiple of the 16 B // that wmma requires for integer ldm) spreads them over 8 banks. - constexpr int SM_STRIDE = K_CHUNK + 16; + constexpr int SM_STRIDE = K_CHUNK+16; const int tx = (int)threadIdx.x; // 0..7 const int ty = (int)threadIdx.y; // 0..15 - const int tid = ty * 8 + tx; // 0..127 + const int tid = ty*8+tx; // 0..127 const int warp = tid >> 5; // 0..3 const int lane = tid & 31; - const int m_base = (int)blockIdx.y * M_ROWS; + const int m_base = (int)blockIdx.y*M_ROWS; __shared__ signed char A_smem[M_ROWS][SM_STRIDE]; __shared__ signed char W_smem[N_TILES][16][SM_STRIDE]; @@ -322,9 +323,9 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel_m_wide( // Column tile this warp owns. N is not always a multiple of 16*N_TILES // (the ViT's 4304 is 269 tiles), so tiles past the end are skipped rather // than clamped -- clamping would double-write real columns. - const int my_tile = (int)blockIdx.x * N_TILES + (warp % N_TILES); + const int my_tile = (int)blockIdx.x*N_TILES+(warp%N_TILES); const bool my_tile_valid = my_tile < N_BLOCKS; - const int m_tile_base = (warp / N_TILES) * M_PER_WARP; + const int m_tile_base = (warp/N_TILES)*M_PER_WARP; int B_reshape_local[1]; signed char B_decode_local[K_per_loop]; @@ -333,45 +334,46 @@ __global__ void __launch_bounds__(128) ladder_int8xint2_kernel_m_wide( wmma::fragment b_frag; wmma::fragment acc[M_PER_WARP]; #pragma unroll - for (int t = 0; t < M_PER_WARP; ++t) wmma::fill_fragment(acc[t], 0); + for (int t=0; t> 3; - const int kk = (idx & 7) * 16; - const int mrow = m_base + m; - const int8_t* aptr = A + ((mrow < M ? mrow : 0) * K) + k_0 * K_CHUNK + kk; + const int kk = (idx & 7)*16; + const int mrow = m_base+m; + const int8_t* aptr = A+((mrow < M ? mrow : 0)*K)+k_0*K_CHUNK+kk; *(int4*)(&A_smem[m][kk]) = *(const int4*)aptr; } // All 128 threads cooperate on one weight tile at a time, reproducing the // pack's swizzle exactly; only the tile base changes per j. #pragma unroll - for (int j = 0; j < N_TILES; ++j) { - const int tile = (int)blockIdx.x * N_TILES + j; + for (int j=0; j> 1) * wmma_K * wmma_N / 4) + - ((ty >> 3) * (wmma_K * wmma_N / 2) / 4) + - ((tx & 1) * (wmma_K * wmma_N / 4) / 4) + - ((ty & 7) * (wmma_K / 2) / 4)); + ((size_t)tile*16*K/4) + + (k_0*8*K_per_loop * wmma_N/4) + + ((tx >> 1)*wmma_K * wmma_N/4) + + ((ty >> 3)*(wmma_K * wmma_N/2)/4) + + ((tx & 1)*(wmma_K * wmma_N/4)/4) + + ((ty & 7)*(wmma_K/2)/4)); decode_i2s_to_i8s(B_reshape_local, B_decode_local, 16); - *(int4*)(&W_smem[j][ty][tx * 16]) = *(int4*)(&B_decode_local[0]); + *(int4*)(&W_smem[j][ty][tx*16]) = *(int4*)(&B_decode_local[0]); } } __syncthreads(); if (my_tile_valid) { #pragma unroll - for (int k16 = 0; k16 < K_CHUNK / 16; ++k16) { - wmma::load_matrix_sync(b_frag, &W_smem[warp % N_TILES][0][k16 * 16], SM_STRIDE); + for (int k16=0; k16> 4; const int col = lin & 15; - const int m = m_base + (m_tile_base + t) * 16 + ml; + const int m = m_base+(m_tile_base+t)*16+ml; if (m < M) - out[(size_t)m * N + n_base + col] = - __float2bfloat16(((float)C_smem[warp][ml][col]) / s[m] * wsv); + out[(size_t)m * N+n_base+col] = + __float2bfloat16(((float)C_smem[warp][ml][col])/s[m]*wsv); } __syncwarp(); } @@ -438,9 +441,9 @@ template static inline void launch_ladder_int8xint2_m_wide( int8_t* A, int8_t* B, __nv_bfloat16* out, float* s, float* ws, int M, cudaStream_t stream) { - constexpr int N_BLOCKS = N / 16; + constexpr int N_BLOCKS = N/16; ladder_int8xint2_kernel_m_wide - <<>>(A, B, out, s, ws, M); } @@ -469,42 +472,50 @@ __global__ void act_quant_kernel( const int m = (int)blockIdx.x; const int tid = (int)threadIdx.x; const __nv_bfloat16* row_in = in + m * K; - int8_t* row_out = out + m * K; + int8_t* row_out = out+m * K; float local_max = 0.0f; - for (int k = tid; k < K; k += BLOCK_THREADS) { + for (int k=tid; k local_max) local_max = v; + if (v > local_max) + local_max = v; } - for (int off = 16; off > 0; off >>= 1) { + for (int off=16; off>0; off >>= 1) { float other = __shfl_down_sync(0xffffffff, local_max, off); - if (other > local_max) local_max = other; + if (other > local_max) + local_max = other; } __shared__ float smem[32]; const int warp_id = tid >> 5; const int lane = tid & 31; - if (lane == 0) smem[warp_id] = local_max; + if (lane == 0) + smem[warp_id] = local_max; __syncthreads(); if (warp_id == 0) { - float v = (tid < (BLOCK_THREADS + 31) / 32) ? smem[lane] : 0.0f; - for (int off = 16; off > 0; off >>= 1) { + float v = (tid < (BLOCK_THREADS+31)/32) ? smem[lane] : 0.0f; + for (int off=16; off>0; off >>= 1) { float other = __shfl_down_sync(0xffffffff, v, off); - if (other > v) v = other; + if (other > v) + v = other; } - if (lane == 0) smem[0] = v; + if (lane == 0) + smem[0] = v; } __syncthreads(); const float amax = smem[0] < 1e-5f ? 1e-5f : smem[0]; - const float scale = 127.0f / amax; - if (tid == 0) scales[m] = scale; + const float scale = 127.0f/amax; + if (tid == 0) + scales[m] = scale; - for (int k = tid; k < K; k += BLOCK_THREADS) { - float v = __bfloat162float(row_in[k]) * scale; + for (int k=tid; k 127.0f) q = 127.0f; - if (q < -128.0f) q = -128.0f; + if (q > 127.0f) + q = 127.0f; + if (q < -128.0f) + q = -128.0f; row_out[k] = (int8_t)q; } } diff --git a/src/kernels/bitvla/bitvla_fp32head_cuda.cu b/src/kernels/bitvla/bitvla_fp32head_cuda.cu index 38108c4..3bd1c24 100644 --- a/src/kernels/bitvla/bitvla_fp32head_cuda.cu +++ b/src/kernels/bitvla/bitvla_fp32head_cuda.cu @@ -58,24 +58,34 @@ struct bitvla_fp32head_cuda_ctx { static float* upload_f32(const float* h, size_t n) { float* d = nullptr; cudaError_t e = cudaMalloc(&d, n * sizeof(float)); - if (e != cudaSuccess) { std::fprintf(stderr, "vla(bitvla_fp32head): cudaMalloc failed (%zu)\n", n); return nullptr; } + if (e != cudaSuccess) { + std::fprintf(stderr, "vla(bitvla_fp32head): cudaMalloc failed (%zu)\n", n); + return nullptr; + } e = cudaMemcpy(d, h, n * sizeof(float), cudaMemcpyHostToDevice); - if (e != cudaSuccess) { std::fprintf(stderr, "vla(bitvla_fp32head): cudaMemcpy H2D failed (%zu)\n", n); cudaFree(d); return nullptr; } + if (e != cudaSuccess) { + std::fprintf(stderr, "vla(bitvla_fp32head): cudaMemcpy H2D failed (%zu)\n", n); + cudaFree(d); + return nullptr; + } return d; } __global__ void gelu_erf_fp32_kernel(const float* __restrict__ in, float* __restrict__ out, int N) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= N) return; + const int i = blockIdx.x*blockDim.x+threadIdx.x; + if (i >= N) + return; float x = in[i]; - out[i] = 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f)); + out[i] = 0.5f * x * (1.0f+erff(x*0.70710678118654752440f)); } __global__ void relu_fp32_kernel(float* __restrict__ inout, int N) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= N) return; + const int i = blockIdx.x*blockDim.x+threadIdx.x; + if (i >= N) + return; float v = inout[i]; - if (v < 0.0f) inout[i] = 0.0f; + if (v < 0.0f) + inout[i] = 0.0f; } template @@ -86,57 +96,68 @@ __global__ void layernorm_fp32_kernel(const float* __restrict__ x, float eps, int K) { const int m = blockIdx.x; const int tid = threadIdx.x; - const float* row = x + (size_t)m * K; - float* o = out + (size_t)m * K; + const float* row = x+(size_t)m * K; + float* o = out+(size_t)m * K; float sum = 0.0f; - for (int k = tid; k < K; k += BLOCK) sum += row[k]; - for (int off = 16; off > 0; off >>= 1) sum += __shfl_down_sync(0xffffffff, sum, off); + for (int k=tid; k0; off >>= 1) + sum += __shfl_down_sync(0xffffffff, sum, off); __shared__ float smem[32]; - if ((tid & 31) == 0) smem[tid >> 5] = sum; + if ((tid & 31) == 0) + smem[tid >> 5] = sum; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); - if (tid == 0) smem[0] = v; + float v = (tid < (BLOCK+31)/32) ? smem[tid] : 0.0f; + for (int off=16; off>0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + if (tid == 0) + smem[0] = v; } __syncthreads(); - const float mean = smem[0] / (float)K; + const float mean = smem[0]/(float)K; float vsum = 0.0f; - for (int k = tid; k < K; k += BLOCK) { - float v = row[k] - mean; + for (int k=tid; k 0; off >>= 1) vsum += __shfl_down_sync(0xffffffff, vsum, off); - if ((tid & 31) == 0) smem[tid >> 5] = vsum; + for (int off=16; off>0; off >>= 1) + vsum += __shfl_down_sync(0xffffffff, vsum, off); + if ((tid & 31) == 0) + smem[tid >> 5] = vsum; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); - if (tid == 0) smem[0] = v; + float v = (tid < (BLOCK+31)/32) ? smem[tid] : 0.0f; + for (int off=16; off>0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + if (tid == 0) + smem[0] = v; } __syncthreads(); - const float inv_std = rsqrtf(smem[0] / (float)K + eps); + const float inv_std = rsqrtf(smem[0]/(float)K+eps); - for (int k = tid; k < K; k += BLOCK) { - o[k] = (row[k] - mean) * inv_std * w[k] + b[k]; + for (int k=tid; k= K) return; - const size_t i = (size_t)m * K + k; - out[i] = x[i] + bias[k]; + const int k = blockIdx.y*blockDim.x+threadIdx.x; + if (k >= K) + return; + const size_t i = (size_t)m * K+k; + out[i] = x[i]+bias[k]; } __global__ void add_fp32_kernel(const float* a, const float* b, float* out, int N) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= N) return; - out[i] = a[i] + b[i]; + const int i = blockIdx.x*blockDim.x+threadIdx.x; + if (i >= N) + return; + out[i] = a[i]+b[i]; } static int linear_bias_fp32( @@ -160,7 +181,7 @@ static int linear_bias_fp32( } if (bias) { const int B = 256; - const int n_kb = (N_out + B - 1) / B; + const int n_kb = (N_out+B-1)/B; add_bias_fp32_kernel<<>>(out, bias, out, M, N_out); } return 0; @@ -225,8 +246,8 @@ extern "C" bitvla_fp32head_cuda_ctx* bitvla_fp32head_cuda_init( CUDA_OK_NULL(cudaMalloc(&ctx->d_pp_h1, (size_t)lm_hidden * sizeof(float))); CUDA_OK_NULL(cudaMalloc(&ctx->d_pp_out, (size_t)lm_hidden * sizeof(float))); - CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_in, (size_t)chunk * ctx->ah_in_dim * sizeof(float))); - CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_norm_big, (size_t)chunk * ctx->ah_in_dim * sizeof(float))); + CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_in, (size_t)chunk * ctx->ah_in_dim*sizeof(float))); + CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_norm_big, (size_t)chunk * ctx->ah_in_dim*sizeof(float))); CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_h, (size_t)chunk * lm_hidden * sizeof(float))); CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_tmp, (size_t)chunk * lm_hidden * sizeof(float))); CUDA_OK_NULL(cudaMalloc(&ctx->d_ah_tmp2, (size_t)chunk * lm_hidden * sizeof(float))); @@ -254,7 +275,7 @@ extern "C" int bitvla_fp32head_proprio_forward( cudaStream_t stream) { - CUDA_OK_RET(cudaMemcpyAsync(ctx->d_state, host_state, (size_t)ctx->proprio_dim * sizeof(float), + CUDA_OK_RET(cudaMemcpyAsync(ctx->d_state, host_state, (size_t)ctx->proprio_dim*sizeof(float), cudaMemcpyHostToDevice, stream)); if (linear_bias_fp32(ctx->cublas, stream, ctx->pp_fc1_w, ctx->pp_fc1_b, @@ -263,13 +284,13 @@ extern "C" int bitvla_fp32head_proprio_forward( { const int B = 256; const int N = ctx->lm_hidden; - gelu_erf_fp32_kernel<<>>(ctx->d_pp_h1, ctx->d_pp_h1, N); + gelu_erf_fp32_kernel<<>>(ctx->d_pp_h1, ctx->d_pp_h1, N); } if (linear_bias_fp32(ctx->cublas, stream, ctx->pp_fc2_w, ctx->pp_fc2_b, ctx->d_pp_h1, ctx->d_pp_out, 1, ctx->lm_hidden, ctx->lm_hidden) != 0) return -1; - CUDA_OK_RET(cudaMemcpyAsync(host_out, ctx->d_pp_out, (size_t)ctx->lm_hidden * sizeof(float), + CUDA_OK_RET(cudaMemcpyAsync(host_out, ctx->d_pp_out, (size_t)ctx->lm_hidden*sizeof(float), cudaMemcpyDeviceToHost, stream)); CUDA_OK_RET(cudaStreamSynchronize(stream)); return 0; @@ -301,7 +322,7 @@ extern "C" int bitvla_fp32head_action_forward( { const int B = 256; const int N = M * H; - relu_fp32_kernel<<>>(ctx->d_ah_h, N); + relu_fp32_kernel<<>>(ctx->d_ah_h, N); } { @@ -316,13 +337,13 @@ extern "C" int bitvla_fp32head_action_forward( { const int B = 256; const int N = M * H; - relu_fp32_kernel<<>>(ctx->d_ah_tmp2, N); + relu_fp32_kernel<<>>(ctx->d_ah_tmp2, N); } { const int B = 256; const int N = M * H; - add_fp32_kernel<<>>(ctx->d_ah_h, ctx->d_ah_tmp2, ctx->d_ah_h, N); + add_fp32_kernel<<>>(ctx->d_ah_h, ctx->d_ah_tmp2, ctx->d_ah_h, N); } { @@ -335,12 +356,12 @@ extern "C" int bitvla_fp32head_action_forward( { const int B = 256; const int N = M * H; - relu_fp32_kernel<<>>(ctx->d_ah_tmp2, N); + relu_fp32_kernel<<>>(ctx->d_ah_tmp2, N); } { const int B = 256; const int N = M * H; - add_fp32_kernel<<>>(ctx->d_ah_h, ctx->d_ah_tmp2, ctx->d_ah_h, N); + add_fp32_kernel<<>>(ctx->d_ah_h, ctx->d_ah_tmp2, ctx->d_ah_h, N); } { @@ -358,7 +379,8 @@ extern "C" int bitvla_fp32head_action_forward( } extern "C" void bitvla_fp32head_cuda_free(bitvla_fp32head_cuda_ctx* ctx) { - if (!ctx) return; + if (!ctx) + return; float* ws[] = { ctx->pp_fc1_w, ctx->pp_fc1_b, ctx->pp_fc2_w, ctx->pp_fc2_b, ctx->ah_ln1_w, ctx->ah_ln1_b, ctx->ah_fc1_w, ctx->ah_fc1_b, @@ -368,7 +390,10 @@ extern "C" void bitvla_fp32head_cuda_free(bitvla_fp32head_cuda_ctx* ctx) { ctx->d_state, ctx->d_pp_h1, ctx->d_pp_out, ctx->d_ah_in, ctx->d_ah_norm_big, ctx->d_ah_h, ctx->d_ah_tmp, ctx->d_ah_tmp2, ctx->d_ah_out, }; - for (float* p : ws) if (p) cudaFree(p); - if (ctx->cublas) cublasDestroy(ctx->cublas); + for (float* p : ws) + if (p) + cudaFree(p); + if (ctx->cublas) + cublasDestroy(ctx->cublas); delete ctx; } diff --git a/src/kernels/bitvla/bitvla_fp32head_cuda.h b/src/kernels/bitvla/bitvla_fp32head_cuda.h index 5f5dd58..dc6ccbc 100644 --- a/src/kernels/bitvla/bitvla_fp32head_cuda.h +++ b/src/kernels/bitvla/bitvla_fp32head_cuda.h @@ -14,7 +14,7 @@ /** * @file bitvla_fp32head_cuda.h - * @brief FP32 BitVLA action head + proprioception projection. + * @brief FP32 BitVLA action head+proprioception projection. * * BitVLA's LM and ViT are 1.58-bit ternary; its small action head and * proprio projector are kept in FP32 (cuBLAS GEMM) to preserve regression @@ -97,7 +97,7 @@ bitvla_fp32head_cuda_ctx* bitvla_fp32head_cuda_init( * @param ctx Context returned by @ref bitvla_fp32head_cuda_init. * @param host_state Length-@c proprio_dim FP32 input on the host. * @param host_out Length-@c lm_hidden FP32 output on the host. - * @param stream CUDA stream used for the H2D / D2H transfers. + * @param stream CUDA stream used for the H2D/D2H transfers. * @return 0 on success, non-zero on dispatch failure. */ int bitvla_fp32head_proprio_forward( @@ -113,7 +113,7 @@ int bitvla_fp32head_proprio_forward( * @param host_norm_actions Length-@c chunk*action_dim FP32 output on the * host, normalised to training statistics. The caller un-normalises * to world units. - * @param stream CUDA stream used for the H2D / D2H transfers. + * @param stream CUDA stream used for the H2D/D2H transfers. * @return 0 on success, non-zero on dispatch failure. */ int bitvla_fp32head_action_forward( diff --git a/src/kernels/bitvla/bitvla_lm_cuda.cu b/src/kernels/bitvla/bitvla_lm_cuda.cu index 0e4c455..7605a0e 100644 --- a/src/kernels/bitvla/bitvla_lm_cuda.cu +++ b/src/kernels/bitvla/bitvla_lm_cuda.cu @@ -42,29 +42,33 @@ __global__ void rmsnorm_bf16_kernel(const __nv_bfloat16* __restrict__ x, { const int m = (int)blockIdx.x; const int tid = (int)threadIdx.x; - const __nv_bfloat16* row = x + m * K; - __nv_bfloat16* o = out + m * K; + const __nv_bfloat16* row = x+m * K; + __nv_bfloat16* o = out+m * K; float ss = 0.0f; - for (int k = tid; k < K; k += BLOCK) { + for (int k=tid; k 0; off >>= 1) ss += __shfl_down_sync(0xffffffff, ss, off); + for (int off=16; off>0; off >>= 1) + ss += __shfl_down_sync(0xffffffff, ss, off); __shared__ float smem[32]; - if ((tid & 31) == 0) smem[tid >> 5] = ss; + if ((tid & 31) == 0) + smem[tid >> 5] = ss; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); - if (tid == 0) smem[0] = v; + float v = (tid < (BLOCK+31)/32) ? smem[tid] : 0.0f; + for (int off=16; off>0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + if (tid == 0) + smem[0] = v; } __syncthreads(); - const float mean = smem[0] / (float)K; - const float scale = rsqrtf(mean + eps); + const float mean = smem[0]/(float)K; + const float scale = rsqrtf(mean+eps); - for (int k = tid; k < K; k += BLOCK) { - float v = __bfloat162float(row[k]) * scale; + for (int k=tid; k mx) mx = v; + for (int i=tid; i mx) + mx = v; } - for (int off = 16; off > 0; off >>= 1) { + for (int off=16; off>0; off >>= 1) { float other = __shfl_down_sync(0xffffffff, mx, off); - if (other > mx) mx = other; + if (other > mx) + mx = other; } __shared__ float smem[32]; - if ((tid & 31) == 0) smem[tid >> 5] = mx; + if ((tid & 31) == 0) + smem[tid >> 5] = mx; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : -INFINITY; - for (int off = 16; off > 0; off >>= 1) { + float v = (tid < (BLOCK+31)/32) ? smem[tid] : -INFINITY; + for (int off=16; off>0; off >>= 1) { float other = __shfl_down_sync(0xffffffff, v, off); - if (other > v) v = other; + if (other > v) + v = other; } - if (tid == 0) smem[0] = v; + if (tid == 0) + smem[0] = v; } __syncthreads(); const float max_v = smem[0]; float s_sum = 0.0f; - for (int i = tid; i < S; i += BLOCK) { - s_sum += expf(__bfloat162float(r[i]) * scale - max_v); + for (int i=tid; i 0; off >>= 1) s_sum += __shfl_down_sync(0xffffffff, s_sum, off); - if ((tid & 31) == 0) smem[tid >> 5] = s_sum; + for (int off=16; off>0; off >>= 1) + s_sum += __shfl_down_sync(0xffffffff, s_sum, off); + if ((tid & 31) == 0) + smem[tid >> 5] = s_sum; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); - if (tid == 0) smem[0] = v; + float v = (tid < (BLOCK+31)/32) ? smem[tid] : 0.0f; + for (int off=16; off>0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + if (tid == 0) + smem[0] = v; } __syncthreads(); - const float inv_sum = 1.0f / smem[0]; + const float inv_sum = 1.0f/smem[0]; - for (int i = tid; i < S; i += BLOCK) { - float v = expf(__bfloat162float(r[i]) * scale - max_v) * inv_sum; + for (int i=tid; i= N) return; + const int i = (int)(blockIdx.x*blockDim.x+threadIdx.x); + if (i >= N) + return; float gv = __bfloat162float(g[i]); - if (gv < 0.0f) gv = 0.0f; + if (gv < 0.0f) + gv = 0.0f; out[i] = __float2bfloat16(gv * gv * __bfloat162float(u[i])); } __global__ void add_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b, __nv_bfloat16* out, int N) { - const int i = (int)(blockIdx.x * blockDim.x + threadIdx.x); - if (i >= N) return; - out[i] = __float2bfloat16(__bfloat162float(a[i]) + __bfloat162float(b[i])); + const int i = (int)(blockIdx.x*blockDim.x+threadIdx.x); + if (i >= N) + return; + out[i] = __float2bfloat16(__bfloat162float(a[i])+__bfloat162float(b[i])); } __global__ void repeat_kv_bf16_kernel(const __nv_bfloat16* in, __nv_bfloat16* out, @@ -167,9 +183,9 @@ __global__ void repeat_kv_bf16_kernel(const __nv_bfloat16* in, __nv_bfloat16* ou const int q_h = (int)blockIdx.x; const int s = (int)blockIdx.y; const int tid = (int)threadIdx.x; - const int kv_h = q_h * n_kv / n_q; - for (int k = tid; k < hd; k += blockDim.x) { - out[((size_t)q_h * seq + s) * hd + k] = in[((size_t)kv_h * seq + s) * hd + k]; + const int kv_h = q_h * n_kv/n_q; + for (int k=tid; k>>(g, u, out, N); + squared_relu_mul_bf16_kernel<<>>(g, u, out, N); } extern "C" void bitvla_add_bf16(const __nv_bfloat16* a, const __nv_bfloat16* b, __nv_bfloat16* out, int N, cudaStream_t stream) { constexpr int B = 256; - add_bf16_kernel<<>>(a, b, out, N); + add_bf16_kernel<<>>(a, b, out, N); } extern "C" void bitvla_repeat_kv_bf16(const __nv_bfloat16* in, __nv_bfloat16* out, int n_q, int n_kv, int seq, int hd, cudaStream_t stream) { @@ -308,29 +324,33 @@ extern "C" bitvla_lm_cuda_ctx* bitvla_lm_cuda_init(int hidden, int n_q, int n_kv ctx->layers.resize(n_layers); cublasStatus_t cbs = cublasCreate(&ctx->cublas); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_lm_cuda): cublasCreate failed\n"); delete ctx; return nullptr; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_lm_cuda): cublasCreate failed\n"); + delete ctx; + return nullptr; + } - const int half = head_dim / 2; + const int half = head_dim/2; std::vector h_cos((size_t)max_seq * half), h_sin((size_t)max_seq * half); - for (int s = 0; s < max_seq; ++s) { - for (int k = 0; k < half; ++k) { - float freq = 1.0f / std::pow(rope_base, (float)(2 * k) / (float)head_dim); + for (int s=0; sd_cos, (size_t)max_seq * half * sizeof(float))); CUDA_OKV(cudaMalloc(&ctx->d_sin, (size_t)max_seq * half * sizeof(float))); - CUDA_OKV(cudaMemcpy(ctx->d_cos, h_cos.data(), h_cos.size() * sizeof(float), cudaMemcpyHostToDevice)); - CUDA_OKV(cudaMemcpy(ctx->d_sin, h_sin.data(), h_sin.size() * sizeof(float), cudaMemcpyHostToDevice)); + CUDA_OKV(cudaMemcpy(ctx->d_cos, h_cos.data(), h_cos.size()*sizeof(float), cudaMemcpyHostToDevice)); + CUDA_OKV(cudaMemcpy(ctx->d_sin, h_sin.data(), h_sin.size()*sizeof(float), cudaMemcpyHostToDevice)); const size_t bf16 = sizeof(__nv_bfloat16); CUDA_OKV(cudaMalloc(&ctx->d_h, (size_t)max_seq * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_h_norm, (size_t)max_seq * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_act_int8_h, (size_t)max_seq * hidden)); CUDA_OKV(cudaMalloc(&ctx->d_act_s, (size_t)max_seq * sizeof(float))); - CUDA_OKV(cudaMalloc(&ctx->d_qkv, (size_t)max_seq * (ctx->hidden_q + 2 * ctx->hidden_kv) * bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_qkv, (size_t)max_seq * (ctx->hidden_q+2*ctx->hidden_kv)*bf16)); CUDA_OKV(cudaMalloc(&ctx->d_q_HShd, (size_t)n_q * max_seq * head_dim * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_k_HShd, (size_t)n_kv * max_seq * head_dim * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_v_HShd, (size_t)n_kv * max_seq * head_dim * bf16)); @@ -341,14 +361,15 @@ extern "C" bitvla_lm_cuda_ctx* bitvla_lm_cuda_init(int hidden, int n_q, int n_kv CUDA_OKV(cudaMalloc(&ctx->d_attn_merged, (size_t)max_seq * ctx->hidden_q * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_o_out, (size_t)max_seq * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_act_int8_ffn,(size_t)max_seq * ffn)); - CUDA_OKV(cudaMalloc(&ctx->d_gate_up, (size_t)max_seq * 2 * ffn * bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_gate_up, (size_t)max_seq*2*ffn * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_gate_sq_up, (size_t)max_seq * ffn * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_down_out, (size_t)max_seq * hidden * bf16)); return ctx; } extern "C" void bitvla_lm_cuda_free(bitvla_lm_cuda_ctx* ctx) { - if (!ctx) return; + if (!ctx) + return; cublasDestroy(ctx->cublas); cudaFree(ctx->d_cos); cudaFree(ctx->d_sin); cudaFree(ctx->d_h); cudaFree(ctx->d_h_norm); @@ -378,12 +399,13 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea const char* l0_dir = (L == 0) ? std::getenv("VLA_BITVLA_DUMP_L0") : nullptr; auto l0_dump = [&](const char* name, const __nv_bfloat16* d_ptr, size_t n) { - if (!l0_dir) return; + if (!l0_dir) + return; cudaStreamSynchronize(stream); std::vector<__nv_bfloat16> tmp(n); std::vector f32(n); cudaMemcpy(tmp.data(), d_ptr, n * sizeof(__nv_bfloat16), cudaMemcpyDeviceToHost); - for (size_t i = 0; i < n; ++i) { + for (size_t i=0; i(&tmp[i]); uint32_t b = ((uint32_t) u) << 16; float f; std::memcpy(&f, &b, 4); @@ -391,7 +413,10 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea } std::string path = std::string(l0_dir) + "/" + name + ".bin"; FILE* f = std::fopen(path.c_str(), "wb"); - if (f) { std::fwrite(f32.data(), sizeof(float), n, f); std::fclose(f); } + if (f) { + std::fwrite(f32.data(), sizeof(float), n, f); + std::fclose(f); + } }; bitvla_rmsnorm_bf16(ctx->d_h, lr.attn_norm_w, ctx->d_h_norm, ctx->rms_eps, seq, hidden, stream); @@ -400,15 +425,15 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea bitvla_act_quant_cuda(ctx->d_h_norm, ctx->d_act_int8_h, ctx->d_act_s, seq, hidden, stream); __nv_bfloat16* q_dense = ctx->d_qkv; - __nv_bfloat16* k_dense = ctx->d_qkv + (size_t)seq * hq; - __nv_bfloat16* v_dense = ctx->d_qkv + (size_t)seq * (hq + hkv); + __nv_bfloat16* k_dense = ctx->d_qkv+(size_t)seq * hq; + __nv_bfloat16* v_dense = ctx->d_qkv+(size_t)seq * (hq+hkv); bitlinear_int8xint2_m(ctx->d_act_int8_h, lr.q_packed, q_dense, ctx->d_act_s, lr.q_ws, seq, hq, hidden, stream); bitlinear_int8xint2_m(ctx->d_act_int8_h, lr.k_packed, k_dense, ctx->d_act_s, lr.k_ws, seq, hkv, hidden, stream); bitlinear_int8xint2_m(ctx->d_act_int8_h, lr.v_packed, v_dense, ctx->d_act_s, lr.v_ws, seq, hkv, hidden, stream); - l0_dump("L0_02_qkv_proj", ctx->d_qkv, (size_t)seq * (hq + 2*hkv)); + l0_dump("L0_02_qkv_proj", ctx->d_qkv, (size_t)seq * (hq+2*hkv)); // Split the interleaved [seq, H*hd] projections into head-major [H, seq, hd] // with one kernel per tensor instead of a cudaMemcpy2DAsync per head (30 tiny @@ -435,9 +460,12 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea ctx->d_scores,CUDA_R_16BF, seq, (long long)seq * seq, n_q, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_lm_cuda): QK^T gemm failed @L%d (%d)\n", L, cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_lm_cuda): QK^T gemm failed @L%d (%d)\n", L, cbs); + return -1; + } - const float scl = 1.0f / std::sqrt((float)hd); + const float scl = 1.0f/std::sqrt((float)hd); bitvla_softmax_scaled_bf16(ctx->d_scores, scl, n_q * seq, seq, stream); cbs = cublasGemmStridedBatchedEx( @@ -450,7 +478,10 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea ctx->d_attn_out, CUDA_R_16BF, hd, (long long)seq * hd, n_q, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_lm_cuda): attn@V gemm failed @L%d (%d)\n", L, cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_lm_cuda): attn@V gemm failed @L%d (%d)\n", L, cbs); + return -1; + } bitvla_transpose_NshHd_to_sNhd_bf16(ctx->d_attn_out, ctx->d_attn_merged, n_q, seq, hd, stream); @@ -473,7 +504,7 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea bitvla_act_quant_cuda(ctx->d_h_norm, ctx->d_act_int8_h, ctx->d_act_s, seq, hidden, stream); bitlinear_int8xint2_m(ctx->d_act_int8_h, lr.gate_up_packed, ctx->d_gate_up, - ctx->d_act_s, lr.gate_up_ws, seq, 2 * ffn, hidden, stream); + ctx->d_act_s, lr.gate_up_ws, seq, 2*ffn, hidden, stream); gate_up_fused_sqrelu_mul_bf16(ctx->d_gate_up, ctx->d_gate_sq_up, seq, ffn, stream); @@ -496,22 +527,24 @@ static int run_layer(bitvla_lm_cuda_ctx* ctx, int L, int seq, cudaStream_t strea __global__ void gate_up_fused_sqrelu_mul_bf16_kernel(const __nv_bfloat16* __restrict__ gu, __nv_bfloat16* __restrict__ out, int seq, int ffn) { - const int idx = (int)(blockIdx.x * blockDim.x + threadIdx.x); + const int idx = (int)(blockIdx.x*blockDim.x+threadIdx.x); const int total = seq * ffn; - if (idx >= total) return; - const int s = idx / ffn; - const int k = idx % ffn; - const size_t row_base = (size_t)s * 2 * ffn; - float g = __bfloat162float(gu[row_base + k]); - float u = __bfloat162float(gu[row_base + ffn + k]); - if (g < 0.0f) g = 0.0f; + if (idx >= total) + return; + const int s = idx/ffn; + const int k = idx%ffn; + const size_t row_base = (size_t)s*2*ffn; + float g = __bfloat162float(gu[row_base+k]); + float u = __bfloat162float(gu[row_base+ffn+k]); + if (g < 0.0f) + g = 0.0f; out[(size_t)idx] = __float2bfloat16(g * g * u); } extern "C" void gate_up_fused_sqrelu_mul_bf16(const __nv_bfloat16* gu, __nv_bfloat16* out, int seq, int ffn, cudaStream_t stream) { const int total = seq * ffn; const int B = 256; - gate_up_fused_sqrelu_mul_bf16_kernel<<>>(gu, out, seq, ffn); } @@ -524,73 +557,85 @@ __global__ void layernorm_bias_bf16_kernel(const __nv_bfloat16* __restrict__ x, { const int m = (int)blockIdx.x; const int tid = (int)threadIdx.x; - const __nv_bfloat16* row = x + (size_t)m * K; - __nv_bfloat16* o = out + (size_t)m * K; + const __nv_bfloat16* row = x+(size_t)m * K; + __nv_bfloat16* o = out+(size_t)m * K; float sum = 0.0f; - for (int k = tid; k < K; k += BLOCK) sum += __bfloat162float(row[k]); - for (int off = 16; off > 0; off >>= 1) sum += __shfl_down_sync(0xffffffff, sum, off); + for (int k=tid; k0; off >>= 1) + sum += __shfl_down_sync(0xffffffff, sum, off); __shared__ float smem[32]; - if ((tid & 31) == 0) smem[tid >> 5] = sum; + if ((tid & 31) == 0) + smem[tid >> 5] = sum; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); - if (tid == 0) smem[0] = v; + float v = (tid < (BLOCK+31)/32) ? smem[tid] : 0.0f; + for (int off=16; off>0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + if (tid == 0) + smem[0] = v; } __syncthreads(); - const float mean = smem[0] / (float)K; + const float mean = smem[0]/(float)K; float vsum = 0.0f; - for (int k = tid; k < K; k += BLOCK) { - float v = __bfloat162float(row[k]) - mean; + for (int k=tid; k 0; off >>= 1) vsum += __shfl_down_sync(0xffffffff, vsum, off); - if ((tid & 31) == 0) smem[tid >> 5] = vsum; + for (int off=16; off>0; off >>= 1) + vsum += __shfl_down_sync(0xffffffff, vsum, off); + if ((tid & 31) == 0) + smem[tid >> 5] = vsum; __syncthreads(); if ((tid >> 5) == 0) { - float v = (tid < (BLOCK + 31) / 32) ? smem[tid] : 0.0f; - for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off); - if (tid == 0) smem[0] = v; + float v = (tid < (BLOCK+31)/32) ? smem[tid] : 0.0f; + for (int off=16; off>0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + if (tid == 0) + smem[0] = v; } __syncthreads(); - const float inv_std = rsqrtf(smem[0] / (float)K + eps); + const float inv_std = rsqrtf(smem[0]/(float)K+eps); - for (int k = tid; k < K; k += BLOCK) { - float v = (__bfloat162float(row[k]) - mean) * inv_std; + for (int k=tid; k= N) return; + const int i = (int)(blockIdx.x*blockDim.x+threadIdx.x); + if (i >= N) + return; float x = __bfloat162float(in[i]); const float kAlpha = 0.7978845608028654f; const float kBeta = 0.044715f; - float u = kAlpha * (x + kBeta * x * x * x); + float u = kAlpha * (x+kBeta * x * x * x); float t = tanhf(u); - out[i] = __float2bfloat16(0.5f * x * (1.0f + t)); + out[i] = __float2bfloat16(0.5f * x * (1.0f+t)); } __global__ void add_bias_bf16_kernel(const __nv_bfloat16* x, const __nv_bfloat16* bias, __nv_bfloat16* out, int M, int K) { const int m = (int)blockIdx.x; - const int k = (int)(blockIdx.y * blockDim.x + threadIdx.x); - if (k >= K) return; - const size_t i = (size_t)m * K + k; - out[i] = __float2bfloat16(__bfloat162float(x[i]) + __bfloat162float(bias[k])); + const int k = (int)(blockIdx.y*blockDim.x+threadIdx.x); + if (k >= K) + return; + const size_t i = (size_t)m * K+k; + out[i] = __float2bfloat16(__bfloat162float(x[i])+__bfloat162float(bias[k])); } __global__ void zero_tail_bf16_kernel(__nv_bfloat16* x, int total_cols, int start_col) { const int m = (int)blockIdx.x; - const int k = (int)(start_col + blockIdx.y * blockDim.x + threadIdx.x); - if (k >= total_cols) return; - x[(size_t)m * total_cols + k] = __float2bfloat16(0.0f); + const int k = (int)(start_col+blockIdx.y*blockDim.x+threadIdx.x); + if (k >= total_cols) + return; + x[(size_t)m * total_cols+k] = __float2bfloat16(0.0f); } extern "C" void bitvla_layernorm_bf16(const __nv_bfloat16* x, const __nv_bfloat16* w, @@ -602,20 +647,21 @@ extern "C" void bitvla_layernorm_bf16(const __nv_bfloat16* x, const __nv_bfloat1 extern "C" void bitvla_gelu_tanh_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int N, cudaStream_t stream) { constexpr int B = 256; - gelu_tanh_bf16_kernel<<>>(x, out, N); + gelu_tanh_bf16_kernel<<>>(x, out, N); } extern "C" void bitvla_add_bias_bf16(const __nv_bfloat16* x, const __nv_bfloat16* bias, __nv_bfloat16* out, int M, int K, cudaStream_t stream) { constexpr int B = 256; - const int n_kb = (K + B - 1) / B; + const int n_kb = (K+B-1)/B; add_bias_bf16_kernel<<>>(x, bias, out, M, K); } extern "C" void bitvla_zero_tail_bf16(__nv_bfloat16* x, int M, int total_cols, int start_col, cudaStream_t stream) { - if (start_col >= total_cols) return; + if (start_col >= total_cols) + return; constexpr int B = 128; - const int len = total_cols - start_col; - const int n_kb = (len + B - 1) / B; + const int len = total_cols-start_col; + const int n_kb = (len+B-1)/B; zero_tail_bf16_kernel<<>>(x, total_cols, start_col); } @@ -633,11 +679,12 @@ extern "C" int bitvla_lm_cuda_forward(bitvla_lm_cuda_ctx* ctx, std::vector<__nv_bfloat16> h_dump; std::vector h_dump_f32; auto dump_to_file = [&](const char* name, const __nv_bfloat16* d_ptr) { - if (!dump_dir) return; + if (!dump_dir) + return; const size_t n = (size_t) seq * ctx->hidden; h_dump.resize(n); h_dump_f32.resize(n); cudaMemcpy(h_dump.data(), d_ptr, n * sizeof(__nv_bfloat16), cudaMemcpyDeviceToHost); - for (size_t i = 0; i < n; ++i) { + for (size_t i=0; i(&h_dump[i]); uint32_t b = ((uint32_t) u) << 16; float f; std::memcpy(&f, &b, 4); @@ -645,18 +692,22 @@ extern "C" int bitvla_lm_cuda_forward(bitvla_lm_cuda_ctx* ctx, } std::string path = std::string(dump_dir) + "/" + name + ".bin"; FILE* f = std::fopen(path.c_str(), "wb"); - if (f) { std::fwrite(h_dump_f32.data(), sizeof(float), n, f); std::fclose(f); } + if (f) { + std::fwrite(h_dump_f32.data(), sizeof(float), n, f); + std::fclose(f); + } }; - CUDA_OK(cudaMemcpyAsync(ctx->d_h, d_in, (size_t)seq * ctx->hidden * sizeof(__nv_bfloat16), + CUDA_OK(cudaMemcpyAsync(ctx->d_h, d_in, (size_t)seq * ctx->hidden*sizeof(__nv_bfloat16), cudaMemcpyDeviceToDevice, stream)); if (dump_dir) { cudaStreamSynchronize(stream); dump_to_file("lm_layer_input", ctx->d_h); } - for (int L = 0; L < ctx->n_layers; ++L) { + for (int L=0; Ln_layers; ++L) { int rc = run_layer(ctx, L, seq, stream); - if (rc != 0) return rc; + if (rc != 0) + return rc; if (dump_dir) { cudaStreamSynchronize(stream); char name[64]; std::snprintf(name, sizeof(name), "lm_layer_%d", L); diff --git a/src/kernels/bitvla/bitvla_lm_cuda.h b/src/kernels/bitvla/bitvla_lm_cuda.h index 60a64b3..ee79799 100644 --- a/src/kernels/bitvla/bitvla_lm_cuda.h +++ b/src/kernels/bitvla/bitvla_lm_cuda.h @@ -20,10 +20,10 @@ * packed format ("ladder int8xint2") alongside FP32 weight scales. This * header exposes: * - * * Standalone bf16 ops (norm / RoPE / softmax / activations) used by + * * Standalone bf16 ops (norm/RoPE/softmax/activations) used by * both the LM and the ViT. * * @ref bitvla_lm_cuda_ctx, an opaque context that owns the device-side - * LM state, plus its init / set-layer / forward / free entry points. + * LM state, plus its init/set-layer/forward/free entry points. * * All functions are @c extern @c "C" so they can be called from C++ or C * driver code. Streams are passed in explicitly; the kernels never call @@ -81,7 +81,7 @@ void bitvla_softmax_scaled_bf16(__nv_bfloat16* inout, float scale, int n_rows, int S, cudaStream_t stream); /** - * @brief Elementwise @c relu(g)^2 * u (BitVLA squared-ReLU FFN gate). + * @brief Elementwise @c relu(g)^2*u (BitVLA squared-ReLU FFN gate). * @param g Gate input (N), bf16 device pointer. * @param u Up input (N), bf16 device pointer. * @param out Output (N), bf16 device pointer. @@ -155,7 +155,7 @@ void bitvla_gather_rows_bf16(const __nv_bfloat16* in, __nv_bfloat16* out, cudaStream_t stream); /** - * @brief Affine LayerNorm in bf16 (mean/variance + scale + bias). + * @brief Affine LayerNorm in bf16 (mean/variance+scale+bias). * @param x Input matrix (M x K), bf16 device pointer. * @param w Per-channel scale (length K). * @param b Per-channel bias (length K). diff --git a/src/kernels/bitvla/bitvla_vit_cuda.cu b/src/kernels/bitvla/bitvla_vit_cuda.cu index 08ae84d..2f48e3e 100644 --- a/src/kernels/bitvla/bitvla_vit_cuda.cu +++ b/src/kernels/bitvla/bitvla_vit_cuda.cu @@ -31,15 +31,16 @@ extern "C" void bitvla_act_quant_cuda(const __nv_bfloat16* in, int8_t* out, int M, int K, cudaStream_t stream); __global__ void gelu_erf_bf16_kernel(const __nv_bfloat16* in, __nv_bfloat16* out, int N) { - const int i = (int)(blockIdx.x * blockDim.x + threadIdx.x); - if (i >= N) return; + const int i = (int)(blockIdx.x*blockDim.x+threadIdx.x); + if (i >= N) + return; const float x = __bfloat162float(in[i]); const float inv_sqrt2 = 0.7071067811865475f; - out[i] = __float2bfloat16(0.5f * x * (1.0f + erff(x * inv_sqrt2))); + out[i] = __float2bfloat16(0.5f * x * (1.0f+erff(x * inv_sqrt2))); } static void gelu_erf_bf16(const __nv_bfloat16* in, __nv_bfloat16* out, int N, cudaStream_t stream) { constexpr int B = 256; - gelu_erf_bf16_kernel<<>>(in, out, N); + gelu_erf_bf16_kernel<<>>(in, out, N); } #define CUDA_OK(call) do { cudaError_t e = (call); if (e != cudaSuccess) { \ @@ -96,14 +97,14 @@ bitvla_vit_cuda_ctx* bitvla_vit_cuda_init(int n_layers, int hidden, int n_heads, ctx->n_layers = n_layers; ctx->hidden = hidden; ctx->n_heads = n_heads; - ctx->head_dim = hidden / n_heads; + ctx->head_dim = hidden/n_heads; ctx->ffn = ffn; ctx->n_patches = n_patches; ctx->patch_flat = patch_flat; ctx->ln_eps = ln_eps; ctx->mm_out = mm_out; - ctx->ffn_pad = ((ffn + 127) / 128) * 128; + ctx->ffn_pad = ((ffn+127)/128)*128; ctx->layers.resize(n_layers); if (cublasCreate(&ctx->cublas) != CUBLAS_STATUS_SUCCESS) { @@ -119,24 +120,25 @@ bitvla_vit_cuda_ctx* bitvla_vit_cuda_init(int n_layers, int hidden, int n_heads, CUDA_OKV(cudaMalloc(&ctx->d_q_proj, (size_t) n_patches * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_k_proj, (size_t) n_patches * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_v_proj, (size_t) n_patches * hidden * bf16)); - CUDA_OKV(cudaMalloc(&ctx->d_q_HShd, (size_t) n_heads * n_patches * ctx->head_dim * bf16)); - CUDA_OKV(cudaMalloc(&ctx->d_k_HShd, (size_t) n_heads * n_patches * ctx->head_dim * bf16)); - CUDA_OKV(cudaMalloc(&ctx->d_v_HShd, (size_t) n_heads * n_patches * ctx->head_dim * bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_q_HShd, (size_t) n_heads * n_patches * ctx->head_dim*bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_k_HShd, (size_t) n_heads * n_patches * ctx->head_dim*bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_v_HShd, (size_t) n_heads * n_patches * ctx->head_dim*bf16)); CUDA_OKV(cudaMalloc(&ctx->d_scores, (size_t) n_heads * n_patches * n_patches * bf16)); - CUDA_OKV(cudaMalloc(&ctx->d_attn_out, (size_t) n_heads * n_patches * ctx->head_dim * bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_attn_out, (size_t) n_heads * n_patches * ctx->head_dim*bf16)); CUDA_OKV(cudaMalloc(&ctx->d_attn_merged, (size_t) n_patches * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_o_out, (size_t) n_patches * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_fc1_dense, (size_t) n_patches * ffn * bf16)); - CUDA_OKV(cudaMalloc(&ctx->d_fc1_padded, (size_t) n_patches * ctx->ffn_pad * bf16)); + CUDA_OKV(cudaMalloc(&ctx->d_fc1_padded, (size_t) n_patches * ctx->ffn_pad*bf16)); - CUDA_OKV(cudaMemset(ctx->d_fc1_padded, 0, (size_t) n_patches * ctx->ffn_pad * bf16)); + CUDA_OKV(cudaMemset(ctx->d_fc1_padded, 0, (size_t) n_patches * ctx->ffn_pad*bf16)); CUDA_OKV(cudaMalloc(&ctx->d_fc2_out, (size_t) n_patches * hidden * bf16)); CUDA_OKV(cudaMalloc(&ctx->d_mm_h1, (size_t) n_patches * mm_out * bf16)); return ctx; } void bitvla_vit_cuda_free(bitvla_vit_cuda_ctx* ctx) { - if (!ctx) return; + if (!ctx) + return; cublasDestroy(ctx->cublas); cudaFree(ctx->d_h); cudaFree(ctx->d_h_norm); cudaFree(ctx->d_act_int8_h); cudaFree(ctx->d_act_int8_ffn); cudaFree(ctx->d_act_s); @@ -200,9 +202,12 @@ static int run_vit_layer(bitvla_vit_cuda_ctx* ctx, int L, cudaStream_t stream) { &beta, ctx->d_scores, CUDA_R_16BF, seq, (long long) seq * seq, n_heads, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_vit_cuda): QK^T gemm @L%d failed (%d)\n", L, cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_vit_cuda): QK^T gemm @L%d failed (%d)\n", L, cbs); + return -1; + } - const float scl = 1.0f / std::sqrt((float) hd); + const float scl = 1.0f/std::sqrt((float) hd); bitvla_softmax_scaled_bf16(ctx->d_scores, scl, n_heads * seq, seq, stream); cbs = cublasGemmStridedBatchedEx( @@ -213,7 +218,10 @@ static int run_vit_layer(bitvla_vit_cuda_ctx* ctx, int L, cudaStream_t stream) { &beta, ctx->d_attn_out, CUDA_R_16BF, hd, (long long) seq * hd, n_heads, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_vit_cuda): attn@V gemm @L%d failed (%d)\n", L, cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_vit_cuda): attn@V gemm @L%d failed (%d)\n", L, cbs); + return -1; + } bitvla_transpose_NshHd_to_sNhd_bf16(ctx->d_attn_out, ctx->d_attn_merged, n_heads, seq, hd, stream); @@ -265,15 +273,19 @@ int bitvla_vit_cuda_forward(bitvla_vit_cuda_ctx* ctx, &beta, ctx->d_h, CUDA_R_16BF, H, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_vit_cuda): patch_embed gemm failed (%d)\n", cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_vit_cuda): patch_embed gemm failed (%d)\n", cbs); + return -1; + } bitvla_add_bias_bf16(ctx->d_h, ctx->patch_b, ctx->d_h, seq, H, stream); bitvla_add_bf16(ctx->d_h, ctx->pos_emb, ctx->d_h, seq * H, stream); - for (int L = 0; L < ctx->n_layers; ++L) { + for (int L=0; Ln_layers; ++L) { int rc = run_vit_layer(ctx, L, stream); - if (rc != 0) return rc; + if (rc != 0) + return rc; } cbs = cublasGemmEx( @@ -284,7 +296,10 @@ int bitvla_vit_cuda_forward(bitvla_vit_cuda_ctx* ctx, &beta, ctx->d_mm_h1, CUDA_R_16BF, M, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_vit_cuda): MM linear_1 gemm failed (%d)\n", cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_vit_cuda): MM linear_1 gemm failed (%d)\n", cbs); + return -1; + } bitvla_add_bias_bf16(ctx->d_mm_h1, ctx->mm_b1, ctx->d_mm_h1, seq, M, stream); gelu_erf_bf16(ctx->d_mm_h1, ctx->d_mm_h1, seq * M, stream); @@ -296,7 +311,10 @@ int bitvla_vit_cuda_forward(bitvla_vit_cuda_ctx* ctx, &beta, d_out, CUDA_R_16BF, M, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); - if (cbs != CUBLAS_STATUS_SUCCESS) { std::fprintf(stderr, "vla(bitvla_vit_cuda): MM linear_2 gemm failed (%d)\n", cbs); return -1; } + if (cbs != CUBLAS_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(bitvla_vit_cuda): MM linear_2 gemm failed (%d)\n", cbs); + return -1; + } bitvla_add_bias_bf16(d_out, ctx->mm_b2, d_out, seq, M, stream); return 0; } diff --git a/src/kernels/bitvla/bitvla_vit_cuda.h b/src/kernels/bitvla/bitvla_vit_cuda.h index 963ba5e..5fee0b6 100644 --- a/src/kernels/bitvla/bitvla_vit_cuda.h +++ b/src/kernels/bitvla/bitvla_vit_cuda.h @@ -14,7 +14,7 @@ /** * @file bitvla_vit_cuda.h - * @brief CUDA forward path for the BitVLA vision tower (ViT + mmproj). + * @brief CUDA forward path for the BitVLA vision tower (ViT+mmproj). * * Mirrors the layout of @ref bitvla_lm_cuda.h: an opaque * @ref bitvla_vit_cuda_ctx owns device buffers; per-layer weights are diff --git a/src/layers/attn.h b/src/layers/attn.h new file mode 100644 index 0000000..bef6af2 --- /dev/null +++ b/src/layers/attn.h @@ -0,0 +1,59 @@ +// 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. + + +// nv == 1 emits the same nodes as the 2d/3d spelling, since ggml pads every +// shape to four dimensions. + +#pragma once + +#include "ggml.h" + +#include + +namespace vla { + +inline ggml_tensor * to_heads(ggml_context * C, ggml_tensor * p, int64_t hd, int64_t heads, + int64_t T, int64_t nv = 1) { + return ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, p, hd, heads, T, nv), 0, 2, 1, 3)); +} + +// V is pre-transposed so that mul_mat(V, aw) lands the right way round. +inline ggml_tensor * to_heads_v(ggml_context * C, ggml_tensor * p, int64_t hd, int64_t heads, + int64_t T, int64_t nv = 1) { + return ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, p, hd, heads, T, nv), 1, 2, 0, 3)); +} + +inline ggml_tensor * attention(ggml_context * C, ggml_tensor * Q, ggml_tensor * K, ggml_tensor * V, + ggml_tensor * mask, float scale, int64_t dim, int64_t T, int64_t nv = 1) { + 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, mask, scale, 0.0f); + ggml_tensor * kqv = ggml_mul_mat(C, V, aw); + return ggml_reshape_3d(C, ggml_cont(C, ggml_permute(C, kqv, 0, 2, 1, 3)), dim, T, nv); +} + +// Takes V laid out like Q/K, not the transposed to_heads_v form. +inline ggml_tensor * flash_attention(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]); +} + +} diff --git a/src/layers/embed.h b/src/layers/embed.h new file mode 100644 index 0000000..4439020 --- /dev/null +++ b/src/layers/embed.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. + +// The sin/cos order differs per family and each matches its reference; +// tests/test_dit_common.cpp pins all three. + +#pragma once + +#include +#include +#include +#include + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +namespace vla { + +// cos first, then sin. +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 & 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 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 & out) { + const float NEG = -std::numeric_limits::infinity(); + + out.assign((size_t)seq*seq, 0.0f); + for (int64_t q=0; q + +namespace vla { + +inline ggml_tensor * linear(ggml_context * C, ggml_tensor * W, ggml_tensor * b, ggml_tensor * x) { + ggml_tensor * y = ggml_mul_mat(C, W, x); + return b ? ggml_add(C, y, b) : y; +} + +// One row of a stacked [out, in, n_embodiment] weight. +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]; + const int64_t 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])); +} + +// Block `blk` of a fused [nblk*E, T] projection, laid out as heads. +inline 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); +} + +} diff --git a/src/layers/norm.h b/src/layers/norm.h new file mode 100644 index 0000000..a10acf7 --- /dev/null +++ b/src/layers/norm.h @@ -0,0 +1,47 @@ +// 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. + + +#pragma once + +#include "layers/linear.h" + +#include "ggml.h" + +#include +#include + +namespace vla { + +inline ggml_tensor * layer_norm(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 * rms_norm(ggml_context * C, ggml_tensor * x, ggml_tensor * w, float eps) { + ggml_tensor * n = ggml_rms_norm(C, x, eps); + return w ? ggml_mul(C, n, w) : n; +} + +// cond is (scale, shift) here; the DiT final projection uses (shift, scale). +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 = linear(C, lw, lb, ggml_silu(C, temb)); + ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0); + ggml_tensor * 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); +} + +} diff --git a/src/layers/rope.h b/src/layers/rope.h new file mode 100644 index 0000000..ca06156 --- /dev/null +++ b/src/layers/rope.h @@ -0,0 +1,78 @@ +// 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. + +// Three conventions, not interchangeable; pinned by +// tests/test_rope_conventions.cpp. rope_pairwise deliberately pairs a +// half-split frequency table with an interleaved rotation, as the VLA-Adapter +// reference does. + +#pragma once + +#include "ggml.h" + +#include + +namespace vla { + +struct RopeSpec { + int type = GGML_ROPE_TYPE_NEOX; + int n_dims = 0; + int sections[4]= {0, 0, 0, 0}; + float freq_base = 10000.0f; + float freq_scale = 1.0f; + float ext_factor = 0.0f; + float attn_factor= 1.0f; + float beta_fast = 32.0f; + float beta_slow = 1.0f; +}; + +inline ggml_tensor * rope(ggml_context * C, const RopeSpec & r, ggml_tensor * x, ggml_tensor * pos) { + if (r.type == GGML_ROPE_TYPE_IMROPE) { + int sect[4] = { r.sections[0], r.sections[1], r.sections[2], r.sections[3] }; + return ggml_rope_multi(C, x, pos, nullptr, r.n_dims, sect, r.type, 0, + r.freq_base, r.freq_scale, r.ext_factor, r.attn_factor, r.beta_fast, r.beta_slow); + } + return ggml_rope_ext(C, x, pos, nullptr, r.n_dims, r.type, 0, + r.freq_base, r.freq_scale, r.ext_factor, r.attn_factor, r.beta_fast, r.beta_slow); +} + +inline ggml_tensor * rope_2d(ggml_context * C, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_t) { + const int64_t hd = x->ne[0]; + const int64_t S = x->ne[1]; + const int64_t 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 ggml_tensor * rope_pairwise_rot(ggml_context * C, ggml_tensor * x, int64_t HD) { + const int64_t L = x->ne[1]; + const int64_t H = x->ne[2]; + + ggml_tensor * xp = ggml_reshape_4d(C, x, 2, HD/2, L, H); + ggml_tensor * ev = ggml_cont(C, ggml_view_4d(C, xp, 1, HD/2, L, H, xp->nb[1], xp->nb[2], xp->nb[3], 0)); + ggml_tensor * od = ggml_cont(C, ggml_view_4d(C, xp, 1, HD/2, L, H, xp->nb[1], xp->nb[2], xp->nb[3], xp->nb[0])); + return ggml_reshape_3d(C, ggml_concat(C, ggml_scale(C, od, -1.0f), ev, 0), HD, L, H); +} + +inline ggml_tensor * rope_pairwise(ggml_context * C, ggml_tensor * x, ggml_tensor * cs, ggml_tensor * sn, int64_t HD) { + ggml_tensor * c = ggml_reshape_3d(C, cs, HD, x->ne[1], 1); + ggml_tensor * s = ggml_reshape_3d(C, sn, HD, x->ne[1], 1); + return ggml_add(C, ggml_mul(C, x, c), ggml_mul(C, rope_pairwise_rot(C, x, HD), s)); +} + +} diff --git a/src/loader.cpp b/src/loader.cpp new file mode 100644 index 0000000..9a67f05 --- /dev/null +++ b/src/loader.cpp @@ -0,0 +1,182 @@ +// 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. + +#include "loader.h" + +#include +#include +#include + +namespace vla { + +namespace { + +constexpr size_t NAME_CAP = 256; + +} + +ggml_tensor * WeightLoader::declare(ggml_type want, bool required, bool gemma_norm, + const char * fmt, va_list ap) { + char name[NAME_CAP]; + const int n = std::vsnprintf(name, sizeof(name), fmt, ap); + if (n < 0 || (size_t)n >= sizeof(name)) { + std::fprintf(stderr, "vla(%s): tensor name too long for a %zu-byte buffer\n", arch_, sizeof(name)); + ok_ = false; + return nullptr; + } + + const ggml_tensor * src = g_.meta(name); + if (!src) { + if (required) { + std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch_, name); + ok_ = false; + } + return nullptr; + } + + ggml_tensor * t = ggml_new_tensor(ctx_, g_.resident_type(src, want), ggml_n_dims(src), src->ne); + if (!t) { + std::fprintf(stderr, "vla(%s): ggml_new_tensor failed for %s (weight context too small?)\n", arch_, name); + ok_ = false; + return nullptr; + } + ggml_set_name(t, name); + + if (gemma_norm) + gemma_norms_.push_back(name); + return t; +} + +#define VLA_DECLARE_FN(fn, type, required, gemma) \ + ggml_tensor * WeightLoader::fn(const char * fmt, ...) { \ + va_list ap; \ + va_start(ap, fmt); \ + ggml_tensor * t = declare(type, required, gemma, fmt, ap); \ + va_end(ap); \ + return t; \ + } + +VLA_DECLARE_FN(gemm, gemm_, true, false) +VLA_DECLARE_FN(f32, GGML_TYPE_F32, true, false) +VLA_DECLARE_FN(opt_gemm, gemm_, false, false) +VLA_DECLARE_FN(opt_f32, GGML_TYPE_F32, false, false) +VLA_DECLARE_FN(f32_gemma_norm, GGML_TYPE_F32, true, true) + +#undef VLA_DECLARE_FN + +ggml_tensor * WeightLoader::typed(ggml_type want, const char * fmt, ...) { + va_list ap; + va_start(ap, fmt); + ggml_tensor * t = declare(want, true, false, fmt, ap); + va_end(ap); + return t; +} + +ggml_tensor * WeightLoader::fuse_gemm(const char * out_name, const std::vector & srcs) { + return fuse(gemm_, out_name, srcs); +} + +ggml_tensor * WeightLoader::fuse_f32(const char * out_name, const std::vector & srcs) { + return fuse(GGML_TYPE_F32, out_name, srcs); +} + +ggml_tensor * WeightLoader::fuse(ggml_type want, const char * out_name, const std::vector & srcs) { + if (srcs.empty()) { + ok_ = false; + return nullptr; + } + + const ggml_tensor * first = g_.meta(srcs[0].c_str()); + if (!first) { + std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch_, srcs[0].c_str()); + ok_ = false; + return nullptr; + } + + const bool is1d = ggml_n_dims(first) == 1; + int64_t rows = 0; + for (const std::string & s : srcs) { + const ggml_tensor * gs = g_.meta(s.c_str()); + if (!gs) { + std::fprintf(stderr, "vla(%s): missing tensor %s\n", arch_, s.c_str()); + ok_ = false; + return nullptr; + } + rows += is1d ? gs->ne[0] : gs->ne[1]; + } + + ggml_tensor * t = is1d ? ggml_new_tensor_1d(ctx_, want, rows) + : ggml_new_tensor_2d(ctx_, want, first->ne[0], rows); + if (!t) { + std::fprintf(stderr, "vla(%s): ggml_new_tensor failed for %s\n", arch_, out_name); + ok_ = false; + return nullptr; + } + ggml_set_name(t, out_name); + fused_.push_back(Fused{t, srcs}); + return t; +} + +bool WeightLoader::upload(ggml_backend_t backend, ggml_backend_buffer_t * out_buf) { + if (!ok_) { + std::fprintf(stderr, "vla(%s): weight tensor setup failed\n", arch_); + return false; + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx_, backend); + if (!buf) { + std::fprintf(stderr, "vla(%s): ggml_backend_alloc_ctx_tensors failed (OOM?)\n", arch_); + return false; + } + *out_buf = buf; + + for (ggml_tensor * t=ggml_get_first_tensor(ctx_); t; t=ggml_get_next_tensor(ctx_, t)) { + const char * name = ggml_get_name(t); + const bool fused = std::any_of(fused_.begin(), fused_.end(), + [&](const Fused & f) { return f.dst == t; }); + if (fused) + continue; + + const bool gn = std::find(gemma_norms_.begin(), gemma_norms_.end(), name) != gemma_norms_.end(); + + std::vector bytes = g_.read_convert(name, t->type, gn); + if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { + std::fprintf(stderr, "vla(%s): failed to load %s (%zu vs %zu bytes)\n", + arch_, name, bytes.size(), ggml_nbytes(t)); + return false; + } + ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); + } + + for (const Fused & f : fused_) { + std::vector buf; + for (const std::string & s : f.srcs) { + std::vector b = g_.read_convert(s.c_str(), f.dst->type); + if (b.empty()) { + std::fprintf(stderr, "vla(%s): fused fill: read %s failed\n", arch_, s.c_str()); + return false; + } + buf.insert(buf.end(), b.begin(), b.end()); + } + if (buf.size() != ggml_nbytes(f.dst)) { + std::fprintf(stderr, "vla(%s): fused fill: %s size %zu vs %zu\n", + arch_, ggml_get_name(f.dst), buf.size(), ggml_nbytes(f.dst)); + return false; + } + ggml_backend_tensor_set(f.dst, buf.data(), 0, buf.size()); + } + return true; +} + +} diff --git a/src/loader.h b/src/loader.h new file mode 100644 index 0000000..649c9e9 --- /dev/null +++ b/src/loader.h @@ -0,0 +1,87 @@ +// 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. + +// gemm() lands in the model's matmul type, unless the GGUF holds the tensor +// quantized, in which case it stays packed and ggml dequantizes at compute. +// A miss is recorded and surfaces once, at ok(). + +#pragma once + +#include "gguf_reader.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include + +namespace vla { + +class WeightLoader { +public: + WeightLoader(const char * arch, gguf_reader & g, ggml_context * ctx, ggml_type gemm_type) + : arch_(arch), g_(g), ctx_(ctx), gemm_(gemm_type) {} + + WeightLoader(const WeightLoader &) = delete; + WeightLoader & operator=(const WeightLoader &) = delete; + + ggml_tensor * gemm(const char * fmt, ...) __attribute__((format(printf, 2, 3))); + ggml_tensor * f32 (const char * fmt, ...) __attribute__((format(printf, 2, 3))); + + // Explicit resident type, for weights that are neither a plain GEMM input + // nor F32 (BitVLA's int2-packed ternary blocks). + ggml_tensor * typed(ggml_type want, const char * fmt, ...) __attribute__((format(printf, 3, 4))); + + // A miss is not an error. + ggml_tensor * opt_gemm(const char * fmt, ...) __attribute__((format(printf, 2, 3))); + ggml_tensor * opt_f32 (const char * fmt, ...) __attribute__((format(printf, 2, 3))); + + // Gemma norms are centred on zero and add 1 at use. + ggml_tensor * f32_gemma_norm(const char * fmt, ...) __attribute__((format(printf, 2, 3))); + + // Several GGUF tensors concatenated into one resident tensor; out_name is + // synthetic and need not exist in the file. + ggml_tensor * fuse_gemm(const char * out_name, const std::vector & srcs); + ggml_tensor * fuse_f32 (const char * out_name, const std::vector & srcs); + + ggml_type gemm_type() const { + return gemm_; + } + bool ok() const { + return ok_; + } + + bool upload(ggml_backend_t backend, ggml_backend_buffer_t * out_buf); + +private: + ggml_tensor * declare(ggml_type want, bool required, bool gemma_norm, const char * fmt, va_list ap); + ggml_tensor * fuse(ggml_type want, const char * out_name, const std::vector & srcs); + + const char * arch_; + gguf_reader & g_; + ggml_context * ctx_; + ggml_type gemm_; + bool ok_ = true; + + struct Fused { + ggml_tensor * dst; + std::vector srcs; + }; + + std::vector gemma_norms_; + std::vector fused_; +}; + +} diff --git a/src/model.cpp b/src/model.cpp index 3b09ceb..b700ec3 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -14,6 +14,7 @@ #include "arch.h" #include "model.h" +#include "options.h" #include "gguf.h" @@ -33,8 +34,9 @@ struct Model { namespace { bool ends_with_gguf(const std::string& p) { - if (p.size() < 5) return false; - return std::strcmp(p.c_str() + p.size() - 5, ".gguf") == 0; + if (p.size() < 5) + return false; + return std::strcmp(p.c_str()+p.size()-5, ".gguf") == 0; } bool detect_arch_gguf(const std::string& path, Arch* out) { @@ -42,16 +44,20 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { p.no_alloc = true; p.ctx = nullptr; gguf_context * gctx = gguf_init_from_file(path.c_str(), p); - if (!gctx) return false; + if (!gctx) + return false; auto try_str = [&](const char * key, std::string& val) -> bool { const int64_t kid = gguf_find_key(gctx, key); - if (kid < 0) return false; + if (kid < 0) + return false; // gguf_get_val_str asserts (aborts) if the key is not a string, so a // malformed GGUF would kill the process here. Fail the probe instead. - if (gguf_get_kv_type(gctx, kid) != GGUF_TYPE_STRING) return false; + if (gguf_get_kv_type(gctx, kid) != GGUF_TYPE_STRING) + return false; const char * s = gguf_get_val_str(gctx, kid); - if (!s) return false; + if (!s) + return false; val = s; return true; }; @@ -70,17 +76,50 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { try_str("openvla_oft.architecture", arch_str) || try_str("vla_jepa.architecture", arch_str) || try_str("vla_adapter.architecture", arch_str)) { - if (arch_str == "smolvla") { *out = Arch::SMOLVLA; ok = true; } - else if (arch_str == "pi0") { *out = Arch::PI0; ok = true; } - else if (arch_str == "pi05") { *out = Arch::PI05; ok = true; } - else if (arch_str == "evo1") { *out = Arch::EVO1; ok = true; } - else if (arch_str == "gr00t_n1_5") { *out = Arch::GR00T_N1_5; ok = true; } - else if (arch_str == "gr00t_n1_6") { *out = Arch::GR00T_N1_6; ok = true; } - else if (arch_str == "gr00t_n1_7") { *out = Arch::GR00T_N1_7; ok = true; } - else if (arch_str == "bitvla") { *out = Arch::BITVLA; ok = true; } - else if (arch_str == "vla_adapter"){ *out = Arch::VLA_ADAPTER;ok = true; } - else if (arch_str == "openvla_oft"){ *out = Arch::OPENVLA_OFT;ok = true; } - else if (arch_str == "vla_jepa") { *out = Arch::VLA_JEPA; ok = true; } + if (arch_str == "smolvla") { + *out = Arch::SMOLVLA; + ok = true; + } + else if (arch_str == "pi0") { + *out = Arch::PI0; + ok = true; + } + else if (arch_str == "pi05") { + *out = Arch::PI05; + ok = true; + } + else if (arch_str == "evo1") { + *out = Arch::EVO1; + ok = true; + } + else if (arch_str == "gr00t_n1_5") { + *out = Arch::GR00T_N1_5; + ok = true; + } + else if (arch_str == "gr00t_n1_6") { + *out = Arch::GR00T_N1_6; + ok = true; + } + else if (arch_str == "gr00t_n1_7") { + *out = Arch::GR00T_N1_7; + ok = true; + } + else if (arch_str == "bitvla") { + *out = Arch::BITVLA; + ok = true; + } + else if (arch_str == "vla_adapter"){ + *out = Arch::VLA_ADAPTER; + ok = true; + } + else if (arch_str == "openvla_oft"){ + *out = Arch::OPENVLA_OFT; + ok = true; + } + else if (arch_str == "vla_jepa") { + *out = Arch::VLA_JEPA; + ok = true; + } } gguf_free(gctx); @@ -113,13 +152,16 @@ namespace { bool detect_arch_safetensors(const std::string& path, Arch* out) { std::ifstream f(path, std::ios::binary); - if (!f) return false; + if (!f) + return false; uint64_t header_size = 0; f.read(reinterpret_cast(&header_size), sizeof(header_size)); - if (!f || header_size == 0 || header_size > (1u << 28)) return false; + if (!f || header_size == 0 || header_size > (1u << 28)) + return false; std::string header(header_size, '\0'); f.read(header.data(), header_size); - if (!f) return false; + if (!f) + return false; if (header.find("vlm_with_expert.vlm.") != std::string::npos) { *out = Arch::SMOLVLA; @@ -143,13 +185,20 @@ bool detect_arch_safetensors(const std::string& path, Arch* out) { } bool detect_arch_from_ckpt(const std::string& ckpt_path, Arch* out) { - if (!out) return false; - if (ends_with_gguf(ckpt_path)) return detect_arch_gguf(ckpt_path, out); + if (!out) + return false; + if (ends_with_gguf(ckpt_path)) + return detect_arch_gguf(ckpt_path, out); return detect_arch_safetensors(ckpt_path, out); } Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, const std::string& config_path) { + return model_load(mmproj_path, ckpt_path, config_path, Options{}); +} + +Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, + const std::string& config_path, const Options& opts) { Arch arch; if (!detect_arch_from_ckpt(ckpt_path, &arch)) { std::fprintf(stderr, @@ -160,53 +209,65 @@ Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, } std::unique_ptr impl; + { + std::string err; + if (!Options::reject_retired_env(err)) { + std::fprintf(stderr, "vla: %s\n", err.c_str()); + return nullptr; + } + } + + set_flash_attn(opts.flash_attn.value_or(false)); + set_mm_prec_f32(opts.mm_prec_f32.value_or(true)); + switch (arch) { case Arch::SMOLVLA: std::printf("vla: arch = smolvla\n"); - impl = smolvla_create(mmproj_path, ckpt_path, config_path); + impl = smolvla_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::PI0: std::printf("vla: arch = pi0\n"); - impl = pi0_create(mmproj_path, ckpt_path, config_path); + impl = pi0_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::PI05: std::printf("vla: arch = pi05\n"); - impl = pi05_create(mmproj_path, ckpt_path, config_path); + impl = pi05_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::EVO1: std::printf("vla: arch = evo1\n"); - impl = evo1_create(mmproj_path, ckpt_path, config_path); + impl = evo1_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::GR00T_N1_5: std::printf("vla: arch = gr00t_n1_5\n"); - impl = gr00t_n1_5_create(mmproj_path, ckpt_path, config_path); + impl = gr00t_n1_5_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::GR00T_N1_6: std::printf("vla: arch = gr00t_n1_6\n"); - impl = gr00t_n1_6_create(mmproj_path, ckpt_path, config_path); + impl = gr00t_n1_6_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::GR00T_N1_7: std::printf("vla: arch = gr00t_n1_7\n"); - impl = gr00t_n1_7_create(mmproj_path, ckpt_path, config_path); + impl = gr00t_n1_7_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::BITVLA: std::printf("vla: arch = bitvla\n"); - impl = bitvla_create(mmproj_path, ckpt_path, config_path); + impl = bitvla_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::VLA_ADAPTER: std::printf("vla: arch = vla_adapter\n"); - impl = vla_adapter_create(mmproj_path, ckpt_path, config_path); + impl = vla_adapter_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::OPENVLA_OFT: std::printf("vla: arch = openvla_oft\n"); - impl = openvla_oft_create(mmproj_path, ckpt_path, config_path); + impl = openvla_oft_create(mmproj_path, ckpt_path, config_path, opts); break; case Arch::VLA_JEPA: std::printf("vla: arch = vla_jepa\n"); - impl = vla_jepa_create(mmproj_path, ckpt_path, config_path); + impl = vla_jepa_create(mmproj_path, ckpt_path, config_path, opts); break; } - if (!impl) return nullptr; + 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; diff --git a/src/model.h b/src/model.h index 6265f34..dfdb170 100644 --- a/src/model.h +++ b/src/model.h @@ -32,6 +32,10 @@ namespace vla { +// Declared in options.h. Kept incomplete here so this header stays free of +// ggml, which tests/test_vision_common.cpp and the pure CI job rely on. +struct Options; + /** * @brief Resolved hyper-parameters of a loaded model. * @@ -150,7 +154,7 @@ struct Inputs { }; /** - * @brief Load a model from one (vision-baked) or two (mmproj + ckpt) GGUFs. + * @brief Load a model from one (vision-baked) or two (mmproj+ckpt) GGUFs. * * The architecture is detected from the checkpoint via * @ref detect_arch_from_ckpt. Fails loud: a missing file, unknown @@ -166,6 +170,13 @@ struct Inputs { Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, const std::string& config_path = ""); +/** + * @brief Load with explicit runtime options; include options.h to use it. + * @copydetails model_load + */ +Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, + const std::string& config_path, const Options& opts); + /** * @brief Release a model handle returned by @ref model_load. * @param m Handle to free; may be @c nullptr. diff --git a/src/models/bitvla.cpp b/src/models/bitvla.cpp index b8a99e2..76e7e66 100644 --- a/src/models/bitvla.cpp +++ b/src/models/bitvla.cpp @@ -13,6 +13,8 @@ // limitations under the License. #include "arch.h" +#include "loader.h" +#include "options.h" #include "model.h" #include "ggml.h" @@ -22,13 +24,14 @@ #include "ggml-cuda.h" #endif #include "gguf.h" -#include "models/gguf_reader.h" -#include "models/scratch_ctx.h" +#include "gguf_reader.h" +#include "scratch_ctx.h" #ifdef VLA_BITVLA_CUDA_KERNELS #include "kernels/bitvla/bitvla_lm_cuda.h" #include "kernels/bitvla/bitvla_vit_cuda.h" #include "kernels/bitvla/bitvla_fp32head_cuda.h" +#include "env_flag.h" #ifdef __GLIBC__ # include #endif @@ -66,23 +69,27 @@ struct LmLayerW { void bitvla_act_quant_op(ggml_tensor * dst, const ggml_tensor * a, int ith, int nth, void * ) { const int64_t cols = a->ne[0]; const int64_t rows = ggml_nrows(a); - const int64_t per = (rows + nth - 1) / nth; + const int64_t per = (rows+nth-1)/nth; const int64_t r0 = ith * per; - const int64_t r1 = std::min(rows, r0 + per); + const int64_t r1 = std::min(rows, r0+per); const float * src = (const float *) a->data; float * out = (float *) dst->data; - for (int64_t r = r0; r < r1; ++r) { - const float * row_in = src + r * cols; - float * row_out = out + r * cols; + for (int64_t r=r0; r 127.0f) q = 127.0f; - if (q < -128.0f) q = -128.0f; + for (int64_t c=0; c 127.0f) + q = 127.0f; + if (q < -128.0f) + q = -128.0f; row_out[c] = q * inv_s; } } @@ -176,7 +183,7 @@ ggml_tensor * rmsnorm(ggml_context * C, ggml_tensor * x, ggml_tensor * w, float ggml_tensor * build_vit_layer(ggml_context * C, const VitLayerW & 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); + const float scale = 1.0f/std::sqrt((float) head_dim); ggml_tensor * x1 = layernorm(C, x, w.ln1w, w.ln1b, ln_eps); ggml_tensor * q = bit_linear(C, w.Wq, w.bq, x1); ggml_tensor * k = bit_linear(C, w.Wk, w.bk, x1); @@ -198,7 +205,7 @@ ggml_tensor * build_vit_layer(ggml_context * C, const VitLayerW & w, ggml_tensor ggml_tensor * build_lm_layer(ggml_context * C, const BitvlaModelArch & m, const LmLayerW & w, ggml_tensor * h, ggml_tensor * positions, int64_t seq) { const int64_t hd = m.lm_head_dim, n_q = m.lm_q, n_kv = m.lm_kv, hq = n_q * hd; - const float scale = 1.0f / std::sqrt((float) hd); + const float scale = 1.0f/std::sqrt((float) hd); ggml_tensor * hn = rmsnorm(C, h, w.attn_norm, m.lm_rms_eps); ggml_tensor * qp = bit_linear(C, w.Wq, nullptr, hn); @@ -235,52 +242,65 @@ bool parse_stats_json(const std::string & js, const char * env_key, std::vector & q01, std::vector & q99, std::vector & mask, std::string & resolved_key) { auto find_obj = [&](const std::string & where, const std::string & key) -> std::string { - const std::string q = std::string("\"") + key + "\""; + const std::string q = std::string("\"")+key + "\""; size_t p = where.find(q); if (p == std::string::npos) return ""; - p = where.find(':', p + q.size()); if (p == std::string::npos) return ""; + p = where.find(':', p+q.size()); if (p == std::string::npos) return ""; size_t s = where.find('{', p); if (s == std::string::npos) return ""; - int depth = 1; size_t i = s + 1; + int depth = 1; size_t i = s+1; while (i < where.size() && depth > 0) { if (where[i] == '{') depth++; else if (where[i] == '}') depth--; - if (depth == 0) break; + if (depth == 0) + break; i++; } - return where.substr(s, i - s + 1); + return where.substr(s, i-s+1); }; auto parse_array_floats = [](const std::string & body, const char * key, std::vector & out) { - const std::string q = std::string("\"") + key + "\""; + const std::string q = std::string("\"")+key + "\""; size_t p = body.find(q); if (p == std::string::npos) return false; p = body.find('[', p); if (p == std::string::npos) return false; size_t e = body.find(']', p); if (e == std::string::npos) return false; - std::string inner = body.substr(p + 1, e - p - 1); + std::string inner = body.substr(p+1, e-p-1); out.clear(); const char * s = inner.c_str(); while (*s) { - while (*s == ' ' || *s == '\t' || *s == '\n' || *s == ',') s++; - if (!*s) break; + while (*s == ' ' || *s == '\t' || *s == '\n' || *s == ',') + s++; + if (!*s) + break; char * end = nullptr; float v = std::strtof(s, &end); - if (end == s) break; + if (end == s) + break; out.push_back(v); s = end; } return true; }; auto parse_array_bools = [](const std::string & body, const char * key, std::vector & out) { - const std::string q = std::string("\"") + key + "\""; + const std::string q = std::string("\"")+key + "\""; size_t p = body.find(q); if (p == std::string::npos) return false; p = body.find('[', p); if (p == std::string::npos) return false; size_t e = body.find(']', p); if (e == std::string::npos) return false; - std::string inner = body.substr(p + 1, e - p - 1); + std::string inner = body.substr(p+1, e-p-1); out.clear(); size_t i = 0; while (i < inner.size()) { - while (i < inner.size() && (inner[i] == ' ' || inner[i] == ',' || inner[i] == '\n' || inner[i] == '\t')) i++; - if (i >= inner.size()) break; - if (inner.compare(i, 4, "true") == 0) { out.push_back(1); i += 4; } - else if (inner.compare(i, 5, "false") == 0) { out.push_back(0); i += 5; } - else i++; + while (i < inner.size() && (inner[i] == ' ' || inner[i] == ',' || inner[i] == '\n' || inner[i] == '\t')) + i++; + if (i >= inner.size()) + break; + if (inner.compare(i, 4, "true") == 0) { + out.push_back(1); + i += 4; + } + else if (inner.compare(i, 5, "false") == 0) { + out.push_back(0); + i += 5; + } + else + i++; } return true; }; @@ -291,20 +311,30 @@ bool parse_stats_json(const std::string & js, const char * env_key, } else { size_t p = js.find('"'); - if (p == std::string::npos) return false; - size_t q = js.find('"', p + 1); - if (q == std::string::npos) return false; - suite = js.substr(p + 1, q - p - 1); + if (p == std::string::npos) + return false; + size_t q = js.find('"', p+1); + if (q == std::string::npos) + return false; + suite = js.substr(p+1, q-p-1); } resolved_key = suite; const std::string suite_obj = find_obj(js, suite); - if (suite_obj.empty()) { std::fprintf(stderr, "vla(bitvla): suite key '%s' not found in statistics_json\n", suite.c_str()); return false; } + if (suite_obj.empty()) { + std::fprintf(stderr, "vla(bitvla): suite key '%s' not found in statistics_json\n", suite.c_str()); + return false; + } const std::string action_obj = find_obj(suite_obj, "action"); - if (action_obj.empty()) { std::fprintf(stderr, "vla(bitvla): no action stats under '%s'\n", suite.c_str()); return false; } + if (action_obj.empty()) { + std::fprintf(stderr, "vla(bitvla): no action stats under '%s'\n", suite.c_str()); + return false; + } - if (!parse_array_floats(action_obj, "q01", q01)) return false; - if (!parse_array_floats(action_obj, "q99", q99)) return false; + if (!parse_array_floats(action_obj, "q01", q01)) + return false; + if (!parse_array_floats(action_obj, "q99", q99)) + return false; if (!parse_array_bools (action_obj, "mask", mask)) { mask.assign(q01.size(), 1); @@ -344,15 +374,15 @@ bool load_config(const gguf_reader & g, BitvlaModelArch & m, Config & cfg) { // 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)) { + 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) { + 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; @@ -402,16 +432,20 @@ static void recover_ternary_and_scale(const float* W, int64_t n, // Per-tensor absmean scale (1/mean|W|), matching scripts/bitvla_int2_pack.py; // the int2-packed path bakes the same scale. double s = 0.0; - for (int64_t i = 0; i < n; ++i) s += std::fabs((double) W[i]); - float mean = n > 0 ? (float) (s / (double) n) : 0.0f; - if (mean < 1e-5f) mean = 1e-5f; + for (int64_t i=0; i 0 ? (float) (s/(double) n) : 0.0f; + if (mean < 1e-5f) + mean = 1e-5f; absmean = mean; - const float inv = 1.0f / mean; + const float inv = 1.0f/mean; ternary.resize(n); - for (int64_t i = 0; i < n; ++i) { - float q = std::nearbyintf(W[i] * inv); - if (q > 1.0f) q = 1.0f; - if (q < -1.0f) q = -1.0f; + for (int64_t i=0; i 1.0f) + q = 1.0f; + if (q < -1.0f) + q = -1.0f; ternary[i] = (int8_t) q; } } @@ -419,30 +453,30 @@ static void recover_ternary_and_scale(const float* W, int64_t n, static std::vector pack_ladder_int2(const int8_t* W, int64_t N, int64_t K) { constexpr int N_BLOCK = 16, K_BLOCK = 8, K_PER_LOOP = 16; constexpr int WMMA_K = 32, K_PER_ITER = K_PER_LOOP * K_BLOCK; - const int64_t n_slots = N * K / 16; - std::vector out(N * K / 4, 0); - for (int64_t s = 0; s < n_slots; ++s) { - const int64_t slots_per_block = (N_BLOCK * K) / 16; - const int64_t n_block = s / slots_per_block; - const int64_t in_block = s % slots_per_block; - const int64_t k_0 = in_block / 128; - const int64_t in_k0 = in_block % 128; - const int64_t major_k = in_k0 / 32; - const int64_t in_major = in_k0 % 32; - const int64_t y_half = in_major / 16; - const int64_t in_yhalf = in_major % 16; - const int64_t sub_k = in_yhalf / 8; - const int64_t y_in_h = in_yhalf % 8; - const int64_t n_global = n_block * N_BLOCK + y_half * 8 + y_in_h; - const int64_t k_sub = k_0 * K_PER_ITER + major_k * WMMA_K + sub_k * K_PER_LOOP; - for (int byte_i = 0; byte_i < 4; ++byte_i) { + const int64_t n_slots = N * K/16; + std::vector out(N * K/4, 0); + for (int64_t s=0; s& out_ptrs) { std::vector tmp(n); - for (size_t i = 0; i < n; ++i) tmp[i] = f32_to_bf16_u16(h[i]); + for (size_t i=0; i& wptrs, std::vector& out_scales, std::vector& out_ptrs) { int64_t N_total = 0; - for (int64_t n : Ns) N_total += n; + for (int64_t n : Ns) + N_total += n; std::vector stacked(N_total * K); out_scales.clear(); int64_t row_off = 0; - for (size_t i = 0; i < wptrs.size(); ++i) { + for (size_t i=0; i tern; float sc; - recover_ternary_and_scale(wptrs[i], Ns[i] * K, tern, sc); - std::memcpy(stacked.data() + row_off * K, tern.data(), (size_t) Ns[i] * K); + recover_ternary_and_scale(wptrs[i], Ns[i]*K, tern, sc); + std::memcpy(stacked.data()+row_off * K, tern.data(), (size_t) Ns[i]*K); out_scales.push_back(sc); row_off += Ns[i]; } @@ -511,44 +547,64 @@ static int8_t* pack_and_upload_fused(const std::vector& wptrs, BitvlaModelArch::~BitvlaModelArch() { #ifdef VLA_BITVLA_CUDA_KERNELS - if (lm_cuda_ctx) bitvla_lm_cuda_free(lm_cuda_ctx); - if (vit_cuda_ctx) bitvla_vit_cuda_free(vit_cuda_ctx); - if (fp32head_cuda_ctx) bitvla_fp32head_cuda_free(fp32head_cuda_ctx); - for (void* p : cuda_devptrs) if (p) cudaFree(p); - if (d_inputs_embeds) cudaFree(d_inputs_embeds); - if (d_last_hidden) cudaFree(d_last_hidden); - if (d_action_hidden) cudaFree(d_action_hidden); - if (d_action_ids) cudaFree(d_action_ids); - if (d_vit_patches) cudaFree(d_vit_patches); - if (d_vit_img_embeds) cudaFree(d_vit_img_embeds); - for (void* p : cpu_kept_ptrs) if (p) std::free(p); + if (lm_cuda_ctx) + bitvla_lm_cuda_free(lm_cuda_ctx); + if (vit_cuda_ctx) + bitvla_vit_cuda_free(vit_cuda_ctx); + if (fp32head_cuda_ctx) + bitvla_fp32head_cuda_free(fp32head_cuda_ctx); + for (void* p : cuda_devptrs) + if (p) + cudaFree(p); + if (d_inputs_embeds) + cudaFree(d_inputs_embeds); + if (d_last_hidden) + cudaFree(d_last_hidden); + if (d_action_hidden) + cudaFree(d_action_hidden); + if (d_action_ids) + cudaFree(d_action_ids); + if (d_vit_patches) + cudaFree(d_vit_patches); + if (d_vit_img_embeds) + cudaFree(d_vit_img_embeds); + for (void* p : cpu_kept_ptrs) + if (p) + std::free(p); #endif - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr bitvla_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(bitvla): note - mmproj '%s' is ignored (the BitSigLIP-L vision tower is bundled in the combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); m->gguf_path = ckpt_path; - m->matmul_type = std::getenv("VLA_BITVLA_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_F32); gguf_reader g("bitvla"); - if (!g.open(ckpt_path)) return nullptr; + if (!g.open(ckpt_path)) + return nullptr; if (!g.has("bitvla.architecture") && !g.has("general.architecture")) { std::fprintf(stderr, "vla(bitvla): %s is not a bitvla GGUF\n", ckpt_path.c_str()); return nullptr; } - if (!load_config(g, *m, m->cfg)) return nullptr; + if (!load_config(g, *m, m->cfg)) + return nullptr; // Keep one reader open for the per-step token-embedding fetches (token_embd // stays on disk under int2 packing) and cache the constant stop-token row, // so predict() no longer re-opens and re-parses the GGUF twice per call. - if (!m->emb_reader.open(ckpt_path)) return nullptr; + if (!m->emb_reader.open(ckpt_path)) + return nullptr; m->stop_embed.resize((size_t) m->lm_hidden); { std::vector sid{ m->stop_id }; @@ -566,25 +622,23 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, // 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; } + 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); std::printf("vla(bitvla): ggml backend = CPU (%d threads) - CUDA LM module activates below if available\n", m->n_threads); - ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; + ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(bitvla): ggml_init(ctx_weights) failed\n"); return nullptr; } - ggml_context * W = m->ctx_weights; - auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { - const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(bitvla): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); - ggml_set_name(t, name); - return t; - }; - auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; - auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; - - auto mk_bit = [&](const char * name) { return mk(name, m->packed_int2 ? GGML_TYPE_I8 : m->matmul_type); }; + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(bitvla): ggml_init(ctx_weights) failed\n"); + return nullptr; + } + WeightLoader L("bitvla", g, m->ctx_weights, m->matmul_type); + auto mk_mm = [&](const char * name) { return L.gemm("%s", name); }; + auto mk_f32 = [&](const char * name) { return L.f32 ("%s", name); }; + auto mk_bit = [&](const char * name) { return L.typed(m->packed_int2 ? GGML_TYPE_I8 : m->matmul_type, "%s", name); }; bool ok = true; @@ -592,7 +646,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, m->vit_patch_b = mk_f32("vit.patch_embd.bias"); m->vit_pos = mk_f32("vit.pos_embd.weight"); m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers && ok; ++i) { + for (int64_t i=0; ivit_layers && ok; ++i) { char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return p; }; auto & w = m->vit[i]; w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); @@ -615,7 +669,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, m->embed_tokens = m->packed_int2 ? nullptr : mk_mm("token_embd.weight"); m->lm_output_norm = mk_f32("lm.output_norm.weight"); m->lm.resize(m->lm_layers); - for (int64_t i = 0; i < m->lm_layers && ok; ++i) { + for (int64_t i=0; ilm_layers && ok; ++i) { char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "lm.blk.%lld.%s", (long long) i, s); return p; }; auto & w = m->lm[i]; w.attn_norm = mk_f32(N("attn_norm.weight")); @@ -649,19 +703,16 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, m->ah_b0_ln_w&&m->ah_b0_ln_b&&m->ah_b0_w&&m->ah_b0_b&& m->ah_b1_ln_w&&m->ah_b1_ln_b&&m->ah_b1_w&&m->ah_b1_b&& m->ah_ln2_w&&m->ah_ln2_b&&m->ah_fc2_w&&m->ah_fc2_b; - if (!ok) { std::fprintf(stderr, "vla(bitvla): weight tensor setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(bitvla): ggml_backend_alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for (ggml_tensor * t = ggml_get_first_tensor(W); t; t = ggml_get_next_tensor(W, t)) { - std::vector bytes = g.read_convert(ggml_get_name(t), t->type); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(bitvla): failed to load %s (%zu vs %zu bytes)\n", ggml_get_name(t), bytes.size(), ggml_nbytes(t)); return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); + if (!ok) { + std::fprintf(stderr, "vla(bitvla): weight tensor setup failed\n"); + return nullptr; } + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + std::printf("vla(bitvla): weights resident in %.2f GiB (%s); image_id=%d proprio_id=%d action_begin_id=%d stop_id=%d\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0), + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), m->packed_int2 ? "int2-packed + F32 sidecars" : (m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"), m->image_token_id, m->proprio_pad_id, m->action_begin_id, m->stop_id); @@ -676,7 +727,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, #ifdef VLA_BITVLA_CUDA_KERNELS - if ((m->packed_int2 || m->matmul_type == GGML_TYPE_F32) && !std::getenv("VLA_BITVLA_NO_CUDA_LM")) { + if ((m->packed_int2 || m->matmul_type == GGML_TYPE_F32) && !vla::env_flag("VLA_BITVLA_NO_CUDA_LM")) { int dev_count = 0; if (cudaGetDeviceCount(&dev_count) == cudaSuccess && dev_count > 0) { cudaSetDevice(0); @@ -689,7 +740,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, 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::string sn = nm.substr(0, nm.size()-7) + ".scale"; std::vector sc = g.read_f32(sn.c_str()); if (sc.empty()) { std::fprintf(stderr, "vla(bitvla): int2 tensor %s has no %s sidecar\n", @@ -711,7 +762,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 && scales_ok; ++L) { + for (int64_t L=0; Llm_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); @@ -720,12 +771,28 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, lyr.ffn_sub_norm_w = upload_bf16_from_f32((const float*) m->lm[L].ffn_sub_norm->data, m->lm_inter, m->cuda_devptrs); const int64_t hq_dim = m->lm_q * m->lm_head_dim; - const int64_t hkv_dim = m->lm_kv * m->lm_head_dim; - { auto r = load_bit(m->lm[L].Wq, hq_dim, m->lm_hidden); lyr.q_packed = r.first; lyr.q_ws = r.second; } - { auto r = load_bit(m->lm[L].Wk, hkv_dim, m->lm_hidden); lyr.k_packed = r.first; lyr.k_ws = r.second; } - { auto r = load_bit(m->lm[L].Wv, hkv_dim, m->lm_hidden); lyr.v_packed = r.first; lyr.v_ws = r.second; } + const int64_t hkv_dim = m->lm_kv*m->lm_head_dim; + { + auto r = load_bit(m->lm[L].Wq, hq_dim, m->lm_hidden); + lyr.q_packed = r.first; + lyr.q_ws = r.second; + } + { + auto r = load_bit(m->lm[L].Wk, hkv_dim, m->lm_hidden); + lyr.k_packed = r.first; + lyr.k_ws = r.second; + } + { + auto r = load_bit(m->lm[L].Wv, hkv_dim, m->lm_hidden); + lyr.v_packed = r.first; + lyr.v_ws = r.second; + } - { auto r = load_bit(m->lm[L].Wo, m->lm_hidden, m->lm_hidden); lyr.o_packed = r.first; lyr.o_ws = r.second; } + { + auto r = load_bit(m->lm[L].Wo, m->lm_hidden, m->lm_hidden); + lyr.o_packed = r.first; + lyr.o_ws = r.second; + } 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); @@ -746,7 +813,11 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, lyr.gate_up_ws = upload_f32_scales(ws2.data(), 2, m->cuda_devptrs); } - { auto r = load_bit(m->lm[L].Wdown, m->lm_hidden, m->lm_inter); lyr.down_packed = r.first; lyr.down_ws = r.second; } + { + 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 (!scales_ok) { @@ -756,10 +827,13 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, __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); - cudaError_t lm_ce = cudaMalloc(&m->d_inputs_embeds, (size_t) max_seq * m->lm_hidden * sizeof(__nv_bfloat16)); - if (lm_ce == cudaSuccess) lm_ce = cudaMalloc(&m->d_last_hidden, (size_t) max_seq * m->lm_hidden * sizeof(__nv_bfloat16)); - if (lm_ce == cudaSuccess) lm_ce = cudaMalloc(&m->d_action_hidden, (size_t) (m->num_actions_chunk * m->action_dim) * m->lm_hidden * sizeof(__nv_bfloat16)); - if (lm_ce == cudaSuccess) lm_ce = cudaMalloc(&m->d_action_ids, (size_t) (m->num_actions_chunk * m->action_dim) * sizeof(int32_t)); + cudaError_t lm_ce = cudaMalloc(&m->d_inputs_embeds, (size_t) max_seq * m->lm_hidden*sizeof(__nv_bfloat16)); + if (lm_ce == cudaSuccess) + lm_ce = cudaMalloc(&m->d_last_hidden, (size_t) max_seq * m->lm_hidden*sizeof(__nv_bfloat16)); + if (lm_ce == cudaSuccess) + lm_ce = cudaMalloc(&m->d_action_hidden, (size_t) (m->num_actions_chunk*m->action_dim)*m->lm_hidden*sizeof(__nv_bfloat16)); + if (lm_ce == cudaSuccess) + lm_ce = cudaMalloc(&m->d_action_ids, (size_t) (m->num_actions_chunk*m->action_dim)*sizeof(int32_t)); // only enable the CUDA LM once every work buffer is really allocated. if (lm_ce != cudaSuccess) { std::fprintf(stderr, "vla(bitvla): CUDA LM buffer alloc failed (%s); using CPU LM\n", @@ -770,25 +844,25 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, m->d_action_hidden = nullptr; m->d_action_ids = nullptr; } else { m->cuda_lm_ready = true; - const size_t packed_bytes = (size_t) m->lm_layers * ( - (size_t)(m->lm_q + 2*m->lm_kv) * m->lm_head_dim * m->lm_hidden / 4 + - (size_t) m->lm_hidden * m->lm_hidden / 4 + - (size_t) 2 * m->lm_inter * m->lm_hidden / 4 + - (size_t) m->lm_hidden * m->lm_inter / 4); + const size_t packed_bytes = (size_t) m->lm_layers*( + (size_t)(m->lm_q+2*m->lm_kv)*m->lm_head_dim*m->lm_hidden/4 + + (size_t) m->lm_hidden*m->lm_hidden/4 + + (size_t) 2*m->lm_inter*m->lm_hidden/4 + + (size_t) m->lm_hidden*m->lm_inter/4); std::printf("vla(bitvla): CUDA LM forward ENABLED - packed int2 weights = %.2f MiB, max_seq=%d\n", - packed_bytes / (1024.0 * 1024.0), max_seq); + packed_bytes/(1024.0*1024.0), max_seq); } - const int patch_flat = 3 * m->patch_size * m->patch_size; + const int patch_flat = 3*m->patch_size*m->patch_size; const int mm_out = (int) m->lm_hidden; - const int ffn_pad = ((m->vit_inter + 127) / 128) * 128; + const int ffn_pad = ((m->vit_inter+127)/128)*128; m->vit_cuda_ctx = bitvla_vit_cuda_init( (int) m->vit_layers, (int) m->vit_hidden, (int) m->vit_heads, (int) m->vit_inter, (int) m->n_patches, patch_flat, m->vit_ln_eps, mm_out); if (m->vit_cuda_ctx) { bool vit_ok = true; - for (int64_t L = 0; L < m->vit_layers && vit_ok; ++L) { + for (int64_t L=0; Lvit_layers && vit_ok; ++L) { bitvla_vit_layer_cuda vl{}; vl.ln1_w = upload_bf16_from_f32((const float*) m->vit[L].ln1w->data, m->vit_hidden, m->cuda_devptrs); @@ -796,31 +870,55 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, vl.ln2_w = upload_bf16_from_f32((const float*) m->vit[L].ln2w->data, m->vit_hidden, m->cuda_devptrs); vl.ln2_b = upload_bf16_from_f32((const float*) m->vit[L].ln2b->data, m->vit_hidden, m->cuda_devptrs); - { auto r = load_bit(m->vit[L].Wq, m->vit_hidden, m->vit_hidden); vl.q_packed = r.first; vl.q_ws = r.second; } + { + auto r = load_bit(m->vit[L].Wq, m->vit_hidden, m->vit_hidden); + vl.q_packed = r.first; + vl.q_ws = r.second; + } vl.q_b = upload_bf16_from_f32((const float*) m->vit[L].bq->data, m->vit_hidden, m->cuda_devptrs); - { auto r = load_bit(m->vit[L].Wk, m->vit_hidden, m->vit_hidden); vl.k_packed = r.first; vl.k_ws = r.second; } + { + auto r = load_bit(m->vit[L].Wk, m->vit_hidden, m->vit_hidden); + vl.k_packed = r.first; + vl.k_ws = r.second; + } vl.k_b = upload_bf16_from_f32((const float*) m->vit[L].bk->data, m->vit_hidden, m->cuda_devptrs); - { auto r = load_bit(m->vit[L].Wv, m->vit_hidden, m->vit_hidden); vl.v_packed = r.first; vl.v_ws = r.second; } + { + auto r = load_bit(m->vit[L].Wv, m->vit_hidden, m->vit_hidden); + vl.v_packed = r.first; + vl.v_ws = r.second; + } vl.v_b = upload_bf16_from_f32((const float*) m->vit[L].bv->data, m->vit_hidden, m->cuda_devptrs); - { auto r = load_bit(m->vit[L].Wo, m->vit_hidden, m->vit_hidden); vl.o_packed = r.first; vl.o_ws = r.second; } + { + auto r = load_bit(m->vit[L].Wo, m->vit_hidden, m->vit_hidden); + vl.o_packed = r.first; + vl.o_ws = r.second; + } vl.o_b = upload_bf16_from_f32((const float*) m->vit[L].bo->data, m->vit_hidden, m->cuda_devptrs); - { auto r = load_bit(m->vit[L].Wfc1, m->vit_inter, m->vit_hidden); vl.fc1_packed = r.first; vl.fc1_ws = r.second; } + { + auto r = load_bit(m->vit[L].Wfc1, m->vit_inter, m->vit_hidden); + vl.fc1_packed = r.first; + vl.fc1_ws = r.second; + } vl.fc1_b = upload_bf16_from_f32((const float*) m->vit[L].bfc1->data, m->vit_inter, m->cuda_devptrs); if (m->packed_int2) { - { auto r = load_bit(m->vit[L].Wfc2, m->vit_hidden, ffn_pad); vl.fc2_packed = r.first; vl.fc2_ws = r.second; } + { + auto r = load_bit(m->vit[L].Wfc2, m->vit_hidden, ffn_pad); + vl.fc2_packed = r.first; + vl.fc2_ws = r.second; + } vl.fc2_b = upload_bf16_from_f32((const float*) m->vit[L].bfc2->data, m->vit_hidden, m->cuda_devptrs); } else { const float* W = (const float*) m->vit[L].Wfc2->data; std::vector tern; float scale; - recover_ternary_and_scale(W, m->vit_hidden * m->vit_inter, tern, scale); + recover_ternary_and_scale(W, m->vit_hidden*m->vit_inter, tern, scale); - std::vector padded((size_t) m->vit_hidden * ffn_pad, 0); - for (int64_t n = 0; n < m->vit_hidden; ++n) { - std::memcpy(padded.data() + n * ffn_pad, - tern.data() + n * m->vit_inter, + std::vector padded((size_t) m->vit_hidden*ffn_pad, 0); + for (int64_t n=0; nvit_hidden; ++n) { + std::memcpy(padded.data()+n * ffn_pad, + tern.data()+n * m->vit_inter, (size_t) m->vit_inter); } @@ -833,9 +931,9 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, } if (vit_ok) { - __nv_bfloat16* pe_w = upload_bf16_from_f32((const float*) m->vit_patch_w->data, m->vit_hidden * patch_flat, m->cuda_devptrs); + __nv_bfloat16* pe_w = upload_bf16_from_f32((const float*) m->vit_patch_w->data, m->vit_hidden*patch_flat, m->cuda_devptrs); __nv_bfloat16* pe_b = upload_bf16_from_f32((const float*) m->vit_patch_b->data, m->vit_hidden, m->cuda_devptrs); - __nv_bfloat16* pos_e = upload_bf16_from_f32((const float*) m->vit_pos->data, m->n_patches * m->vit_hidden, m->cuda_devptrs); + __nv_bfloat16* pos_e = upload_bf16_from_f32((const float*) m->vit_pos->data, m->n_patches*m->vit_hidden, m->cuda_devptrs); bitvla_vit_cuda_set_embed(m->vit_cuda_ctx, pe_w, pe_b, pos_e); __nv_bfloat16* mm_W1 = upload_bf16_from_f32((const float*) m->mm_l1_w->data, mm_out * m->vit_hidden, m->cuda_devptrs); @@ -844,15 +942,15 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, __nv_bfloat16* mm_b2 = upload_bf16_from_f32((const float*) m->mm_l2_b->data, mm_out, m->cuda_devptrs); bitvla_vit_cuda_set_mmproj(m->vit_cuda_ctx, mm_W1, mm_b1, mm_W2, mm_b2); - cudaMalloc(&m->d_vit_patches, (size_t) m->n_patches * patch_flat * sizeof(__nv_bfloat16)); - cudaMalloc(&m->d_vit_img_embeds, (size_t) m->n_patches * mm_out * sizeof(__nv_bfloat16)); + cudaMalloc(&m->d_vit_patches, (size_t) m->n_patches*patch_flat * sizeof(__nv_bfloat16)); + cudaMalloc(&m->d_vit_img_embeds, (size_t) m->n_patches*mm_out * sizeof(__nv_bfloat16)); m->cuda_vit_ready = true; - const size_t vit_packed_bytes = (size_t) m->vit_layers * ( - 4 * (size_t) m->vit_hidden * m->vit_hidden / 4 + - (size_t) m->vit_inter * m->vit_hidden / 4 + - (size_t) m->vit_hidden * ffn_pad / 4); + const size_t vit_packed_bytes = (size_t) m->vit_layers*( + 4*(size_t) m->vit_hidden*m->vit_hidden/4 + + (size_t) m->vit_inter*m->vit_hidden/4 + + (size_t) m->vit_hidden*ffn_pad/4); std::printf("vla(bitvla): CUDA ViT forward ENABLED - packed int2 weights = %.2f MiB, ffn_pad=%d\n", - vit_packed_bytes / (1024.0 * 1024.0), ffn_pad); + vit_packed_bytes/(1024.0*1024.0), ffn_pad); } else { bitvla_vit_cuda_free(m->vit_cuda_ctx); m->vit_cuda_ctx = nullptr; @@ -881,7 +979,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, } if (m->cuda_lm_ready && m->cuda_vit_ready && - std::getenv("VLA_BITVLA_CPU_HEAD") == nullptr) { + !vla::env_flag("VLA_BITVLA_CPU_HEAD")) { m->fp32head_cuda_ctx = bitvla_fp32head_cuda_init( (int) m->proprio_dim, (int) m->lm_hidden, (int) m->num_actions_chunk, (int) m->action_dim, @@ -902,19 +1000,19 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, const int64_t h = m->lm_hidden, p = m->proprio_dim, a = m->action_dim; const int64_t adh = a * h; const size_t bytes = (size_t)( - h*p + h + h*h + h + - 2*adh + h*adh + h + - 2 * (2*h + h*h + h) + - 2*h + a*h + a) * sizeof(float); + h*p+h+h*h+h + + 2*adh+h*adh+h + + 2*(2*h+h*h+h) + + 2*h+a*h+a)*sizeof(float); std::printf("vla(bitvla): CUDA fp32head (ProprioProj+ActionHead) ENABLED - weights = %.2f MiB on GPU\n", - bytes / (1024.0 * 1024.0)); + bytes/(1024.0*1024.0)); } else { std::fprintf(stderr, "vla(bitvla): bitvla_fp32head_cuda_init failed; falling back to CPU ggml head\n"); } } if (m->cuda_lm_ready && m->cuda_vit_ready && m->weight_buf && - std::getenv("VLA_BITVLA_KEEP_CPU_WEIGHTS") == nullptr) { + !vla::env_flag("VLA_BITVLA_KEEP_CPU_WEIGHTS")) { size_t bytes_kept = 0; if (!m->cuda_fp32head_ready) { @@ -926,7 +1024,8 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, m->ah_ln2_w, m->ah_ln2_b, m->ah_fc2_w, m->ah_fc2_b, }; for (ggml_tensor* t : keep_tensors) { - if (!t) continue; + if (!t) + continue; const size_t nb = ggml_nbytes(t); void* copy = std::malloc(nb); if (!copy) { @@ -939,7 +1038,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, bytes_kept += nb; } } - const double before_gb = ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0); + const double before_gb = ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0); ggml_backend_buffer_free(m->weight_buf); m->weight_buf = nullptr; @@ -952,7 +1051,7 @@ std::unique_ptr bitvla_create(const std::string& mmproj_path, } else { std::printf("vla(bitvla): freed weight_buf (%.2f GiB); kept %.2f MiB of CPU-resident " "ProprioProj+ActionHead weights as standalone copies\n", - before_gb, bytes_kept / (1024.0 * 1024.0)); + before_gb, bytes_kept/(1024.0*1024.0)); } } #endif @@ -968,18 +1067,22 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { const char* _dump_dir = std::getenv("VLA_BITVLA_DUMP_DIR"); auto _dump_bin = [&](const char* name, const float* data, size_t nelem) { - if (!_dump_dir) return; + if (!_dump_dir) + return; std::string path = std::string(_dump_dir) + "/" + name + ".bin"; FILE* f = std::fopen(path.c_str(), "wb"); - if (!f) return; + if (!f) + return; std::fwrite(data, sizeof(float), nelem, f); std::fclose(f); }; auto _dump_manifest = [&](const std::string& line) { - if (!_dump_dir) return; + if (!_dump_dir) + return; std::string path = std::string(_dump_dir) + "/manifest.txt"; FILE* f = std::fopen(path.c_str(), "a"); - if (!f) return; + if (!f) + return; std::fwrite(line.data(), 1, line.size(), f); std::fputc('\n', f); std::fclose(f); @@ -992,7 +1095,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { const int64_t H = image_size, P = patch_size, N = n_patches; const int64_t hidden_v = vit_hidden, hidden_l = lm_hidden; - const int64_t patch_flat = 3 * P * P; + const int64_t patch_flat = 3*P * P; const bool use_precomp_img = (in.precomputed_img_emb != nullptr) && in.n_img_views > 0; const int64_t n_views = use_precomp_img ? (int64_t) in.n_img_views : (int64_t) in.n_images; @@ -1008,7 +1111,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { img_embeds_host.assign((size_t) n_views * N * hidden_l, 0.0f); std::vector patches((size_t) N * patch_flat); const auto t_v0 = clk::now(); - for (int64_t v = 0; v < n_views; ++v) { + for (int64_t v=0; v BitvlaModelArch::predict(const Inputs& in) { return {}; } - for (int64_t pi = 0; pi < H / P; ++pi) - for (int64_t pj = 0; pj < H / P; ++pj) { - const int64_t p_idx = pi * (H / P) + pj; - float * p_dst = patches.data() + p_idx * patch_flat; + for (int64_t pi=0; pi BitvlaModelArch::predict(const Inputs& in) { if (cuda_vit_ready) { std::vector patches_bf16((size_t) N * patch_flat); - for (size_t i = 0; i < patches_bf16.size(); ++i) patches_bf16[i] = f32_to_bf16_u16(patches[i]); - cudaMemcpy(d_vit_patches, patches_bf16.data(), patches_bf16.size() * sizeof(uint16_t), cudaMemcpyHostToDevice); + for (size_t i=0; i img_bf16((size_t) N * hidden_l); - cudaMemcpy(img_bf16.data(), d_vit_img_embeds, img_bf16.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost); - float* dst = img_embeds_host.data() + (size_t) v * N * hidden_l; - for (size_t i = 0; i < img_bf16.size(); ++i) { + cudaMemcpy(img_bf16.data(), d_vit_img_embeds, img_bf16.size()*sizeof(uint16_t), cudaMemcpyDeviceToHost); + float* dst = img_embeds_host.data()+(size_t) v * N * hidden_l; + for (size_t i=0; i BitvlaModelArch::predict(const Inputs& in) { #endif { - ggml_context * ctx = vision_scratch.reset((size_t) 24 * 1024 * 1024); + 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"); ggml_tensor * pe = ggml_add(ctx, ggml_mul_mat(ctx, vit_patch_w, x_in), vit_patch_b); ggml_tensor * h = ggml_add(ctx, pe, vit_pos); - for (int64_t L = 0; L < vit_layers; ++L) { + for (int64_t L=0; L BitvlaModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(bitvla): vision graph compute failed (view %lld)\n", (long long) v); 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_backend_tensor_get(mm2, img_embeds_host.data()+(size_t) v * N * hidden_l, 0, (size_t) N * hidden_l * sizeof(float)); } } - stats.ms_vision = std::chrono::duration(clk::now() - t_v0).count(); + stats.ms_vision = std::chrono::duration(clk::now()-t_v0).count(); } else { - img_embeds_host.assign(in.precomputed_img_emb, in.precomputed_img_emb + (size_t) n_views * N * hidden_l); + img_embeds_host.assign(in.precomputed_img_emb, in.precomputed_img_emb+(size_t) n_views * N * hidden_l); } _dump_bin("mm_proj_out", img_embeds_host.data(), img_embeds_host.size()); - _dump_manifest(std::string("mm_proj_out fp32 ") + std::to_string(n_views) + " " + std::to_string(N) + " " + std::to_string(hidden_l)); + _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)); + 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, state_host.data(), proprio_embed_host.data(), 0) != 0) { @@ -1107,7 +1212,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } else #endif { - ggml_context * ctx = proprio_scratch.reset((size_t) 4 * 1024 * 1024); + 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); @@ -1124,15 +1229,16 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } _dump_bin("proprio_features", proprio_embed_host.data(), proprio_embed_host.size()); - _dump_manifest(std::string("proprio_features fp32 1 ") + std::to_string(hidden_l)); + _dump_manifest(std::string("proprio_features fp32 1 ")+std::to_string(hidden_l)); const int64_t n_lang_in = (int64_t) in.n_lang; const int64_t n_img_tok = n_views * N; const int64_t n_action = num_actions_chunk * action_dim; int64_t n_image_markers = 0, n_proprio_markers = 0; - for (int64_t i = 0; i < n_lang_in; ++i) { - if (in.lang_tokens[i] == image_token_id) n_image_markers++; + for (int64_t i=0; i BitvlaModelArch::predict(const Inputs& in) { } const int64_t seq = full_prefix - ? (n_lang_in + n_action + 1) - : (n_lang_in + n_img_tok + 1 + n_action + 1); + ? (n_lang_in+n_action+1) + : (n_lang_in+n_img_tok+1+n_action+1); if (seq > lm_max_pos) { std::fprintf(stderr, "vla(bitvla): seq=%lld > lm_max_pos=%lld\n", (long long) seq, (long long) lm_max_pos); return {}; @@ -1164,7 +1270,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { if (full_prefix) { - std::vector ids(in.lang_tokens, in.lang_tokens + n_lang_in); + std::vector ids(in.lang_tokens, in.lang_tokens+n_lang_in); for (int32_t id : ids) { if (id < 0 || id >= vocab_size) { std::fprintf(stderr, "vla(bitvla): prompt token %d out of vocab\n", id); return {}; @@ -1173,27 +1279,27 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { if (!emb_reader.fetch_rows_f32("token_embd.weight", ids, inputs_embeds.data(), hidden_l)) return {}; int64_t k_img = 0; - for (int64_t i = 0; i < n_lang_in; ++i) { + for (int64_t i=0; i 0) { - std::vector ids(in.lang_tokens, in.lang_tokens + n_prompt); + std::vector ids(in.lang_tokens, in.lang_tokens+n_prompt); for (int32_t id : ids) { if (id < 0 || id >= vocab_size) { std::fprintf(stderr, "vla(bitvla): prompt token %d out of vocab\n", id); return {}; @@ -1201,16 +1307,16 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } if (!emb_reader.fetch_rows_f32("token_embd.weight", ids, inputs_embeds.data(), hidden_l)) return {}; } - std::memcpy(inputs_embeds.data() + (size_t) n_prompt * hidden_l, + std::memcpy(inputs_embeds.data()+(size_t) n_prompt * hidden_l, img_embeds_host.data(), (size_t) n_img_tok * hidden_l * sizeof(float)); - std::memcpy(inputs_embeds.data() + (size_t) (n_prompt + n_img_tok) * hidden_l, + std::memcpy(inputs_embeds.data()+(size_t) (n_prompt+n_img_tok)*hidden_l, proprio_embed_host.data(), (size_t) hidden_l * sizeof(float)); - std::memcpy(inputs_embeds.data() + (size_t) (seq - 1) * hidden_l, + std::memcpy(inputs_embeds.data()+(size_t) (seq-1)*hidden_l, stop_embed.data(), (size_t) hidden_l * sizeof(float)); } _dump_bin("inputs_embeds", inputs_embeds.data(), inputs_embeds.size()); - _dump_manifest(std::string("inputs_embeds fp32 1 ") + std::to_string(seq) + + _dump_manifest(std::string("inputs_embeds fp32 1 ")+std::to_string(seq) + " " + std::to_string(hidden_l)); std::vector last_hidden_at_actions((size_t) n_action * hidden_l); @@ -1220,20 +1326,22 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { if (cuda_lm_ready && seq <= cuda_max_seq) { std::vector in_bf16((size_t) seq * hidden_l); - for (size_t i = 0; i < in_bf16.size(); ++i) in_bf16[i] = f32_to_bf16_u16(inputs_embeds[i]); - cudaMemcpy(d_inputs_embeds, in_bf16.data(), in_bf16.size() * sizeof(uint16_t), cudaMemcpyHostToDevice); + for (size_t i=0; i aids(n_action); - for (int64_t i = 0; i < n_action; ++i) aids[i] = (int32_t) (seq - 2 - n_action + i); + for (int64_t i=0; i out_bf16((size_t) n_action * hidden_l); - cudaMemcpy(out_bf16.data(), d_action_hidden, out_bf16.size() * sizeof(uint16_t), cudaMemcpyDeviceToHost); - for (size_t i = 0; i < out_bf16.size(); ++i) { + cudaMemcpy(out_bf16.data(), d_action_hidden, out_bf16.size()*sizeof(uint16_t), cudaMemcpyDeviceToHost); + for (size_t i=0; i BitvlaModelArch::predict(const Inputs& in) { } else #endif { - ggml_context * ctx = lm_scratch.reset((size_t) 64 * 1024 * 1024); + 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); ggml_set_name(positions, "positions"); ggml_tensor * h = x_in; - for (int64_t L = 0; L < lm_layers; ++L) { + for (int64_t L=0; L BitvlaModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(x_in, inputs_embeds.data(), 0, ggml_nbytes(x_in)); std::vector pos_v(seq); - for (int64_t i = 0; i < seq; ++i) pos_v[i] = (int32_t) i; + for (int64_t i=0; i aids(n_action); - for (int64_t i = 0; i < n_action; ++i) aids[i] = (int32_t) (seq - 2 - n_action + i); + for (int64_t i=0; i(clk::now() - t_p0).count(); + if (timing_phase) + stats.ms_prefill = std::chrono::duration(clk::now()-t_p0).count(); _dump_bin("ah_input", last_hidden_at_actions.data(), last_hidden_at_actions.size()); - _dump_manifest(std::string("ah_input fp32 1 ") + std::to_string(num_actions_chunk) + + _dump_manifest(std::string("ah_input fp32 1 ")+std::to_string(num_actions_chunk) + " " + std::to_string(action_dim * hidden_l)); const int64_t chunk = num_actions_chunk; @@ -1294,7 +1405,7 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { } else #endif { - ggml_context * ctx = head_scratch.reset((size_t) 8 * 1024 * 1024); + 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"); @@ -1322,27 +1433,28 @@ std::vector BitvlaModelArch::predict(const Inputs& in) { 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)); } - if (timing_phase) stats.ms_denoise = std::chrono::duration(clk::now() - t_d0).count(); + if (timing_phase) + stats.ms_denoise = std::chrono::duration(clk::now()-t_d0).count(); _dump_bin("ah_norm_actions", normalized_actions.data(), normalized_actions.size()); - _dump_manifest(std::string("ah_norm_actions fp32 1 ") + std::to_string(num_actions_chunk) + + _dump_manifest(std::string("ah_norm_actions fp32 1 ")+std::to_string(num_actions_chunk) + " " + std::to_string(action_dim)); std::vector actions = std::move(normalized_actions); - for (int64_t t = 0; t < chunk; ++t) { - for (int64_t d = 0; d < action_dim; ++d) { - const float a = actions[t * action_dim + d]; + for (int64_t t=0; t(clk::now() - t_start).count() - stats.ms_vision); - stats.ms_total = std::chrono::duration(clk::now() - t_start).count(); + ? stats.ms_prefill+stats.ms_denoise + : std::chrono::duration(clk::now()-t_start).count()-stats.ms_vision); + stats.ms_total = std::chrono::duration(clk::now()-t_start).count(); return actions; } diff --git a/src/models/dit_common.h b/src/models/dit_common.h index 83f6cc7..ab6047b 100644 --- a/src/models/dit_common.h +++ b/src/models/dit_common.h @@ -48,27 +48,35 @@ inline ggml_tensor * adaln(ggml_context * C, ggml_tensor * x, ggml_tensor * temb 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); } + for (int64_t i=0; i & out) { - const int64_t half = dim / 2; const float step = std::log(10000.0f) / (float) half; const float t = (float) bucket; + 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); } + for (int64_t tk=0; tk sinusoidal_time_emb(double t, int64_t dim, double min_p, double max_p) { - const int64_t half = dim / 2; + 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; + for (int64_t i=0; i sinusoidal_time_emb(double t, int64_t dim, double min_ 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; + for (int64_t q=0; q #include @@ -70,7 +73,9 @@ struct Evo1ModelArch : public ModelArchBase { struct MainKey { int64_t seq=-1, nsteps=-1; - bool operator==(const MainKey & o) const { return seq==o.seq && nsteps==o.nsteps; } + 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; @@ -126,7 +131,7 @@ ggml_tensor * build_qwen2_layer(ggml_context * C, const Evo1ModelArch & m, const ggml_tensor * qmask = nullptr) { const int64_t hd = m.lm_head_dim, n_q = m.n_q, n_kv = m.n_kv, hq = n_q * hd; const ggml_type at = m.act_type; - const float scale = 1.0f / std::sqrt((float) hd); + const float scale = 1.0f/std::sqrt((float) hd); ggml_tensor * h_n1 = ggml_mul(C, ggml_rms_norm(C, h, m.lm_rms_eps), w.attn_norm); ggml_tensor * qp = ggml_add(C, mm_act(C, w.Wq, h_n1, at), w.bq); ggml_tensor * kp = ggml_add(C, mm_act(C, w.Wk, h_n1, at), w.bk); @@ -146,7 +151,8 @@ ggml_tensor * build_qwen2_layer(ggml_context * C, const Evo1ModelArch & m, const 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)), hq, seq); ggml_tensor * attn_out = mm_act(C, w.Wo, as_type(C, att, at), at); - if (qmask) attn_out = ggml_mul(C, attn_out, qmask); + if (qmask) + attn_out = ggml_mul(C, attn_out, qmask); ggml_tensor * h_attn = ggml_add(C, h, attn_out); ggml_tensor * h_n2 = ggml_mul(C, ggml_rms_norm(C, h_attn, m.lm_rms_eps), w.ffn_norm); ggml_tensor * gate = ggml_silu(C, mm_act(C, w.Wgate, h_n2, at)); @@ -161,17 +167,17 @@ bool preprocess_image_chw(const ImageView & v, int64_t side, std::vector std::fprintf(stderr, "vla(evo1): 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) { + out.assign((size_t) 3*side * side, 0.0f); + for (int64_t h=0; h // (n=100, same binary). That is inside sampling noise at ~1.6 SE, but the drop // concentrated in the two tasks the control aced, so the default stays on the // accuracy-preserving path and the speedup is opt-in. -inline bool evo1_vit_fa_enabled() { - static const bool enabled = (std::getenv("VLA_EVO1_FA") != nullptr); - return enabled; -} ggml_tensor * evo1_flash_attn(ggml_context * C, ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, float scale, int64_t hidden, int64_t N) { @@ -218,20 +220,20 @@ ggml_tensor * evo1_flash_attn(ggml_context * C, ggml_tensor * q, ggml_tensor * k ggml_tensor * build_internvit_layer(ggml_context * C, const Evo1ModelArch & m, const ViTLayerW & w, ggml_tensor * x, int64_t N) { - const int64_t H = m.vit_hidden, n_heads = m.vit_heads, hd = H / n_heads; + const int64_t H = m.vit_hidden, n_heads = m.vit_heads, hd = H/n_heads; const ggml_type at = m.act_type; - const float scale = 1.0f / std::sqrt((float) hd); + const float scale = 1.0f/std::sqrt((float) hd); ggml_tensor * x_n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, m.vit_ln_eps), w.n1w), w.n1b); ggml_tensor * qkv = ggml_add(C, mm_act(C, w.Wqkv, x_n1, at), w.bqkv); // one cast of the packed QKV rather than three of its slices qkv = as_type(C, qkv, GGML_TYPE_F32); - ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv, H, N, qkv->nb[1], 0 * H * ggml_element_size(qkv))); - ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv, H, N, qkv->nb[1], 1 * H * ggml_element_size(qkv))); - ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv, H, N, qkv->nb[1], 2 * H * ggml_element_size(qkv))); + ggml_tensor * q = ggml_cont(C, ggml_view_2d(C, qkv, H, N, qkv->nb[1], 0*H * ggml_element_size(qkv))); + ggml_tensor * k = ggml_cont(C, ggml_view_2d(C, qkv, H, N, qkv->nb[1], 1*H * ggml_element_size(qkv))); + ggml_tensor * v = ggml_cont(C, ggml_view_2d(C, qkv, H, N, qkv->nb[1], 2*H * ggml_element_size(qkv))); ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, n_heads, N), 0, 2, 1, 3)); ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, n_heads, N), 0, 2, 1, 3)); ggml_tensor * att; - if (evo1_vit_fa_enabled()) { + if (vla::flash_attn_enabled()) { ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, n_heads, N), 0, 2, 1, 3)); att = evo1_flash_attn(C, Q, K, V, scale, H, N); } else { @@ -251,8 +253,8 @@ ggml_tensor * build_internvit_layer(ggml_context * C, const Evo1ModelArch & m, c } ggml_tensor * build_internvit_view(ggml_context * C, const Evo1ModelArch & m, ggml_tensor * pixels) { - const int64_t H = m.vit_hidden, grid = m.image_size / m.patch_size, n_patches = grid * grid, n_tok = n_patches + 1; - const int64_t shuf_c = H * 4, sgrid = grid / 2; + const int64_t H = m.vit_hidden, grid = m.image_size/m.patch_size, n_patches = grid * grid, n_tok = n_patches+1; + const int64_t shuf_c = H*4, sgrid = grid/2; ggml_tensor * conv = ggml_conv_2d(C, m.vit_patch_w, pixels, (int) m.patch_size, (int) m.patch_size, 0, 0, 1, 1); ggml_tensor * patches = ggml_add(C, ggml_cont(C, ggml_transpose(C, ggml_reshape_2d(C, conv, n_patches, H))), m.vit_patch_b); @@ -261,10 +263,11 @@ ggml_tensor * build_internvit_view(ggml_context * C, const Evo1ModelArch & m, gg // the activation dtype from here ggml_tensor * x = as_type(C, ggml_add(C, ggml_concat(C, cls2d, patches, 1), m.vit_pos), m.act_type); - for (int64_t i = 0; i < m.vit_layers; ++i) x = build_internvit_layer(C, m, m.vit[i], x, n_tok); + for (int64_t i=0; inb[1], x->nb[1])); - ggml_tensor * s1 = ggml_reshape_3d(C, pnc, 2 * H, sgrid, grid); + ggml_tensor * s1 = ggml_reshape_3d(C, pnc, 2*H, sgrid, grid); ggml_tensor * s2 = ggml_cont(C, ggml_permute(C, s1, 0, 2, 1, 3)); ggml_tensor * s3 = ggml_reshape_3d(C, s2, shuf_c, sgrid, sgrid); ggml_tensor * s4 = ggml_cont(C, ggml_permute(C, s3, 0, 2, 1, 3)); @@ -284,7 +287,7 @@ ggml_tensor * inproj_split_b(ggml_context * C, ggml_tensor * bin, int64_t E, int } bool load_config(const gguf_reader & g, Evo1ModelArch & m, Config & cfg) { - auto u = [&](const char * k, int64_t & dst) { if (g.has((std::string("evo1.") + k).c_str())) dst = g.u32((std::string("evo1.") + k).c_str()); }; + auto u = [&](const char * k, int64_t & dst) { if (g.has((std::string("evo1.")+k).c_str())) dst = g.u32((std::string("evo1.")+k).c_str()); }; u("lm_hidden", m.lm_hidden); u("lm_layers_used", m.lm_layers); u("lm_q_heads", m.n_q); u("lm_kv_heads", m.n_kv); u("lm_head_dim", m.lm_head_dim); u("lm_inter", m.lm_inter); u("embed_dim", m.embed_dim); u("dit_layers", m.dit_layers); u("dit_heads", m.dit_heads); u("mlp_head_hidden", m.mlp_head_hidden); u("horizon", m.horizon); u("per_action_dim", m.per_a); @@ -295,18 +298,23 @@ bool load_config(const gguf_reader & g, Evo1ModelArch & m, Config & cfg) { u("real_state_dim", m.real_state_dim); u("real_action_dim", m.real_action_dim); u("vit_hidden", m.vit_hidden); u("vit_layers", m.vit_layers); u("vit_heads", m.vit_heads); u("vit_inter", m.vit_inter); u("image_size", m.image_size); u("patch_size", m.patch_size); - if (g.has("evo1.lm_rms_eps")) m.lm_rms_eps = g.f32("evo1.lm_rms_eps"); - if (g.has("evo1.proj_ln_eps")) m.proj_ln_eps = g.f32("evo1.proj_ln_eps"); - if (g.has("evo1.norm_eps")) m.norm_eps_denom = g.f32("evo1.norm_eps"); - if (g.has("evo1.vit_ln_eps")) m.vit_ln_eps = g.f32("evo1.vit_ln_eps"); - if (g.has("evo1.lm_rope_theta")) m.lm_rope_base = (float) g.f64("evo1.lm_rope_theta"); + if (g.has("evo1.lm_rms_eps")) + m.lm_rms_eps = g.f32("evo1.lm_rms_eps"); + if (g.has("evo1.proj_ln_eps")) + m.proj_ln_eps = g.f32("evo1.proj_ln_eps"); + if (g.has("evo1.norm_eps")) + m.norm_eps_denom = g.f32("evo1.norm_eps"); + if (g.has("evo1.vit_ln_eps")) + m.vit_ln_eps = g.f32("evo1.vit_ln_eps"); + if (g.has("evo1.lm_rope_theta")) + m.lm_rope_base = (float) g.f64("evo1.lm_rope_theta"); if (m.embed_dim != m.lm_hidden) { std::fprintf(stderr, "vla(evo1): embed_dim (%lld) != lm_hidden (%lld) - not handled\n", (long long) m.embed_dim, (long long) m.lm_hidden); return false; } // predict() reads action_dim noise floats; the server sizes noise as // horizon*per_action_dim, so they must agree or a client noise buffer underruns. - if (m.action_dim != m.horizon * m.per_a) { + if (m.action_dim != m.horizon*m.per_a) { std::fprintf(stderr, "vla(evo1): action_dim (%lld) != horizon (%lld) * per_action_dim (%lld)\n", (long long) m.action_dim, (long long) m.horizon, (long long) m.per_a); return false; } @@ -340,28 +348,34 @@ bool load_config(const gguf_reader & g, Evo1ModelArch & m, Config & cfg) { } Evo1ModelArch::~Evo1ModelArch() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr evo1_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(evo1): note - mmproj '%s' is ignored (the vision tower is bundled in the combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); m->gguf_path = ckpt_path; - m->matmul_type = std::getenv("VLA_EVO1_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); - if (!m->io.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; } - if (!load_config(g, *m, m->cfg)) return nullptr; + if (!load_config(g, *m, m->cfg)) + return nullptr; std::printf("vla(evo1): lm=%lldd×%lldL (%lldq/%lldkv×%lld) inter=%lld embed=%lld dit=%lldL×%lldh " "horizon=%lld per_a=%lld N_steps=%lld resident matmul=%s\n", (long long) m->lm_hidden, (long long) m->lm_layers, (long long) m->n_q, (long long) m->n_kv, @@ -371,11 +385,13 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, { const Backend b = backend_init("vla(evo1)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; // BF16 activations need BF16-resident weights and the CUDA BF16 GEMM path. - if (std::getenv("VLA_EVO1_BF16_ACT")) { + if (opts.act_dtype.value_or(GGML_TYPE_F32) == GGML_TYPE_BF16) { if (b.is_cuda && m->matmul_type == GGML_TYPE_BF16) { m->act_type = GGML_TYPE_BF16; cuda_register_bf16_ops(); // installs the in-tree BF16 CUDA kernels @@ -386,26 +402,23 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, } } - ggml_init_params wp = { (size_t) 32 * 1024 * 1024, + ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(evo1): ggml_init(ctx_weights) failed\n"); return nullptr; } - ggml_context * W = m->ctx_weights; - - auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { - const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(evo1): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); - ggml_set_name(t, name); - return t; - }; - auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; - auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(evo1): ggml_init(ctx_weights) failed\n"); + return nullptr; + } + // The vision tower is optional here, so misses are reported by the ok chain + // below rather than by the loader. + WeightLoader L("evo1", g, m->ctx_weights, m->matmul_type); + auto mk_mm = [&](const char * name) { return L.opt_gemm("%s", name); }; + auto mk_f32 = [&](const char * name) { return L.opt_f32 ("%s", name); }; bool ok = true; m->lm_output_norm = mk_f32("vlm.output_norm.weight"); ok &= (m->lm_output_norm != nullptr); m->lm.resize(m->lm_layers); - for (int64_t i = 0; i < m->lm_layers && ok; ++i) { + for (int64_t i=0; ilm_layers && ok; ++i) { char p[64]; auto N = [&](const char * suf) { std::snprintf(p, sizeof(p), "vlm.blk.%lld.%s", (long long) i, suf); return p; }; auto & w = m->lm[i]; w.attn_norm = mk_f32(N("attn_norm.weight")); @@ -423,7 +436,7 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, m->ae_W3 = mk_mm("aex.ae.W3.weight"); m->ae_b3 = mk_f32("aex.ae.W3.bias"); m->ae_pos = mk_f32("aex.ae.pos_enc"); m->dit.resize(m->dit_layers); - for (int64_t i = 0; i < m->dit_layers && ok; ++i) { + for (int64_t i=0; idit_layers && ok; ++i) { char p[64]; auto N = [&](const char * suf) { std::snprintf(p, sizeof(p), "aex.blk.%lld.%s", (long long) i, suf); return p; }; auto & w = m->dit[i]; w.n1w = mk_f32(N("norm1.weight")); w.n1b = mk_f32(N("norm1.bias")); @@ -446,12 +459,12 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, m->head_W2 && m->head_b2 && m->time_pos && m->state_W1 && m->state_b1 && m->state_W2 && m->state_b2; if (g.meta("vit.patch_embd.weight") && ok) { - m->vit_patch_w = mk("vit.patch_embd.weight", GGML_TYPE_F32); + m->vit_patch_w = mk_f32("vit.patch_embd.weight"); m->vit_patch_b = mk_f32("vit.patch_embd.bias"); m->vit_cls = mk_f32("vit.class_embd"); m->vit_pos = mk_f32("vit.pos_embd"); m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers && ok; ++i) { + for (int64_t i=0; ivit_layers && ok; ++i) { char p[64]; auto N = [&](const char * suf) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, suf); return p; }; auto & w = m->vit[i]; w.n1w = mk_f32(N("norm1.weight")); w.n1b = mk_f32(N("norm1.bias")); @@ -471,21 +484,17 @@ std::unique_ptr evo1_create(const std::string& mmproj_path, } else { std::printf("vla(evo1): note - no vit.*/mm.* weights in the GGUF; predict() will require Inputs::precomputed_img_emb\n"); } - if (!ok) { std::fprintf(stderr, "vla(evo1): weight tensor setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(evo1): ggml_backend_alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for (ggml_tensor * t = ggml_get_first_tensor(W); t; t = ggml_get_next_tensor(W, t)) { - std::vector bytes = g.read_convert(ggml_get_name(t), t->type); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(evo1): failed to load %s (%zu vs %zu bytes)\n", - ggml_get_name(t), bytes.size(), ggml_nbytes(t)); return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); + if (!ok) { + std::fprintf(stderr, "vla(evo1): weight tensor setup failed\n"); + return nullptr; } + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + std::printf("vla(evo1): weights resident in %.2f GiB (%s)%s\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0), - m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16", + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), + dtype_name(m->matmul_type), m->have_vision ? " - incl. InternViT vision tower" : " - vision tower NOT loaded (precomputed_img_emb required)"); m->state_min = g.read_f32("state_min"); @@ -529,19 +538,23 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { // // The branches are independent, so the arithmetic per view is unchanged // - only the submission pattern differs. - const size_t want_arena = (size_t) 32 * 1024 * 1024 * (size_t) std::max(n_views, 1); - if (want_arena > vision_arena) { vision_scratch.release(); vision_arena = want_arena; } + const size_t want_arena = (size_t) 32*1024*1024*(size_t) std::max(n_views, 1); + if (want_arena > vision_arena) { + vision_scratch.release(); + vision_arena = want_arena; + } ggml_context * VC = vision_scratch.reset(vision_arena); if (!VC) { std::fprintf(stderr, "vla(evo1): ggml_init(vision ctx) failed\n"); return {}; } std::vector t_px((size_t) n_views), t_ie((size_t) n_views); - for (int64_t v = 0; v < n_views; ++v) { + for (int64_t v=0; v(n_views, 1), false); - for (int64_t v = 0; v < n_views; ++v) ggml_build_forward_expand(vg, t_ie[v]); + ggml_cgraph * vg = ggml_new_graph_custom(VC, (size_t) 8192*std::max(n_views, 1), false); + for (int64_t v=0; v Evo1ModelArch::predict(const Inputs& in) { 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) { + for (int64_t v=0; v Evo1ModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(evo1): vision graph compute failed (%lld views)\n", (long long) n_views); return {}; } - for (int64_t v = 0; v < n_views; ++v) { - ggml_backend_tensor_get(t_ie[v], img_emb_host.data() + v * num_image_token * lm_hidden, + for (int64_t v=0; v(std::chrono::steady_clock::now() - tv0).count(); + stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now()-tv0).count(); img_emb_ptr = img_emb_host.data(); } else { std::fprintf(stderr, "vla(evo1): no images and no precomputed_img_emb in the request\n"); @@ -572,19 +585,25 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(evo1): note - %lld image views (model n_images=%lld); prompt adapts\n", (long long) n_views, (long long) n_images); bool pre_built = false; - for (int j = 0; j < in.n_lang; ++j) - if (in.lang_tokens[j] == (int32_t) img_ctx_id) { pre_built = true; break; } + for (int j=0; j input_ids; input_ids.reserve(max_text_length); if (pre_built) { - for (int j = 0; j < in.n_lang; ++j) input_ids.push_back(in.lang_tokens[j]); + for (int j=0; j max_text_length) { @@ -599,11 +618,11 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { 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) { + for (int64_t p=0; p= n_img_tokens) { std::fprintf(stderr, "vla(evo1): more IMG_CTX tokens than ViT embeds\n"); return {}; } - std::memcpy(inputs_embeds.data() + p * lm_hidden, - img_emb_ptr + img_idx * lm_hidden, lm_hidden * sizeof(float)); + std::memcpy(inputs_embeds.data()+p * lm_hidden, + img_emb_ptr+img_idx * lm_hidden, lm_hidden * sizeof(float)); ++img_idx; } } @@ -620,43 +639,50 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { in.attention_mask_n, (long long) SEQ); return {}; } - for (int64_t p = 0; p < SEQ; ++p) attn_ok[p] = in.attention_mask[p] ? 1 : 0; + for (int64_t p=0; p state_norm(per_a, 0.0f); - for (int64_t i = 0; i < per_a; ++i) { + for (int64_t i=0; i 1.0f) xn = 1.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; } std::vector x_init((size_t) action_dim); if (in.noise) { - std::memcpy(x_init.data(), in.noise, x_init.size() * sizeof(float)); + std::memcpy(x_init.data(), in.noise, x_init.size()*sizeof(float)); } else { // Evo1 is trained with uniform[-1,1] noise uint64_t s = 0xE701ACE5ULL ^ (uint64_t) std::chrono::steady_clock::now().time_since_epoch().count(); for (auto & v : x_init) { - s = s * 6364136223846793005ULL + 1442695040888963407ULL; - v = ((float) (uint32_t) (s >> 32) / 2147483648.0f) - 1.0f; // uniform[-1,1) + s = s*6364136223846793005ULL+1442695040888963407ULL; + v = ((float) (uint32_t) (s >> 32)/2147483648.0f)-1.0f; // uniform[-1,1) } } // 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, + 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; + 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; ggml_tensor * t_embeds = ggml_new_tensor_2d(C, GGML_TYPE_F32, lm_hidden, SEQ); ggml_set_input(t_embeds); ggml_tensor * t_pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ); ggml_set_input(t_pos); @@ -670,7 +696,8 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { const ggml_type at = act_type; ggml_tensor * h = as_type(C, t_embeds, at); - for (int64_t i = 0; i < lm_layers; ++i) h = build_qwen2_layer(C, *this, lm[i], h, t_pos, t_lmmask, SEQ, t_qmask); + for (int64_t i=0; i Evo1ModelArch::predict(const Inputs& in) { struct DC { ggml_tensor *Wq, *bq, *K, *V; }; std::vector dc(dit_layers); - for (int64_t i = 0; i < dit_layers; ++i) { + for (int64_t i=0; i Evo1ModelArch::predict(const Inputs& in) { ae = ggml_add(C, ae, ae_pos); ae = ggml_relu(C, ggml_add(C, mm_act(C, ae_W2, ae, at), ae_b2)); ggml_tensor * x = ggml_add(C, mm_act(C, ae_W3, ae, at), ae_b3); - for (int64_t i = 0; i < dit_layers; ++i) { + for (int64_t i=0; i Evo1ModelArch::predict(const Inputs& in) { return as_type(C, ggml_add(C, mm_act(C, head_W2, mh, at), head_b2), GGML_TYPE_F32); }; - const float dt = 1.0f / (float) num_steps; + const float dt = 1.0f/(float) num_steps; ggml_tensor * x_action = t_x; - for (int64_t step = 0; step < num_steps; ++step) { - const int64_t time_index = (int64_t) ((double) step / (double) num_steps * 1000.0); + for (int64_t step=0; step Evo1ModelArch::predict(const Inputs& in) { ggml_tensor * t_amask = gio.t_amask, * x_action = gio.x_action; 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 pp(SEQ); + for (int64_t i=0; i mk((size_t) SEQ * SEQ); const float NEG = -std::numeric_limits::infinity(); - for (int64_t q = 0; q < SEQ; ++q) for (int64_t kv = 0; kv < SEQ; ++kv) mk[q * SEQ + kv] = (kv <= q && attn_ok[kv]) ? 0.0f : NEG; + for (int64_t q=0; q am(per_a, 0.0f); for (int64_t i = 0; i < real_action_dim && i < per_a; ++i) am[i] = 1.0f; + { std::vector am(per_a, 0.0f); for (int64_t i=0; i qm(SEQ, 0.0f); for (int64_t p = 0; p < SEQ; ++p) qm[p] = attn_ok[p] ? 1.0f : 0.0f; + { std::vector qm(SEQ, 0.0f); for (int64_t p=0; p Evo1ModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(evo1): ggml_backend_graph_compute failed (%d)\n", (int) st); return {}; } - stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); + 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_backend_tensor_get(x_action, x_final.data(), 0, x_final.size()*sizeof(float)); std::vector out((size_t) horizon * per_a); - for (int64_t hstep = 0; hstep < horizon; ++hstep) - for (int64_t c = 0; c < per_a; ++c) { - const double a = (double) x_final[hstep * per_a + c]; - out[hstep * per_a + c] = (float) ((a + 1.0) / 2.0 * ((double) action_max[c] - (double) action_min[c] + (double) norm_eps_denom) + (double) action_min[c]); + for (int64_t hstep=0; hstep(std::chrono::steady_clock::now() - t0).count(); + stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now()-t0).count(); return out; } diff --git a/src/models/gr00tn1d5.cpp b/src/models/gr00tn1d5.cpp index 45c946f..a9228c8 100644 --- a/src/models/gr00tn1d5.cpp +++ b/src/models/gr00tn1d5.cpp @@ -13,360 +13,259 @@ // limitations under the License. #include "arch.h" +#include "options.h" +#include "backend.h" +#include "env_flag.h" +#include "gguf_reader.h" +#include "layers/embed.h" +#include "layers/linear.h" +#include "layers/norm.h" #include "model.h" +#include "modules/action_expert.h" +#include "modules/dit_head.h" +#include "modules/encoder.h" +#include "modules/preprocess.h" +#include "modules/prompt.h" +#include "modules/qwen3_lm.h" +#include "modules/siglip_vit.h" +#include "scratch_ctx.h" #include "ggml.h" -#include "ggml-cpu.h" #include "ggml-backend.h" -#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 #include #include #include #include -#include #include #include #include #include namespace vla { -namespace { - - -struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; -struct Qwen3LayerW { ggml_tensor *attn_norm,*Wq,*Wk,*Wv,*Wo,*q_norm,*k_norm,*ffn_norm,*Wgate,*Wup,*Wdown; }; -struct VlsaLayerW { ggml_tensor *n1w,*n1b,*n3w,*n3b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wff0,*bff0,*Wff2,*bff2; }; -struct DitLayerW { ggml_tensor *adaln_w,*adaln_b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wff0,*bff0,*Wff2,*bff2; }; - -} struct Gr00tN1d5ModelArch : public ModelArchBase { Gr00tN1d5ModelArch() : ModelArchBase(Arch::GR00T_N1_5) {} ~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; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + ggml_backend_buffer_t weight_buf = nullptr; + ggml_type matmul_type = GGML_TYPE_F32; 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; } + 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; - int64_t vit_hidden=1152, vit_layers=27, vit_heads=16, vit_inter=4304, image_size=224, patch_size=14, n_img_tokens=256; - int64_t lm_hidden=2048, lm_layers=12, n_q=16, n_kv=8, lm_head_dim=128, lm_inter=6144, vocab=151680, image_token_index=151669; - int64_t bb_embed_dim=2048, in_embed_dim=1536, dit_hidden=1536, dit_heads=32, dit_head_dim=48, dit_layers=16, dit_interleave=1; - int64_t vlsa_layers=4, vlsa_heads=32, vlsa_head_dim=64, vlsa_inter=8192; + SigLipTower vit; + Qwen3LM lm; + EncStack vlsa; + ActionExpert aex; + DitHead dit; + ggml_tensor *mm_W=nullptr, *mm_b=nullptr; + ggml_tensor *vlln_w=nullptr, *vlln_b=nullptr; + ggml_tensor *future_tokens=nullptr; + + int64_t vit_layers=27, vit_inter=4304, image_size=224, patch_size=14, n_img_tokens=256; + int64_t lm_inter=6144, vocab=151680, image_token_index=151669; + int64_t bb_embed_dim=2048, in_embed_dim=1536, dit_interleave=1, vlsa_layers=4; int64_t num_future=32, action_horizon=16, action_dim=32, max_state_dim=64; int64_t num_steps=4, num_buckets=1000, max_embodiments=32, max_seq_len=1024; - float vit_ln_eps=1e-6f, lm_rms_eps=1e-6f, ln_eps=1e-5f, norm_out_eps=1e-6f, vlln_eps=1e-5f, lm_rope_base=1000000.0f; - int64_t embodiment_id = 24; - - ggml_tensor *vit_patch_w=nullptr,*vit_patch_b=nullptr,*vit_pos=nullptr,*vit_post_ln_w=nullptr,*vit_post_ln_b=nullptr; - std::vector vit; - ggml_tensor *mm_W=nullptr,*mm_b=nullptr; - - ggml_tensor *lm_output_norm=nullptr; - std::vector lm; - - ggml_tensor *vlln_w=nullptr,*vlln_b=nullptr; - std::vector vlsa; - ggml_tensor *se_l1W=nullptr,*se_l1b=nullptr,*se_l2W=nullptr,*se_l2b=nullptr; - ggml_tensor *ae_W1W=nullptr,*ae_W1b=nullptr,*ae_W2W=nullptr,*ae_W2b=nullptr,*ae_W3W=nullptr,*ae_W3b=nullptr; - ggml_tensor *ad_l1W=nullptr,*ad_l1b=nullptr,*ad_l2W=nullptr,*ad_l2b=nullptr; - ggml_tensor *future_tokens=nullptr,*pos_embd=nullptr; - ggml_tensor *te_l1W=nullptr,*te_l1b=nullptr,*te_l2W=nullptr,*te_l2b=nullptr; - std::vector dit; - ggml_tensor *po1W=nullptr,*po1b=nullptr,*po2W=nullptr,*po2b=nullptr; + float vlln_eps=1e-5f; std::vector predict(const Inputs& in) override; }; namespace { -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); - ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, 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_qwen3_layer(ggml_context * C, const Gr00tN1d5ModelArch & 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; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * hn = ggml_mul(C, ggml_rms_norm(C, h, m.lm_rms_eps), w.attn_norm); - ggml_tensor * qp = ggml_mul_mat(C, w.Wq, hn); - ggml_tensor * kp = ggml_mul_mat(C, w.Wk, hn); - ggml_tensor * vp = ggml_mul_mat(C, w.Wv, hn); - ggml_tensor * qh = ggml_reshape_3d(C, qp, hd, n_q, seq); - ggml_tensor * kh = ggml_reshape_3d(C, kp, hd, n_kv, seq); - ggml_tensor * vh = ggml_reshape_3d(C, vp, hd, n_kv, seq); - ggml_tensor * qn = ggml_mul(C, ggml_rms_norm(C, qh, m.lm_rms_eps), w.q_norm); - ggml_tensor * kn = ggml_mul(C, ggml_rms_norm(C, kh, m.lm_rms_eps), w.k_norm); - ggml_tensor * qr = ggml_rope_ext(C, qn, positions, nullptr, (int) hd, GGML_ROPE_TYPE_NEOX, 0, m.lm_rope_base, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - ggml_tensor * kr = ggml_rope_ext(C, kn, positions, nullptr, (int) hd, GGML_ROPE_TYPE_NEOX, 0, m.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)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, kr, 0, 2, 1, 3)); - 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); - ggml_tensor * aw = ggml_soft_max_ext(C, kq, mask, 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)), hq, seq); - ggml_tensor * h_attn = ggml_add(C, h, ggml_mul_mat(C, w.Wo, att)); - ggml_tensor * hn2 = ggml_mul(C, ggml_rms_norm(C, h_attn, m.lm_rms_eps), w.ffn_norm); - ggml_tensor * gate = ggml_silu(C, ggml_mul_mat(C, w.Wgate, hn2)); - ggml_tensor * up = ggml_mul_mat(C, w.Wup, hn2); - return ggml_add(C, h_attn, ggml_mul_mat(C, w.Wdown, ggml_mul(C, gate, up))); -} - -ggml_tensor * build_vlsa_block(ggml_context * C, const VlsaLayerW & 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); - ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.n1w), w.n1b); - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, 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 * n3 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.n3w), w.n3b); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wff2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wff0, n3), w.bff0))), w.bff2); - return ggml_add(C, h1, ff); -} - -void dit_kv(ggml_context * C, const Gr00tN1d5ModelArch & m, const DitLayerW & w, ggml_tensor * kv, - ggml_tensor ** K_out, ggml_tensor ** V_out) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, Tkv = kv->ne[1]; - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, kv), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, kv), w.bv); - *K_out = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, Tkv), 0, 2, 1, 3)); - *V_out = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, Tkv), 1, 2, 0, 3)); -} - -ggml_tensor * build_dit_block(ggml_context * C, const Gr00tN1d5ModelArch & m, const DitLayerW & w, - ggml_tensor * h, ggml_tensor * temb, ggml_tensor * enc , - ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, dim = m.dit_hidden, Tk = h->ne[1]; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * n = adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, m.ln_eps); - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n), w.bq); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, Tk), 0, 2, 1, 3)); - ggml_tensor * K, * V; - if (K_pre) { K = K_pre; V = V_pre; } - else { dit_kv(C, m, w, enc ? enc : n, &K, &V); } - 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)), dim, Tk); - ggml_tensor * h1 = ggml_add(C, h, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); - ggml_tensor * n3 = ggml_norm(C, h1, m.ln_eps); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wff2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wff0, n3), w.bff0))), w.bff2); - return ggml_add(C, h1, ff); -} - - 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); }; - auto F = [&](const char * k, float & dst) { if (g.has(k)) dst = g.f32(k); }; + 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); }; auto fk = [&](const char * s) { static char b[64]; std::snprintf(b, sizeof(b), "gr00t_n1_5.%s", s); return b; }; - U(fk("vit_hidden"), m.vit_hidden); U(fk("vit_layers"), m.vit_layers); U(fk("vit_heads"), m.vit_heads); U(fk("vit_inter"), m.vit_inter); - U(fk("image_size"), m.image_size); U(fk("patch_size"), m.patch_size); U(fk("n_img_tokens"), m.n_img_tokens); - U(fk("lm_hidden"), m.lm_hidden); U(fk("lm_layers_used"), m.lm_layers); U(fk("lm_q_heads"), m.n_q); U(fk("lm_kv_heads"), m.n_kv); - U(fk("lm_head_dim"), m.lm_head_dim); U(fk("lm_inter"), m.lm_inter); U(fk("vocab_size"), m.vocab); U(fk("image_token_index"), m.image_token_index); - U(fk("backbone_embedding_dim"), m.bb_embed_dim); U(fk("input_embedding_dim"), m.in_embed_dim); - U(fk("dit_hidden"), m.dit_hidden); U(fk("dit_heads"), m.dit_heads); U(fk("dit_head_dim"), m.dit_head_dim); U(fk("dit_layers"), m.dit_layers); U(fk("dit_interleave"), m.dit_interleave); - U(fk("vlsa_layers"), m.vlsa_layers); U(fk("vlsa_heads"), m.vlsa_heads); U(fk("vlsa_head_dim"), m.vlsa_head_dim); U(fk("vlsa_inter"), m.vlsa_inter); - U(fk("num_target_vision_tokens"), m.num_future); U(fk("action_horizon"), m.action_horizon); U(fk("action_dim"), m.action_dim); U(fk("max_state_dim"), m.max_state_dim); - 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); - F(fk("vit_ln_eps"), m.vit_ln_eps); F(fk("lm_rms_eps"), m.lm_rms_eps); F(fk("ln_eps"), m.ln_eps); F(fk("norm_out_eps"), m.norm_out_eps); F(fk("vlln_eps"), m.vlln_eps); - if (g.has(fk("lm_rope_theta"))) m.lm_rope_base = (float) g.f64(fk("lm_rope_theta")); - - m.embodiment_id = 24; + + U(fk("vit_hidden" ), m.vit.enc.cfg.hidden); + U(fk("vit_layers" ), m.vit_layers); + U(fk("vit_heads" ), m.vit.enc.cfg.heads); + U(fk("vit_inter" ), m.vit_inter); + U(fk("image_size" ), m.image_size); + U(fk("patch_size" ), m.patch_size); + U(fk("n_img_tokens" ), m.n_img_tokens); + U(fk("lm_hidden" ), m.lm.cfg.hidden); + U(fk("lm_layers_used" ), m.lm.cfg.layers); + U(fk("lm_q_heads" ), m.lm.cfg.n_q); + U(fk("lm_kv_heads" ), m.lm.cfg.n_kv); + U(fk("lm_head_dim" ), m.lm.cfg.head_dim); + U(fk("lm_inter" ), m.lm_inter); + U(fk("vocab_size" ), m.vocab); + U(fk("image_token_index"), m.image_token_index); + U(fk("backbone_embedding_dim"), m.bb_embed_dim); + U(fk("input_embedding_dim" ), m.in_embed_dim); + U(fk("dit_hidden" ), m.dit.cfg.hidden); + U(fk("dit_heads" ), m.dit.cfg.heads); + U(fk("dit_head_dim" ), m.dit.cfg.head_dim); + U(fk("dit_layers" ), m.dit.cfg.layers); + U(fk("dit_interleave" ), m.dit_interleave); + U(fk("vlsa_layers" ), m.vlsa_layers); + U(fk("vlsa_heads" ), m.vlsa.cfg.heads); + U(fk("vlsa_head_dim" ), m.vlsa.cfg.head_dim); + U(fk("num_target_vision_tokens"), m.num_future); + U(fk("action_horizon" ), m.action_horizon); + U(fk("action_dim" ), m.action_dim); + U(fk("max_state_dim" ), m.max_state_dim); + 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); + + F(fk("vit_ln_eps" ), m.vit.enc.cfg.ln_eps); + F(fk("lm_rms_eps" ), m.lm.cfg.rms_eps); + F(fk("ln_eps" ), m.dit.cfg.ln_eps); + F(fk("norm_out_eps" ), m.dit.cfg.norm_out_eps); + F(fk("vlln_eps" ), m.vlln_eps); + + if (g.has(fk("lm_rope_theta"))) + m.lm.cfg.rope.freq_base = (float) g.f64(fk("lm_rope_theta")); + + m.vit.enc.cfg.head_dim = m.vit.enc.cfg.hidden/m.vit.enc.cfg.heads; + m.vlsa.cfg.hidden = m.bb_embed_dim; + m.vlsa.cfg.ln_eps = m.dit.cfg.ln_eps; + m.lm.cfg.rope.n_dims = (int) m.lm.cfg.head_dim; + + m.aex.embodiment_id = 24; if (const char * e = std::getenv("VLA_GR00T_EMBODIMENT")) { - char * end = nullptr; long v = std::strtol(e, &end, 10); - if (end && *end == '\0') { m.embodiment_id = (int64_t) v; } - else { - const std::string js = g.str(fk("embodiment_tag_mapping")); - const std::string key = std::string("\"") + e + "\":"; + char * end = nullptr; + const long v = std::strtol(e, &end, 10); + if (end && *end == '\0') { + m.aex.embodiment_id = (int64_t) v; + } else { + const std::string js = g.str(fk("embodiment_tag_mapping")); + const std::string key = std::string("\"")+e+"\":"; const size_t p = js.find(key); - if (p != std::string::npos) m.embodiment_id = std::strtol(js.c_str() + p + key.size(), nullptr, 10); - else std::fprintf(stderr, "vla(gr00tn1d5): embodiment tag '%s' not in embodiment_tag_mapping; using id %lld\n", e, (long long) m.embodiment_id); + if (p != std::string::npos) + m.aex.embodiment_id = std::strtol(js.c_str()+p+key.size(), nullptr, 10); + else std::fprintf(stderr, "vla(gr00tn1d5): embodiment tag '%s' not in embodiment_tag_mapping; using id %lld\n", e, (long long) m.aex.embodiment_id); } } - if (m.embodiment_id < 0 || m.embodiment_id >= m.max_embodiments) { std::fprintf(stderr, "vla(gr00tn1d5): embodiment id %lld out of range [0,%lld)\n", (long long) m.embodiment_id, (long long) m.max_embodiments); return false; } + if (m.aex.embodiment_id < 0 || m.aex.embodiment_id >= m.max_embodiments) { + std::fprintf(stderr, "vla(gr00tn1d5): embodiment id %lld out of range [0,%lld)\n", + (long long) m.aex.embodiment_id, (long long) m.max_embodiments); + 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; - cfg.real_state_dim = m.max_state_dim; cfg.real_action_dim = m.action_dim; - 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; + 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; + cfg.real_state_dim = m.max_state_dim; + cfg.real_action_dim = m.action_dim; + cfg.hidden = m.lm.cfg.hidden; + cfg.n_q_heads = m.lm.cfg.n_q; + cfg.n_kv_heads = m.lm.cfg.n_kv; + cfg.head_dim = m.lm.cfg.head_dim; + cfg.n_layers = m.lm.cfg.layers; + cfg.num_steps = (int) m.num_steps; + cfg.rms_eps = m.lm.cfg.rms_eps; + cfg.rope_n_dims = (int) m.lm.cfg.head_dim; + cfg.rope_mode = GGML_ROPE_TYPE_NEOX; + cfg.rope_freq_base = m.lm.cfg.rope.freq_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; + cfg.denormalized = false; + cfg.norm_eps = 1e-8f; return true; } } Gr00tN1d5ModelArch::~Gr00tN1d5ModelArch() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(gr00tn1d5): note - mmproj '%s' is ignored (the vision tower is bundled in the combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); - m->gguf_path = ckpt_path; - m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); + m->lm.cfg.rope.freq_base = 1000000.0f; - if (!m->io.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; + 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) " "dit=%lldL×%lldh×%lld(inner %lld) interleave=%lld vlsa=%lldL×%lldh×%lld in_emb=%lld horizon=%lld action_dim=%lld N_steps=%lld embodiment=%lld resident=%s\n", - (long long) m->vit_hidden, (long long) m->vit_layers, (long long) m->vit_heads, (long long) m->n_img_tokens, - (long long) m->lm_hidden, (long long) m->lm_layers, (long long) m->n_q, (long long) m->n_kv, (long long) m->lm_head_dim, - (long long) m->dit_layers, (long long) m->dit_heads, (long long) m->dit_head_dim, (long long) m->dit_hidden, (long long) m->dit_interleave, - (long long) m->vlsa_layers, (long long) m->vlsa_heads, (long long) m->vlsa_head_dim, (long long) m->in_embed_dim, - (long long) m->action_horizon, (long long) m->action_dim, (long long) m->num_steps, (long long) m->embodiment_id, + (long long) m->vit.enc.cfg.hidden, (long long) m->vit_layers, (long long) m->vit.enc.cfg.heads, (long long) m->n_img_tokens, + (long long) m->lm.cfg.hidden, (long long) m->lm.cfg.layers, (long long) m->lm.cfg.n_q, (long long) m->lm.cfg.n_kv, (long long) m->lm.cfg.head_dim, + (long long) m->dit.cfg.layers, (long long) m->dit.cfg.heads, (long long) m->dit.cfg.head_dim, (long long) m->dit.cfg.hidden, (long long) m->dit_interleave, + (long long) m->vlsa_layers, (long long) m->vlsa.cfg.heads, (long long) m->vlsa.cfg.head_dim, (long long) m->in_embed_dim, + (long long) m->action_horizon, (long long) m->action_dim, (long long) m->num_steps, (long long) m->aex.embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); - { - const Backend b = backend_init("vla(gr00tn1d5)", m->n_threads); - if (!b.handle) { return nullptr; } - m->backend = b.handle; - } + 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 }; + ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(gr00tn1d5): ggml_init(ctx_weights) failed\n"); return nullptr; } - ggml_context * W = m->ctx_weights; - auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { - const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(gr00tn1d5): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); - ggml_set_name(t, name); return t; - }; - auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; - auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; - - bool ok = true; - - m->vit_patch_w = mk("vit.patch_embd.weight", GGML_TYPE_F32); - m->vit_patch_b = mk_f32("vit.patch_embd.bias"); - m->vit_pos = mk_f32("vit.pos_embd"); - m->vit_post_ln_w = mk_f32("vit.post_ln.weight"); m->vit_post_ln_b = mk_f32("vit.post_ln.bias"); - m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->vit[i]; - w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); - ok &= w.ln1w&&w.ln1b&&w.ln2w&&w.ln2b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wfc1&&w.bfc1&&w.Wfc2&&w.bfc2; - } - m->mm_W = mk_mm("mm.fc.weight"); m->mm_b = mk_f32("mm.fc.bias"); - - m->lm_output_norm = mk_f32("vlm.output_norm.weight"); - m->lm.resize(m->lm_layers); - for (int64_t i = 0; i < m->lm_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vlm.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->lm[i]; - w.attn_norm=mk_f32(N("attn_norm.weight")); - w.Wq=mk_mm(N("attn_q.weight")); w.Wk=mk_mm(N("attn_k.weight")); w.Wv=mk_mm(N("attn_v.weight")); w.Wo=mk_mm(N("attn_o.weight")); - w.q_norm=mk_f32(N("attn_q_norm.weight")); w.k_norm=mk_f32(N("attn_k_norm.weight")); w.ffn_norm=mk_f32(N("ffn_norm.weight")); - w.Wgate=mk_mm(N("ffn_gate.weight")); w.Wup=mk_mm(N("ffn_up.weight")); w.Wdown=mk_mm(N("ffn_down.weight")); - ok &= w.attn_norm&&w.Wq&&w.Wk&&w.Wv&&w.Wo&&w.q_norm&&w.k_norm&&w.ffn_norm&&w.Wgate&&w.Wup&&w.Wdown; + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(gr00tn1d5): ggml_init(ctx_weights) failed\n"); + return nullptr; } - m->vlln_w=mk_f32("aex.vlln.weight"); m->vlln_b=mk_f32("aex.vlln.bias"); - m->vlsa.resize(m->vlsa_layers); - for (int64_t i = 0; i < m->vlsa_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "aex.vlsa.%lld.%s", (long long) i, s); return p; }; - auto & w = m->vlsa[i]; - w.n1w=mk_f32(N("norm1.weight")); w.n1b=mk_f32(N("norm1.bias")); w.n3w=mk_f32(N("norm3.weight")); w.n3b=mk_f32(N("norm3.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wff0=mk_mm(N("ff0.weight")); w.bff0=mk_f32(N("ff0.bias")); w.Wff2=mk_mm(N("ff2.weight")); w.bff2=mk_f32(N("ff2.bias")); - ok &= w.n1w&&w.n1b&&w.n3w&&w.n3b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wff0&&w.bff0&&w.Wff2&&w.bff2; - } + WeightLoader L("gr00tn1d5", g, m->ctx_weights, m->matmul_type); + + m->vit.declare(L, "vit", m->vit_layers); + m->mm_W = L.gemm("mm.fc.weight"); + m->mm_b = L.f32 ("mm.fc.bias"); + + m->lm.declare(L, "vlm"); + + m->vlln_w = L.f32("aex.vlln.weight"); + m->vlln_b = L.f32("aex.vlln.bias"); + m->vlsa.declare(L, "aex.vlsa", m->vlsa_layers, EncNames{"norm1", "norm3", "ff0", "ff2"}); + + m->aex.declare(L, "aex"); + m->future_tokens = L.f32("aex.future_tokens"); + m->dit.declare(L, "aex.dit"); + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; - m->se_l1W=mk_f32("aex.state_enc.l1.W"); m->se_l1b=mk_f32("aex.state_enc.l1.b"); m->se_l2W=mk_f32("aex.state_enc.l2.W"); m->se_l2b=mk_f32("aex.state_enc.l2.b"); - m->ae_W1W=mk_f32("aex.act_enc.W1.W"); m->ae_W1b=mk_f32("aex.act_enc.W1.b"); m->ae_W2W=mk_f32("aex.act_enc.W2.W"); m->ae_W2b=mk_f32("aex.act_enc.W2.b"); m->ae_W3W=mk_f32("aex.act_enc.W3.W"); m->ae_W3b=mk_f32("aex.act_enc.W3.b"); - m->ad_l1W=mk_f32("aex.act_dec.l1.W"); m->ad_l1b=mk_f32("aex.act_dec.l1.b"); m->ad_l2W=mk_f32("aex.act_dec.l2.W"); m->ad_l2b=mk_f32("aex.act_dec.l2.b"); - m->future_tokens=mk_f32("aex.future_tokens"); m->pos_embd=mk_f32("aex.pos_embd"); - m->te_l1W=mk_mm("aex.dit.time_emb.l1.weight"); m->te_l1b=mk_f32("aex.dit.time_emb.l1.bias"); m->te_l2W=mk_mm("aex.dit.time_emb.l2.weight"); m->te_l2b=mk_f32("aex.dit.time_emb.l2.bias"); - m->dit.resize(m->dit_layers); - for (int64_t i = 0; i < m->dit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "aex.dit.%lld.%s", (long long) i, s); return p; }; - auto & w = m->dit[i]; - w.adaln_w=mk_mm(N("adaln.weight")); w.adaln_b=mk_f32(N("adaln.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wff0=mk_mm(N("ff0.weight")); w.bff0=mk_f32(N("ff0.bias")); w.Wff2=mk_mm(N("ff2.weight")); w.bff2=mk_f32(N("ff2.bias")); - ok &= w.adaln_w&&w.adaln_b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wff0&&w.bff0&&w.Wff2&&w.bff2; - } - m->po1W=mk_mm("aex.dit.proj_out1.weight"); m->po1b=mk_f32("aex.dit.proj_out1.bias"); m->po2W=mk_mm("aex.dit.proj_out2.weight"); m->po2b=mk_f32("aex.dit.proj_out2.bias"); - ok &= m->vit_patch_w&&m->vit_patch_b&&m->vit_pos&&m->vit_post_ln_w&&m->vit_post_ln_b&&m->mm_W&&m->mm_b&&m->lm_output_norm&& - m->vlln_w&&m->vlln_b&&m->se_l1W&&m->se_l1b&&m->se_l2W&&m->se_l2b&&m->ae_W1W&&m->ae_W1b&&m->ae_W2W&&m->ae_W2b&&m->ae_W3W&&m->ae_W3b&& - m->ad_l1W&&m->ad_l1b&&m->ad_l2W&&m->ad_l2b&&m->future_tokens&&m->pos_embd&&m->te_l1W&&m->te_l1b&&m->te_l2W&&m->te_l2b&&m->po1W&&m->po1b&&m->po2W&&m->po2b; - if (!ok) { std::fprintf(stderr, "vla(gr00tn1d5): weight tensor setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(gr00tn1d5): ggml_backend_alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for (ggml_tensor * t = ggml_get_first_tensor(W); t; t = ggml_get_next_tensor(W, t)) { - std::vector bytes = g.read_convert(ggml_get_name(t), t->type); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(gr00tn1d5): failed to load %s (%zu vs %zu bytes)\n", ggml_get_name(t), bytes.size(), ggml_nbytes(t)); return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); - } std::printf("vla(gr00tn1d5): weights resident in %.2f GiB (%s) - incl. SigLIP vision tower; embodiment id %lld\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0), m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16", (long long) m->embodiment_id); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), + m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16", (long long) m->aex.embodiment_id); return m; } @@ -374,178 +273,165 @@ std::vector Gr00tN1d5ModelArch::predict(const Inputs& in) { const auto t0 = std::chrono::steady_clock::now(); stats = Stats{}; - const int64_t H = lm_hidden, K = n_img_tokens, E = in_embed_dim; - const int64_t Nsa = 1 + num_future + action_horizon; + const int64_t H = lm.cfg.hidden; + const int64_t K = n_img_tokens; + const int64_t E = in_embed_dim; + const int64_t AD = action_dim; + const int64_t AH = action_horizon; + const int64_t Nsa = 1+num_future+AH; int64_t n_views = 0; std::vector img_emb_host; const float * img_emb_ptr = nullptr; + if (in.precomputed_img_emb && in.n_img_views > 0) { - n_views = in.n_img_views; + n_views = in.n_img_views; img_emb_ptr = in.precomputed_img_emb; } 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_context * VC = vision_scratch.reset((size_t) 64 * 1024 * 1024); + img_emb_host.assign((size_t) n_views*K*H, 0.0f); + + 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); - ggml_tensor * conv = ggml_conv_2d(VC, vit_patch_w, t_px, (int) patch_size, (int) patch_size, 0, 0, 1, 1); - ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, grid * grid, vit_hidden))); - ggml_tensor * h = ggml_add(VC, ggml_add(VC, patches, vit_patch_b), vit_pos); - for (int64_t i = 0; i < vit_layers; ++i) h = build_siglip_layer(VC, vit[i], h, K, vit_heads, vit_hidden / vit_heads, vit_hidden, vit_ln_eps); - h = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, h, vit_ln_eps), vit_post_ln_w), vit_post_ln_b); - ggml_tensor * vit_emb = ggml_add(VC, ggml_mul_mat(VC, mm_W, h), mm_b); + + 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); + + ggml_tensor * h = vit.build(VC, vit.embed_conv(VC, t_px, patch_size, grid), K); + ggml_tensor * vit_emb = linear(VC, mm_W, mm_b, h); ggml_set_output(vit_emb); + ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vg, vit_emb); 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("gr00tn1d5", in.images[v], image_size, chw)) { return {}; } + for (int64_t v=0; v(std::chrono::steady_clock::now() - tv0).count(); + stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now()-tv0).count(); 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 {}; - } - const int64_t n_img = n_views * K; - - std::vector input_ids; - int64_t n_lang_img = 0; - for (int j = 0; j < in.n_lang; ++j) if (in.lang_tokens[j] == (int32_t) image_token_index) ++n_lang_img; - if (n_lang_img > 0) { - if (n_lang_img != n_img) { - std::fprintf(stderr, "vla(gr00tn1d5): lang_tokens has %lld placeholders but %lld ViT embeds " - "(n_views=%lld × K=%lld)\n", (long long) n_lang_img, (long long) n_img, (long long) n_views, (long long) K); - return {}; - } - input_ids.assign(in.lang_tokens, in.lang_tokens + in.n_lang); - } else { - input_ids.reserve(n_img + in.n_lang); - for (int64_t i = 0; i < n_img; ++i) input_ids.push_back((int32_t) image_token_index); - for (int j = 0; j < in.n_lang; ++j) input_ids.push_back(in.lang_tokens[j]); - } - 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 {}; } - - std::vector inputs_embeds((size_t) SEQ * H); - 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 {}; } - std::memcpy(inputs_embeds.data() + p * H, img_emb_ptr + k * H, H * sizeof(float)); ++k; - } + std::fprintf(stderr, "vla(gr00tn1d5): no images and no precomputed_img_emb in the request\n"); + return {}; } + const int64_t n_img = n_views*K; + + Prompt prompt; + if (!build_prompt("gr00tn1d5", in, n_img, (int32_t) image_token_index, max_seq_len, prompt)) return {}; + const int64_t SEQ = prompt.len(); - const int64_t AD = action_dim, AH = action_horizon; - std::vector x_init((size_t) AH * AD); - 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); } + std::vector inputs_embeds; + if (!fetch_embeds("gr00tn1d5", io, prompt, img_emb_ptr, H, inputs_embeds)) return {}; + + std::vector x_init; + init_noise(in, (size_t) AH*AD, x_init); - // 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, + 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); - ggml_tensor * t_state = ggml_new_tensor_2d(C, GGML_TYPE_F32, max_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); - std::vector t_tau(num_steps), t_tproj(num_steps); - for (int64_t s = 0; s < num_steps; ++s) { - t_tau[s] = ggml_new_tensor_2d(C, GGML_TYPE_F32, E, AH); ggml_set_input(t_tau[s]); - t_tproj[s] = ggml_new_tensor_1d(C, GGML_TYPE_F32, 256); ggml_set_input(t_tproj[s]); - } - - ggml_tensor * h = t_embeds; - for (int64_t i = 0; i < lm_layers; ++i) h = build_qwen3_layer(C, *this, lm[i], h, t_pos, t_lmmask, SEQ); - ggml_tensor * eagle = ggml_mul(C, ggml_rms_norm(C, h, lm_rms_eps), lm_output_norm); + 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); + ggml_tensor * t_state = ggml_new_tensor_2d(C, GGML_TYPE_F32, max_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); + + std::vector t_tau(num_steps), t_tproj(num_steps); + for (int64_t s=0; s Kc(dit_layers, nullptr), Vc(dit_layers, nullptr); - for (int64_t i = 0; i < dit_layers; ++i) { - if (dit_interleave && (i % 2 == 1)) continue; - dit_kv(C, *this, dit[i], vl_embs, &Kc[i], &Vc[i]); - } + std::vector Kc(dit.cfg.layers, nullptr), Vc(dit.cfg.layers, nullptr); + for (int64_t i=0; inb[1], (size_t)(Nsa-AH)*pred->nb[1])); + actions = ggml_add(C, actions, ggml_scale(C, vel, dt)); + } + ggml_set_name(actions, "action_pred"); + ggml_set_output(actions); - ggml_tensor * temb = ggml_add(C, ggml_mul_mat(C, te_l2W, ggml_silu(C, ggml_add(C, ggml_mul_mat(C, te_l1W, t_tproj[s]), te_l1b))), te_l2b); + 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_tensor * a_emb = cat_linear(C, ae_W1W, ae_W1b, embodiment_id, actions); - ggml_tensor * x_w2 = ggml_silu(C, cat_linear(C, ae_W2W, ae_W2b, embodiment_id, ggml_concat(C, a_emb, t_tau[s], 0))); - ggml_tensor * af = ggml_add(C, cat_linear(C, ae_W3W, ae_W3b, embodiment_id, x_w2), ggml_view_2d(C, pos_embd, E, AH, pos_embd->nb[1], 0)); + 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_tensor * sa = ggml_concat(C, ggml_concat(C, state_features, future_tokens, 1), af, 1); + MainIO & gio = main_graph.io(); - ggml_tensor * hh = sa; - for (int64_t i = 0; i < dit_layers; ++i) { - ggml_tensor * enc = (dit_interleave && (i % 2 == 1)) ? nullptr : vl_embs; - hh = build_dit_block(C, *this, dit[i], hh, temb, enc, Kc[i], Vc[i]); - } + ggml_backend_tensor_set(gio.t_embeds, inputs_embeds.data(), 0, ggml_nbytes(gio.t_embeds)); - ggml_tensor * po = ggml_add(C, ggml_mul_mat(C, po1W, ggml_silu(C, temb)), po1b); - ggml_tensor * sh = ggml_view_1d(C, po, dit_hidden, 0), * sc = ggml_view_1d(C, po, dit_hidden, (size_t) dit_hidden * sizeof(float)); - ggml_tensor * hn = ggml_norm(C, hh, norm_out_eps); - ggml_tensor * h_mod = ggml_add(C, ggml_add(C, hn, ggml_mul(C, hn, sc)), sh); - ggml_tensor * model_output = ggml_add(C, ggml_mul_mat(C, po2W, h_mod), po2b); + std::vector pp(SEQ); + for (int64_t i=0; inb[1], (size_t) (Nsa - AH) * pred->nb[1])); - actions = ggml_add(C, actions, ggml_scale(C, vel, dt)); - } - ggml_set_name(actions, "action_pred"); ggml_set_output(actions); + std::vector mask; + build_causal_mask(SEQ, mask); + ggml_backend_tensor_set(gio.t_lmmask, mask.data(), 0, ggml_nbytes(gio.t_lmmask)); - 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; + std::vector st(max_state_dim, 0.0f); + for (int64_t i=0; i & 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)); } - { std::vector mk((size_t) SEQ * SEQ); const float NEG = -std::numeric_limits::infinity(); - for (int64_t q = 0; q < SEQ; ++q) for (int64_t kv = 0; kv < SEQ; ++kv) mk[q * SEQ + kv] = (kv <= q) ? 0.0f : NEG; - ggml_backend_tensor_set(t_lmmask, mk.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)); - for (int64_t s = 0; s < num_steps; ++s) { - const int64_t bucket = (int64_t) ((double) s / (double) num_steps * (double) num_buckets); - std::vector tau, tpr; action_sinusoid(bucket, E, AH, tau); timesteps_proj(bucket, tpr); - ggml_backend_tensor_set(t_tau[s], tau.data(), 0, ggml_nbytes(t_tau[s])); - ggml_backend_tensor_set(t_tproj[s], tpr.data(), 0, ggml_nbytes(t_tproj[s])); + for (int64_t s=0; s tau, tpr; + action_sinusoid(bucket, E, AH, tau); + timesteps_proj(bucket, tpr); + ggml_backend_tensor_set(gio.t_tau[s], tau.data(), 0, ggml_nbytes(gio.t_tau[s])); + ggml_backend_tensor_set(gio.t_tproj[s], tpr.data(), 0, ggml_nbytes(gio.t_tproj[s])); } const auto tc0 = std::chrono::steady_clock::now(); - const ggml_status st = ggml_backend_graph_compute(backend, gf); + const ggml_status status = ggml_backend_graph_compute(backend, main_graph.graph()); 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); return {}; } - stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(gr00tn1d5): graph compute failed (%d)\n", (int) status); + 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)); - stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::vector out((size_t) AH*AD); + ggml_backend_tensor_get(gio.actions, out.data(), 0, out.size()*sizeof(float)); + 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 aff2f29..52f14fe 100644 --- a/src/models/gr00tn1d6.cpp +++ b/src/models/gr00tn1d6.cpp @@ -13,49 +13,49 @@ // limitations under the License. #include "arch.h" +#include "options.h" +#include "backend.h" +#include "env_flag.h" +#include "gguf_reader.h" +#include "layers/embed.h" +#include "layers/ffn.h" +#include "layers/linear.h" +#include "layers/norm.h" #include "model.h" +#include "modules/action_expert.h" +#include "modules/dit_head.h" +#include "modules/preprocess.h" +#include "modules/prompt.h" +#include "modules/qwen3_lm.h" +#include "modules/siglip_vit.h" +#include "scratch_ctx.h" #include "ggml.h" -#include "ggml-cpu.h" #include "ggml-backend.h" -#include "backend.h" -#include "gguf.h" -#include "models/gguf_reader.h" -#include "models/scratch_ctx.h" -#include "models/dit_common.h" #include -#include #include #include #include #include -#include #include #include #include #include namespace vla { -namespace { - - -struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; -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; }; - -} struct Gr00tN1d6ModelArch : public ModelArchBase { Gr00tN1d6ModelArch() : ModelArchBase(Arch::GR00T_N1_6) {} ~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; int n_threads = default_cpu_threads(); ggml_context * ctx_weights = nullptr; + ggml_backend_buffer_t weight_buf = nullptr; + ggml_type matmul_type = GGML_TYPE_F32; scratch_ctx vision_scratch; scratch_ctx merge_scratch; @@ -71,198 +71,128 @@ struct Gr00tN1d6ModelArch : public ModelArchBase { std::vector t_tau, t_tproj; }; graph_cache main_graph; - ggml_backend_buffer_t weight_buf = nullptr; - ggml_type matmul_type = GGML_TYPE_F32; - int64_t vit_hidden=1152, vit_layers=27, vit_heads=16, vit_inter=4304, image_size=224, patch_size=14; + SigLipTower vit; + Qwen3LM lm; + ActionExpert aex; + DitHead dit; + ggml_tensor *mm_ln_w=nullptr, *mm_ln_b=nullptr; + ggml_tensor *mm_fc1_w=nullptr, *mm_fc1_b=nullptr, *mm_fc2_w=nullptr, *mm_fc2_b=nullptr; + ggml_tensor *vlln_w=nullptr, *vlln_b=nullptr; + + int64_t vit_layers=27, vit_inter=4304, image_size=224, patch_size=14; int64_t vit_num_patches=256, n_img_tokens=64, vit_pixel_shuffle=2, mlp_inner=4608; - int64_t lm_hidden=2048, lm_layers=16, n_q=16, n_kv=8, lm_head_dim=128, lm_inter=6144, vocab=151680, image_token_index=151669; - int64_t bb_embed_dim=2048, in_embed_dim=1536, dit_hidden=1536, dit_heads=32, dit_head_dim=48, dit_layers=32, dit_interleave=1, attend_text_every_n=2; + int64_t lm_inter=6144, vocab=151680, image_token_index=151669; + int64_t bb_embed_dim=2048, in_embed_dim=1536, dit_interleave=1, attend_text_every_n=2; int64_t action_horizon=50, action_dim=128, max_state_dim=128; int64_t num_steps=4, num_buckets=1000, max_embodiments=32, max_seq_len=1024; - float vit_ln_eps=1e-6f, lm_rms_eps=1e-6f, ln_eps=1e-5f, norm_out_eps=1e-6f, vlln_eps=1e-5f, connector_ln_eps=1e-5f, lm_rope_base=1000000.0f; - int64_t embodiment_id = 20; - - ggml_tensor *vit_patch_w=nullptr,*vit_patch_b=nullptr,*vit_pos=nullptr,*vit_post_ln_w=nullptr,*vit_post_ln_b=nullptr; - std::vector vit; - ggml_tensor *mm_ln_w=nullptr,*mm_ln_b=nullptr,*mm_fc1_w=nullptr,*mm_fc1_b=nullptr,*mm_fc2_w=nullptr,*mm_fc2_b=nullptr; - - ggml_tensor *lm_output_norm=nullptr; - std::vector lm; - - ggml_tensor *vlln_w=nullptr,*vlln_b=nullptr; - ggml_tensor *se_l1W=nullptr,*se_l1b=nullptr,*se_l2W=nullptr,*se_l2b=nullptr; - ggml_tensor *ae_W1W=nullptr,*ae_W1b=nullptr,*ae_W2W=nullptr,*ae_W2b=nullptr,*ae_W3W=nullptr,*ae_W3b=nullptr; - ggml_tensor *ad_l1W=nullptr,*ad_l1b=nullptr,*ad_l2W=nullptr,*ad_l2b=nullptr; - ggml_tensor *pos_embd=nullptr; - ggml_tensor *te_l1W=nullptr,*te_l1b=nullptr,*te_l2W=nullptr,*te_l2b=nullptr; - std::vector dit; - ggml_tensor *po1W=nullptr,*po1b=nullptr,*po2W=nullptr,*po2b=nullptr; + float vlln_eps=1e-5f, connector_ln_eps=1e-5f; std::vector predict(const Inputs& in) override; }; namespace { -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]; - const float scale = 1.0f / std::sqrt((float) head_dim); - ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, q, head_dim, heads, seq, nv), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, k, head_dim, heads, seq, nv), 0, 2, 1, 3)); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_4d(C, v, head_dim, heads, seq, nv), 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_3d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hidden, seq, nv); - 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_qwen3_layer(ggml_context * C, const Gr00tN1d6ModelArch & 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; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * hn = ggml_mul(C, ggml_rms_norm(C, h, m.lm_rms_eps), w.attn_norm); - ggml_tensor * qp = ggml_mul_mat(C, w.Wq, hn); - ggml_tensor * kp = ggml_mul_mat(C, w.Wk, hn); - ggml_tensor * vp = ggml_mul_mat(C, w.Wv, hn); - ggml_tensor * qh = ggml_reshape_3d(C, qp, hd, n_q, seq); - ggml_tensor * kh = ggml_reshape_3d(C, kp, hd, n_kv, seq); - ggml_tensor * vh = ggml_reshape_3d(C, vp, hd, n_kv, seq); - ggml_tensor * qn = ggml_mul(C, ggml_rms_norm(C, qh, m.lm_rms_eps), w.q_norm); - ggml_tensor * kn = ggml_mul(C, ggml_rms_norm(C, kh, m.lm_rms_eps), w.k_norm); - ggml_tensor * qr = ggml_rope_ext(C, qn, positions, nullptr, (int) hd, GGML_ROPE_TYPE_NEOX, 0, m.lm_rope_base, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - ggml_tensor * kr = ggml_rope_ext(C, kn, positions, nullptr, (int) hd, GGML_ROPE_TYPE_NEOX, 0, m.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)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, kr, 0, 2, 1, 3)); - 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); - ggml_tensor * aw = ggml_soft_max_ext(C, kq, mask, 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)), hq, seq); - ggml_tensor * h_attn = ggml_add(C, h, ggml_mul_mat(C, w.Wo, att)); - ggml_tensor * hn2 = ggml_mul(C, ggml_rms_norm(C, h_attn, m.lm_rms_eps), w.ffn_norm); - ggml_tensor * gate = ggml_silu(C, ggml_mul_mat(C, w.Wgate, hn2)); - ggml_tensor * up = ggml_mul_mat(C, w.Wup, hn2); - return ggml_add(C, h_attn, ggml_mul_mat(C, w.Wdown, ggml_mul(C, gate, up))); -} - -void dit_kv(ggml_context * C, const Gr00tN1d6ModelArch & m, const DitLayerW & w, ggml_tensor * kv, - ggml_tensor ** K_out, ggml_tensor ** V_out) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, Tkv = kv->ne[1]; - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, kv), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, kv), w.bv); - *K_out = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, Tkv), 0, 2, 1, 3)); - *V_out = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, Tkv), 1, 2, 0, 3)); -} - -ggml_tensor * build_dit_block(ggml_context * C, const Gr00tN1d6ModelArch & m, const DitLayerW & w, - ggml_tensor * h, ggml_tensor * temb, ggml_tensor * enc , - ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, dim = m.dit_hidden, Tk = h->ne[1]; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * n = adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, m.ln_eps); - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n), w.bq); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, Tk), 0, 2, 1, 3)); - ggml_tensor * K, * V; - if (K_pre) { K = K_pre; V = V_pre; } - else { dit_kv(C, m, w, enc ? enc : n, &K, &V); } - 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)), dim, Tk); - ggml_tensor * h1 = ggml_add(C, h, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); - ggml_tensor * n3 = ggml_norm(C, h1, m.ln_eps); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wff2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wff0, n3), w.bff0))), w.bff2); - return ggml_add(C, h1, ff); -} - -bool preprocess_image_patches(const ImageView & v, int64_t side, int64_t ps, std::vector & out) { - if (v.w != (int) side || v.h != (int) side || !v.data) { - std::fprintf(stderr, "vla(gr00tn1d6): image view is %dx%d, expected %lldx%lld\n", v.w, v.h, (long long) side, (long long) side); return false; - } - const int64_t grid = side / ps, pd = 3 * ps * ps, np = grid * grid; - out.assign((size_t) pd * np, 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 row = 0; row < grid; ++row) - for (int64_t col = 0; col < grid; ++col) { - const int64_t t = row * grid + col; - for (int64_t ph = 0; ph < ps; ++ph) - for (int64_t pw = 0; pw < ps; ++pw) - for (int64_t ch = 0; ch < 3; ++ch) - out[t * pd + ph * ps * 3 + pw * 3 + ch] = px(row * ps + ph, col * ps + pw, ch) * 2.0f - 1.0f; - } - return true; -} - -void pixel_shuffle_back(const float * src, int64_t grid, int64_t hidden, int64_t r, float * dst) { - const int64_t g2 = grid / r, c4 = hidden * r * r; - for (int64_t y = 0; y < g2; ++y) - for (int64_t x = 0; x < g2; ++x) { - const int64_t t = y * g2 + x; - for (int64_t c = 0; c < hidden; ++c) - for (int64_t i = 0; i < r; ++i) - for (int64_t j = 0; j < r; ++j) { - const int64_t pp = (r * y + i) * grid + (r * x + j); - const int64_t cp = c * r * r + i * r + j; - dst[t * c4 + cp] = src[pp * hidden + c]; - } - } -} - 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); }; + 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); }; auto fk = [&](const char * s) { static char b[64]; std::snprintf(b, sizeof(b), "gr00t_n1_6.%s", s); return b; }; - U(fk("vit_hidden"), m.vit_hidden); U(fk("vit_layers"), m.vit_layers); U(fk("vit_heads"), m.vit_heads); U(fk("vit_inter"), m.vit_inter); - U(fk("image_size"), m.image_size); U(fk("patch_size"), m.patch_size); - U(fk("vit_num_patches"), m.vit_num_patches); U(fk("n_img_tokens"), m.n_img_tokens); U(fk("vit_pixel_shuffle"), m.vit_pixel_shuffle); U(fk("mlp_connector_inner"), m.mlp_inner); - U(fk("lm_hidden"), m.lm_hidden); U(fk("lm_layers_used"), m.lm_layers); U(fk("lm_q_heads"), m.n_q); U(fk("lm_kv_heads"), m.n_kv); - U(fk("lm_head_dim"), m.lm_head_dim); U(fk("lm_inter"), m.lm_inter); U(fk("vocab_size"), m.vocab); U(fk("image_token_index"), m.image_token_index); - U(fk("backbone_embedding_dim"), m.bb_embed_dim); U(fk("input_embedding_dim"), m.in_embed_dim); - U(fk("dit_hidden"), m.dit_hidden); U(fk("dit_heads"), m.dit_heads); U(fk("dit_head_dim"), m.dit_head_dim); U(fk("dit_layers"), m.dit_layers); U(fk("dit_interleave"), m.dit_interleave); - U(fk("attend_text_every_n_blocks"), m.attend_text_every_n); - U(fk("action_horizon"), m.action_horizon); U(fk("action_dim"), m.action_dim); U(fk("max_state_dim"), m.max_state_dim); - 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); - F(fk("vit_ln_eps"), m.vit_ln_eps); F(fk("lm_rms_eps"), m.lm_rms_eps); F(fk("ln_eps"), m.ln_eps); F(fk("norm_out_eps"), m.norm_out_eps); F(fk("vlln_eps"), m.vlln_eps); F(fk("connector_ln_eps"), m.connector_ln_eps); - if (g.has(fk("lm_rope_theta"))) m.lm_rope_base = (float) g.f64(fk("lm_rope_theta")); - m.embodiment_id = 20; + U(fk("vit_hidden" ), m.vit.enc.cfg.hidden); + U(fk("vit_layers" ), m.vit_layers); + U(fk("vit_heads" ), m.vit.enc.cfg.heads); + U(fk("vit_inter" ), m.vit_inter); + U(fk("image_size" ), m.image_size); + U(fk("patch_size" ), m.patch_size); + U(fk("vit_num_patches" ), m.vit_num_patches); + U(fk("n_img_tokens" ), m.n_img_tokens); + U(fk("vit_pixel_shuffle"), m.vit_pixel_shuffle); + U(fk("mlp_connector_inner"), m.mlp_inner); + U(fk("lm_hidden" ), m.lm.cfg.hidden); + U(fk("lm_layers_used" ), m.lm.cfg.layers); + U(fk("lm_q_heads" ), m.lm.cfg.n_q); + U(fk("lm_kv_heads" ), m.lm.cfg.n_kv); + U(fk("lm_head_dim" ), m.lm.cfg.head_dim); + U(fk("lm_inter" ), m.lm_inter); + U(fk("vocab_size" ), m.vocab); + U(fk("image_token_index"), m.image_token_index); + U(fk("backbone_embedding_dim"), m.bb_embed_dim); + U(fk("input_embedding_dim" ), m.in_embed_dim); + U(fk("dit_hidden" ), m.dit.cfg.hidden); + U(fk("dit_heads" ), m.dit.cfg.heads); + U(fk("dit_head_dim" ), m.dit.cfg.head_dim); + U(fk("dit_layers" ), m.dit.cfg.layers); + U(fk("dit_interleave" ), m.dit_interleave); + U(fk("attend_text_every_n_blocks"), m.attend_text_every_n); + U(fk("action_horizon" ), m.action_horizon); + U(fk("action_dim" ), m.action_dim); + U(fk("max_state_dim" ), m.max_state_dim); + 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); + + F(fk("vit_ln_eps" ), m.vit.enc.cfg.ln_eps); + F(fk("lm_rms_eps" ), m.lm.cfg.rms_eps); + F(fk("ln_eps" ), m.dit.cfg.ln_eps); + F(fk("norm_out_eps" ), m.dit.cfg.norm_out_eps); + F(fk("vlln_eps" ), m.vlln_eps); + F(fk("connector_ln_eps" ), m.connector_ln_eps); + + if (g.has(fk("lm_rope_theta"))) + m.lm.cfg.rope.freq_base = (float) g.f64(fk("lm_rope_theta")); + + m.vit.enc.cfg.head_dim = m.vit.enc.cfg.hidden/m.vit.enc.cfg.heads; + m.lm.cfg.rope.n_dims = (int) m.lm.cfg.head_dim; + + m.aex.embodiment_id = 20; { const std::string js = g.str(fk("embodiment_id_mapping")); auto lookup = [&](const char * key) -> long { - const std::string k = std::string("\"") + key + "\""; - size_t p = js.find(k); if (p == std::string::npos) return -1; - p = js.find(':', p + k.size()); if (p == std::string::npos) return -1; - return std::strtol(js.c_str() + p + 1, nullptr, 10); + const std::string k = std::string("\"")+key+"\""; + size_t p = js.find(k); + if (p == std::string::npos) + return -1; + p = js.find(':', p+k.size()); + if (p == std::string::npos) + return -1; + return std::strtol(js.c_str()+p+1, nullptr, 10); }; - long gr1 = lookup("gr1"); if (gr1 >= 0) m.embodiment_id = gr1; + + const long gr1 = lookup("gr1"); + if (gr1 >= 0) + m.aex.embodiment_id = gr1; + if (const char * e = std::getenv("VLA_GR00T_EMBODIMENT")) { - char * end = nullptr; long v = std::strtol(e, &end, 10); - if (end && *end == '\0') m.embodiment_id = v; - else { long id = lookup(e); if (id >= 0) m.embodiment_id = id; else std::fprintf(stderr, "vla(gr00tn1d6): embodiment tag '%s' not in embodiment_id_mapping; using id %lld\n", e, (long long) m.embodiment_id); } + char * end = nullptr; + const long v = std::strtol(e, &end, 10); + if (end && *end == '\0') { + m.aex.embodiment_id = v; + } else { + const long id = lookup(e); + if (id >= 0) + m.aex.embodiment_id = id; + else std::fprintf(stderr, "vla(gr00tn1d6): embodiment tag '%s' not in embodiment_id_mapping; using id %lld\n", e, (long long) m.aex.embodiment_id); + } } } - 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; } + if (m.aex.embodiment_id < 0 || m.aex.embodiment_id >= m.max_embodiments) { + std::fprintf(stderr, "vla(gr00tn1d6): embodiment id %lld out of range [0,%lld)\n", + (long long) m.aex.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) { + 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) { + 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; @@ -270,136 +200,109 @@ bool load_config(const gguf_reader & g, Gr00tN1d6ModelArch & m, Config & cfg) { } 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; - cfg.real_state_dim = m.max_state_dim; cfg.real_action_dim = m.action_dim; - 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; + 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; + cfg.real_state_dim = m.max_state_dim; + cfg.real_action_dim = m.action_dim; + cfg.hidden = m.lm.cfg.hidden; + cfg.n_q_heads = m.lm.cfg.n_q; + cfg.n_kv_heads = m.lm.cfg.n_kv; + cfg.head_dim = m.lm.cfg.head_dim; + cfg.n_layers = m.lm.cfg.layers; + cfg.num_steps = (int) m.num_steps; + cfg.rms_eps = m.lm.cfg.rms_eps; + cfg.rope_n_dims = (int) m.lm.cfg.head_dim; + cfg.rope_mode = GGML_ROPE_TYPE_NEOX; + cfg.rope_freq_base = m.lm.cfg.rope.freq_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; + cfg.denormalized = false; + cfg.norm_eps = 1e-8f; return true; } } Gr00tN1d6ModelArch::~Gr00tN1d6ModelArch() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(gr00tn1d6): note - mmproj '%s' is ignored (the vision tower is bundled in the combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); - m->gguf_path = ckpt_path; - m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); + m->lm.cfg.rope.freq_base = 1000000.0f; - if (!m->io.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; + 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 " "lm=Qwen3 %lldd×%lldL (%lldq/%lldkv×%lld) dit=AlternateVLDiT %lldL×%lldh×%lld(inner %lld) attend_text_every_n=%lld in_emb=%lld " "horizon=%lld action_dim=%lld max_state=%lld N_steps=%lld embodiment=%lld resident=%s\n", - (long long) m->vit_hidden, (long long) m->vit_layers, (long long) m->vit_heads, (long long) m->vit_pixel_shuffle, (long long) m->n_img_tokens, (long long) m->mlp_inner, - (long long) m->lm_hidden, (long long) m->lm_layers, (long long) m->n_q, (long long) m->n_kv, (long long) m->lm_head_dim, - (long long) m->dit_layers, (long long) m->dit_heads, (long long) m->dit_head_dim, (long long) m->dit_hidden, (long long) m->attend_text_every_n, (long long) m->in_embed_dim, - (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, + (long long) m->vit.enc.cfg.hidden, (long long) m->vit_layers, (long long) m->vit.enc.cfg.heads, (long long) m->vit_pixel_shuffle, (long long) m->n_img_tokens, (long long) m->mlp_inner, + (long long) m->lm.cfg.hidden, (long long) m->lm.cfg.layers, (long long) m->lm.cfg.n_q, (long long) m->lm.cfg.n_kv, (long long) m->lm.cfg.head_dim, + (long long) m->dit.cfg.layers, (long long) m->dit.cfg.heads, (long long) m->dit.cfg.head_dim, (long long) m->dit.cfg.hidden, (long long) m->attend_text_every_n, (long long) m->in_embed_dim, + (long long) m->action_horizon, (long long) m->action_dim, (long long) m->max_state_dim, (long long) m->num_steps, (long long) m->aex.embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); - { - const Backend b = backend_init("vla(gr00tn1d6)", m->n_threads); - if (!b.handle) { return nullptr; } - m->backend = b.handle; - } + 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 }; + ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_init(ctx_weights) failed\n"); return nullptr; } - ggml_context * W = m->ctx_weights; - auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { - const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(gr00tn1d6): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); - ggml_set_name(t, name); return t; - }; - auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; - auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; - - bool ok = true; - - m->vit_patch_w = mk_mm("vit.patch_embd.weight"); - m->vit_patch_b = mk_f32("vit.patch_embd.bias"); - m->vit_pos = mk_f32("vit.pos_embd"); - m->vit_post_ln_w = mk_f32("vit.post_ln.weight"); m->vit_post_ln_b = mk_f32("vit.post_ln.bias"); - m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->vit[i]; - w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); - ok &= w.ln1w&&w.ln1b&&w.ln2w&&w.ln2b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wfc1&&w.bfc1&&w.Wfc2&&w.bfc2; + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(gr00tn1d6): ggml_init(ctx_weights) failed\n"); + return nullptr; } - m->mm_ln_w = mk_f32("mm.ln.weight"); m->mm_ln_b = mk_f32("mm.ln.bias"); - m->mm_fc1_w = mk_mm("mm.fc1.weight"); m->mm_fc1_b = mk_f32("mm.fc1.bias"); - m->mm_fc2_w = mk_mm("mm.fc2.weight"); m->mm_fc2_b = mk_f32("mm.fc2.bias"); - - m->lm_output_norm = mk_f32("vlm.output_norm.weight"); - m->lm.resize(m->lm_layers); - for (int64_t i = 0; i < m->lm_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vlm.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->lm[i]; - w.attn_norm=mk_f32(N("attn_norm.weight")); - w.Wq=mk_mm(N("attn_q.weight")); w.Wk=mk_mm(N("attn_k.weight")); w.Wv=mk_mm(N("attn_v.weight")); w.Wo=mk_mm(N("attn_o.weight")); - w.q_norm=mk_f32(N("attn_q_norm.weight")); w.k_norm=mk_f32(N("attn_k_norm.weight")); w.ffn_norm=mk_f32(N("ffn_norm.weight")); - w.Wgate=mk_mm(N("ffn_gate.weight")); w.Wup=mk_mm(N("ffn_up.weight")); w.Wdown=mk_mm(N("ffn_down.weight")); - ok &= w.attn_norm&&w.Wq&&w.Wk&&w.Wv&&w.Wo&&w.q_norm&&w.k_norm&&w.ffn_norm&&w.Wgate&&w.Wup&&w.Wdown; - } + WeightLoader L("gr00tn1d6", g, m->ctx_weights, m->matmul_type); + + m->vit.declare(L, "vit", m->vit_layers, /*patch_embd_is_gemm=*/true); + + m->mm_ln_w = L.f32 ("mm.ln.weight"); + m->mm_ln_b = L.f32 ("mm.ln.bias"); + m->mm_fc1_w = L.gemm("mm.fc1.weight"); + m->mm_fc1_b = L.f32 ("mm.fc1.bias"); + m->mm_fc2_w = L.gemm("mm.fc2.weight"); + m->mm_fc2_b = L.f32 ("mm.fc2.bias"); + + m->lm.declare(L, "vlm"); + + m->vlln_w = L.f32("aex.vlln.weight"); + m->vlln_b = L.f32("aex.vlln.bias"); + + m->aex.declare(L, "aex"); + m->dit.declare(L, "aex.dit"); + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; - m->vlln_w=mk_f32("aex.vlln.weight"); m->vlln_b=mk_f32("aex.vlln.bias"); - - m->se_l1W=mk_f32("aex.state_enc.l1.W"); m->se_l1b=mk_f32("aex.state_enc.l1.b"); m->se_l2W=mk_f32("aex.state_enc.l2.W"); m->se_l2b=mk_f32("aex.state_enc.l2.b"); - m->ae_W1W=mk_f32("aex.act_enc.W1.W"); m->ae_W1b=mk_f32("aex.act_enc.W1.b"); m->ae_W2W=mk_f32("aex.act_enc.W2.W"); m->ae_W2b=mk_f32("aex.act_enc.W2.b"); m->ae_W3W=mk_f32("aex.act_enc.W3.W"); m->ae_W3b=mk_f32("aex.act_enc.W3.b"); - m->ad_l1W=mk_f32("aex.act_dec.l1.W"); m->ad_l1b=mk_f32("aex.act_dec.l1.b"); m->ad_l2W=mk_f32("aex.act_dec.l2.W"); m->ad_l2b=mk_f32("aex.act_dec.l2.b"); - m->pos_embd=mk_f32("aex.pos_embd"); - m->te_l1W=mk_mm("aex.dit.time_emb.l1.weight"); m->te_l1b=mk_f32("aex.dit.time_emb.l1.bias"); m->te_l2W=mk_mm("aex.dit.time_emb.l2.weight"); m->te_l2b=mk_f32("aex.dit.time_emb.l2.bias"); - m->dit.resize(m->dit_layers); - for (int64_t i = 0; i < m->dit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "aex.dit.%lld.%s", (long long) i, s); return p; }; - auto & w = m->dit[i]; - w.adaln_w=mk_mm(N("adaln.weight")); w.adaln_b=mk_f32(N("adaln.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wff0=mk_mm(N("ff0.weight")); w.bff0=mk_f32(N("ff0.bias")); w.Wff2=mk_mm(N("ff2.weight")); w.bff2=mk_f32(N("ff2.bias")); - ok &= w.adaln_w&&w.adaln_b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wff0&&w.bff0&&w.Wff2&&w.bff2; - } - m->po1W=mk_mm("aex.dit.proj_out1.weight"); m->po1b=mk_f32("aex.dit.proj_out1.bias"); m->po2W=mk_mm("aex.dit.proj_out2.weight"); m->po2b=mk_f32("aex.dit.proj_out2.bias"); - ok &= m->vit_patch_w&&m->vit_patch_b&&m->vit_pos&&m->vit_post_ln_w&&m->vit_post_ln_b&&m->mm_ln_w&&m->mm_ln_b&&m->mm_fc1_w&&m->mm_fc1_b&&m->mm_fc2_w&&m->mm_fc2_b&&m->lm_output_norm&& - m->vlln_w&&m->vlln_b&&m->se_l1W&&m->se_l1b&&m->se_l2W&&m->se_l2b&&m->ae_W1W&&m->ae_W1b&&m->ae_W2W&&m->ae_W2b&&m->ae_W3W&&m->ae_W3b&& - m->ad_l1W&&m->ad_l1b&&m->ad_l2W&&m->ad_l2b&&m->pos_embd&&m->te_l1W&&m->te_l1b&&m->te_l2W&&m->te_l2b&&m->po1W&&m->po1b&&m->po2W&&m->po2b; - if (!ok) { std::fprintf(stderr, "vla(gr00tn1d6): weight tensor setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(gr00tn1d6): ggml_backend_alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for (ggml_tensor * t = ggml_get_first_tensor(W); t; t = ggml_get_next_tensor(W, t)) { - std::vector bytes = g.read_convert(ggml_get_name(t), t->type); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(gr00tn1d6): failed to load %s (%zu vs %zu bytes)\n", ggml_get_name(t), bytes.size(), ggml_nbytes(t)); return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); - } std::printf("vla(gr00tn1d6): weights resident in %.2f GiB (%s) - incl. SigLIP2 vision tower; embodiment id %lld\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0), m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16", (long long) m->embodiment_id); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), + m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16", (long long) m->aex.embodiment_id); return m; } @@ -407,230 +310,223 @@ std::vector Gr00tN1d6ModelArch::predict(const Inputs& in) { const auto t0 = std::chrono::steady_clock::now(); stats = Stats{}; - const int64_t H = lm_hidden, K = n_img_tokens, E = in_embed_dim; - const int64_t grid = image_size / patch_size; - const int64_t r = vit_pixel_shuffle; - const int64_t n_patches = grid * grid; - const int64_t patch_dim = 3 * patch_size * patch_size; - const int64_t c4 = vit_hidden * r * r; - const int64_t Nsa = 1 + action_horizon; + const int64_t H = lm.cfg.hidden; + const int64_t K = n_img_tokens; + const int64_t E = in_embed_dim; + const int64_t grid = image_size/patch_size; + const int64_t r = vit_pixel_shuffle; + const int64_t n_patches = grid*grid; + const int64_t patch_dim = 3*patch_size*patch_size; + const int64_t c4 = vit.enc.cfg.hidden*r*r; + const int64_t AD = action_dim; + const int64_t AH = action_horizon; + const int64_t Nsa = 1+AH; int64_t n_views = 0; std::vector img_emb_host; const float * img_emb_ptr = nullptr; + if (in.precomputed_img_emb && in.n_img_views > 0) { - n_views = in.n_img_views; + n_views = in.n_img_views; img_emb_ptr = in.precomputed_img_emb; } 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); + img_emb_host.assign((size_t) n_views*K*H, 0.0f); - ggml_context * VC = vision_scratch.reset((size_t) 64 * 1024 * 1024); + 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); - for (int64_t i = 0; i < vit_layers; ++i) h = build_siglip_layer(VC, vit[i], h, n_patches, vit_heads, vit_hidden / vit_heads, vit_hidden, vit_ln_eps); - ggml_tensor * post_ln = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, h, vit_ln_eps), vit_post_ln_w), vit_post_ln_b); + + 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 * post_ln = vit.build(VC, vit.embed_patches(VC, t_patches), n_patches, n_views); ggml_set_output(post_ln); + ggml_cgraph * vgA = ggml_new_graph_custom(VC, 8192, false); ggml_build_forward_expand(vgA, post_ln); if (!vision_scratch.alloc(backend, vgA)) { std::fprintf(stderr, "vla(gr00tn1d6): vision gallocr A alloc failed\n"); return {}; } - ggml_context * MC = merge_scratch.reset((size_t) 16 * 1024 * 1024); + 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); - ggml_tensor * vit_embeds = ggml_add(MC, ggml_mul_mat(MC, mm_fc2_w, ggml_gelu_erf(MC, mz1)), mm_fc2_b); + + ggml_tensor * t_shuf = ggml_new_tensor_3d(MC, GGML_TYPE_F32, c4, K, n_views); + ggml_set_input(t_shuf); + ggml_tensor * mln = layer_norm(MC, t_shuf, mm_ln_w, mm_ln_b, connector_ln_eps); + ggml_tensor * vit_embeds = ffn_gelu_erf(MC, mm_fc1_w, mm_fc1_b, mm_fc2_w, mm_fc2_b, mln); ggml_set_output(vit_embeds); + ggml_cgraph * vgB = ggml_new_graph(MC); ggml_build_forward_expand(vgB, vit_embeds); 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, - patches_all((size_t) patch_dim * n_patches * n_views), - post_ln_host((size_t) vit_hidden * n_patches * n_views), - shuf_host((size_t) c4 * K * n_views); + std::vector patches; + std::vector patches_all((size_t) patch_dim*n_patches*n_views); + std::vector post_ln_host((size_t) vit.enc.cfg.hidden*n_patches*n_views); + std::vector shuf_host((size_t) c4*K*n_views); + bool vok = true; - for (int64_t v = 0; v < n_views && vok; ++v) { - if (!preprocess_image_patches(in.images[v], image_size, patch_size, patches)) { vok = false; break; } - std::memcpy(patches_all.data() + (size_t) v * patch_dim * n_patches, patches.data(), patches.size() * sizeof(float)); + for (int64_t v=0; v(std::chrono::steady_clock::now() - tv0).count(); + 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(); if (!vok) return {}; img_emb_ptr = img_emb_host.data(); } else { - std::fprintf(stderr, "vla(gr00tn1d6): no images and no precomputed_img_emb in the request\n"); return {}; - } - const int64_t n_img = n_views * K; - - std::vector input_ids; - int64_t n_img_slots = 0; - for (int j = 0; j < in.n_lang; ++j) if (in.lang_tokens[j] == (int32_t) image_token_index) ++n_img_slots; - if (n_img_slots == n_img) { - input_ids.assign(in.lang_tokens, in.lang_tokens + in.n_lang); - } else if (n_img_slots == 0) { - input_ids.reserve(n_img + in.n_lang); - for (int64_t i = 0; i < n_img; ++i) input_ids.push_back((int32_t) image_token_index); - for (int j = 0; j < in.n_lang; ++j) input_ids.push_back(in.lang_tokens[j]); - } else { - std::fprintf(stderr, "vla(gr00tn1d6): lang_tokens has %lld image-token slots but n_img=%lld; expected 0 (v1 fallback) or %lld (chat-template path)\n", - (long long) n_img_slots, (long long) n_img, (long long) n_img); + std::fprintf(stderr, "vla(gr00tn1d6): no images and no precomputed_img_emb in the request\n"); return {}; } - 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 {}; } - - std::vector inputs_embeds((size_t) SEQ * H); - 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 {}; } - std::memcpy(inputs_embeds.data() + p * H, img_emb_ptr + k * H, H * sizeof(float)); ++k; - } - } + const int64_t n_img = n_views*K; - std::vector image_pos_idx, text_pos_idx; - image_pos_idx.reserve((size_t) n_img); text_pos_idx.reserve((size_t) (SEQ - n_img)); - for (int64_t p = 0; p < SEQ; ++p) { - if (input_ids[p] == (int32_t) image_token_index) image_pos_idx.push_back((int32_t) p); - else text_pos_idx.push_back((int32_t) p); - } - const int64_t SEQ_TXT = (int64_t) text_pos_idx.size(); - if ((int64_t) image_pos_idx.size() != n_img) { - std::fprintf(stderr, "vla(gr00tn1d6): internal: built %zu image positions, expected %lld\n", image_pos_idx.size(), (long long) n_img); return {}; - } + Prompt prompt; + if (!build_prompt("gr00tn1d6", in, n_img, (int32_t) image_token_index, max_seq_len, prompt)) return {}; + const int64_t SEQ = prompt.len(); + const int64_t SEQ_TXT = prompt.n_text(); + + std::vector inputs_embeds; + if (!fetch_embeds("gr00tn1d6", io, prompt, img_emb_ptr, H, inputs_embeds)) return {}; - const int64_t AD = action_dim, AH = action_horizon; - std::vector x_init((size_t) AH * AD); - 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); } + std::vector x_init; + init_noise(in, (size_t) AH*AD, x_init); - // 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, + 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); - ggml_tensor * t_state = ggml_new_tensor_2d(C, GGML_TYPE_F32, max_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); - - ggml_tensor * t_img_idx = ggml_new_tensor_1d(C, GGML_TYPE_I32, n_img); ggml_set_input(t_img_idx); - ggml_tensor * t_txt_idx = (SEQ_TXT > 0) ? ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ_TXT) : nullptr; - if (t_txt_idx) ggml_set_input(t_txt_idx); - std::vector t_tau(num_steps), t_tproj(num_steps); - for (int64_t s = 0; s < num_steps; ++s) { - t_tau[s] = ggml_new_tensor_2d(C, GGML_TYPE_F32, E, AH); ggml_set_input(t_tau[s]); - t_tproj[s] = ggml_new_tensor_1d(C, GGML_TYPE_F32, 256); ggml_set_input(t_tproj[s]); - } - - ggml_tensor * h = t_embeds; - for (int64_t i = 0; i < lm_layers; ++i) h = build_qwen3_layer(C, *this, lm[i], h, t_pos, t_lmmask, SEQ); - ggml_tensor * eagle = ggml_mul(C, ggml_rms_norm(C, h, lm_rms_eps), lm_output_norm); - - ggml_tensor * vl_embs = ggml_add(C, ggml_mul(C, ggml_norm(C, eagle, vlln_eps), vlln_w), vlln_b); - - ggml_tensor * vl_img = ggml_get_rows(C, vl_embs, t_img_idx); - ggml_tensor * vl_txt = (t_txt_idx ? ggml_get_rows(C, vl_embs, t_txt_idx) : vl_img); - - ggml_tensor * state_features = cat_linear(C, se_l2W, se_l2b, embodiment_id, ggml_relu(C, cat_linear(C, se_l1W, se_l1b, embodiment_id, t_state))); - - const float dt = 1.0f / (float) num_steps; - const int64_t every2 = 2 * attend_text_every_n; - - std::vector Kc(dit_layers, nullptr), Vc(dit_layers, nullptr); - for (int64_t i = 0; i < dit_layers; ++i) { - if (dit_interleave && (i % 2 == 1)) continue; - ggml_tensor * enc = (i % every2 == 0) ? vl_txt : vl_img; - dit_kv(C, *this, dit[i], enc, &Kc[i], &Vc[i]); - } - ggml_tensor * actions = t_x0; - for (int64_t s = 0; s < num_steps; ++s) { + 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); + ggml_tensor * t_state = ggml_new_tensor_2d(C, GGML_TYPE_F32, max_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); + + ggml_tensor * t_img_idx = ggml_new_tensor_1d(C, GGML_TYPE_I32, n_img); + ggml_set_input(t_img_idx); + ggml_tensor * t_txt_idx = (SEQ_TXT > 0) ? ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ_TXT) : nullptr; + if (t_txt_idx) + ggml_set_input(t_txt_idx); + + std::vector t_tau(num_steps), t_tproj(num_steps); + for (int64_t s=0; snb[1], 0)); + ggml_tensor * state_features = aex.encode_state(C, t_state); - ggml_tensor * sa = ggml_concat(C, state_features, af, 1); + const float dt = 1.0f/(float) num_steps; + const int64_t every2 = 2*attend_text_every_n; - ggml_tensor * hh = sa; - for (int64_t i = 0; i < dit_layers; ++i) { - ggml_tensor * enc; - if (dit_interleave && (i % 2 == 1)) enc = nullptr; - else if (i % every2 == 0) enc = vl_txt; - else enc = vl_img; - hh = build_dit_block(C, *this, dit[i], hh, temb, enc, Kc[i], Vc[i]); + std::vector Kc(dit.cfg.layers, nullptr), Vc(dit.cfg.layers, nullptr); + for (int64_t i=0; inb[1], (size_t) (Nsa - AH) * pred->nb[1])); - actions = ggml_add(C, actions, ggml_scale(C, vel, dt)); - } - ggml_set_name(actions, "action_pred"); ggml_set_output(actions); + ggml_tensor * actions = t_x0; + for (int64_t s=0; snb[1], (size_t)(Nsa-AH)*pred->nb[1])); + actions = ggml_add(C, actions, ggml_scale(C, vel, dt)); + } + 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; + 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; + 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 {}; } 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)); } - { std::vector mk((size_t) SEQ * SEQ); const float NEG = -std::numeric_limits::infinity(); - for (int64_t q = 0; q < SEQ; ++q) for (int64_t kv = 0; kv < SEQ; ++kv) mk[q * SEQ + kv] = (kv <= q) ? 0.0f : NEG; - ggml_backend_tensor_set(t_lmmask, mk.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)); - ggml_backend_tensor_set(t_img_idx, image_pos_idx.data(), 0, ggml_nbytes(t_img_idx)); - if (t_txt_idx) ggml_backend_tensor_set(t_txt_idx, text_pos_idx.data(), 0, ggml_nbytes(t_txt_idx)); - for (int64_t s = 0; s < num_steps; ++s) { - const int64_t bucket = (int64_t) ((double) s / (double) num_steps * (double) num_buckets); - std::vector tau, tpr; action_sinusoid(bucket, E, AH, tau); timesteps_proj(bucket, tpr); - ggml_backend_tensor_set(t_tau[s], tau.data(), 0, ggml_nbytes(t_tau[s])); - ggml_backend_tensor_set(t_tproj[s], tpr.data(), 0, ggml_nbytes(t_tproj[s])); + + ggml_backend_tensor_set(gio.t_embeds, inputs_embeds.data(), 0, ggml_nbytes(gio.t_embeds)); + + std::vector pp(SEQ); + for (int64_t i=0; i mask; + build_causal_mask(SEQ, mask); + ggml_backend_tensor_set(gio.t_lmmask, mask.data(), 0, ggml_nbytes(gio.t_lmmask)); + + std::vector st(max_state_dim, 0.0f); + for (int64_t i=0; i tau, tpr; + action_sinusoid(bucket, E, AH, tau); + timesteps_proj(bucket, tpr); + ggml_backend_tensor_set(gio.t_tau[s], tau.data(), 0, ggml_nbytes(gio.t_tau[s])); + ggml_backend_tensor_set(gio.t_tproj[s], tpr.data(), 0, ggml_nbytes(gio.t_tproj[s])); } const auto tc0 = std::chrono::steady_clock::now(); - const ggml_status st = ggml_backend_graph_compute(backend, gf); + const ggml_status status = ggml_backend_graph_compute(backend, main_graph.graph()); 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); return {}; } - stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(gr00tn1d6): graph compute failed (%d)\n", (int) status); + 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)); - stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::vector out((size_t) AH*AD); + ggml_backend_tensor_get(gio.actions, out.data(), 0, out.size()*sizeof(float)); + 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 d9742a7..ce19844 100644 --- a/src/models/gr00tn1d7.cpp +++ b/src/models/gr00tn1d7.cpp @@ -13,6 +13,14 @@ // limitations under the License. #include "arch.h" +#include "layers/attn.h" +#include "layers/linear.h" +#include "layers/norm.h" +#include "modules/action_expert.h" +#include "modules/dit_head.h" +#include "modules/encoder.h" +#include "modules/qwen3_lm.h" +#include "options.h" #include "model.h" #include "ggml.h" @@ -20,10 +28,11 @@ #include "ggml-backend.h" #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 "gguf_reader.h" +#include "scratch_ctx.h" +#include "layers/embed.h" +#include "modules/qwen3vl_vit.h" +#include "env_flag.h" #include #include @@ -43,9 +52,6 @@ namespace { 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; - ggml_tensor *Wqkv=nullptr,*bqkv=nullptr,*Wkv=nullptr,*bkv=nullptr; }; } @@ -74,21 +80,12 @@ struct Gr00tN1d7ModelArch : public ModelArchBase { float vlln_eps=1e-5f, vlsa_ln_eps=1e-5f, ln_eps=1e-5f, norm_out_eps=1e-6f, connector_ln_eps=1e-6f; int64_t embodiment_id = 2; - ggml_tensor *vit_patch_w=nullptr,*vit_patch_b=nullptr,*vit_pos=nullptr; - std::vector vit; - MergerW deepstack[3]; - MergerW merger; - ggml_tensor *lm_output_norm=nullptr; - std::vector lm; + Qwen3VLTower vit; + Qwen3LM lm; + EncStack vlsa; + ActionExpert aex; + DitHead dit; ggml_tensor *vlln_w=nullptr,*vlln_b=nullptr; - std::vector vlsa; - ggml_tensor *se_l1W=nullptr,*se_l1b=nullptr,*se_l2W=nullptr,*se_l2b=nullptr; - ggml_tensor *ae_W1W=nullptr,*ae_W1b=nullptr,*ae_W2W=nullptr,*ae_W2b=nullptr,*ae_W3W=nullptr,*ae_W3b=nullptr; - ggml_tensor *ad_l1W=nullptr,*ad_l1b=nullptr,*ad_l2W=nullptr,*ad_l2b=nullptr; - ggml_tensor *pos_embd=nullptr; - ggml_tensor *te_l1W=nullptr,*te_l1b=nullptr,*te_l2W=nullptr,*te_l2b=nullptr; - std::vector dit; - ggml_tensor *po1W=nullptr,*po1b=nullptr,*po2W=nullptr,*po2b=nullptr; bool caches_ready = false; std::vector c_grow, c_gcol; @@ -119,113 +116,9 @@ struct Gr00tN1d7ModelArch : public ModelArchBase { namespace { -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 * 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); - ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.n1w), w.n1b); - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, n1), w.bv); - 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)); - 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 * n3 = ggml_add(C, ggml_mul(C, ggml_norm(C, h1, ln_eps), w.n3w), w.n3b); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wff2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wff0, n3), w.bff0))), w.bff2); - return ggml_add(C, h1, ff); -} -ggml_tensor * build_qwen3_layer(ggml_context * C, const Gr00tN1d7ModelArch & 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; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * hn = ggml_mul(C, ggml_rms_norm(C, h, m.lm_rms_eps), w.attn_norm); - ggml_tensor * qp = ggml_mul_mat(C, w.Wq, hn); - ggml_tensor * kp = ggml_mul_mat(C, w.Wk, hn); - ggml_tensor * vp = ggml_mul_mat(C, w.Wv, hn); - ggml_tensor * qh = ggml_reshape_3d(C, qp, hd, n_q, seq); - ggml_tensor * kh = ggml_reshape_3d(C, kp, hd, n_kv, seq); - ggml_tensor * vh = ggml_reshape_3d(C, vp, hd, n_kv, seq); - ggml_tensor * qn = ggml_mul(C, ggml_rms_norm(C, qh, m.lm_rms_eps), w.q_norm); - ggml_tensor * kn = ggml_mul(C, ggml_rms_norm(C, kh, m.lm_rms_eps), w.k_norm); - int sections[4] = { 24, 20, 20, 0 }; - ggml_tensor * qr = ggml_rope_multi(C, qn, positions, nullptr, (int) hd, sections, GGML_ROPE_TYPE_IMROPE, 0, m.lm_rope_base, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - ggml_tensor * kr = ggml_rope_multi(C, kn, positions, nullptr, (int) hd, sections, GGML_ROPE_TYPE_IMROPE, 0, m.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)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, kr, 0, 2, 1, 3)); - 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); - } 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); - ggml_tensor * aw = ggml_soft_max_ext(C, kq, mask, scale, 0.0f); - att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, ggml_mul_mat(C, V, aw), 0, 2, 1, 3)), hq, seq); - } - ggml_tensor * h_attn = ggml_add(C, h, ggml_mul_mat(C, w.Wo, att)); - ggml_tensor * hn2 = ggml_mul(C, ggml_rms_norm(C, h_attn, m.lm_rms_eps), w.ffn_norm); - ggml_tensor * gate = ggml_silu(C, ggml_mul_mat(C, w.Wgate, hn2)); - ggml_tensor * up = ggml_mul_mat(C, w.Wup, hn2); - return ggml_add(C, h_attn, ggml_mul_mat(C, w.Wdown, ggml_mul(C, gate, up))); -} -void dit_kv(ggml_context * C, const Gr00tN1d7ModelArch & m, const DitLayerW & w, ggml_tensor * kv, - ggml_tensor ** K_out, ggml_tensor ** V_out) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, dim = m.dit_hidden, Tkv = kv->ne[1]; - if (w.Wkv) { - ggml_tensor * kvp = ggml_add(C, ggml_mul_mat(C, w.Wkv, kv), w.bkv); - *K_out = ggml_cont(C, ggml_permute(C, head_view(C, kvp, hd, heads, Tkv, dim, 2, 0), 0, 2, 1, 3)); - *V_out = ggml_cont(C, ggml_permute(C, head_view(C, kvp, hd, heads, Tkv, dim, 2, 1), 1, 2, 0, 3)); - return; - } - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, kv), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, kv), w.bv); - *K_out = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, Tkv), 0, 2, 1, 3)); - *V_out = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, Tkv), 1, 2, 0, 3)); -} - -ggml_tensor * build_dit_block(ggml_context * C, const Gr00tN1d7ModelArch & m, const DitLayerW & w, - ggml_tensor * h, ggml_tensor * temb, ggml_tensor * enc , - ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, dim = m.dit_hidden, Tk = h->ne[1]; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * n = adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, m.ln_eps); - ggml_tensor * K, * V, * Q; - if (!enc && w.Wqkv) { - ggml_tensor * qkv = ggml_add(C, ggml_mul_mat(C, w.Wqkv, n), w.bqkv); - Q = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 0), 0, 2, 1, 3)); - K = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 1), 0, 2, 1, 3)); - V = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 2), 1, 2, 0, 3)); - } else { - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n), w.bq); - Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, Tk), 0, 2, 1, 3)); - if (K_pre) { K = K_pre; V = V_pre; } - else { dit_kv(C, m, w, enc ? enc : n, &K, &V); } - } - 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)), dim, Tk); - ggml_tensor * h1 = ggml_add(C, h, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); - ggml_tensor * n3 = ggml_norm(C, h1, m.ln_eps); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wff2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wff0, n3), w.bff0))), w.bff2); - return ggml_add(C, h1, ff); -} 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); }; @@ -247,8 +140,8 @@ bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { // 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) { + 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; @@ -256,29 +149,66 @@ bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { 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); } + 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); + } } F(fk("vit_ln_eps"), m.vit_ln_eps); F(fk("lm_rms_eps"), m.lm_rms_eps); F(fk("ln_eps"), m.ln_eps); F(fk("norm_out_eps"), m.norm_out_eps); F(fk("vlln_eps"), m.vlln_eps); F(fk("vlsa_ln_eps"), m.vlsa_ln_eps); F(fk("connector_ln_eps"), m.connector_ln_eps); F(fk("vit_rope_theta"), m.vit_rope_base); - if (g.has(fk("lm_rope_theta"))) m.lm_rope_base = (float) g.f64(fk("lm_rope_theta")); - - m.embodiment_id = 2; + if (g.has(fk("lm_rope_theta"))) + m.lm_rope_base = (float) g.f64(fk("lm_rope_theta")); + + m.lm.cfg.hidden = m.lm_hidden; + m.lm.cfg.layers = m.lm_layers; + m.lm.cfg.n_q = m.n_q; + m.lm.cfg.n_kv = m.n_kv; + m.lm.cfg.head_dim = m.lm_head_dim; + m.lm.cfg.inter = m.lm_inter; + m.lm.cfg.rms_eps = m.lm_rms_eps; + m.lm.cfg.flash_attn = flash_attn_enabled(); + m.lm.cfg.rope.type = GGML_ROPE_TYPE_IMROPE; + m.lm.cfg.rope.n_dims = (int) m.lm_head_dim; + m.lm.cfg.rope.freq_base = m.lm_rope_base; + m.lm.cfg.rope.sections[0]= 24; + m.lm.cfg.rope.sections[1]= 20; + m.lm.cfg.rope.sections[2]= 20; + m.lm.cfg.rope.sections[3]= 0; + + m.vlsa.cfg.hidden = m.bb_embed_dim; + m.vlsa.cfg.heads = m.vlsa_heads; + m.vlsa.cfg.head_dim = m.vlsa_head_dim; + m.vlsa.cfg.ln_eps = m.vlsa_ln_eps; + m.vlsa.cfg.flash_attn = flash_attn_enabled(); + + m.dit.cfg.hidden = m.dit_hidden; + m.dit.cfg.heads = m.dit_heads; + m.dit.cfg.head_dim = m.dit_head_dim; + m.dit.cfg.layers = m.dit_layers; + m.dit.cfg.ln_eps = m.ln_eps; + m.dit.cfg.norm_out_eps = m.norm_out_eps; + + m.aex.embodiment_id = 2; { const std::string js = g.str(fk("embodiment_id_mapping")); auto lookup = [&](const char * key) -> long { - const std::string k = std::string("\"") + key + "\""; + const std::string k = std::string("\"")+key + "\""; size_t p = js.find(k); if (p == std::string::npos) return -1; - p = js.find(':', p + k.size()); if (p == std::string::npos) return -1; - return std::strtol(js.c_str() + p + 1, nullptr, 10); + p = js.find(':', p+k.size()); if (p == std::string::npos) return -1; + return std::strtol(js.c_str()+p+1, nullptr, 10); }; - long ls = lookup("libero_sim"); if (ls >= 0) m.embodiment_id = ls; + long ls = lookup("libero_sim"); if (ls >= 0) m.aex.embodiment_id = ls; if (const char * e = std::getenv("VLA_GR00T_EMBODIMENT")) { char * end = nullptr; long v = std::strtol(e, &end, 10); - if (end && *end == '\0') m.embodiment_id = v; - else { long id = lookup(e); if (id >= 0) m.embodiment_id = id; else std::fprintf(stderr, "vla(gr00tn1d7): embodiment tag '%s' not in embodiment_id_mapping; using id %lld\n", e, (long long) m.embodiment_id); } + if (end && *end == '\0') + m.aex.embodiment_id = v; + else { long id = lookup(e); if (id >= 0) m.aex.embodiment_id = id; else std::fprintf(stderr, "vla(gr00tn1d7): embodiment tag '%s' not in embodiment_id_mapping; using id %lld\n", e, (long long) m.aex.embodiment_id); } } } - if (m.embodiment_id < 0 || m.embodiment_id >= m.max_embodiments) { std::fprintf(stderr, "vla(gr00tn1d7): embodiment id %lld out of range [0,%lld)\n", (long long) m.embodiment_id, (long long) m.max_embodiments); return false; } + if (m.aex.embodiment_id < 0 || m.aex.embodiment_id >= m.max_embodiments) { + std::fprintf(stderr, "vla(gr00tn1d7): embodiment id %lld out of range [0,%lld)\n", (long long) m.aex.embodiment_id, (long long) m.max_embodiments); + return false; + } cfg = Config{}; cfg.n_img = 64; cfg.n_lang = m.max_seq_len; cfg.n_state = 1; @@ -298,25 +228,34 @@ bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { Gr00tN1d7ModelArch::~Gr00tN1d7ModelArch() { mg.release(); - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(gr00tn1d7): note - mmproj '%s' is ignored (the vision tower is bundled in the combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); m->gguf_path = ckpt_path; - m->matmul_type = std::getenv("VLA_GR00T_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); gguf_reader g("gr00tn1d7"); - if (!g.open(ckpt_path)) return nullptr; - if (!g.has("gr00t_n1_7.architecture")) { std::fprintf(stderr, "vla(gr00tn1d7): %s is not a gr00t_n1_7 GGUF\n", ckpt_path.c_str()); return nullptr; } - if (!load_config(g, *m, m->cfg)) return nullptr; + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("gr00t_n1_7.architecture")) { + std::fprintf(stderr, "vla(gr00tn1d7): %s is not a gr00t_n1_7 GGUF\n", ckpt_path.c_str()); + return nullptr; + } + if (!load_config(g, *m, m->cfg)) + return nullptr; std::printf("vla(gr00tn1d7): vit=Qwen3-VL %lldd×%lldL×%lldh (Conv3d patch %lld², temporal %lld; learned pos %lld + 2D rope; deepstack@{%lld,%lld,%lld}; merge÷%lld) " "lm=Qwen3-VL %lldd×%lldL (%lldq/%lldkv×%lld, θ=%g) vlsa=%lldL×%lldh×%lld dit=AlternateVLDiT %lldL×%lldh×%lld(inner %lld) attend_text_every_n=%lld " "in_emb=%lld horizon=%lld action_dim=%lld max_state=%lld N_steps=%lld embodiment=%lld resident=%s\n", @@ -325,171 +264,66 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, (long long) m->lm_hidden, (long long) m->lm_layers, (long long) m->n_q, (long long) m->n_kv, (long long) m->lm_head_dim, (double) m->lm_rope_base, (long long) m->vlsa_layers, (long long) m->vlsa_heads, (long long) m->vlsa_head_dim, (long long) m->dit_layers, (long long) m->dit_heads, (long long) m->dit_head_dim, (long long) m->dit_hidden, (long long) m->attend_text_every_n, (long long) m->in_embed_dim, - (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, + (long long) m->action_horizon, (long long) m->action_dim, (long long) m->max_state_dim, (long long) m->num_steps, (long long) m->aex.embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); { const Backend b = backend_init("vla(gr00tn1d7)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; } - ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; + ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(gr00tn1d7): ggml_init(ctx_weights) failed\n"); return nullptr; } - ggml_context * W = m->ctx_weights; - auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { - const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(gr00tn1d7): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); - ggml_set_name(t, name); return t; - }; - auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; - auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; - - struct FusedSpec { ggml_tensor * dst; std::vector srcs; }; - std::vector fused; - const bool fuse = true; - auto mk_fused = [&](const char * out_name, std::vector srcs, ggml_type type) -> ggml_tensor * { - const ggml_tensor * g0 = g.meta(srcs[0]); - if (!g0) { std::fprintf(stderr, "vla(gr00tn1d7): fused src missing %s\n", srcs[0]); return nullptr; } - const bool is1d = ggml_n_dims(g0) == 1; - int64_t ne0 = g0->ne[0], acc = 0; - for (const char * s : srcs) { const ggml_tensor * gs = g.meta(s); if (!gs) { std::fprintf(stderr, "vla(gr00tn1d7): fused src missing %s\n", s); return nullptr; } acc += is1d ? gs->ne[0] : gs->ne[1]; } - ggml_tensor * t = is1d ? ggml_new_tensor_1d(W, type, acc) : ggml_new_tensor_2d(W, type, ne0, acc); - ggml_set_name(t, out_name); - FusedSpec fs{t, {}}; for (const char * s : srcs) fs.srcs.emplace_back(s); fused.push_back(std::move(fs)); - return t; - }; - - bool ok = true; - - m->vit_patch_w = mk_mm("vit.patch_embd.weight"); m->vit_patch_b = mk_f32("vit.patch_embd.bias"); m->vit_pos = mk_f32("vit.pos_embd"); - m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->vit[i]; - w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); - w.Wqkv=mk_mm(N("attn_qkv.weight")); w.bqkv=mk_f32(N("attn_qkv.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); - ok &= w.ln1w&&w.ln1b&&w.ln2w&&w.ln2b&&w.Wqkv&&w.bqkv&&w.Wo&&w.bo&&w.Wfc1&&w.bfc1&&w.Wfc2&&w.bfc2; - } - for (int j = 0; j < 3; ++j) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.deepstack.%d.%s", j, s); return p; }; - auto & w = m->deepstack[j]; - w.nw=mk_f32(N("norm.weight")); w.nb=mk_f32(N("norm.bias")); w.fc1w=mk_mm(N("fc1.weight")); w.fc1b=mk_f32(N("fc1.bias")); w.fc2w=mk_mm(N("fc2.weight")); w.fc2b=mk_f32(N("fc2.bias")); - ok &= w.nw&&w.nb&&w.fc1w&&w.fc1b&&w.fc2w&&w.fc2b; - } - { auto & w = m->merger; - w.nw=mk_f32("vit.merger.norm.weight"); w.nb=mk_f32("vit.merger.norm.bias"); w.fc1w=mk_mm("vit.merger.fc1.weight"); w.fc1b=mk_f32("vit.merger.fc1.bias"); w.fc2w=mk_mm("vit.merger.fc2.weight"); w.fc2b=mk_f32("vit.merger.fc2.bias"); - ok &= w.nw&&w.nb&&w.fc1w&&w.fc1b&&w.fc2w&&w.fc2b; } - - m->lm_output_norm = mk_f32("vlm.output_norm.weight"); - m->lm.resize(m->lm_layers); - for (int64_t i = 0; i < m->lm_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vlm.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->lm[i]; - w.attn_norm=mk_f32(N("attn_norm.weight")); - w.Wq=mk_mm(N("attn_q.weight")); w.Wk=mk_mm(N("attn_k.weight")); w.Wv=mk_mm(N("attn_v.weight")); w.Wo=mk_mm(N("attn_o.weight")); - w.q_norm=mk_f32(N("attn_q_norm.weight")); w.k_norm=mk_f32(N("attn_k_norm.weight")); w.ffn_norm=mk_f32(N("ffn_norm.weight")); - w.Wgate=mk_mm(N("ffn_gate.weight")); w.Wup=mk_mm(N("ffn_up.weight")); w.Wdown=mk_mm(N("ffn_down.weight")); - ok &= w.attn_norm&&w.Wq&&w.Wk&&w.Wv&&w.Wo&&w.q_norm&&w.k_norm&&w.ffn_norm&&w.Wgate&&w.Wup&&w.Wdown; + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(gr00tn1d7): ggml_init(ctx_weights) failed\n"); + return nullptr; } - m->vlln_w=mk_f32("aex.vlln.weight"); m->vlln_b=mk_f32("aex.vlln.bias"); - m->vlsa.resize(m->vlsa_layers); - for (int64_t i = 0; i < m->vlsa_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "aex.vlsa.%lld.%s", (long long) i, s); return p; }; - auto & w = m->vlsa[i]; - w.n1w=mk_f32(N("norm1.weight")); w.n1b=mk_f32(N("norm1.bias")); w.n3w=mk_f32(N("norm3.weight")); w.n3b=mk_f32(N("norm3.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wff0=mk_mm(N("ff0.weight")); w.bff0=mk_f32(N("ff0.bias")); w.Wff2=mk_mm(N("ff2.weight")); w.bff2=mk_f32(N("ff2.bias")); - ok &= w.n1w&&w.n1b&&w.n3w&&w.n3b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wff0&&w.bff0&&w.Wff2&&w.bff2; - } - m->se_l1W=mk_f32("aex.state_enc.l1.W"); m->se_l1b=mk_f32("aex.state_enc.l1.b"); m->se_l2W=mk_f32("aex.state_enc.l2.W"); m->se_l2b=mk_f32("aex.state_enc.l2.b"); - m->ae_W1W=mk_f32("aex.act_enc.W1.W"); m->ae_W1b=mk_f32("aex.act_enc.W1.b"); m->ae_W2W=mk_f32("aex.act_enc.W2.W"); m->ae_W2b=mk_f32("aex.act_enc.W2.b"); m->ae_W3W=mk_f32("aex.act_enc.W3.W"); m->ae_W3b=mk_f32("aex.act_enc.W3.b"); - m->ad_l1W=mk_f32("aex.act_dec.l1.W"); m->ad_l1b=mk_f32("aex.act_dec.l1.b"); m->ad_l2W=mk_f32("aex.act_dec.l2.W"); m->ad_l2b=mk_f32("aex.act_dec.l2.b"); - m->pos_embd=mk_f32("aex.pos_embd"); - m->te_l1W=mk_mm("aex.dit.time_emb.l1.weight"); m->te_l1b=mk_f32("aex.dit.time_emb.l1.bias"); m->te_l2W=mk_mm("aex.dit.time_emb.l2.weight"); m->te_l2b=mk_f32("aex.dit.time_emb.l2.bias"); - m->dit.resize(m->dit_layers); - for (int64_t i = 0; i < m->dit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "aex.dit.%lld.%s", (long long) i, s); return p; }; - auto & w = m->dit[i]; - w.adaln_w=mk_mm(N("adaln.weight")); w.adaln_b=mk_f32(N("adaln.bias")); - if (fuse) { - const std::string pre = "aex.dit." + std::to_string((long long) i) + "."; - const std::string qn=pre+"attn_q.weight", kn=pre+"attn_k.weight", vn=pre+"attn_v.weight"; - const std::string qb=pre+"attn_q.bias", kb=pre+"attn_k.bias", vb=pre+"attn_v.bias"; - if (m->dit_interleave && (i % 2 == 1)) { - const std::string ow=pre+"attn_qkv.fused.w", ob=pre+"attn_qkv.fused.b"; - w.Wqkv=mk_fused(ow.c_str(), {qn.c_str(),kn.c_str(),vn.c_str()}, m->matmul_type); - w.bqkv=mk_fused(ob.c_str(), {qb.c_str(),kb.c_str(),vb.c_str()}, GGML_TYPE_F32); - ok &= w.Wqkv&&w.bqkv; - } else { - const std::string ow=pre+"attn_kv.fused.w", ob=pre+"attn_kv.fused.b"; - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); - w.Wkv=mk_fused(ow.c_str(), {kn.c_str(),vn.c_str()}, m->matmul_type); - w.bkv=mk_fused(ob.c_str(), {kb.c_str(),vb.c_str()}, GGML_TYPE_F32); - ok &= w.Wq&&w.bq&&w.Wkv&&w.bkv; - } - } else { - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); - ok &= w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv; - } - w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wff0=mk_mm(N("ff0.weight")); w.bff0=mk_f32(N("ff0.bias")); w.Wff2=mk_mm(N("ff2.weight")); w.bff2=mk_f32(N("ff2.bias")); - ok &= w.adaln_w&&w.adaln_b&&w.Wo&&w.bo&&w.Wff0&&w.bff0&&w.Wff2&&w.bff2; - } - m->po1W=mk_mm("aex.dit.proj_out1.weight"); m->po1b=mk_f32("aex.dit.proj_out1.bias"); m->po2W=mk_mm("aex.dit.proj_out2.weight"); m->po2b=mk_f32("aex.dit.proj_out2.bias"); - ok &= m->vit_patch_w&&m->vit_patch_b&&m->vit_pos&&m->lm_output_norm&&m->vlln_w&&m->vlln_b&&m->se_l1W&&m->se_l1b&&m->se_l2W&&m->se_l2b&& - m->ae_W1W&&m->ae_W1b&&m->ae_W2W&&m->ae_W2b&&m->ae_W3W&&m->ae_W3b&&m->ad_l1W&&m->ad_l1b&&m->ad_l2W&&m->ad_l2b&&m->pos_embd&&m->te_l1W&&m->te_l1b&&m->te_l2W&&m->te_l2b&&m->po1W&&m->po1b&&m->po2W&&m->po2b; - if (!ok) { std::fprintf(stderr, "vla(gr00tn1d7): weight tensor setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(gr00tn1d7): ggml_backend_alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - std::set fused_dst; - for (const auto & fs : fused) fused_dst.insert(fs.dst); - for (ggml_tensor * t = ggml_get_first_tensor(W); t; t = ggml_get_next_tensor(W, t)) { - if (fused_dst.count(t)) continue; - std::vector bytes = g.read_convert(ggml_get_name(t), t->type); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(gr00tn1d7): failed to load %s (%zu vs %zu bytes)\n", ggml_get_name(t), bytes.size(), ggml_nbytes(t)); return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); - } - for (const auto & fs : fused) { - std::vector buf; - for (const std::string & s : fs.srcs) { - std::vector b = g.read_convert(s.c_str(), fs.dst->type); - if (b.empty()) { std::fprintf(stderr, "vla(gr00tn1d7): fused fill: read %s failed\n", s.c_str()); return nullptr; } - buf.insert(buf.end(), b.begin(), b.end()); - } - if (buf.size() != ggml_nbytes(fs.dst)) { - std::fprintf(stderr, "vla(gr00tn1d7): fused fill: %s size %zu vs %zu\n", ggml_get_name(fs.dst), buf.size(), ggml_nbytes(fs.dst)); return nullptr; - } - ggml_backend_tensor_set(fs.dst, buf.data(), 0, buf.size()); - } - if (fuse) std::printf("vla(gr00tn1d7): QKV-fused DiT (self Wqkv / cross Wkv) - %zu fused tensors\n", fused.size()); + WeightLoader L("gr00tn1d7", g, m->ctx_weights, m->matmul_type); + + m->vit.declare(L, "vit", m->vit_layers); + m->lm.declare(L, "vlm"); + + m->vlln_w = L.f32("aex.vlln.weight"); + m->vlln_b = L.f32("aex.vlln.bias"); + m->vlsa.declare(L, "aex.vlsa", m->vlsa_layers, EncNames{"norm1", "norm3", "ff0", "ff2"}); + + m->aex.declare(L, "aex"); + m->dit.declare(L, "aex.dit", true, m->dit_interleave != 0); + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + + std::printf("vla(gr00tn1d7): QKV-fused DiT (self Wqkv / cross Wkv)\n"); std::printf("vla(gr00tn1d7): weights resident in %.2f GiB (%s) - incl. Qwen3-VL vision tower + deepstack + vl_self_attention; embodiment id %lld\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0), m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16", (long long) m->embodiment_id); - if (!m->build_caches()) { std::fprintf(stderr, "vla(gr00tn1d7): build_caches failed\n"); return nullptr; } + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), + dtype_name(m->matmul_type), (long long) m->aex.embodiment_id); + if (!m->build_caches()) { + std::fprintf(stderr, "vla(gr00tn1d7): build_caches failed\n"); + return nullptr; + } return m; } bool Gr00tN1d7ModelArch::build_caches() { - if (caches_ready) return true; + if (caches_ready) + return true; const int64_t side = image_target_size, ps = patch_size, m2 = spatial_merge; - const int64_t grid = side / ps; - const int64_t hd_vit = vit_hidden / vit_heads; + const int64_t grid = side/ps; + const int64_t hd_vit = vit_hidden/vit_heads; const int64_t num_side = (int64_t) std::lround(std::sqrt((double) vit_num_pos)); const int64_t E = in_embed_dim, AH = action_horizon; merge_block_coords(grid, grid, m2, c_grow, c_gcol); vit_rope_tables(c_grow, c_gcol, hd_vit, (double) vit_rope_base, c_rope_cos, c_rope_sin); - if (!io.open(gguf_path)) { std::fprintf(stderr, "vla(gr00tn1d7): build_caches: io.open(%s) failed\n", gguf_path.c_str()); return false; } + if (!io.open(gguf_path)) { + std::fprintf(stderr, "vla(gr00tn1d7): build_caches: io.open(%s) failed\n", gguf_path.c_str()); + return false; + } std::vector pos_table = io.read_f32("vit.pos_embd"); if (pos_table.empty() || (int64_t) pos_table.size() != vit_num_pos * vit_hidden) { std::fprintf(stderr, "vla(gr00tn1d7): build_caches: vit.pos_embd unreadable\n"); return false; @@ -497,8 +331,8 @@ bool Gr00tN1d7ModelArch::build_caches() { interp_pos_embed(pos_table, num_side, vit_hidden, c_grow, c_gcol, grid, grid, c_pos_interp); c_tau.assign((size_t) num_steps, {}); c_tproj.assign((size_t) num_steps, {}); - for (int64_t s = 0; s < num_steps; ++s) { - const int64_t bucket = (int64_t) ((double) s / (double) num_steps * (double) num_buckets); + for (int64_t s=0; s Gr00tN1d7ModelArch::predict(const Inputs& in) { const int64_t H = lm_hidden, E = in_embed_dim; const int64_t side = image_target_size; const int64_t ps = patch_size, m2 = spatial_merge; - const int64_t grid = side / ps; + const int64_t grid = side/ps; const int64_t n_patches = grid * grid; - const int64_t K = (grid / m2) * (grid / m2); - const int64_t hd_vit = vit_hidden / vit_heads; - const int64_t AD = action_dim, AH = action_horizon, Nsa = 1 + AH; + const int64_t K = (grid/m2)*(grid/m2); + const int64_t hd_vit = vit_hidden/vit_heads; + const int64_t AD = action_dim, AH = action_horizon, Nsa = 1+AH; const bool do_dump = (std::getenv("VLA_GR00T_N17_DUMP") != nullptr); if (!caches_ready) { std::fprintf(stderr, "vla(gr00tn1d7): caches not ready\n"); return {}; } @@ -530,50 +364,66 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { if (in.precomputed_img_emb && in.n_img_views > 0) { n_views = in.n_img_views; img_emb_ptr = in.precomputed_img_emb; - for (int j = 0; j < 3; ++j) ds_host[j].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); } 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); - for (int j = 0; j < 3; ++j) ds_host[j].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_context * VC = vision_scratch.reset((size_t) 512 * 1024 * 1024); + 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); ggml_tensor * t_cos = ggml_new_tensor_2d(VC, GGML_TYPE_F32, hd_vit, n_patches); ggml_set_input(t_cos); ggml_tensor * t_sin = ggml_new_tensor_2d(VC, GGML_TYPE_F32, hd_vit, n_patches); ggml_set_input(t_sin); - ggml_tensor * h = ggml_add(VC, ggml_add(VC, ggml_mul_mat(VC, vit_patch_w, t_patches), vit_patch_b), t_pos); + ggml_tensor * h = ggml_add(VC, ggml_add(VC, ggml_mul_mat(VC, vit.patch_w, t_patches), vit.patch_b), t_pos); ggml_set_output(h); ggml_tensor * stash[3] = {nullptr, nullptr, nullptr}; - for (int64_t i = 0; i < vit_layers; ++i) { - h = build_vit_layer(VC, vit[i], h, t_cos, t_sin, n_patches, vit_heads, hd_vit, vit_hidden, vit_ln_eps); + for (int64_t i=0; i patches; bool vok = true; - for (int64_t v = 0; v < n_views && vok; ++v) { - if (!preprocess_image_patches("gr00tn1d7", in.images[v], side, ps, temporal_patch, grow, gcol, patches)) { vok = false; break; } + for (int64_t v=0; v(std::chrono::steady_clock::now() - tv0).count(); + stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now()-tv0).count(); if (!vok) return {}; img_emb_ptr = img_emb_host.data(); } else { @@ -583,13 +433,17 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { std::vector input_ids; int64_t n_img_slots = 0; - for (int j = 0; j < in.n_lang; ++j) if (in.lang_tokens[j] == (int32_t) image_token_index) ++n_img_slots; + for (int j=0; j Gr00tN1d7ModelArch::predict(const Inputs& in) { std::vector inputs_embeds((size_t) SEQ * H); 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) { + for (int64_t p=0; p= n_img) { std::fprintf(stderr, "vla(gr00tn1d7): more tokens than ViT embeds\n"); return {}; } - std::memcpy(inputs_embeds.data() + p * H, img_emb_ptr + k * H, H * sizeof(float)); ++k; + std::memcpy(inputs_embeds.data()+p * H, img_emb_ptr+k * H, H * sizeof(float)); ++k; } } std::vector image_pos_idx, text_pos_idx; - image_pos_idx.reserve((size_t) n_img); text_pos_idx.reserve((size_t) (SEQ - n_img)); - for (int64_t p = 0; p < SEQ; ++p) { - if (input_ids[p] == (int32_t) image_token_index) image_pos_idx.push_back((int32_t) p); - else text_pos_idx.push_back((int32_t) p); + image_pos_idx.reserve((size_t) n_img); text_pos_idx.reserve((size_t) (SEQ-n_img)); + for (int64_t p=0; p Gr00tN1d7ModelArch::predict(const Inputs& in) { std::vector> ds_pad(3); const bool inject_deepstack = (in.images && in.n_images > 0); - if (inject_deepstack) for (int j = 0; j < 3; ++j) { + if (inject_deepstack) 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)); + for (int64_t k=0; k x_init((size_t) AH * AD); - 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); } + 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); + } // 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(); + if (!use_cache) + mg.release(); ggml_tensor * eagle = nullptr, * vl_embs = nullptr; std::vector lm_h_dump, vlsa_dump; const MainKey mkey{ SEQ, n_img, SEQ_TXT, num_steps, inject_deepstack }; - const bool built = mg.ensure(backend, mkey, (size_t) 256 * 1024 * 1024, + 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_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); ggml_tensor * t_state = ggml_new_tensor_2d(C, GGML_TYPE_F32, max_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); ggml_tensor * t_ds[3] = {nullptr,nullptr,nullptr}; - if (inject_deepstack) for (int j = 0; j < 3; ++j) { t_ds[j] = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); ggml_set_input(t_ds[j]); } + if (inject_deepstack) for (int j=0; j<3; ++j) { + t_ds[j] = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); + ggml_set_input(t_ds[j]); + } ggml_tensor * t_img_idx = ggml_new_tensor_1d(C, GGML_TYPE_I32, n_img); ggml_set_input(t_img_idx); ggml_tensor * t_txt_idx = (SEQ_TXT > 0) ? ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ_TXT) : nullptr; - if (t_txt_idx) ggml_set_input(t_txt_idx); + if (t_txt_idx) + ggml_set_input(t_txt_idx); std::vector t_tau(num_steps), t_tproj(num_steps); - for (int64_t s = 0; s < num_steps; ++s) { + for (int64_t s=0; s Kc(dit_layers, nullptr), Vc(dit_layers, nullptr); - for (int64_t i = 0; i < dit_layers; ++i) { - if (dit_interleave && (i % 2 == 1)) continue; - ggml_tensor * enc = (i % every2 == 0) ? vl_txt : vl_img; - dit_kv(C, *this, dit[i], enc, &Kc[i], &Vc[i]); + for (int64_t i=0; inb[1], 0)); + for (int64_t s=0; snb[1], 0)); ggml_tensor * sa = ggml_concat(C, state_features, af, 1); ggml_tensor * hh = sa; - for (int64_t i = 0; i < dit_layers; ++i) { + for (int64_t i=0; inb[1], (size_t) (Nsa - AH) * pred->nb[1])); + ggml_tensor * model_output = ggml_add(C, ggml_mul_mat(C, dit.po2W, h_mod), dit.po2b); + ggml_tensor * pred = cat_linear(C, aex.ad_l2W, aex.ad_l2b, aex.embodiment_id, ggml_relu(C, cat_linear(C, aex.ad_l1W, aex.ad_l1b, aex.embodiment_id, model_output))); + ggml_tensor * vel = ggml_cont(C, ggml_view_2d(C, pred, AD, AH, pred->nb[1], (size_t) (Nsa-AH)*pred->nb[1])); actions = ggml_add(C, actions, ggml_scale(C, vel, dt)); } ggml_set_name(actions, "action_pred"); ggml_set_output(actions); @@ -740,62 +620,81 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { - const int64_t llm_grid_h = image_target_size / patch_size / spatial_merge; + const int64_t llm_grid_h = image_target_size/patch_size/spatial_merge; const int64_t llm_grid_w = llm_grid_h; - std::vector pp((size_t) 4 * SEQ, 0); + std::vector pp((size_t) 4*SEQ, 0); int64_t st = 0, st_idx = 0; while (st < SEQ) { int64_t img_start = -1; - for (int64_t i = st; i < SEQ; ++i) if (input_ids[i] == (int32_t) image_token_index) { img_start = i; break; } + for (int64_t i=st; i max_image_pos) max_image_pos = llm_grid_h - 1; - if (llm_grid_w - 1 > max_image_pos) max_image_pos = llm_grid_w - 1; - st_idx = image_offset + max_image_pos + 1; + int64_t max_image_pos = this_t-1; + if (llm_grid_h-1 > max_image_pos) + max_image_pos = llm_grid_h-1; + if (llm_grid_w-1 > max_image_pos) + max_image_pos = llm_grid_w-1; + st_idx = image_offset+max_image_pos+1; st = img_end; } - std::memcpy(pp.data() + (size_t) 3 * SEQ, pp.data() + (size_t) 0 * SEQ, (size_t) SEQ * sizeof(int32_t)); + 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) { build_causal_mask(SEQ, c_mask); 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)); } + { + std::vector st(max_state_dim, 0.0f); + for (int64_t i=0; i Gr00tN1d7ModelArch::predict(const Inputs& in) { 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(gr00tn1d7): graph compute failed (%d)\n", (int) st); mg.release(); return {}; } - stats.ms_inference = std::chrono::duration(tc1 - tc0).count(); + 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_backend_tensor_get(actions, out.data(), 0, out.size()*sizeof(float)); if (const char * dump = std::getenv("VLA_GR00T_N17_DUMP")) { auto dump_t = [&](const char * name, ggml_tensor * t) { const int64_t n0 = t->ne[0], n1 = t->ne[1]; - std::vector buf((size_t) n0 * n1); - ggml_backend_tensor_get(t, buf.data(), 0, buf.size() * sizeof(float)); + std::vector buf((size_t) n0*n1); + ggml_backend_tensor_get(t, buf.data(), 0, buf.size()*sizeof(float)); char path[1024]; std::snprintf(path, sizeof(path), "%s_%s_%lldx%lld.f32", dump, name, (long long) n0, (long long) n1); FILE * fp = std::fopen(path, "wb"); - if (fp) { std::fwrite(buf.data(), sizeof(float), buf.size(), fp); std::fclose(fp); std::fprintf(stderr, "vla(gr00tn1d7): dumped %s shape=(%lld,%lld) to %s\n", name, (long long) n1, (long long) n0, path); } + if (fp) { + std::fwrite(buf.data(), sizeof(float), buf.size(), fp); + std::fclose(fp); + std::fprintf(stderr, "vla(gr00tn1d7): dumped %s shape=(%lld,%lld) to %s\n", name, (long long) n1, (long long) n0, path); + } }; dump_t("eagle", eagle); dump_t("vl_embs", vl_embs); - for (size_t li = 0; li < lm_h_dump.size(); ++li) { char nm[32]; std::snprintf(nm, sizeof(nm), "lm_h_%02zu", li); dump_t(nm, lm_h_dump[li]); } + for (size_t li=0; li(std::chrono::steady_clock::now() - t0).count(); + if (!use_cache) + mg.release(); + stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now()-t0).count(); return out; } diff --git a/src/models/openvla_oft.cpp b/src/models/openvla_oft.cpp index 4650272..ed107bf 100644 --- a/src/models/openvla_oft.cpp +++ b/src/models/openvla_oft.cpp @@ -13,17 +13,19 @@ // limitations under the License. #include "arch.h" +#include "options.h" #include "model.h" -#include "vision_common.h" -#include "models/dual_tower.h" +#include "modules/preprocess.h" +#include "modules/dual_tower.h" #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" #include "backend.h" #include "gguf.h" -#include "models/gguf_reader.h" -#include "models/scratch_ctx.h" +#include "gguf_reader.h" +#include "scratch_ctx.h" +#include "env_flag.h" #include #include @@ -47,33 +49,52 @@ bool parse_stats(const std::string & js, int64_t want, std::vector & q01, }; const char * env = std::getenv("VLA_OPENVLA_OFT_UNNORM_KEY"); size_t suite_pos; - if (env) { suite = env; suite_pos = find_key(0, suite); } + if (env) { + suite = env; + suite_pos = find_key(0, suite); + } else { size_t b = js.find('{'); size_t q = js.find('"', b); - size_t qe = js.find('"', q + 1); - suite = js.substr(q + 1, qe - q - 1); suite_pos = q; + size_t qe = js.find('"', q+1); + suite = js.substr(q+1, qe-q-1); suite_pos = q; + } + if (suite_pos == std::string::npos) { + std::fprintf(stderr, "vla(openvla_oft): suite '%s' not in stats\n", suite.c_str()); + return false; } - if (suite_pos == std::string::npos) { std::fprintf(stderr, "vla(openvla_oft): suite '%s' not in stats\n", suite.c_str()); return false; } size_t act = find_key(suite_pos, "action"); - if (act == std::string::npos) return false; + if (act == std::string::npos) + return false; auto read_arr = [&](const std::string & key, std::vector & out) -> bool { size_t k = find_key(act, key); if (k == std::string::npos) return false; size_t lb = js.find('[', k); size_t rb = js.find(']', lb); - if (lb == std::string::npos || rb == std::string::npos) return false; - out.clear(); size_t p = lb + 1; + if (lb == std::string::npos || rb == std::string::npos) + return false; + out.clear(); size_t p = lb+1; while (p < rb) { - while (p < rb && (js[p] == ',' || js[p] == ' ' || js[p] == '\n' || js[p] == '\t' || js[p] == '\r')) ++p; - if (p >= rb) break; + while (p < rb && (js[p] == ',' || js[p] == ' ' || js[p] == '\n' || js[p] == '\t' || js[p] == '\r')) + ++p; + if (p >= rb) + break; bool t = (js.compare(p, 4, "true") == 0), f = (js.compare(p, 5, "false") == 0); - if (t || f) { out.push_back(t ? 1.0f : 0.0f); p += t ? 4 : 5; } - else { out.push_back(std::strtof(js.c_str() + p, nullptr)); while (p < rb && js[p] != ',') ++p; } + if (t || f) { + out.push_back(t ? 1.0f : 0.0f); + p += t ? 4 : 5; + } + else { + out.push_back(std::strtof(js.c_str()+p, nullptr)); + while (p < rb && js[p] != ',') + ++p; + } } return true; }; std::vector mk; - if (!read_arr("q01", q01) || !read_arr("q99", q99)) return false; - if (!read_arr("mask", mk)) mk.assign(want, 1.0f); - mask.assign(mk.size(), 1); for (size_t i = 0; i < mk.size(); ++i) mask[i] = mk[i] != 0.0f ? 1 : 0; + if (!read_arr("q01", q01) || !read_arr("q99", q99)) + return false; + if (!read_arr("mask", mk)) + mk.assign(want, 1.0f); + mask.assign(mk.size(), 1); for (size_t i=0; i dvit; - ggml_tensor *s_patch_w,*s_patch_b,*s_pos; std::vector svit; - ggml_tensor *pj_fc1w,*pj_fc1b,*pj_fc2w,*pj_fc2b,*pj_fc3w,*pj_fc3b; + DualTower vis; ggml_tensor *token_embd,*lm_out_norm; std::vector lm; ggml_tensor *pp_fc1w,*pp_fc1b,*pp_fc2w,*pp_fc2b; ggml_tensor *h_ln1w,*h_ln1b,*h_fc1w,*h_fc1b,*h_ln2w,*h_ln2b,*h_fc2w,*h_fc2b; std::vector hblk; @@ -130,15 +154,20 @@ struct OpenVlaOftModelArch : public ModelArchBase { std::unique_ptr openvla_oft_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(openvla_oft): note - mmproj '%s' ignored (vision baked into combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); - m->mt = std::getenv("VLA_OPENVLA_OFT_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; + m->mt = opts.weight_dtype.value_or(GGML_TYPE_BF16); gguf_reader g("openvla_oft"); - if (!g.open(ckpt_path)) return nullptr; - if (!g.has("openvla_oft.architecture")) { std::fprintf(stderr, "vla(openvla_oft): not an openvla_oft GGUF\n"); return nullptr; } + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("openvla_oft.architecture")) { + std::fprintf(stderr, "vla(openvla_oft): not an openvla_oft GGUF\n"); + return nullptr; + } auto U=[&](const char*k,int64_t&d){ if(g.has(k)) d=(int64_t)g.u32(k); }; auto F=[&](const char*k,float&d){ if(g.has(k)) d=g.f32(k); }; @@ -159,52 +188,34 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path // 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 (m->lm_head_dim==0) + m->lm_head_dim = m->lm_hidden/m->n_q; if (g.has("openvla_oft.statistics_json")) { if (!parse_stats(g.str("openvla_oft.statistics_json"), m->action_dim, m->q01, m->q99, m->unnorm_mask, m->suite)) - { std::fprintf(stderr, "vla(openvla_oft): failed to parse statistics_json\n"); return nullptr; } + { + std::fprintf(stderr, "vla(openvla_oft): failed to parse statistics_json\n"); + return nullptr; + } std::printf("vla(openvla_oft): unnorm suite = %s (q99 dim %zu)\n", m->suite.c_str(), m->q99.size()); } { const Backend b = backend_init("vla(openvla_oft)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; } ggml_init_params wp = { (size_t)64*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - ggml_context * W = m->ctx_weights; bool ok = true; - auto mk=[&](const char*name, ggml_type ty)->ggml_tensor*{ const ggml_tensor*gt=g.meta(name); - if(!gt){ std::fprintf(stderr,"vla(openvla_oft): missing %s\n",name); ok=false; return nullptr; } - ggml_tensor*t=ggml_new_tensor(W,g.resident_type(gt,ty),ggml_n_dims(gt),gt->ne); ggml_set_name(t,name); return t; }; - auto mm=[&](const char*n){ return mk(n,m->mt); }; - auto f32=[&](const char*n){ return mk(n,GGML_TYPE_F32); }; - - m->d_patch_w=mk("vis.d.patch.weight",GGML_TYPE_F32); m->d_patch_b=f32("vis.d.patch.bias"); - m->d_cls=f32("vis.d.cls"); m->d_reg=f32("vis.d.reg"); m->d_pos=f32("vis.d.pos"); - m->dvit.resize(m->d_layers); - for(int i=0;id_layers;++i){ auto&w=m->dvit[i]; char b[64]; - auto N=[&](const char*s){ std::snprintf(b,sizeof(b),"vis.d.blk.%d.%s",i,s); return (const char*)b; }; - w.n1w=f32(N("ln1.weight")); w.n1b=f32(N("ln1.bias")); w.n2w=f32(N("ln2.weight")); w.n2b=f32(N("ln2.bias")); - w.ls1=f32(N("ls1")); w.ls2=f32(N("ls2")); w.Wqkv=mm(N("qkv.weight")); w.bqkv=f32(N("qkv.bias")); - w.Wproj=mm(N("proj.weight")); w.bproj=f32(N("proj.bias")); w.Wfc1=mm(N("fc1.weight")); w.bfc1=f32(N("fc1.bias")); - w.Wfc2=mm(N("fc2.weight")); w.bfc2=f32(N("fc2.bias")); } - - m->s_patch_w=mk("vis.s.patch.weight",GGML_TYPE_F32); m->s_patch_b=f32("vis.s.patch.bias"); m->s_pos=f32("vis.s.pos"); - m->svit.resize(m->s_layers); - for(int i=0;is_layers;++i){ auto&w=m->svit[i]; char b[64]; - auto N=[&](const char*s){ std::snprintf(b,sizeof(b),"vis.s.blk.%d.%s",i,s); return (const char*)b; }; - w.n1w=f32(N("ln1.weight")); w.n1b=f32(N("ln1.bias")); w.n2w=f32(N("ln2.weight")); w.n2b=f32(N("ln2.bias")); - w.ls1=nullptr; w.ls2=nullptr; w.Wqkv=mm(N("qkv.weight")); w.bqkv=f32(N("qkv.bias")); - w.Wproj=mm(N("proj.weight")); w.bproj=f32(N("proj.bias")); w.Wfc1=mm(N("fc1.weight")); w.bfc1=f32(N("fc1.bias")); - w.Wfc2=mm(N("fc2.weight")); w.bfc2=f32(N("fc2.bias")); } - - m->pj_fc1w=mm("vis.proj.fc1.weight"); m->pj_fc1b=f32("vis.proj.fc1.bias"); - m->pj_fc2w=mm("vis.proj.fc2.weight"); m->pj_fc2b=f32("vis.proj.fc2.bias"); - m->pj_fc3w=mm("vis.proj.fc3.weight"); m->pj_fc3b=f32("vis.proj.fc3.bias"); + WeightLoader L("openvla_oft", g, m->ctx_weights, m->mt); + auto mm = [&](const char * n) { return L.gemm("%s", n); }; + auto f32 = [&](const char * n) { return L.f32("%s", n); }; + + m->vis.declare(L, m->d_layers, m->s_layers); m->token_embd=mm("token_embd.weight"); m->lm_out_norm=f32("lm.output_norm.weight"); m->lm.resize(m->lm_layers); @@ -225,17 +236,16 @@ std::unique_ptr openvla_oft_create(const std::string& mmproj_path for(int i=0;ihead_blocks;++i){ auto&w=m->hblk[i]; char b[64]; auto N=[&](const char*s){ std::snprintf(b,sizeof(b),"aex.head.blk.%d.%s",i,s); return (const char*)b; }; w.lnw=f32(N("ln.weight")); w.lnb=f32(N("ln.bias")); w.linw=mm(N("lin.weight")); w.linb=f32(N("lin.bias")); } - if(!ok){ std::fprintf(stderr,"vla(openvla_oft): weight setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if(!m->weight_buf){ std::fprintf(stderr,"vla(openvla_oft): alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for(ggml_tensor*t=ggml_get_first_tensor(W); t; t=ggml_get_next_tensor(W,t)){ - std::vector bytes=g.read_convert(ggml_get_name(t),t->type); - if(bytes.empty()||bytes.size()!=ggml_nbytes(t)){ std::fprintf(stderr,"vla(openvla_oft): load %s (%zu vs %zu)\n",ggml_get_name(t),bytes.size(),ggml_nbytes(t)); return nullptr; } - ggml_backend_tensor_set(t,bytes.data(),0,bytes.size()); + if(!ok){ + std::fprintf(stderr,"vla(openvla_oft): weight setup failed\n"); + return nullptr; } + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + std::printf("vla(openvla_oft): weights resident %.2f GiB (%s) - DINOv2+SigLIP towers + Llama-2-7B + MLPResNet L1 head\n", - ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), m->mt==GGML_TYPE_F32?"F32":"BF16"); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), dtype_name(m->mt)); m->cfg.n_suffix = m->chunk; m->cfg.max_action_dim = m->action_dim; m->cfg.real_action_dim = m->action_dim; m->cfg.real_state_dim = m->proprio_dim; @@ -254,7 +264,7 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { if (n_views < 1) { std::fprintf(stderr, "vla(openvla_oft): need >=1 image view\n"); return {}; } if (!in.images) { std::fprintf(stderr, "vla(openvla_oft): n_images=%d but the images pointer is null\n", in.n_images); return {}; } // towers read S*S*3 per view; reject any view that is not exactly SxS. - for (int64_t v = 0; v < n_views; ++v) { + for (int64_t v=0; v OpenVlaOftModelArch::predict(const Inputs& in) { for(int v=0; v dbuf, sbuf; @@ -296,7 +306,7 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { const int64_t L = in.n_lang; // ggml_get_rows does not bound-check, so reject out-of-range tokens here. - for (int64_t i = 0; i < L; ++i) + for (int64_t i=0; i= vocab) { std::fprintf(stderr, "vla(openvla_oft): token %d out of vocab\n", in.lang_tokens[i]); return {}; @@ -306,10 +316,10 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { return {}; } const int64_t n_act = chunk * action_dim; - const int64_t NUM_PATCHES = NPATCH + 1; - const int64_t NUM_PROMPT_TOKENS = L - 1; - const int64_t ACT_START = NUM_PATCHES + NUM_PROMPT_TOKENS; - const int64_t SEQ = 1 + NUM_PATCHES + (L-1) + n_act + 1; + const int64_t NUM_PATCHES = NPATCH+1; + const int64_t NUM_PROMPT_TOKENS = L-1; + 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(); // LM + action head graph depends only on the sequence layout. const MainKey mkey{ SEQ, n_views, L }; @@ -317,7 +327,8 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { [&](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); + if(emb->type!=GGML_TYPE_F32) + emb=ggml_cast(C,emb,GGML_TYPE_F32); ggml_tensor*bos =ggml_cont(C,ggml_view_2d(C,emb,HC,1,emb->nb[1],0)); ggml_tensor*rest=ggml_cont(C,ggml_view_2d(C,emb,HC,L-1,emb->nb[1],emb->nb[1])); ggml_tensor*stop=ggml_cont(C,ggml_view_2d(C,emb,HC,1,emb->nb[1],L*emb->nb[1])); @@ -386,14 +397,23 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { 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 pp(SEQ); for(int64_t i=0;i pp(SEQ); + for(int64_t i=0;i sv(proprio_dim,0.0f); for(int64_t i=0;i z((size_t)HC*n_act,0.0f); ggml_backend_tensor_set(act0,z.data(),0,ggml_nbytes(act0)); } + { + 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"); return {}; } std::vector na((size_t)action_dim*chunk); diff --git a/src/models/pi0.cpp b/src/models/pi0.cpp index 51f2f28..c70dcca 100644 --- a/src/models/pi0.cpp +++ b/src/models/pi0.cpp @@ -13,6 +13,9 @@ // limitations under the License. #include "arch.h" +#include "modules/gemma_expert.h" +#include "modules/siglip_vit.h" +#include "options.h" #include "model.h" #include "ggml.h" @@ -21,12 +24,13 @@ #include "ggml-alloc.h" #include "backend.h" #include "gguf.h" -#include "models/gguf_reader.h" -#include "models/scratch_ctx.h" +#include "gguf_reader.h" +#include "scratch_ctx.h" #include "models/dit_common.h" -#include "models/vision_common.h" -#include "models/act_dtype.h" +#include "modules/preprocess.h" +#include "act_dtype.h" #include "cuda/vla_cuda_ops.h" +#include "env_flag.h" #include #include @@ -46,31 +50,9 @@ namespace vla { namespace { -struct GemmaLayerW { - ggml_tensor * ln_in = nullptr; - ggml_tensor * Wq = nullptr; - ggml_tensor * Wk = nullptr; - ggml_tensor * Wv = nullptr; - ggml_tensor * Wo = nullptr; - ggml_tensor * ln_post = nullptr; - ggml_tensor * Wgate = nullptr; - ggml_tensor * Wup = nullptr; - ggml_tensor * Wdown = nullptr; -}; - -// 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; }; - -// The +1 RMSNorm fixup is a Gemma quirk that must only touch the language towers, -// never the SigLIP LayerNorms (whose names would otherwise never match anyway). -bool is_gemma_norm(const std::string & name) { - const bool lm = name.rfind("vlm.", 0) == 0 || name.rfind("aex.", 0) == 0; - return lm && name.find("norm.weight") != std::string::npos; -} - 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; + return s.size() >= n && s.compare(s.size()-n, n, sfx) == 0; } } @@ -88,7 +70,9 @@ struct Pi0ModelArch : public ModelArchBase { 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; } + 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; @@ -109,15 +93,11 @@ struct Pi0ModelArch : public ModelArchBase { int64_t vit_hidden = 1152, vit_layers = 27, vit_heads = 16; int64_t vit_image_size = 224, vit_patch_size = 14, vit_n_tokens = 256; float vit_ln_eps = 1e-6f; - ggml_tensor * vit_patch_w = nullptr, * vit_patch_b = nullptr, * vit_pos = nullptr; - ggml_tensor * vit_post_ln_w = nullptr, * vit_post_ln_b = nullptr; - std::vector vit; + SigLipTower vit; ggml_tensor * mm_proj_w = nullptr, * mm_proj_b = nullptr; - std::vector pl_layers; - - std::vector ex_layers; - ggml_tensor * ex_final_norm = nullptr; + GemmaStack pl; + GemmaStack ex; ggml_tensor * W_sp = nullptr, * b_sp = nullptr; ggml_tensor * W_ain = nullptr, * b_ain = nullptr; @@ -143,15 +123,11 @@ namespace { // measured, so it stays opt-in on an unquantified risk rather than a measured // cost. (The evo1 SR drop this used to cite did not reproduce.) // VLA_PI0_BF16_ACT is the better lever here: 9.1%, and its SR was measured. -static inline bool pi0_fa_enabled() { - static const bool enabled = (std::getenv("VLA_PI0_FA") != nullptr); - return enabled; -} -ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, +ggml_tensor * build_siglip_layer(ggml_context * C, const EncBlockW & w, ggml_tensor * x, int64_t seq, int64_t heads, int64_t head_dim, int64_t hidden, float ln_eps, ggml_type at) { - const float scale = 1.0f / std::sqrt((float) head_dim); + const float scale = 1.0f/std::sqrt((float) head_dim); ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); ggml_tensor * q = as_type(C, ggml_add(C, mm_act(C, w.Wq, n1, at), w.bq), GGML_TYPE_F32); ggml_tensor * k = as_type(C, ggml_add(C, mm_act(C, w.Wk, n1, at), w.bk), GGML_TYPE_F32); @@ -159,7 +135,7 @@ ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_ ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); ggml_tensor * att; - if (pi0_fa_enabled()) { + if (vla::flash_attn_enabled()) { // Avoids materialising the per-head score matrix; K/V stay F32 so the // numerics track the explicit path below. ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, head_dim, heads, seq), 0, 2, 1, 3)); @@ -212,8 +188,10 @@ ggml_tensor * build_gemma_layer( ggml_tensor * q_rope = rope_call(q_h); ggml_tensor * k_rope = rope_call(k_h); - if (k_out) *k_out = k_rope; - if (v_out) *v_out = v_h; + if (k_out) + *k_out = k_rope; + if (v_out) + *v_out = v_h; ggml_tensor * K_full = k_rope; ggml_tensor * V_full = v_h; @@ -222,11 +200,11 @@ ggml_tensor * build_gemma_layer( V_full = ggml_concat(ctx, cached_V, v_h, 2); } - const float scale = 1.f / std::sqrt((float) hd); + const float scale = 1.f/std::sqrt((float) hd); ggml_tensor * Q = ggml_cont(ctx, ggml_permute(ctx, q_rope, 0, 2, 1, 3)); ggml_tensor * K = ggml_cont(ctx, ggml_permute(ctx, K_full, 0, 2, 1, 3)); ggml_tensor * att_pre; - if (pi0_fa_enabled()) { + if (vla::flash_attn_enabled()) { ggml_tensor * V = ggml_cont(ctx, ggml_permute(ctx, V_full, 0, 2, 1, 3)); // ggml_flash_attn_ext asserts an F16 mask. The mask holds only 0 and // -inf, both exactly representable in F16, so the cast is lossless. @@ -271,7 +249,10 @@ ggml_tensor * build_embed_suffix(ggml_context * ctx, const Pi0ModelArch & m, bool load_config(const gguf_reader & g, Config & cfg) { auto need = [&](const char * k) { - if (!g.has(k)) { std::fprintf(stderr, "vla(pi0): gguf missing key %s\n", k); return false; } + if (!g.has(k)) { + std::fprintf(stderr, "vla(pi0): gguf missing key %s\n", k); + return false; + } return true; }; for (const char * k : {"pi0.hidden", "pi0.intermediate", "pi0.n_q_heads", "pi0.n_kv_heads", @@ -279,7 +260,8 @@ bool load_config(const gguf_reader & g, Config & cfg) { "pi0.chunk_size", "pi0.num_steps", "pi0.max_state_dim", "pi0.max_action_dim", "pi0.real_state_dim", "pi0.real_action_dim", "pi0.tokenizer_max_length", "pi0.min_period", "pi0.max_period"}) { - if (!need(k)) return false; + if (!need(k)) + return false; } cfg = Config{}; cfg.hidden = g.u32("pi0.hidden"); @@ -303,7 +285,7 @@ bool load_config(const gguf_reader & g, Config & cfg) { cfg.n_state = 1; cfg.n_img = 256; cfg.q_full_dim = cfg.n_q_heads * cfg.head_dim; - cfg.kv_full_dim = cfg.n_kv_heads * cfg.head_dim; + cfg.kv_full_dim = cfg.n_kv_heads*cfg.head_dim; cfg.self_attn_every_n = 0; cfg.rms_eps = g.has("pi0.rms_norm_eps") ? g.f32("pi0.rms_norm_eps") : 1e-6f; cfg.norm_eps = g.has("pi0.norm_eps") ? g.f32("pi0.norm_eps") : 1e-8f; @@ -323,10 +305,16 @@ bool load_stats(gguf_reader & g, Pi0ModelArch & m) { m.action_std .assign(cfg.real_action_dim, 1.f); auto read1d = [&](const char * name, std::vector & dst) { 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 (!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; + } const std::vector identity = dst; - if (!g.read_raw(name, dst.data(), dst.size() * sizeof(float))) { + 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); @@ -342,14 +330,18 @@ bool load_stats(gguf_reader & g, Pi0ModelArch & m) { } Pi0ModelArch::~Pi0ModelArch() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr pi0_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path) { + const std::string& config_path, + const Options& opts) { (void) config_path; if (!ends_with(ckpt_path, ".gguf")) { @@ -362,16 +354,18 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, auto m = std::make_unique(); m->ckpt_path_ = ckpt_path; - m->matmul_type = std::getenv("VLA_PI0_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); - if (!m->io.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()); return nullptr; } - if (!load_config(g, m->cfg)) return nullptr; + if (!load_config(g, m->cfg)) + return nullptr; const Config & cfg = m->cfg; std::printf("vla(pi0): hidden=%lld inter=%lld heads=%lldq/%lldkv x%lld n_layers=%lld " "expert_h=%lld expert_inter=%lld chunk=%lld steps=%d real_state=%lld real_action=%lld " @@ -385,11 +379,13 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, m->n_threads = default_cpu_threads(); { const Backend b = backend_init("vla(pi0)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; // BF16 activations need BF16-resident weights and the CUDA BF16 GEMM path. - if (std::getenv("VLA_PI0_BF16_ACT")) { + if (opts.act_dtype.value_or(GGML_TYPE_F32) == GGML_TYPE_BF16) { if (b.is_cuda && m->matmul_type == GGML_TYPE_BF16) { m->act_type = GGML_TYPE_BF16; cuda_register_bf16_ops(); // installs the in-tree BF16 CUDA kernels @@ -407,8 +403,9 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, vu("pi0.vit_hidden", m->vit_hidden); vu("pi0.vit_layers", m->vit_layers); vu("pi0.vit_heads", m->vit_heads); vu("pi0.image_size", m->vit_image_size); vu("pi0.patch_size", m->vit_patch_size); vu("pi0.n_img_tokens", m->vit_n_tokens); - if (g.has("pi0.vit_ln_eps")) m->vit_ln_eps = g.f32("pi0.vit_ln_eps"); - const int64_t grid = m->vit_image_size / m->vit_patch_size; + if (g.has("pi0.vit_ln_eps")) + m->vit_ln_eps = g.f32("pi0.vit_ln_eps"); + const int64_t grid = m->vit_image_size/m->vit_patch_size; if (grid * grid != m->vit_n_tokens || m->vit_n_tokens != cfg.n_img) { std::fprintf(stderr, "vla(pi0): vit geometry mismatch (grid^2=%lld n_img_tokens=%lld cfg.n_img=%lld)\n", (long long) (grid * grid), (long long) m->vit_n_tokens, (long long) cfg.n_img); @@ -417,103 +414,36 @@ std::unique_ptr pi0_create(const std::string& mmproj_path, } { - ggml_init_params wp = { (size_t) 16 * 1024 * 1024, nullptr, true }; + ggml_init_params wp = { (size_t) 16*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(pi0): ggml_init(ctx_weights) failed\n"); return nullptr; } + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(pi0): ggml_init(ctx_weights) failed\n"); + return nullptr; + } } - 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); 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); - return t; - }; + WeightLoader L("pi0", g, m->ctx_weights, m->matmul_type); - 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); 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); missing = true; return nullptr; } - return mk(name, GGML_TYPE_F32, GGML_MAX_DIMS, gt->ne); - }; + m->vit.declare(L, "vit", m->vit_layers); + m->mm_proj_w = L.gemm ("mm.proj.weight"); + m->mm_proj_b = L.opt_f32("mm.proj.bias"); - auto load_layer = [&](const char * tower, int i, GemmaLayerW & lw) -> bool { - char b[256]; - auto suf = [&](const char * s) { std::snprintf(b, sizeof(b), "%s.blk.%d.%s", tower, i, s); return b; }; - lw.ln_in = mk_f32(suf("attn_norm.weight")); - lw.Wq = mk_mm (suf("attn_q.weight")); - lw.Wk = mk_mm (suf("attn_k.weight")); - lw.Wv = mk_mm (suf("attn_v.weight")); - lw.Wo = mk_mm (suf("attn_o.weight")); - lw.ln_post = mk_f32(suf("ffn_norm.weight")); - lw.Wgate = mk_mm (suf("ffn_gate.weight")); - lw.Wup = mk_mm (suf("ffn_up.weight")); - lw.Wdown = mk_mm (suf("ffn_down.weight")); - return lw.ln_in && lw.Wq && lw.Wk && lw.Wv && lw.Wo && lw.ln_post && lw.Wgate && lw.Wup && lw.Wdown; - }; + m->pl.declare(L, "vlm", cfg.n_layers, false); + m->ex.declare(L, "aex", cfg.n_layers, true); - // Vision tower weights (SigLIP-So400m + PaliGemma projector), bundled in the ckpt GGUF. - m->vit_patch_w = mk_f32("vit.patch_embd.weight"); - m->vit_patch_b = mk_f32("vit.patch_embd.bias"); - m->vit_pos = mk_f32("vit.pos_embd"); - m->vit_post_ln_w = mk_f32("vit.post_ln.weight"); - m->vit_post_ln_b = mk_f32("vit.post_ln.bias"); - m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers; ++i) { - char p[64]; - auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return (const char *) p; }; - auto & w = m->vit[i]; - w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); - } - m->mm_proj_w = mk_mm("mm.proj.weight"); - m->mm_proj_b = g.meta("mm.proj.bias") ? mk_f32("mm.proj.bias") : nullptr; // PaliGemma projector bias (optional) - - m->pl_layers.resize(cfg.n_layers); - m->ex_layers.resize(cfg.n_layers); - for (int64_t i = 0; i < cfg.n_layers; ++i) { - if (!load_layer("vlm", (int) i, m->pl_layers[i])) return nullptr; - if (!load_layer("aex", (int) i, m->ex_layers[i])) return nullptr; - } - m->ex_final_norm = mk_f32("aex.output_norm.weight"); - m->W_sp = mk_f32("state_proj.weight"); m->b_sp = mk_f32("state_proj.bias"); - m->W_ain = mk_f32("action_in_proj.weight"); m->b_ain = mk_f32("action_in_proj.bias"); - 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) { - std::fprintf(stderr, "vla(pi0): failed to wire projection / norm tensors\n"); return nullptr; - } + m->W_sp = L.f32("state_proj.weight"); m->b_sp = L.f32("state_proj.bias"); + m->W_ain = L.f32("action_in_proj.weight"); m->b_ain = L.f32("action_in_proj.bias"); + m->W_at1 = L.f32("action_time_mlp_in.weight"); m->b_at1 = L.f32("action_time_mlp_in.bias"); + m->W_at2 = L.f32("action_time_mlp_out.weight"); m->b_at2 = L.f32("action_time_mlp_out.bias"); + m->W_aout = L.f32("action_out_proj.weight"); m->b_aout = L.f32("action_out_proj.bias"); + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(pi0): ggml_backend_alloc_ctx_tensors failed (out of memory?)\n"); return nullptr; } - for (ggml_tensor * t : weights) { - std::vector bytes = g.read_convert(t->name, t->type, is_gemma_norm(t->name)); - if (bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(pi0): upload size mismatch for %s (%zu vs %zu)\n", - t->name, bytes.size(), ggml_nbytes(t)); - return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); - } std::printf("vla(pi0): resident weights = %.2f GiB\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0)); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0)); - if (!load_stats(g, *m)) return nullptr; + if (!load_stats(g, *m)) + return nullptr; std::printf("vla(pi0): model loaded (n_threads=%d)\n", m->n_threads); return m; } @@ -527,44 +457,45 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { const int64_t hidden_pl = cfg.hidden; const int64_t hidden_ex = cfg.expert_h; const int64_t chunk = cfg.n_suffix; - const int64_t n_suf = 1 + chunk; + const int64_t n_suf = 1+chunk; const int64_t n_layers = cfg.n_layers; const int64_t max_sd = cfg.max_state_dim; const int64_t max_ad = cfg.max_action_dim; const int num_steps = cfg.num_steps; - const float dt = -1.0f / (float) num_steps; + const float dt = -1.0f/(float) num_steps; const float rope_base = cfg.rope_freq_base; std::vector img_emb_host; int64_t n_img_tokens = 0; if (in.precomputed_img_emb) { - n_img_tokens = (int64_t) in.n_img_views * cfg.n_img; + n_img_tokens = (int64_t) in.n_img_views*cfg.n_img; img_emb_host.assign(in.precomputed_img_emb, - in.precomputed_img_emb + (size_t) n_img_tokens * hidden_pl); + in.precomputed_img_emb+(size_t) n_img_tokens * hidden_pl); } else { if (in.n_images < 1 || !in.images) { std::fprintf(stderr, "vla(pi0): predict: no images and no precomputed_img_emb\n"); return {}; } - const int64_t K = vit_n_tokens, H = hidden_pl, grid = vit_image_size / vit_patch_size; - n_img_tokens = (int64_t) in.n_images * K; - img_emb_host.assign((size_t) in.n_images * K * H, 0.0f); + const int64_t K = vit_n_tokens, H = hidden_pl, grid = vit_image_size/vit_patch_size; + n_img_tokens = (int64_t) in.n_images*K; + img_emb_host.assign((size_t) in.n_images*K * H, 0.0f); - ggml_context * VC = vision_scratch.reset((size_t) 128 * 1024 * 1024); + 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); + ggml_tensor * conv = ggml_conv_2d(VC, vit.patch_w, t_px, (int) vit_patch_size, (int) vit_patch_size, 0, 0, 1, 1); ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, grid * grid, vit_hidden))); // patch embed (conv_2d) stays F32; the tower runs in the activation dtype - ggml_tensor * h = as_type(VC, ggml_add(VC, ggml_add(VC, patches, vit_patch_b), vit_pos), act_type); - for (int64_t i = 0; i < vit_layers; ++i) - h = build_siglip_layer(VC, vit[i], h, K, vit_heads, vit_hidden / vit_heads, vit_hidden, vit_ln_eps, act_type); - h = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, h, vit_ln_eps), vit_post_ln_w), vit_post_ln_b); + ggml_tensor * h = as_type(VC, ggml_add(VC, ggml_add(VC, patches, vit.patch_b), vit.pos), act_type); + for (int64_t i=0; ine[0])), GGML_TYPE_F32); + ggml_tensor * vit_emb = as_type(VC, ggml_scale(VC, proj, 1.0f/std::sqrt((float) proj->ne[0])), GGML_TYPE_F32); ggml_set_output(vit_emb); ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); @@ -576,16 +507,16 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { } const auto tv0 = clk::now(); std::vector chw; - for (int v = 0; v < in.n_images; ++v) { + for (int v=0; v(clk::now() - tv0).count(); + stats.ms_vision = std::chrono::duration(clk::now()-tv0).count(); } if (in.n_lang < 1 || !in.lang_tokens) { @@ -593,10 +524,10 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { return {}; } const int64_t n_lang = in.n_lang; - const int64_t n_prefix = n_img_tokens + n_lang; - const int64_t n_total = n_prefix + n_suf; + const int64_t n_prefix = n_img_tokens+n_lang; + const int64_t n_total = n_prefix+n_suf; - std::vector lang_ids(in.lang_tokens, in.lang_tokens + n_lang); + std::vector lang_ids(in.lang_tokens, in.lang_tokens+n_lang); std::vector lang_rows((size_t) n_lang * hidden_pl); { if (!io.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; @@ -604,7 +535,7 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { // 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, + 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); @@ -614,7 +545,7 @@ std::vector Pi0ModelArch::predict(const Inputs& in) { ggml_tensor * t_suffix_pos= ggml_new_tensor_1d(C, GGML_TYPE_I32, n_suf); ggml_set_input(t_suffix_pos); ggml_tensor * t_full_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, n_total, n_suf); ggml_set_input(t_full_mask); std::vector t_time(num_steps); - for (int s = 0; s < num_steps; ++s) { + for (int s=0; s Pi0ModelArch::predict(const Inputs& in) { std::vector cK(n_layers), cV(n_layers); { ggml_tensor * h = prefix_embs; - for (int64_t i = 0; i < n_layers; ++i) { - h = build_gemma_layer(C, pl_layers[i], h, t_prefix_pos, cfg, n_prefix, rope_base, + for (int64_t i=0; i Pi0ModelArch::predict(const Inputs& in) { // do not accumulate in 8 mantissa bits. ggml_tensor * x_t = t_x0; std::vector v_steps(num_steps); - for (int step = 0; step < num_steps; ++step) { + for (int step=0; step Pi0ModelArch::predict(const Inputs& in) { 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)); { - std::vector pp(n_prefix); for (int64_t i = 0; i < n_prefix; ++i) pp[i] = (int32_t) i; + std::vector pp(n_prefix); for (int64_t i=0; i sp(n_suf); for (int64_t i = 0; i < n_suf; ++i) sp[i] = (int32_t) (n_prefix + i); + std::vector sp(n_suf); for (int64_t i=0; i sh(max_sd, 0.f); - for (int64_t i = 0; i < max_sd; ++i) sh[i] = in.state ? in.state[i] : 0.f; - for (int64_t i = 0; i < cfg.real_state_dim && i < max_sd; ++i) - sh[i] = (sh[i] - state_mean[i]) / (state_std[i] + cfg.norm_eps); + for (int64_t i=0; i x0h((size_t) max_ad * chunk); - if (in.noise) std::memcpy(x0h.data(), in.noise, x0h.size() * sizeof(float)); - else { std::normal_distribution nd(0.f, 1.f); for (auto & v : x0h) v = nd(rng); } + if (in.noise) + std::memcpy(x0h.data(), in.noise, x0h.size()*sizeof(float)); + else { + std::normal_distribution nd(0.f, 1.f); + for (auto & v : x0h) + v = nd(rng); + } ggml_backend_tensor_set(t_x0, x0h.data(), 0, ggml_nbytes(t_x0)); } { std::vector mk((size_t) n_total * n_suf); - for (int64_t i = 0; i < n_suf; ++i) - for (int64_t j = 0; j < n_total; ++j) { + for (int64_t i=0; i tv = sinusoidal_time_emb(timestep, hidden_ex, cfg.min_period, cfg.max_period); std::vector tile((size_t) hidden_ex * chunk); - for (int64_t c = 0; c < chunk; ++c) std::memcpy(tile.data() + c * hidden_ex, tv.data(), hidden_ex * sizeof(float)); + for (int64_t c=0; c(clk::now() - ti0).count(); + 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); return {}; } std::vector out((size_t) chunk * max_ad); - ggml_backend_tensor_get(x_final, out.data(), 0, out.size() * sizeof(float)); - 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) - row[j] = j < cfg.real_action_dim ? row[j] * (action_std[j] + cfg.norm_eps) + action_mean[j] : 0.0f; + ggml_backend_tensor_get(x_final, out.data(), 0, out.size()*sizeof(float)); + for (int64_t t=0; t(clk::now() - t0).count(); + stats.ms_total = std::chrono::duration(clk::now()-t0).count(); return out; } diff --git a/src/models/pi05.cpp b/src/models/pi05.cpp index 7e8c877..f16bc56 100644 --- a/src/models/pi05.cpp +++ b/src/models/pi05.cpp @@ -13,6 +13,9 @@ // limitations under the License. #include "arch.h" +#include "modules/gemma_expert.h" +#include "modules/siglip_vit.h" +#include "options.h" #include "model.h" #include "ggml.h" @@ -21,10 +24,11 @@ #include "ggml-alloc.h" #include "backend.h" #include "gguf.h" -#include "models/gguf_reader.h" -#include "models/scratch_ctx.h" +#include "gguf_reader.h" +#include "scratch_ctx.h" #include "models/dit_common.h" -#include "models/vision_common.h" +#include "modules/preprocess.h" +#include "env_flag.h" #include #include @@ -44,18 +48,6 @@ namespace vla { namespace { -struct VlmLayerW { - ggml_tensor * ln_in = nullptr; - ggml_tensor * Wq = nullptr; - ggml_tensor * Wk = nullptr; - ggml_tensor * Wv = nullptr; - ggml_tensor * Wo = nullptr; - ggml_tensor * ln_post = nullptr; - ggml_tensor * Wgate = nullptr; - ggml_tensor * Wup = nullptr; - ggml_tensor * Wdown = nullptr; -}; - struct ExpertLayerW { ggml_tensor * ada_in_w = nullptr; ggml_tensor * ada_in_b = nullptr; @@ -70,22 +62,10 @@ struct ExpertLayerW { ggml_tensor * Wdown = nullptr; }; -// 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; }; - 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; -} -bool starts_with(const std::string & s, const char * pfx) { - const size_t n = std::strlen(pfx); - return s.size() >= n && s.compare(0, n, pfx) == 0; + return s.size() >= n && s.compare(s.size()-n, n, sfx) == 0; } - -bool is_gemma_norm_pi05(const std::string & name) { - return starts_with(name, "vlm.") && name.find("norm.weight") != std::string::npos; -} - } struct Pi05ModelArch : public ModelArchBase { @@ -101,7 +81,9 @@ struct Pi05ModelArch : public ModelArchBase { 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; } + 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; @@ -119,12 +101,10 @@ struct Pi05ModelArch : public ModelArchBase { int64_t vit_hidden = 1152, vit_layers = 27, vit_heads = 16; int64_t vit_image_size = 224, vit_patch_size = 14, vit_n_tokens = 256; float vit_ln_eps = 1e-6f; - ggml_tensor * vit_patch_w = nullptr, * vit_patch_b = nullptr, * vit_pos = nullptr; - ggml_tensor * vit_post_ln_w = nullptr, * vit_post_ln_b = nullptr; - std::vector vit; + SigLipTower vit; ggml_tensor * mm_proj_w = nullptr, * mm_proj_b = nullptr; - std::vector pl_layers; + GemmaStack pl; std::vector ex_layers; ggml_tensor * ex_final_w = nullptr; @@ -148,9 +128,9 @@ namespace { // One pre-norm SigLIP encoder block, identical to gr00tn1d5's in-tree tower // (the PaliGemma vision tower is the same SigLIP-So400m/14). Bidirectional // attention (nullptr mask), F32 score accumulation, tanh GELU FFN. -ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, +ggml_tensor * build_siglip_layer(ggml_context * C, const EncBlockW & 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); + const float scale = 1.0f/std::sqrt((float) head_dim); ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); @@ -170,7 +150,7 @@ 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). ggml_tensor * build_vlm_layer( - ggml_context * ctx, const VlmLayerW & w, + ggml_context * ctx, const GemmaLayerW & w, ggml_tensor * x_in, ggml_tensor * positions, const Config & cfg, int64_t seq, float rope_base, ggml_tensor ** k_out, ggml_tensor ** v_out) { @@ -197,8 +177,10 @@ ggml_tensor * build_vlm_layer( ggml_tensor * q_rope = rope_call(q_h); ggml_tensor * k_rope = rope_call(k_h); - if (k_out) *k_out = k_rope; - if (v_out) *v_out = v_h; + if (k_out) + *k_out = k_rope; + if (v_out) + *v_out = v_h; ggml_tensor * Q = ggml_cont(ctx, ggml_permute(ctx, q_rope, 0, 2, 1, 3)); ggml_tensor * K = ggml_cont(ctx, ggml_permute(ctx, k_rope, 0, 2, 1, 3)); @@ -206,7 +188,7 @@ ggml_tensor * build_vlm_layer( ggml_tensor * kq = ggml_mul_mat(ctx, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - const float scale = 1.f / std::sqrt((float) hd); + const float scale = 1.f/std::sqrt((float) hd); ggml_tensor * attn = ggml_soft_max_ext(ctx, kq, nullptr, scale, 0.f); ggml_tensor * kqv = ggml_mul_mat(ctx, V, attn); @@ -230,12 +212,13 @@ ggml_tensor * build_adarms( ggml_tensor * mod = ggml_add(ctx, ggml_mul_mat(ctx, dense_w, cond), dense_b); ggml_tensor * scale = ggml_view_1d(ctx, mod, h, 0); ggml_tensor * shift = ggml_view_1d(ctx, mod, h, (size_t) h * sizeof(float)); - ggml_tensor * gate = ggml_view_1d(ctx, mod, h, (size_t) 2 * h * sizeof(float)); + ggml_tensor * gate = ggml_view_1d(ctx, mod, h, (size_t) 2*h * sizeof(float)); ggml_tensor * normed = ggml_rms_norm(ctx, x, eps); ggml_tensor * out = ggml_add(ctx, ggml_add(ctx, normed, ggml_mul(ctx, normed, scale)), shift); - if (gate_out) *gate_out = gate; + if (gate_out) + *gate_out = gate; return out; } @@ -278,7 +261,7 @@ ggml_tensor * build_expert_layer( ggml_tensor * kq = ggml_mul_mat(ctx, K, Q); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - const float scale = 1.f / std::sqrt((float) hd); + const float scale = 1.f/std::sqrt((float) hd); ggml_tensor * attn = ggml_soft_max_ext(ctx, kq, nullptr, scale, 0.f); ggml_tensor * kqv = ggml_mul_mat(ctx, V, attn); @@ -300,7 +283,10 @@ ggml_tensor * build_expert_layer( bool load_config(const gguf_reader & g, Config & cfg) { auto need = [&](const char * k) { - if (!g.has(k)) { std::fprintf(stderr, "vla(pi05): gguf missing key %s\n", k); return false; } + if (!g.has(k)) { + std::fprintf(stderr, "vla(pi05): gguf missing key %s\n", k); + return false; + } return true; }; for (const char * k : {"pi05.hidden", "pi05.intermediate", "pi05.n_q_heads", "pi05.n_kv_heads", @@ -308,7 +294,8 @@ bool load_config(const gguf_reader & g, Config & cfg) { "pi05.chunk_size", "pi05.num_steps", "pi05.max_state_dim", "pi05.max_action_dim", "pi05.real_state_dim", "pi05.real_action_dim", "pi05.tokenizer_max_length", "pi05.min_period", "pi05.max_period"}) { - if (!need(k)) return false; + if (!need(k)) + return false; } cfg = Config{}; cfg.hidden = g.u32("pi05.hidden"); @@ -332,7 +319,7 @@ bool load_config(const gguf_reader & g, Config & cfg) { cfg.n_state = 0; cfg.n_img = 256; cfg.q_full_dim = cfg.n_q_heads * cfg.head_dim; - cfg.kv_full_dim = cfg.n_kv_heads * cfg.head_dim; + cfg.kv_full_dim = cfg.n_kv_heads*cfg.head_dim; cfg.self_attn_every_n = 0; cfg.rms_eps = g.has("pi05.rms_norm_eps") ? g.f32("pi05.rms_norm_eps") : 1e-6f; cfg.norm_eps = g.has("pi05.norm_eps") ? g.f32("pi05.norm_eps") : 1e-8f; @@ -352,10 +339,16 @@ bool load_stats(gguf_reader & g, Pi05ModelArch & m) { m.action_std .assign(cfg.real_action_dim, 1.f); auto read1d = [&](const char * name, std::vector & dst) { 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 (!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; + } const std::vector identity = dst; - if (!g.read_raw(name, dst.data(), dst.size() * sizeof(float))) { + 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); @@ -378,14 +371,18 @@ bool load_stats(gguf_reader & g, Pi05ModelArch & m) { } Pi05ModelArch::~Pi05ModelArch() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr pi05_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path) { + const std::string& config_path, + const Options& opts) { (void) config_path; if (!ends_with(ckpt_path, ".gguf")) { @@ -397,16 +394,18 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, auto m = std::make_unique(); m->ckpt_path_ = ckpt_path; - m->matmul_type = std::getenv("VLA_PI05_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); - if (!m->io.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()); return nullptr; } - if (!load_config(g, m->cfg)) return nullptr; + if (!load_config(g, m->cfg)) + return nullptr; const Config & cfg = m->cfg; m->adarms_cond_dim = g.has("pi05.adarms_cond_dim") ? g.u32("pi05.adarms_cond_dim") : cfg.expert_h; m->quantile_norm = g.has("pi05.norm_mode") && g.str("pi05.norm_mode") == "quantiles"; @@ -423,7 +422,9 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, m->n_threads = default_cpu_threads(); { const Backend b = backend_init("vla(pi05)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; } @@ -434,8 +435,9 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, vu("pi05.vit_hidden", m->vit_hidden); vu("pi05.vit_layers", m->vit_layers); vu("pi05.vit_heads", m->vit_heads); vu("pi05.image_size", m->vit_image_size); vu("pi05.patch_size", m->vit_patch_size); vu("pi05.n_img_tokens", m->vit_n_tokens); - if (g.has("pi05.vit_ln_eps")) m->vit_ln_eps = g.f32("pi05.vit_ln_eps"); - const int64_t grid = m->vit_image_size / m->vit_patch_size; + if (g.has("pi05.vit_ln_eps")) + m->vit_ln_eps = g.f32("pi05.vit_ln_eps"); + const int64_t grid = m->vit_image_size/m->vit_patch_size; if (grid * grid != m->vit_n_tokens || m->vit_n_tokens != cfg.n_img) { std::fprintf(stderr, "vla(pi05): vit geometry mismatch (grid^2=%lld n_img_tokens=%lld cfg.n_img=%lld)\n", (long long) (grid * grid), (long long) m->vit_n_tokens, (long long) cfg.n_img); @@ -444,119 +446,52 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, } { - ggml_init_params wp = { (size_t) 16 * 1024 * 1024, nullptr, true }; + ggml_init_params wp = { (size_t) 16*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(pi05): ggml_init(ctx_weights) failed\n"); return nullptr; } + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(pi05): ggml_init(ctx_weights) failed\n"); + return nullptr; + } } - 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); 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); - return t; - }; - 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); 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); missing = true; return nullptr; } - return mk(name, GGML_TYPE_F32, GGML_MAX_DIMS, gt->ne); - }; + WeightLoader L("pi05", g, m->ctx_weights, m->matmul_type); - auto load_vlm = [&](int i, VlmLayerW & lw) -> bool { - char b[256]; - auto suf = [&](const char * s) { std::snprintf(b, sizeof(b), "vlm.blk.%d.%s", i, s); return b; }; - lw.ln_in = mk_f32(suf("attn_norm.weight")); - lw.Wq = mk_mm (suf("attn_q.weight")); - lw.Wk = mk_mm (suf("attn_k.weight")); - lw.Wv = mk_mm (suf("attn_v.weight")); - lw.Wo = mk_mm (suf("attn_o.weight")); - lw.ln_post = mk_f32(suf("ffn_norm.weight")); - lw.Wgate = mk_mm (suf("ffn_gate.weight")); - lw.Wup = mk_mm (suf("ffn_up.weight")); - lw.Wdown = mk_mm (suf("ffn_down.weight")); - return lw.ln_in && lw.Wq && lw.Wk && lw.Wv && lw.Wo && lw.ln_post && lw.Wgate && lw.Wup && lw.Wdown; - }; - auto load_expert = [&](int i, ExpertLayerW & lw) -> bool { - char b[256]; - auto suf = [&](const char * s) { std::snprintf(b, sizeof(b), "aex.blk.%d.%s", i, s); return b; }; - lw.ada_in_w = mk_f32(suf("attn_norm.weight")); - lw.ada_in_b = mk_f32(suf("attn_norm.bias")); - lw.Wq = mk_mm (suf("attn_q.weight")); - lw.Wk = mk_mm (suf("attn_k.weight")); - lw.Wv = mk_mm (suf("attn_v.weight")); - lw.Wo = mk_mm (suf("attn_o.weight")); - lw.ada_post_w = mk_f32(suf("ffn_norm.weight")); - lw.ada_post_b = mk_f32(suf("ffn_norm.bias")); - lw.Wgate = mk_mm (suf("ffn_gate.weight")); - lw.Wup = mk_mm (suf("ffn_up.weight")); - lw.Wdown = mk_mm (suf("ffn_down.weight")); - return lw.ada_in_w && lw.ada_in_b && lw.Wq && lw.Wk && lw.Wv && lw.Wo && - lw.ada_post_w && lw.ada_post_b && lw.Wgate && lw.Wup && lw.Wdown; - }; + m->vit.declare(L, "vit", m->vit_layers); + m->mm_proj_w = L.gemm ("mm.proj.weight"); + m->mm_proj_b = L.opt_f32("mm.proj.bias"); - // Vision tower weights (SigLIP-So400m + PaliGemma projector), bundled in the ckpt GGUF. - m->vit_patch_w = mk_f32("vit.patch_embd.weight"); - m->vit_patch_b = mk_f32("vit.patch_embd.bias"); - m->vit_pos = mk_f32("vit.pos_embd"); - m->vit_post_ln_w = mk_f32("vit.post_ln.weight"); - m->vit_post_ln_b = mk_f32("vit.post_ln.bias"); - m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers; ++i) { - char p[64]; - auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return (const char *) p; }; - auto & w = m->vit[i]; - w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); - } - m->mm_proj_w = mk_mm("mm.proj.weight"); - m->mm_proj_b = g.meta("mm.proj.bias") ? mk_f32("mm.proj.bias") : nullptr; // PaliGemma projector bias (optional) + m->pl.declare(L, "vlm", cfg.n_layers, false); - m->pl_layers.resize(cfg.n_layers); m->ex_layers.resize(cfg.n_layers); - for (int64_t i = 0; i < cfg.n_layers; ++i) { - if (!load_vlm ((int) i, m->pl_layers[i])) return nullptr; - if (!load_expert((int) i, m->ex_layers[i])) return nullptr; - } - m->ex_final_w = mk_f32("aex.output_norm.weight"); - m->ex_final_b = mk_f32("aex.output_norm.bias"); - m->W_ain = mk_f32("action_in_proj.weight"); m->b_ain = mk_f32("action_in_proj.bias"); - 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) { - std::fprintf(stderr, "vla(pi05): failed to wire projection / norm tensors\n"); return nullptr; + for (int64_t i=0; iex_layers[i]; + w.ada_in_w = L.f32 ("aex.blk.%lld.attn_norm.weight", (long long)i); + w.ada_in_b = L.f32 ("aex.blk.%lld.attn_norm.bias", (long long)i); + w.Wq = L.gemm("aex.blk.%lld.attn_q.weight", (long long)i); + w.Wk = L.gemm("aex.blk.%lld.attn_k.weight", (long long)i); + w.Wv = L.gemm("aex.blk.%lld.attn_v.weight", (long long)i); + w.Wo = L.gemm("aex.blk.%lld.attn_o.weight", (long long)i); + w.ada_post_w = L.f32 ("aex.blk.%lld.ffn_norm.weight", (long long)i); + w.ada_post_b = L.f32 ("aex.blk.%lld.ffn_norm.bias", (long long)i); + w.Wgate = L.gemm("aex.blk.%lld.ffn_gate.weight", (long long)i); + w.Wup = L.gemm("aex.blk.%lld.ffn_up.weight", (long long)i); + w.Wdown = L.gemm("aex.blk.%lld.ffn_down.weight", (long long)i); } - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(pi05): ggml_backend_alloc_ctx_tensors failed (out of memory?)\n"); return nullptr; } - for (ggml_tensor * t : weights) { - std::vector bytes = g.read_convert(t->name, t->type, is_gemma_norm_pi05(t->name)); - if (bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(pi05): upload size mismatch for %s (%zu vs %zu)\n", - t->name, bytes.size(), ggml_nbytes(t)); - return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); - } + m->ex_final_w = L.f32("aex.output_norm.weight"); + m->ex_final_b = L.f32("aex.output_norm.bias"); + m->W_ain = L.f32("action_in_proj.weight"); m->b_ain = L.f32("action_in_proj.bias"); + m->W_tin = L.f32("time_mlp_in.weight"); m->b_tin = L.f32("time_mlp_in.bias"); + m->W_tout = L.f32("time_mlp_out.weight"); m->b_tout = L.f32("time_mlp_out.bias"); + m->W_aout = L.f32("action_out_proj.weight"); m->b_aout = L.f32("action_out_proj.bias"); + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + std::printf("vla(pi05): resident weights = %.2f GiB\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0)); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0)); - if (!load_stats(g, *m)) return nullptr; + if (!load_stats(g, *m)) + return nullptr; std::printf("vla(pi05): model loaded (n_threads=%d)\n", m->n_threads); return m; } @@ -574,37 +509,38 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { const int64_t n_layers = cfg.n_layers; const int64_t max_ad = cfg.max_action_dim; const int num_steps = cfg.num_steps; - const float dt = -1.0f / (float) num_steps; + const float dt = -1.0f/(float) num_steps; const float rope_base = cfg.rope_freq_base; std::vector img_emb_host; int64_t n_img_tokens = 0; if (in.precomputed_img_emb) { - n_img_tokens = (int64_t) in.n_img_views * cfg.n_img; + n_img_tokens = (int64_t) in.n_img_views*cfg.n_img; img_emb_host.assign(in.precomputed_img_emb, - in.precomputed_img_emb + (size_t) n_img_tokens * hidden_pl); + in.precomputed_img_emb+(size_t) n_img_tokens * hidden_pl); } else { if (in.n_images < 1 || !in.images) { std::fprintf(stderr, "vla(pi05): predict: no images and no precomputed_img_emb\n"); return {}; } - const int64_t K = vit_n_tokens, H = hidden_pl, grid = vit_image_size / vit_patch_size; - n_img_tokens = (int64_t) in.n_images * K; - img_emb_host.assign((size_t) in.n_images * K * H, 0.0f); + const int64_t K = vit_n_tokens, H = hidden_pl, grid = vit_image_size/vit_patch_size; + n_img_tokens = (int64_t) in.n_images*K; + img_emb_host.assign((size_t) in.n_images*K * H, 0.0f); - ggml_context * VC = vision_scratch.reset((size_t) 128 * 1024 * 1024); + 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); + ggml_tensor * conv = ggml_conv_2d(VC, vit.patch_w, t_px, (int) vit_patch_size, (int) vit_patch_size, 0, 0, 1, 1); ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, grid * grid, vit_hidden))); - ggml_tensor * h = ggml_add(VC, ggml_add(VC, patches, vit_patch_b), vit_pos); - for (int64_t i = 0; i < vit_layers; ++i) - h = build_siglip_layer(VC, vit[i], h, K, vit_heads, vit_hidden / vit_heads, vit_hidden, vit_ln_eps); - h = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, h, vit_ln_eps), vit_post_ln_w), vit_post_ln_b); + ggml_tensor * h = ggml_add(VC, ggml_add(VC, patches, vit.patch_b), vit.pos); + for (int64_t i=0; ine[0])); + if (mm_proj_b) + proj = ggml_add(VC, proj, mm_proj_b); + ggml_tensor * vit_emb = ggml_scale(VC, proj, 1.0f/std::sqrt((float) proj->ne[0])); ggml_set_output(vit_emb); ggml_cgraph * vg = ggml_new_graph_custom(VC, 8192, false); @@ -616,22 +552,23 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { } const auto tv0 = clk::now(); std::vector chw; - for (int v = 0; v < in.n_images; ++v) { + for (int v=0; v(clk::now() - tv0).count(); + stats.ms_vision = std::chrono::duration(clk::now()-tv0).count(); // 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; + for (float & x : img_emb_host) + x *= img_scale; } if (in.n_lang < 1 || !in.lang_tokens) { @@ -639,9 +576,9 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { return {}; } const int64_t n_lang = in.n_lang; - const int64_t n_prefix = n_img_tokens + n_lang; + const int64_t n_prefix = n_img_tokens+n_lang; - std::vector lang_ids(in.lang_tokens, in.lang_tokens + n_lang); + std::vector lang_ids(in.lang_tokens, in.lang_tokens+n_lang); std::vector lang_rows((size_t) n_lang * hidden_pl); { if (!io.fetch_rows_f32("token_embd.weight", lang_ids, lang_rows.data(), hidden_pl)) return {}; @@ -649,7 +586,7 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { // 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, + 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); @@ -658,7 +595,7 @@ std::vector Pi05ModelArch::predict(const Inputs& in) { ggml_tensor * t_suffix_pos= ggml_new_tensor_1d(C, GGML_TYPE_I32, n_suf); ggml_set_input(t_suffix_pos); std::vector t_time(num_steps); - for (int s = 0; s < num_steps; ++s) { + for (int s=0; s Pi05ModelArch::predict(const Inputs& in) { std::vector cK(n_layers), cV(n_layers); { ggml_tensor * h = prefix_embs; - for (int64_t i = 0; i < n_layers; ++i) { - h = build_vlm_layer(C, pl_layers[i], h, t_prefix_pos, cfg, n_prefix, rope_base, + for (int64_t i=0; i Pi05ModelArch::predict(const Inputs& in) { 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)); { - std::vector pp(n_prefix); for (int64_t i = 0; i < n_prefix; ++i) pp[i] = (int32_t) i; + std::vector pp(n_prefix); for (int64_t i=0; i sp(n_suf); for (int64_t i = 0; i < n_suf; ++i) sp[i] = (int32_t) (n_prefix + i); + std::vector sp(n_suf); for (int64_t i=0; i x0h((size_t) max_ad * chunk); - if (in.noise) std::memcpy(x0h.data(), in.noise, x0h.size() * sizeof(float)); - else { std::normal_distribution nd(0.f, 1.f); for (auto & v : x0h) v = nd(rng); } + if (in.noise) + std::memcpy(x0h.data(), in.noise, x0h.size()*sizeof(float)); + else { + std::normal_distribution nd(0.f, 1.f); + for (auto & v : x0h) + v = nd(rng); + } ggml_backend_tensor_set(t_x0, x0h.data(), 0, ggml_nbytes(t_x0)); } - for (int s = 0; s < num_steps; ++s) { - const float timestep = 1.0f + (float) s * dt; + for (int s=0; s tv = sinusoidal_time_emb(timestep, hidden_ex, cfg.min_period, cfg.max_period); ggml_backend_tensor_set(t_time[s], tv.data(), 0, ggml_nbytes(t_time[s])); } const auto ti0 = clk::now(); const ggml_status st = ggml_backend_graph_compute(backend, gf); - stats.ms_inference = std::chrono::duration(clk::now() - ti0).count(); + 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); return {}; } std::vector out((size_t) chunk * max_ad); - ggml_backend_tensor_get(x_final, out.data(), 0, out.size() * sizeof(float)); + ggml_backend_tensor_get(x_final, out.data(), 0, out.size()*sizeof(float)); - if (!std::getenv("VLA_PI05_SKIP_UNNORM")) { - 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) { + if (!vla::env_flag("VLA_PI05_SKIP_UNNORM")) { + for (int64_t t=0; t= cfg.real_action_dim) row[j] = 0.0f; else if (quantile_norm) - row[j] = (row[j] + 1.0f) * (action_q99[j] - action_q01[j]) * 0.5f + action_q01[j]; + row[j] = (row[j]+1.0f)*(action_q99[j]-action_q01[j])*0.5f+action_q01[j]; else - row[j] = row[j] * (action_std[j] + cfg.norm_eps) + action_mean[j]; + row[j] = row[j]*(action_std[j]+cfg.norm_eps)+action_mean[j]; } } } - stats.ms_total = std::chrono::duration(clk::now() - t0).count(); + stats.ms_total = std::chrono::duration(clk::now()-t0).count(); return out; } diff --git a/src/models/smolvla.cpp b/src/models/smolvla.cpp index 39ad167..62c062c 100644 --- a/src/models/smolvla.cpp +++ b/src/models/smolvla.cpp @@ -16,10 +16,12 @@ // distilled action-expert weights and force num_steps = 1 at the denoise loops. #include "arch.h" +#include "modules/encoder.h" +#include "options.h" #include "model.h" -#include "vision_common.h" +#include "modules/preprocess.h" #include "scratch_ctx.h" -#include "dit_common.h" +#include "models/dit_common.h" #include "ggml.h" #include "ggml-backend.h" @@ -28,6 +30,7 @@ #include "backend.h" #include "nlohmann/json.hpp" +#include "env_flag.h" #include #include @@ -60,15 +63,17 @@ struct safetensors { bool open(const std::string & path) { file.open(path, std::ios::binary); - if (!file) return false; + if (!file) + return false; uint64_t header_size = 0; file.read(reinterpret_cast(&header_size), sizeof(header_size)); std::string header_str(header_size, '\0'); file.read(header_str.data(), header_size); - data_blob_start = sizeof(uint64_t) + header_size; + data_blob_start = sizeof(uint64_t)+header_size; json j = json::parse(header_str); - for (auto it = j.begin(); it != j.end(); ++it) { - if (it.key() == "__metadata__") continue; + for (auto it=j.begin(); it!=j.end(); ++it) { + if (it.key() == "__metadata__") + continue; const auto & v = it.value(); st_tensor_info info; info.dtype = v.at("dtype").get(); @@ -102,17 +107,20 @@ struct safetensors { 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; } + 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) { + 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); + 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)); + 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 { @@ -129,12 +137,12 @@ struct safetensors { return false; } const auto & info = it->second; - if ((info.off_end - info.off_begin) != expected_bytes || + if ((info.off_end-info.off_begin) != expected_bytes || info.dtype != expected_dtype) { std::fprintf(stderr, "vla: bad raw read for %s\n", name.c_str()); return false; } - file.seekg(data_blob_start + info.off_begin, std::ios::beg); + file.seekg(data_blob_start+info.off_begin, std::ios::beg); file.read(static_cast(dst), expected_bytes); return true; } @@ -167,20 +175,26 @@ struct gguf_source { } ~gguf_source() { - if (fp) std::fclose(fp); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); + if (fp) + std::fclose(fp); + if (gctx) + gguf_free(gctx); + if (meta_ctx) + ggml_free(meta_ctx); } static bool shape_matches(const ggml_tensor * t, const std::vector & pt_shape) { const int nd_used = std::max(1, (int) pt_shape.size()); - if (nd_used > GGML_MAX_DIMS) return false; - for (int d = 0; d < (int) pt_shape.size(); ++d) { - const int64_t expected = pt_shape[pt_shape.size() - 1 - d]; - if (t->ne[d] != expected) return false; + if (nd_used > GGML_MAX_DIMS) + return false; + for (int d=0; d<(int) pt_shape.size(); ++d) { + const int64_t expected = pt_shape[pt_shape.size()-1-d]; + if (t->ne[d] != expected) + return false; } - for (int d = (int) pt_shape.size(); d < GGML_MAX_DIMS; ++d) { - if (t->ne[d] != 1) return false; + for (int d=(int) pt_shape.size(); dne[d] != 1) + return false; } return true; } @@ -197,17 +211,19 @@ struct gguf_source { return false; } 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 offset = data_off+gguf_get_tensor_offset(gctx, id); const size_t bytes = gguf_get_tensor_size(gctx, id); if (fseeko(fp, (off_t) offset, SEEK_SET) != 0) { std::fprintf(stderr, "vla: fseek failed for %s\n", name.c_str()); return false; } if (t->type == GGML_TYPE_F32) { - if (std::fread(dst, 1, bytes, fp) != bytes) return false; + if (std::fread(dst, 1, bytes, fp) != bytes) + return false; } else if (t->type == GGML_TYPE_BF16) { - std::vector tmp(bytes / sizeof(ggml_bf16_t)); - if (std::fread(tmp.data(), 1, bytes, fp) != bytes) return false; + std::vector tmp(bytes/sizeof(ggml_bf16_t)); + if (std::fread(tmp.data(), 1, bytes, fp) != bytes) + return false; ggml_bf16_to_fp32_row(tmp.data(), dst, tmp.size()); } else { std::fprintf(stderr, "vla: gguf unsupported dtype %d for %s\n", @@ -232,8 +248,9 @@ struct gguf_source { std::fprintf(stderr, "vla: gguf bad raw read for %s\n", name.c_str()); return false; } - const size_t offset = data_off + gguf_get_tensor_offset(gctx, id); - if (fseeko(fp, (off_t) offset, SEEK_SET) != 0) return false; + const size_t offset = data_off+gguf_get_tensor_offset(gctx, id); + if (fseeko(fp, (off_t) offset, SEEK_SET) != 0) + return false; return std::fread(dst, 1, bytes, fp) == bytes; } @@ -243,11 +260,21 @@ struct gguf_source { bool has_key(const char * key) const { return find_key(key) >= 0; } - uint32_t get_u32(const char * key) const { return gguf_get_val_u32(gctx, find_key(key)); } - int32_t get_i32(const char * key) const { return gguf_get_val_i32(gctx, find_key(key)); } - float get_f32(const char * key) const { return gguf_get_val_f32(gctx, find_key(key)); } - double get_f64(const char * key) const { return gguf_get_val_f64(gctx, find_key(key)); } - std::string get_str(const char * key) const { return gguf_get_val_str(gctx, find_key(key)); } + uint32_t get_u32(const char * key) const { + return gguf_get_val_u32(gctx, find_key(key)); + } + int32_t get_i32(const char * key) const { + return gguf_get_val_i32(gctx, find_key(key)); + } + float get_f32(const char * key) const { + return gguf_get_val_f32(gctx, find_key(key)); + } + double get_f64(const char * key) const { + return gguf_get_val_f64(gctx, find_key(key)); + } + std::string get_str(const char * key) const { + return gguf_get_val_str(gctx, find_key(key)); + } bool has_tensor(const char * name) const { return ggml_get_tensor(meta_ctx, name) != nullptr; @@ -280,7 +307,6 @@ struct ExpertLayerW { }; // SigLIP-B/16 vision block weights (SmolVLM2 tower, built in-tree). -struct SigLipLayerW { ggml_tensor *ln1w,*ln1b,*ln2w,*ln2b,*Wq,*bq,*Wk,*bk,*Wv,*bv,*Wo,*bo,*Wfc1,*bfc1,*Wfc2,*bfc2; }; } @@ -296,7 +322,7 @@ struct SmolVLAModelArch : public ModelArchBase { float vit_ln_eps = 1e-6f; ggml_tensor * vit_patch_w = nullptr, * vit_patch_b = nullptr, * vit_pos = nullptr; ggml_tensor * vit_post_ln_w = nullptr, * vit_post_ln_b = nullptr, * mm_fc = nullptr; - std::vector vit; + std::vector vit; ggml_backend_t backend = nullptr; ggml_backend_buffer_t weight_buf = nullptr; @@ -362,16 +388,12 @@ namespace { // reinterpreting it as F16), and that measured 92/100 on libero_object against // 96/100 for explicit attention. evo1 showed the same ~4-5 pp drop, so the // default stays on the accuracy-preserving path. -static inline bool siglip_fa_enabled() { - static const bool enabled = (std::getenv("VLA_SMOLVLA_FA") != nullptr); - return enabled; -} // One pre-norm SigLIP encoder block (SmolVLM2 tower), same graph as the other // in-tree models. Bidirectional attention, F32 score accumulation, tanh GELU. -ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_tensor * x, +ggml_tensor * build_siglip_layer(ggml_context * C, const EncBlockW & 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); + const float scale = 1.0f/std::sqrt((float) head_dim); ggml_tensor * n1 = ggml_add(C, ggml_mul(C, ggml_norm(C, x, ln_eps), w.ln1w), w.ln1b); ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n1), w.bq); ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, n1), w.bk); @@ -379,7 +401,7 @@ ggml_tensor * build_siglip_layer(ggml_context * C, const SigLipLayerW & w, ggml_ ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, head_dim, heads, seq), 0, 2, 1, 3)); ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, head_dim, heads, seq), 0, 2, 1, 3)); ggml_tensor * att; - if (siglip_fa_enabled()) { + if (vla::flash_attn_enabled()) { // The tower runs 1024 tokens (512/16 grid) over 12 layers, so the // explicit path below materialises a 1024x1024 score matrix per head — // written by the matmul, read and rewritten by the softmax, then read @@ -442,7 +464,7 @@ bool load_config_from_json(const std::string & path, Config & cfg) { cfg.n_layers = j.at("num_vlm_layers").get(); const double mul = j.at("expert_width_multiplier").get(); - cfg.expert_h = static_cast(std::round(double(cfg.hidden) * mul)); + cfg.expert_h = static_cast(std::round(double(cfg.hidden)*mul)); cfg.real_state_dim = j.at("input_features").at("observation.state").at("shape").at(0).get(); cfg.real_action_dim = j.at("output_features").at("action").at("shape").at(0).get(); @@ -453,7 +475,7 @@ bool load_config_from_json(const std::string & path, Config & cfg) { cfg.n_state = 1; cfg.q_full_dim = cfg.n_q_heads * cfg.head_dim; - cfg.kv_full_dim = cfg.n_kv_heads * cfg.head_dim; + cfg.kv_full_dim = cfg.n_kv_heads*cfg.head_dim; cfg.rope_n_dims = static_cast(cfg.head_dim); cfg.norm_eps = 1e-8f; @@ -542,7 +564,7 @@ std::string default_config_path(const std::string & ckpt_path) { bool ends_with_gguf(const std::string & path) { static const std::string sfx = ".gguf"; return path.size() >= sfx.size() - && path.compare(path.size() - sfx.size(), sfx.size(), sfx) == 0; + && path.compare(path.size()-sfx.size(), sfx.size(), sfx) == 0; } bool load_config_from_gguf(const gguf_source & st, Config & cfg) { @@ -575,7 +597,8 @@ bool load_config_from_gguf(const gguf_source & st, Config & cfg) { "smolvla.real_state_dim", "smolvla.real_action_dim", "smolvla.self_attn_every_n_layers", "smolvla.tokenizer_max_length", "smolvla.min_period", "smolvla.max_period"}) { - if (!need(k)) return false; + if (!need(k)) + return false; } cfg.hidden = st.get_u32("smolvla.hidden"); @@ -603,7 +626,7 @@ bool load_config_from_gguf(const gguf_source & st, Config & cfg) { cfg.n_state = 1; cfg.q_full_dim = cfg.n_q_heads * cfg.head_dim; - cfg.kv_full_dim = cfg.n_kv_heads * cfg.head_dim; + cfg.kv_full_dim = cfg.n_kv_heads*cfg.head_dim; cfg.rope_n_dims = static_cast(cfg.head_dim); cfg.norm_eps = st.has_key("smolvla.norm_eps") ? st.get_f32("smolvla.norm_eps") : 1e-8f; @@ -628,15 +651,24 @@ std::string hf_to_gguf(const std::string & n) { static const char * MODEL_PFX = "model."; auto map_suffix = [](const std::string & s) -> std::string { - if (s == "input_layernorm.weight") return "attn_norm.weight"; - if (s == "self_attn.q_proj.weight") return "attn_q.weight"; - if (s == "self_attn.k_proj.weight") return "attn_k.weight"; - if (s == "self_attn.v_proj.weight") return "attn_v.weight"; - if (s == "self_attn.o_proj.weight") return "attn_o.weight"; - if (s == "post_attention_layernorm.weight") return "ffn_norm.weight"; - if (s == "mlp.gate_proj.weight") return "ffn_gate.weight"; - if (s == "mlp.up_proj.weight") return "ffn_up.weight"; - if (s == "mlp.down_proj.weight") return "ffn_down.weight"; + if (s == "input_layernorm.weight") + return "attn_norm.weight"; + if (s == "self_attn.q_proj.weight") + return "attn_q.weight"; + if (s == "self_attn.k_proj.weight") + return "attn_k.weight"; + if (s == "self_attn.v_proj.weight") + return "attn_v.weight"; + if (s == "self_attn.o_proj.weight") + return "attn_o.weight"; + if (s == "post_attention_layernorm.weight") + return "ffn_norm.weight"; + if (s == "mlp.gate_proj.weight") + return "ffn_gate.weight"; + if (s == "mlp.up_proj.weight") + return "ffn_up.weight"; + if (s == "mlp.down_proj.weight") + return "ffn_down.weight"; return s; }; auto starts_with = [](const std::string & s, const char * pfx) -> bool { @@ -653,11 +685,13 @@ std::string hf_to_gguf(const std::string & n) { auto layer_translate = [&](const std::string & rest, const char * dst_blk) -> std::string { - if (!starts_with(rest, "layers.")) return n; + if (!starts_with(rest, "layers.")) + return n; const size_t end_i = rest.find('.', 7); - if (end_i == std::string::npos) return n; - const std::string idx = rest.substr(7, end_i - 7); - const std::string suf = rest.substr(end_i + 1); + if (end_i == std::string::npos) + return n; + const std::string idx = rest.substr(7, end_i-7); + const std::string suf = rest.substr(end_i+1); return std::string(dst_blk) + ".blk." + idx + "." + map_suffix(suf); }; @@ -673,18 +707,25 @@ std::string hf_to_gguf(const std::string & n) { return "mm.fc.weight"; if (starts_with(n, VIS_PFX)) { const std::string rest = n.substr(std::strlen(VIS_PFX)); - if (rest == "embeddings.patch_embedding.weight") return "vit.patch_embd.weight"; - if (rest == "embeddings.patch_embedding.bias") return "vit.patch_embd.bias"; - if (rest == "embeddings.position_embedding.weight") return "vit.pos_embd"; - if (rest == "post_layernorm.weight") return "vit.post_ln.weight"; - if (rest == "post_layernorm.bias") return "vit.post_ln.bias"; + if (rest == "embeddings.patch_embedding.weight") + return "vit.patch_embd.weight"; + if (rest == "embeddings.patch_embedding.bias") + return "vit.patch_embd.bias"; + if (rest == "embeddings.position_embedding.weight") + return "vit.pos_embd"; + if (rest == "post_layernorm.weight") + return "vit.post_ln.weight"; + if (rest == "post_layernorm.bias") + return "vit.post_ln.bias"; if (starts_with(rest, "encoder.layers.")) { const size_t e = rest.find('.', 15); - if (e == std::string::npos) return n; - const std::string idx = rest.substr(15, e - 15); - const std::string suf = rest.substr(e + 1); + if (e == std::string::npos) + return n; + const std::string idx = rest.substr(15, e-15); + const std::string suf = rest.substr(e+1); std::string ds; - if (suf == "layer_norm1.weight") ds = "ln1.weight"; + if (suf == "layer_norm1.weight") + ds = "ln1.weight"; else if (suf == "layer_norm1.bias") ds = "ln1.bias"; else if (suf == "layer_norm2.weight") ds = "ln2.weight"; else if (suf == "layer_norm2.bias") ds = "ln2.bias"; @@ -700,7 +741,8 @@ std::string hf_to_gguf(const std::string & n) { else if (suf == "mlp.fc1.bias") ds = "fc1.bias"; else if (suf == "mlp.fc2.weight") ds = "fc2.weight"; else if (suf == "mlp.fc2.bias") ds = "fc2.bias"; - else return n; + else + return n; return "vit.blk." + idx + "." + ds; } return n; @@ -724,12 +766,12 @@ ggml_tensor * rope_q_or_k(ggml_context * ctx, ggml_tensor * x, } static inline bool tower_mm_f32_prec() { - const char * e = std::getenv("VLA_MM_PREC"); - return !(e && std::strcmp(e, "default") == 0); + return vla::mm_prec_f32_enabled(); } static inline ggml_tensor * mm_w(ggml_context * ctx, ggml_tensor * w, ggml_tensor * x) { ggml_tensor * r = ggml_mul_mat(ctx, w, x); - if (tower_mm_f32_prec()) ggml_mul_mat_set_prec(r, GGML_PREC_F32); + if (tower_mm_f32_prec()) + ggml_mul_mat_set_prec(r, GGML_PREC_F32); return r; } @@ -754,7 +796,7 @@ ggml_tensor * build_vlm_layer(ggml_context * ctx, const VlmLayerW & w, ggml_tensor * Q = ggml_permute(ctx, q_rope, 0, 2, 1, 3); ggml_tensor * K = ggml_permute(ctx, k_rope, 0, 2, 1, 3); ggml_tensor * V = ggml_permute(ctx, v_h, 0, 2, 1, 3); - const float scale = 1.f / std::sqrt(static_cast(cfg.head_dim)); + const float scale = 1.f/std::sqrt(static_cast(cfg.head_dim)); ggml_tensor * fa = ggml_flash_attn_ext(ctx, Q, K, V, mask, scale, 0.f, 0.f); ggml_flash_attn_ext_set_prec(fa, GGML_PREC_F32); @@ -793,7 +835,7 @@ ggml_tensor * build_expert_self_attn_layer( ggml_tensor * Q = ggml_permute(ctx, q_rope, 0, 2, 1, 3); ggml_tensor * Kp = ggml_permute(ctx, K_full, 0, 2, 1, 3); ggml_tensor * Vp = ggml_permute(ctx, V_full, 0, 2, 1, 3); - const float scale = 1.f / std::sqrt(static_cast(cfg.head_dim)); + const float scale = 1.f/std::sqrt(static_cast(cfg.head_dim)); ggml_tensor * fa = ggml_flash_attn_ext(ctx, Q, Kp, Vp, mask_full, scale, 0.f, 0.f); ggml_flash_attn_ext_set_prec(fa, GGML_PREC_F32); @@ -832,7 +874,7 @@ ggml_tensor * build_expert_cross_attn_layer( ggml_tensor * Q = ggml_permute(ctx, q_rope, 0, 2, 1, 3); ggml_tensor * Kp = ggml_permute(ctx, K_repro, 0, 2, 1, 3); ggml_tensor * Vp = ggml_permute(ctx, V_repro, 0, 2, 1, 3); - const float scale = 1.f / std::sqrt(static_cast(cfg.head_dim)); + const float scale = 1.f/std::sqrt(static_cast(cfg.head_dim)); ggml_tensor * fa = ggml_flash_attn_ext(ctx, Q, Kp, Vp, mask_prefix_only, scale, 0.f, 0.f); ggml_flash_attn_ext_set_prec(fa, GGML_PREC_F32); @@ -851,34 +893,26 @@ namespace { static void vram_probe(ggml_backend_t backend, const char * label) { ggml_backend_dev_t dev = ggml_backend_get_device(backend); - if (!dev) return; + if (!dev) + return; size_t free_b = 0, total_b = 0; ggml_backend_dev_memory(dev, &free_b, &total_b); static size_t prev_free = 0; static bool have_prev = false; - const double MiB = 1024.0 * 1024.0; - const long long used = (long long)(total_b - free_b); + const double MiB = 1024.0*1024.0; + const long long used = (long long)(total_b-free_b); if (have_prev) { - const long long delta = (long long)prev_free - (long long)free_b; + const long long delta = (long long)prev_free-(long long)free_b; std::printf("vla: [vram] %-22s used=%.1f MiB free=%.1f MiB (+%.1f MiB)\n", - label, used / MiB, free_b / MiB, delta / MiB); + label, used/MiB, free_b/MiB, delta/MiB); } else { std::printf("vla: [vram] %-22s used=%.1f MiB free=%.1f MiB\n", - label, used / MiB, free_b / MiB); + label, used/MiB, free_b/MiB); } prev_free = free_b; have_prev = true; } -static ggml_type resolve_weight_dtype() { - const char * e = std::getenv("VLA_WEIGHT_DTYPE"); - if (!e) return GGML_TYPE_BF16; - if (std::strcmp(e, "f32") == 0) return GGML_TYPE_F32; - if (std::strcmp(e, "bf16") == 0) return GGML_TYPE_BF16; - if (std::strcmp(e, "f16") == 0) return GGML_TYPE_F16; - std::fprintf(stderr, "vla: unknown VLA_WEIGHT_DTYPE='%s', using bf16\n", e); - return GGML_TYPE_BF16; -} static void backend_set_from_f32(ggml_tensor * t, const float * src, int64_t n) { switch (t->type) { @@ -904,7 +938,8 @@ static void backend_set_from_f32(ggml_tensor * t, const float * src, int64_t n) } } -SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, +SmolVLAModelArch* smolvla_load_impl(ggml_type weight_dtype, + const std::string& mmproj_path, const std::string& ckpt_path, const std::string& config_path) { auto* m = new SmolVLAModelArch(); @@ -935,12 +970,15 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, { const Backend b = backend_init("vla", default_cpu_threads()); - if (!b.handle) { delete m; return nullptr; } + if (!b.handle) { + delete m; + return nullptr; + } m->backend = b.handle; } vram_probe(m->backend, "after backend init"); - m->weight_dtype = resolve_weight_dtype(); + m->weight_dtype = weight_dtype; std::printf("vla: tower weights resident as %s\n", ggml_type_name(m->weight_dtype)); // Vision tower geometry: from gguf KV (self-contained ckpt), else SmolVLM2-500M defaults. @@ -951,11 +989,12 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, vu("smolvla.vit_heads", m->vit_heads); vu("smolvla.patch_size", m->vit_patch); vu("smolvla.image_size", m->vit_image); vu("smolvla.vit_pixel_shuffle", m->vit_scale); vu("smolvla.n_img_tokens", m->vit_n_tokens); vu("smolvla.vit_inter", m->vit_inter); - if (gst.has_key("smolvla.vit_ln_eps")) m->vit_ln_eps = gst.get_f32("smolvla.vit_ln_eps"); + if (gst.has_key("smolvla.vit_ln_eps")) + m->vit_ln_eps = gst.get_f32("smolvla.vit_ln_eps"); } { - const int64_t grid = m->vit_image / m->vit_patch; - const int64_t k = grid / m->vit_scale; + const int64_t grid = m->vit_image/m->vit_patch; + const int64_t k = grid/m->vit_scale; 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); @@ -964,8 +1003,8 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, } m->cfg.n_img = m->vit_n_tokens; } - m->cfg.n_prefix = m->cfg.n_img + m->cfg.n_lang + m->cfg.n_state; - m->cfg.n_full = m->cfg.n_prefix + m->cfg.n_suffix; + m->cfg.n_prefix = m->cfg.n_img+m->cfg.n_lang+m->cfg.n_state; + m->cfg.n_full = m->cfg.n_prefix+m->cfg.n_suffix; if (!use_gguf) { if (!st.open(ckpt_path)) { @@ -978,8 +1017,9 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, int max_layer = -1; for (const auto & kv : st.tensors) { if (kv.first.compare(0, prefix.size(), prefix) == 0) { - const int idx = std::atoi(kv.first.c_str() + prefix.size()); - if (idx > max_layer) max_layer = idx; + const int idx = std::atoi(kv.first.c_str()+prefix.size()); + if (idx > max_layer) + max_layer = idx; } } if (max_layer < 0) { @@ -987,7 +1027,7 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, delete m; return nullptr; } - m->cfg.n_layers = max_layer + 1; + m->cfg.n_layers = max_layer+1; } { const auto it = st.tensors.find("model.vlm_with_expert.lm_expert.layers.0.mlp.gate_proj.weight"); @@ -1034,7 +1074,7 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, } ggml_init_params gparams = { - size_t(32) * 1024 * 1024, + size_t(32)*1024*1024, nullptr, true, }; @@ -1066,8 +1106,8 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, // Vision tower weights (SigLIP-B/16 encoder + single-linear pixel-shuffle connector). { const int64_t H = m->vit_hidden, FF = m->vit_inter, P = m->vit_patch; - const int64_t grid = m->vit_image / P, n_patches = grid * grid; - const int64_t c4 = H * m->vit_scale * m->vit_scale; + const int64_t grid = m->vit_image/P, n_patches = grid * grid; + const int64_t c4 = H * m->vit_scale*m->vit_scale; const char * VP = "model.vlm_with_expert.vlm.model.vision_model."; m->vit_patch_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, P, P, 3, H); m->vit_patch_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); @@ -1080,8 +1120,8 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, pending_f32.push_back({std::string(VP) + "post_layernorm.weight", m->vit_post_ln_w, {H}}); pending_f32.push_back({std::string(VP) + "post_layernorm.bias", m->vit_post_ln_b, {H}}); m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers; ++i) { - SigLipLayerW & w = m->vit[i]; + for (int64_t i=0; ivit_layers; ++i) { + EncBlockW & w = m->vit[i]; char pb[256]; std::snprintf(pb, sizeof(pb), "%sencoder.layers.%lld.", VP, (long long) i); const std::string pf = pb; w.ln1w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); w.ln1b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); @@ -1107,7 +1147,7 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, } m->vlm_layers.resize(cfg.n_layers); - for (int i = 0; i < cfg.n_layers; ++i) { + for (int i=0; ivlm_layers[i]; w.Wln_in = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, cfg.hidden); w.Wq = ggml_new_tensor_2d(ctx, wdt, cfg.hidden, cfg.q_full_dim); @@ -1137,9 +1177,9 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, m->Wnorm_vlm, {cfg.hidden}}); m->expert_layers.resize(cfg.n_layers); - for (int i = 0; i < cfg.n_layers; ++i) { + for (int i=0; iexpert_layers[i]; - w.is_self_attn = (i % cfg.self_attn_every_n == 0); + w.is_self_attn = (i%cfg.self_attn_every_n == 0); w.Wln_in = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, cfg.expert_h); w.Wq = ggml_new_tensor_2d(ctx, wdt, cfg.expert_h, cfg.q_full_dim); if (w.is_self_attn) { @@ -1179,13 +1219,13 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, m->W_ain = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, cfg.max_action_dim, cfg.expert_h); m->b_ain = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, cfg.expert_h); - m->W_at1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2 * cfg.expert_h, cfg.expert_h); + m->W_at1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2*cfg.expert_h, cfg.expert_h); m->b_at1 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, cfg.expert_h); m->W_at2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, cfg.expert_h, cfg.expert_h); m->b_at2 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, cfg.expert_h); pending_f32.push_back({"model.action_in_proj.weight", m->W_ain, {cfg.expert_h, cfg.max_action_dim}}); pending_f32.push_back({"model.action_in_proj.bias", m->b_ain, {cfg.expert_h}}); - pending_f32.push_back({"model.action_time_mlp_in.weight", m->W_at1, {cfg.expert_h, 2 * cfg.expert_h}}); + pending_f32.push_back({"model.action_time_mlp_in.weight", m->W_at1, {cfg.expert_h, 2*cfg.expert_h}}); pending_f32.push_back({"model.action_time_mlp_in.bias", m->b_at1, {cfg.expert_h}}); pending_f32.push_back({"model.action_time_mlp_out.weight", m->W_at2, {cfg.expert_h, cfg.expert_h}}); pending_f32.push_back({"model.action_time_mlp_out.bias", m->b_at2, {cfg.expert_h}}); @@ -1196,7 +1236,7 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, pending_f32.push_back({"model.action_out_proj.bias", m->b_aout, {cfg.max_action_dim}}); m->time_bcasts.assign(cfg.num_steps, nullptr); - for (int step = 0; step < cfg.num_steps; ++step) { + for (int step=0; steptime_bcasts[step] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, cfg.expert_h, cfg.n_suffix); } @@ -1208,7 +1248,7 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, return nullptr; } std::printf("vla: [vram] weight_buf = %.1f MiB\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0)); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0)); vram_probe(m->backend, "after weights alloc"); auto stream_f32 = [&](const std::string & hf_name, ggml_tensor * t, @@ -1238,25 +1278,31 @@ SmolVLAModelArch* smolvla_load_impl(const std::string& mmproj_path, }; for (auto & p : pending_f32) { - if (!stream_f32(p.name, p.t, p.shape)) { delete m; return nullptr; } + if (!stream_f32(p.name, p.t, p.shape)) { + delete m; + return nullptr; + } } for (auto & p : pending_bf16) { - if (!stream_bf16(p.name, p.t)) { delete m; return nullptr; } + if (!stream_bf16(p.name, p.t)) { + delete m; + return nullptr; + } } { - const float dt = -1.f / static_cast(cfg.num_steps); - for (int step = 0; step < cfg.num_steps; ++step) { - const double time = 1.0 + double(step) * double(dt); + const float dt = -1.f/static_cast(cfg.num_steps); + for (int step=0; step tile(cfg.expert_h * cfg.n_suffix); - for (int64_t t = 0; t < cfg.n_suffix; ++t) { - std::memcpy(tile.data() + t * cfg.expert_h, - te.data(), cfg.expert_h * sizeof(float)); + std::vector tile(cfg.expert_h*cfg.n_suffix); + for (int64_t t=0; ttime_bcasts[step], tile.data(), - 0, tile.size() * sizeof(float)); + 0, tile.size()*sizeof(float)); } } @@ -1273,10 +1319,10 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { const Config & cfg_model = m->cfg; Config cfg = cfg_model; - cfg.n_img = cfg_model.n_img * int64_t(n_views); + cfg.n_img = cfg_model.n_img*int64_t(n_views); ggml_init_params gparams = { - size_t(64) * 1024 * 1024, + size_t(64)*1024*1024, nullptr, true, }; @@ -1287,8 +1333,8 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { } const int64_t n_lang_max = cfg.n_lang; - const int64_t n_prefix_max = cfg.n_img + n_lang_max + cfg.n_state; - const int64_t n_full_max = n_prefix_max + cfg.n_suffix; + const int64_t n_prefix_max = cfg.n_img+n_lang_max+cfg.n_state; + const int64_t n_full_max = n_prefix_max+cfg.n_suffix; ggml_tensor * img_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, cfg.hidden, cfg.n_img); ggml_tensor * lang_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_lang_max); @@ -1329,7 +1375,7 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { std::vector v_cache(cfg.n_layers); { ggml_tensor * h = prefix_embs; - for (int i = 0; i < cfg.n_layers; ++i) { + for (int i=0; ivlm_layers[i], h, mask_prefill_f16, pos_prefill, cfg_built, &k_cache[i], &v_cache[i]); } @@ -1338,16 +1384,16 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { // Reproject each cross-attn layer's prefix K/V once; reused every denoise step. std::vector xk_cache(cfg.n_layers, nullptr); std::vector xv_cache(cfg.n_layers, nullptr); - for (int li = 0; li < cfg.n_layers; ++li) { + for (int li=0; liexpert_layers[li].is_self_attn) expert_cross_kv(ctx, m->expert_layers[li], k_cache[li], v_cache[li], cfg_built, &xk_cache[li], &xv_cache[li]); } - const float dt = -1.f / static_cast(cfg.num_steps); + const float dt = -1.f/static_cast(cfg.num_steps); ggml_tensor * x_t = x0; - for (int step = 0; step < cfg.num_steps; ++step) { + for (int step=0; steptime_bcasts[step]; ggml_tensor * action_emb = ggml_add(ctx, ggml_mul_mat(ctx, m->W_ain, x_t), m->b_ain); ggml_tensor * action_time_in = ggml_concat(ctx, action_emb, time_bcast, 0); @@ -1356,7 +1402,7 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { ggml_mul_mat(ctx, m->W_at2, ggml_silu(ctx, mlp1)), m->b_at2); ggml_tensor * h = suffix_embs; - for (int li = 0; li < cfg.n_layers; ++li) { + for (int li=0; liexpert_layers[li].is_self_attn) { h = build_expert_self_attn_layer(ctx, m->expert_layers[li], h, k_cache[li], v_cache[li], @@ -1392,7 +1438,7 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { } std::printf("vla: [vram] gallocr compute buf = %.1f MiB\n", - ggml_gallocr_get_buffer_size(galloc, 0) / (1024.0 * 1024.0)); + ggml_gallocr_get_buffer_size(galloc, 0)/(1024.0*1024.0)); vram_probe(m->backend, "after gallocr reserve"); m->ctx_compute = ctx; @@ -1416,11 +1462,16 @@ bool build_compute_graph(SmolVLAModelArch* m, int n_views) { SmolVLAModelArch::~SmolVLAModelArch() { - if (galloc) ggml_gallocr_free(galloc); - if (ctx_compute) ggml_free(ctx_compute); - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (galloc) + ggml_gallocr_free(galloc); + if (ctx_compute) + ggml_free(ctx_compute); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } namespace { @@ -1432,7 +1483,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { Config cfg = m->cfg; - const size_t per_view_n = size_t(m->cfg.n_img * cfg.hidden); + const size_t per_view_n = size_t(m->cfg.n_img*cfg.hidden); int n_views = 0; size_t img_emb_n = 0; @@ -1446,7 +1497,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { } n_views = in.n_img_views; img_emb_n = per_view_n * size_t(n_views); - img_emb_pre.assign(in.precomputed_img_emb, in.precomputed_img_emb + img_emb_n); + img_emb_pre.assign(in.precomputed_img_emb, in.precomputed_img_emb+img_emb_n); } else { if (in.n_images < 1 || in.images == nullptr) { @@ -1457,19 +1508,19 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { img_emb_n = per_view_n * size_t(n_views); img_emb_pre.resize(img_emb_n); - const int64_t H = m->vit_hidden, grid = m->vit_image / m->vit_patch, n_patches = grid * grid; + const int64_t H = m->vit_hidden, grid = m->vit_image/m->vit_patch, n_patches = grid * grid; const int64_t s = m->vit_scale, c4 = H * s * s, K = m->vit_n_tokens; const auto t_vision_begin = clock::now(); // Graph A: SigLIP ViT (conv patch-embed -> +pos -> layers -> post_ln), plain sequential positions. - ggml_context * VC = m->vision_scratch.reset(size_t(256) * 1024 * 1024); + 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); ggml_tensor * patches = ggml_cont(VC, ggml_transpose(VC, ggml_reshape_2d(VC, conv, n_patches, H))); ggml_tensor * hv = ggml_add(VC, ggml_add(VC, patches, m->vit_patch_b), m->vit_pos); - for (int64_t i = 0; i < m->vit_layers; ++i) - hv = build_siglip_layer(VC, m->vit[i], hv, n_patches, m->vit_heads, H / m->vit_heads, H, m->vit_ln_eps); + for (int64_t i=0; ivit_layers; ++i) + hv = build_siglip_layer(VC, m->vit[i], hv, n_patches, m->vit_heads, H/m->vit_heads, H, m->vit_ln_eps); ggml_tensor * post_ln = ggml_add(VC, ggml_mul(VC, ggml_norm(VC, hv, m->vit_ln_eps), m->vit_post_ln_w), m->vit_post_ln_b); ggml_set_output(post_ln); ggml_cgraph * gA = ggml_new_graph_custom(VC, 8192, false); @@ -1480,7 +1531,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { } // Graph B: pixel-shuffle connector, a single bias-free matmul (c4 -> hidden). - ggml_context * MC = m->connector_scratch.reset(size_t(64) * 1024 * 1024); + 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); @@ -1492,10 +1543,13 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { return {}; } - std::vector chw, post_host((size_t) H * n_patches), shuf_host((size_t) c4 * K); + 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("smolvla", in.images[v], m->vit_image, chw)) { vok = false; break; } + for (int v=0; vvit_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; @@ -1506,15 +1560,15 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { if (ggml_backend_graph_compute(m->backend, gB) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "vla(smolvla): connector compute failed (view %d)\n", v); vok = false; break; } - ggml_backend_tensor_get(img_embeds, img_emb_pre.data() + size_t(v) * per_view_n, 0, ggml_nbytes(img_embeds)); + ggml_backend_tensor_get(img_embeds, img_emb_pre.data()+size_t(v)*per_view_n, 0, ggml_nbytes(img_embeds)); } if (!vok) return {}; - m->stats.ms_vision = std::chrono::duration(clock::now() - t_vision_begin).count(); + m->stats.ms_vision = std::chrono::duration(clock::now()-t_vision_begin).count(); } - cfg.n_img = m->cfg.n_img * int64_t(n_views); - cfg.n_prefix = cfg.n_img + cfg.n_lang + cfg.n_state; - cfg.n_full = cfg.n_prefix + cfg.n_suffix; + cfg.n_img = m->cfg.n_img*int64_t(n_views); + cfg.n_prefix = cfg.n_img+cfg.n_lang+cfg.n_state; + cfg.n_full = cfg.n_prefix+cfg.n_suffix; if (in.n_lang < 1 || in.n_lang > int(cfg.n_lang)) { std::fprintf(stderr, "vla: lang_tokens length %d out of range [1, %lld]\n", @@ -1526,7 +1580,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { // its indices. Reject any token id outside the embedding table before the // gather so an out-of-range id cannot read past the weights. const int64_t vocab_rows = m->E_lang ? m->E_lang->ne[1] : 0; - for (int i = 0; i < in.n_lang; ++i) { + for (int i=0; i= vocab_rows) { std::fprintf(stderr, "vla: lang_tokens[%d]=%d out of vocab range [0, %lld)\n", i, in.lang_tokens[i], (long long) vocab_rows); @@ -1537,8 +1591,10 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { if (in.timing_detail == TimingDetail::NONE) { if (m->gf_cached == nullptr || m->cached_n_views != n_views) { - if (m->galloc) ggml_gallocr_free(m->galloc); - if (m->ctx_compute) ggml_free(m->ctx_compute); + if (m->galloc) + ggml_gallocr_free(m->galloc); + if (m->ctx_compute) + ggml_free(m->ctx_compute); m->galloc = nullptr; m->ctx_compute = nullptr; m->gf_cached = nullptr; @@ -1550,40 +1606,44 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { } const int64_t n_lang_max = m->cfg.n_lang; - const int64_t n_prefix_max = cfg.n_img + n_lang_max + cfg.n_state; - const int64_t n_full_max = n_prefix_max + cfg.n_suffix; - const int64_t pad_start = cfg.n_img + in.n_lang; - const int64_t pad_end = cfg.n_img + n_lang_max; + const int64_t n_prefix_max = cfg.n_img+n_lang_max+cfg.n_state; + const int64_t n_full_max = n_prefix_max+cfg.n_suffix; + 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, 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); + if (in.state) + std::memcpy(state_host.data(), in.state, cfg.max_state_dim*sizeof(float)); + for (int64_t i=0; istate_mean[i])/(m->state_std[i]+cfg.norm_eps); } - std::vector noise_host(cfg.n_suffix * cfg.max_action_dim); + std::vector noise_host(cfg.n_suffix*cfg.max_action_dim); if (in.noise) { - std::memcpy(noise_host.data(), in.noise, noise_host.size() * sizeof(float)); + std::memcpy(noise_host.data(), in.noise, noise_host.size()*sizeof(float)); } else { std::normal_distribution dist(0.f, 1.f); - for (auto & v : noise_host) v = dist(m->rng); + for (auto & v : noise_host) + v = dist(m->rng); } std::vector lang_host(n_lang_max, 0); - std::memcpy(lang_host.data(), in.lang_tokens, in.n_lang * sizeof(int32_t)); + std::memcpy(lang_host.data(), in.lang_tokens, in.n_lang*sizeof(int32_t)); - const int64_t state_pos = cfg.n_img + in.n_lang; - const int64_t suffix_pos_base = state_pos + 1; + const int64_t state_pos = cfg.n_img+in.n_lang; + const int64_t suffix_pos_base = state_pos+1; std::vector mask_prefill_host(n_prefix_max * n_prefix_max); std::vector pos_prefill_host (n_prefix_max); - for (int64_t i = 0; i < n_prefix_max; ++i) { - for (int64_t j = 0; j < n_prefix_max; ++j) { + for (int64_t i=0; i= pad_start && j < pad_end) blocked = true; - mask_prefill_host[i * n_prefix_max + j] = blocked ? -INFINITY : 0.f; + if ((i < n_prefix_max-1) && (j == n_prefix_max-1)) + blocked = true; + if (j >= pad_start && j < pad_end) + blocked = true; + mask_prefill_host[i * n_prefix_max+j] = blocked ? -INFINITY : 0.f; } - pos_prefill_host[i] = (i == n_prefix_max - 1) + pos_prefill_host[i] = (i == n_prefix_max-1) ? static_cast(state_pos) : static_cast(i); } @@ -1592,22 +1652,22 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { std::vector mask_prefix_only_host(n_prefix_max * cfg.n_suffix, 0.f); std::vector pos_full_host (cfg.n_suffix); std::vector pos_rebased_host (cfg.n_suffix); - for (int64_t i = 0; i < cfg.n_suffix; ++i) { - for (int64_t j = 0; j < n_full_max; ++j) { + for (int64_t i=0; i= pad_start && j < pad_end); } else { - blocked = ((j - n_prefix_max) > i); + blocked = ((j-n_prefix_max) > i); } - mask_full_host[i * n_full_max + j] = blocked ? -INFINITY : 0.f; + mask_full_host[i * n_full_max+j] = blocked ? -INFINITY : 0.f; } - for (int64_t j = 0; j < n_prefix_max; ++j) { + for (int64_t j=0; j= pad_start && j < pad_end) { - mask_prefix_only_host[i * n_prefix_max + j] = -INFINITY; + mask_prefix_only_host[i * n_prefix_max+j] = -INFINITY; } } - pos_full_host [i] = static_cast(suffix_pos_base + i); + pos_full_host [i] = static_cast(suffix_pos_base+i); pos_rebased_host[i] = static_cast(i); } @@ -1618,12 +1678,12 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_backend_tensor_set(m->in_img_emb, img_emb_pre.data(), 0, img_emb_n * sizeof(float)); ggml_backend_tensor_set(m->in_lang_ids, lang_host.data(), 0, n_lang_max * sizeof(int32_t)); - ggml_backend_tensor_set(m->in_state, state_host.data(), 0, cfg.max_state_dim * sizeof(float)); - ggml_backend_tensor_set(m->in_x0, noise_host.data(), 0, noise_host.size() * sizeof(float)); + ggml_backend_tensor_set(m->in_state, state_host.data(), 0, cfg.max_state_dim*sizeof(float)); + ggml_backend_tensor_set(m->in_x0, noise_host.data(), 0, noise_host.size()*sizeof(float)); ggml_backend_tensor_set(m->in_mask_prefill, mask_prefill_host.data(), 0, mask_prefill_host.size() * sizeof(float)); ggml_backend_tensor_set(m->in_pos_prefill, pos_prefill_host.data(), 0, pos_prefill_host.size() * sizeof(int32_t)); ggml_backend_tensor_set(m->in_mask_full, mask_full_host.data(), 0, mask_full_host.size() * sizeof(float)); - ggml_backend_tensor_set(m->in_mask_pfx_only, mask_prefix_only_host.data(), 0, mask_prefix_only_host.size() * sizeof(float)); + ggml_backend_tensor_set(m->in_mask_pfx_only, mask_prefix_only_host.data(), 0, mask_prefix_only_host.size()*sizeof(float)); ggml_backend_tensor_set(m->in_pos_full, pos_full_host.data(), 0, pos_full_host.size() * sizeof(int32_t)); ggml_backend_tensor_set(m->in_pos_rebased, pos_rebased_host.data(), 0, pos_rebased_host.size() * sizeof(int32_t)); @@ -1633,24 +1693,24 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { return {}; } m->stats.ms_inference = std::chrono::duration( - clock::now() - t0).count(); - - std::vector out(cfg.n_suffix * cfg.max_action_dim); - ggml_backend_tensor_get(m->out_x_t, out.data(), 0, out.size() * sizeof(float)); - for (int64_t r = 0; r < cfg.n_suffix; ++r) { - float * row = out.data() + r * cfg.max_action_dim; - for (int64_t j = 0; j < cfg.max_action_dim; ++j) { - row[j] = j < cfg.real_action_dim ? row[j] * (m->action_std[j] + cfg.norm_eps) + m->action_mean[j] : 0.0f; + clock::now()-t0).count(); + + std::vector out(cfg.n_suffix*cfg.max_action_dim); + ggml_backend_tensor_get(m->out_x_t, out.data(), 0, out.size()*sizeof(float)); + for (int64_t r=0; raction_std[j]+cfg.norm_eps)+m->action_mean[j] : 0.0f; } } m->stats.ms_total = std::chrono::duration( - clock::now() - t_total_begin).count(); + clock::now()-t_total_begin).count(); return out; } ggml_init_params gparams = { - size_t(64) * 1024 * 1024, + size_t(64)*1024*1024, nullptr, true, }; @@ -1668,8 +1728,8 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { } cfg.n_lang = in.n_lang; - cfg.n_prefix = cfg.n_img + cfg.n_lang + cfg.n_state; - cfg.n_full = cfg.n_prefix + cfg.n_suffix; + cfg.n_prefix = cfg.n_img+cfg.n_lang+cfg.n_state; + cfg.n_full = cfg.n_prefix+cfg.n_suffix; ggml_tensor * img_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, cfg.hidden, cfg.n_img); ggml_tensor * lang_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, cfg.n_lang); ggml_tensor * state_t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, cfg.max_state_dim); @@ -1683,41 +1743,45 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_tensor * pos_rebased = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, cfg.n_suffix); 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); + if (in.state) + std::memcpy(state_host.data(), in.state, cfg.max_state_dim*sizeof(float)); + for (int64_t i=0; istate_mean[i])/(m->state_std[i]+cfg.norm_eps); } - std::vector noise_host(cfg.n_suffix * cfg.max_action_dim); + std::vector noise_host(cfg.n_suffix*cfg.max_action_dim); if (in.noise) { - std::memcpy(noise_host.data(), in.noise, noise_host.size() * sizeof(float)); + std::memcpy(noise_host.data(), in.noise, noise_host.size()*sizeof(float)); } else { std::normal_distribution dist(0.f, 1.f); - for (auto & v : noise_host) v = dist(m->rng); + for (auto & v : noise_host) + v = dist(m->rng); } - std::vector mask_prefill_host(cfg.n_prefix * cfg.n_prefix); + std::vector mask_prefill_host(cfg.n_prefix*cfg.n_prefix); std::vector pos_prefill_host (cfg.n_prefix); - for (int64_t i = 0; i < cfg.n_prefix; ++i) { - for (int64_t j = 0; j < cfg.n_prefix; ++j) { - const bool blocked = (i < cfg.n_prefix - 1) && (j == cfg.n_prefix - 1); - mask_prefill_host[i * cfg.n_prefix + j] = blocked ? -INFINITY : 0.f; + for (int64_t i=0; i(i); } std::vector mask_full_host (cfg.n_full * cfg.n_suffix); - std::vector mask_prefix_only_host(cfg.n_prefix * cfg.n_suffix, 0.f); + std::vector mask_prefix_only_host(cfg.n_prefix*cfg.n_suffix, 0.f); std::vector pos_full_host (cfg.n_suffix); std::vector pos_rebased_host (cfg.n_suffix); - for (int64_t i = 0; i < cfg.n_suffix; ++i) { - for (int64_t j = 0; j < cfg.n_full; ++j) { + for (int64_t i=0; i i); - mask_full_host[i * cfg.n_full + j] = blocked ? -INFINITY : 0.f; + if (j < cfg.n_prefix) + blocked = false; + else + blocked = ((j-cfg.n_prefix) > i); + mask_full_host[i * cfg.n_full+j] = blocked ? -INFINITY : 0.f; } - pos_full_host [i] = static_cast(cfg.n_prefix + i); + pos_full_host [i] = static_cast(cfg.n_prefix+i); pos_rebased_host[i] = static_cast(i); } @@ -1737,7 +1801,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { std::vector v_cache(cfg.n_layers); { ggml_tensor * h = prefix_embs; - for (int i = 0; i < cfg.n_layers; ++i) { + for (int i=0; ivlm_layers[i], h, mask_prefill_f16, pos_prefill, cfg, &k_cache[i], &v_cache[i]); } @@ -1746,7 +1810,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { std::vector K_storage(cfg.n_layers); std::vector V_storage(cfg.n_layers); if (in.timing_detail == TimingDetail::PHASE) { - for (int i = 0; i < cfg.n_layers; ++i) { + for (int i=0; i predict_impl(SmolVLAModelArch* m, const Inputs& in) { } auto ms_since = [](auto t0) { - return std::chrono::duration(clock::now() - t0).count(); + return std::chrono::duration(clock::now()-t0).count(); }; auto & K_ref = (in.timing_detail == TimingDetail::PHASE) ? K_storage : k_cache; auto & V_ref = (in.timing_detail == TimingDetail::PHASE) ? V_storage : v_cache; - const float dt = -1.f / static_cast(cfg.num_steps); + const float dt = -1.f/static_cast(cfg.num_steps); ggml_tensor * x_t = x0; std::vector time_bcasts(cfg.num_steps, nullptr); std::vector> time_host (cfg.num_steps); std::vector xk_cache(cfg.n_layers, nullptr); std::vector xv_cache(cfg.n_layers, nullptr); - for (int li = 0; li < cfg.n_layers; ++li) { + for (int li=0; liexpert_layers[li].is_self_attn) expert_cross_kv(ctx, m->expert_layers[li], K_ref[li], V_ref[li], cfg, &xk_cache[li], &xv_cache[li]); } - for (int step = 0; step < cfg.num_steps; ++step) { - const double time = 1.0 + static_cast(step) * static_cast(dt); + for (int step=0; step(step)*static_cast(dt); time_host[step] = sinusoidal_time_emb(time, cfg.expert_h, cfg.min_period, cfg.max_period); ggml_tensor * time_bcast = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, cfg.expert_h, cfg.n_suffix); @@ -1788,7 +1852,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { m->b_at2); ggml_tensor * h = suffix_embs; - for (int li = 0; li < cfg.n_layers; ++li) { + for (int li=0; liexpert_layers[li].is_self_attn) { h = build_expert_self_attn_layer(ctx, m->expert_layers[li], h, K_ref[li], V_ref[li], @@ -1813,28 +1877,28 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_backend_tensor_set(img_emb_in, img_emb_pre.data(), 0, img_emb_n * sizeof(float)); ggml_backend_tensor_set(lang_ids, in.lang_tokens, 0, cfg.n_lang * sizeof(int32_t)); - ggml_backend_tensor_set(state_t, state_host.data(), 0, cfg.max_state_dim * sizeof(float)); + ggml_backend_tensor_set(state_t, state_host.data(), 0, cfg.max_state_dim*sizeof(float)); ggml_backend_tensor_set(x0, noise_host.data(), 0, noise_host.size() * sizeof(float)); ggml_backend_tensor_set(mask_prefill, mask_prefill_host.data(), 0, mask_prefill_host.size() * sizeof(float)); ggml_backend_tensor_set(pos_prefill, pos_prefill_host.data(), 0, pos_prefill_host.size() * sizeof(int32_t)); ggml_backend_tensor_set(mask_full, mask_full_host.data(), 0, mask_full_host.size() * sizeof(float)); - ggml_backend_tensor_set(mask_prefix_only, mask_prefix_only_host.data(), 0, mask_prefix_only_host.size() * sizeof(float)); + ggml_backend_tensor_set(mask_prefix_only, mask_prefix_only_host.data(), 0, mask_prefix_only_host.size()*sizeof(float)); ggml_backend_tensor_set(pos_full, pos_full_host.data(), 0, pos_full_host.size() * sizeof(int32_t)); ggml_backend_tensor_set(pos_rebased, pos_rebased_host.data(), 0, pos_rebased_host.size() * sizeof(int32_t)); - for (int step = 0; step < cfg.num_steps; ++step) { + for (int step=0; step tile(cfg.expert_h * cfg.n_suffix); - for (int64_t t = 0; t < cfg.n_suffix; ++t) { - std::memcpy(tile.data() + t * cfg.expert_h, - time_host[step].data(), cfg.expert_h * sizeof(float)); + std::vector tile(cfg.expert_h*cfg.n_suffix); + for (int64_t t=0; t predict_impl(SmolVLAModelArch* m, const Inputs& in) { } m->stats.ms_prefill = ms_since(t0); - for (int i = 0; i < cfg.n_layers; ++i) { + for (int i=0; i predict_impl(SmolVLAModelArch* m, const Inputs& in) { m->stats.ms_denoise = ms; } - m->stats.ms_inference = m->stats.ms_prefill + ms; + m->stats.ms_inference = m->stats.ms_prefill+ms; } - std::vector out(cfg.n_suffix * cfg.max_action_dim); - ggml_backend_tensor_get(x_t, out.data(), 0, out.size() * sizeof(float)); + std::vector out(cfg.n_suffix*cfg.max_action_dim); + ggml_backend_tensor_get(x_t, out.data(), 0, out.size()*sizeof(float)); - for (int64_t r = 0; r < cfg.n_suffix; ++r) { - float * row = out.data() + r * cfg.max_action_dim; - for (int64_t j = 0; j < cfg.max_action_dim; ++j) { - row[j] = j < cfg.real_action_dim ? row[j] * (m->action_std[j] + cfg.norm_eps) + m->action_mean[j] : 0.0f; + for (int64_t r=0; raction_std[j]+cfg.norm_eps)+m->action_mean[j] : 0.0f; } } @@ -1885,7 +1949,7 @@ std::vector predict_impl(SmolVLAModelArch* m, const Inputs& in) { ggml_free(ctx); m->stats.ms_total = std::chrono::duration( - clock::now() - t_total_begin).count(); + clock::now()-t_total_begin).count(); return out; } } @@ -1896,9 +1960,12 @@ std::vector SmolVLAModelArch::predict(const Inputs& in) { std::unique_ptr smolvla_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& config_path) { - SmolVLAModelArch* raw = smolvla_load_impl(mmproj_path, ckpt_path, config_path); - if (!raw) return nullptr; + const std::string& config_path, + const Options& opts) { + SmolVLAModelArch* raw = smolvla_load_impl(opts.weight_dtype.value_or(GGML_TYPE_BF16), + mmproj_path, ckpt_path, config_path); + if (!raw) + return nullptr; return std::unique_ptr(raw); } diff --git a/src/models/vision_common.h b/src/models/vision_common.h deleted file mode 100644 index 9e33f7b..0000000 --- a/src/models/vision_common.h +++ /dev/null @@ -1,75 +0,0 @@ -// 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. - -// Small pure vision helpers shared by the in-tree towers, split out so they can -// be unit-tested without a model or a GPU. - -#pragma once - -#include "model.h" - -#include -#include -#include -#include - -namespace vla { - -// A tower that reads side*side*3 from a view needs the view to be exactly that -// size with real data, else it runs past the buffer. -inline bool view_is_side(const void * data, int w, int h, int64_t side) { - return data != nullptr && (int64_t) w == side && (int64_t) h == side; -} - -// IDEFICS3/SmolVLM pixel-shuffle (space-to-depth), c-innermost channel order. -// src [embed, n_patches] row-major (patch p, channel e) -> dst [embed*s^2, (grid/s)^2]. -inline void pixel_shuffle_hf(const float * src, float * dst, - int64_t embed, int64_t grid, int64_t s) { - const int64_t g2 = grid / s, c4 = embed * s * s; - for (int64_t h2 = 0; h2 < g2; ++h2) - for (int64_t w2 = 0; w2 < g2; ++w2) { - const int64_t t = h2 * g2 + w2; - for (int64_t hs = 0; hs < s; ++hs) - for (int64_t ws = 0; ws < s; ++ws) { - const int64_t p = (h2 * s + hs) * grid + (w2 * s + ws); - const int64_t base = (hs * s + ws) * embed; - std::memcpy(dst + t * c4 + base, src + p * embed, - (size_t) embed * sizeof(float)); - } - } -} - -// 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 9351784..af00615 100644 --- a/src/models/vla_adapter.cpp +++ b/src/models/vla_adapter.cpp @@ -13,18 +13,20 @@ // limitations under the License. #include "arch.h" +#include "options.h" #include "model.h" -#include "vision_common.h" -#include "models/dual_tower.h" +#include "modules/preprocess.h" +#include "modules/dual_tower.h" #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" #include "backend.h" #include "gguf.h" -#include "models/gguf_reader.h" -#include "models/scratch_ctx.h" +#include "gguf_reader.h" +#include "scratch_ctx.h" #include "models/dit_common.h" +#include "env_flag.h" #include #include @@ -49,33 +51,52 @@ bool parse_stats(const std::string & js, int64_t want, std::vector & q01, const char * env = std::getenv("VLA_ADAPTER_UNNORM_KEY"); size_t suite_pos; - if (env) { suite = env; suite_pos = find_key(0, suite); } + if (env) { + suite = env; + suite_pos = find_key(0, suite); + } else { size_t b = js.find('{'); size_t q = js.find('"', b); - size_t qe = js.find('"', q + 1); - suite = js.substr(q + 1, qe - q - 1); suite_pos = q; + size_t qe = js.find('"', q+1); + suite = js.substr(q+1, qe-q-1); suite_pos = q; + } + if (suite_pos == std::string::npos) { + std::fprintf(stderr, "vla(vla_adapter): suite '%s' not in stats\n", suite.c_str()); + return false; } - if (suite_pos == std::string::npos) { std::fprintf(stderr, "vla(vla_adapter): suite '%s' not in stats\n", suite.c_str()); return false; } size_t act = find_key(suite_pos, "action"); - if (act == std::string::npos) return false; + if (act == std::string::npos) + return false; auto read_arr = [&](const std::string & key, std::vector & out) -> bool { size_t k = find_key(act, key); if (k == std::string::npos) return false; size_t lb = js.find('[', k); size_t rb = js.find(']', lb); - if (lb == std::string::npos || rb == std::string::npos) return false; - out.clear(); size_t p = lb + 1; + if (lb == std::string::npos || rb == std::string::npos) + return false; + out.clear(); size_t p = lb+1; while (p < rb) { - while (p < rb && (js[p] == ',' || js[p] == ' ' || js[p] == '\n' || js[p] == '\t' || js[p] == '\r')) ++p; - if (p >= rb) break; + while (p < rb && (js[p] == ',' || js[p] == ' ' || js[p] == '\n' || js[p] == '\t' || js[p] == '\r')) + ++p; + if (p >= rb) + break; bool t = (js.compare(p, 4, "true") == 0), f = (js.compare(p, 5, "false") == 0); - if (t || f) { out.push_back(t ? 1.0f : 0.0f); p += t ? 4 : 5; } - else { out.push_back(std::strtof(js.c_str() + p, nullptr)); while (p < rb && js[p] != ',') ++p; } + if (t || f) { + out.push_back(t ? 1.0f : 0.0f); + p += t ? 4 : 5; + } + else { + out.push_back(std::strtof(js.c_str()+p, nullptr)); + while (p < rb && js[p] != ',') + ++p; + } } return true; }; std::vector mk; - if (!read_arr("q01", q01) || !read_arr("q99", q99)) return false; - if (!read_arr("mask", mk)) mk.assign(want, 1.0f); - mask.assign(mk.size(), 1); for (size_t i = 0; i < mk.size(); ++i) mask[i] = mk[i] != 0.0f ? 1 : 0; + if (!read_arr("q01", q01) || !read_arr("q99", q99)) + return false; + if (!read_arr("mask", mk)) + mk.assign(want, 1.0f); + mask.assign(mk.size(), 1); for (size_t i=0; i dvit; - ggml_tensor *s_patch_w,*s_patch_b,*s_pos; std::vector svit; - ggml_tensor *pj_fc1w,*pj_fc1b,*pj_fc2w,*pj_fc2b,*pj_fc3w,*pj_fc3b; + DualTower vis; ggml_tensor *token_embd,*action_queries,*lm_out_norm; std::vector lm; ggml_tensor *h_ln1w,*h_ln1b,*h_fc1w,*h_fc1b,*h_ln2w,*h_ln2b,*h_fc2w,*h_fc2b; ggml_tensor *pp_fc1w,*pp_fc1b,*pp_fc2w,*pp_fc2b; std::vector hblk; @@ -144,7 +168,9 @@ static ggml_tensor* hrot(ggml_context*C, ggml_tensor*x, int64_t HD){ ggml_tensor*od=ggml_cont(C,ggml_view_4d(C,xp,1,HD/2,L,H,xp->nb[1],xp->nb[2],xp->nb[3],xp->nb[0])); return ggml_reshape_3d(C,ggml_concat(C,ggml_scale(C,od,-1.0f),ev,0),HD,L,H); } -static ggml_tensor* hheads(ggml_context*C, ggml_tensor*p, int64_t HD, int64_t NH){ return ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,p,HD,NH,p->ne[1]),0,2,1,3)); } +static ggml_tensor* hheads(ggml_context*C, ggml_tensor*p, int64_t HD, int64_t NH){ + return ggml_cont(C,ggml_permute(C,ggml_reshape_3d(C,p,HD,NH,p->ne[1]),0,2,1,3)); +} static ggml_tensor* hrope(ggml_context*C, ggml_tensor*x, ggml_tensor*cs, ggml_tensor*sn, int64_t HD){ ggml_tensor*c=ggml_reshape_3d(C,cs,HD,x->ne[1],1),*s=ggml_reshape_3d(C,sn,HD,x->ne[1],1); return ggml_add(C,ggml_mul(C,x,c),ggml_mul(C,hrot(C,x,HD),s)); @@ -154,15 +180,20 @@ static ggml_tensor* hrope(ggml_context*C, ggml_tensor*x, ggml_tensor*cs, ggml_te std::unique_ptr vla_adapter_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(vla_adapter): note - mmproj '%s' ignored (vision baked into combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); - m->mt = std::getenv("VLA_ADAPTER_F32_WEIGHTS") ? GGML_TYPE_F32 : GGML_TYPE_BF16; + m->mt = opts.weight_dtype.value_or(GGML_TYPE_BF16); gguf_reader g("vla_adapter"); - if (!g.open(ckpt_path)) return nullptr; - if (!g.has("vla_adapter.architecture")) { std::fprintf(stderr, "vla(vla_adapter): not a vla_adapter GGUF\n"); return nullptr; } + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("vla_adapter.architecture")) { + std::fprintf(stderr, "vla(vla_adapter): not a vla_adapter GGUF\n"); + return nullptr; + } auto U=[&](const char*k,int64_t&d){ if(g.has(k)) d=(int64_t)g.u32(k); }; auto F=[&](const char*k,float&d){ if(g.has(k)) d=g.f32(k); }; @@ -193,50 +224,31 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path 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::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()); } { const Backend b = backend_init("vla(vla_adapter)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; } ggml_init_params wp = { (size_t)64*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - ggml_context * W = m->ctx_weights; bool ok = true; - auto mk=[&](const char*name, ggml_type ty)->ggml_tensor*{ const ggml_tensor*gt=g.meta(name); - if(!gt){ std::fprintf(stderr,"vla(vla_adapter): missing %s\n",name); ok=false; return nullptr; } - ggml_tensor*t=ggml_new_tensor(W,g.resident_type(gt,ty),ggml_n_dims(gt),gt->ne); ggml_set_name(t,name); return t; }; - auto mm=[&](const char*n){ return mk(n,m->mt); }; - auto f32=[&](const char*n){ return mk(n,GGML_TYPE_F32); }; + WeightLoader L("vla_adapter", g, m->ctx_weights, m->mt); + auto mm = [&](const char * n) { return L.gemm("%s", n); }; + auto f32 = [&](const char * n) { return L.f32("%s", n); }; char nm[96]; auto P=[&](const char*fmt,int i){ std::snprintf(nm,sizeof(nm),fmt,i); return (const char*)nm; }; - m->d_patch_w=mk("vis.d.patch.weight",GGML_TYPE_F32); m->d_patch_b=f32("vis.d.patch.bias"); - m->d_cls=f32("vis.d.cls"); m->d_reg=f32("vis.d.reg"); m->d_pos=f32("vis.d.pos"); - m->dvit.resize(m->d_layers); - for(int i=0;id_layers;++i){ auto&w=m->dvit[i]; char b[64]; - auto N=[&](const char*s){ std::snprintf(b,sizeof(b),"vis.d.blk.%d.%s",i,s); return (const char*)b; }; - w.n1w=f32(N("ln1.weight")); w.n1b=f32(N("ln1.bias")); w.n2w=f32(N("ln2.weight")); w.n2b=f32(N("ln2.bias")); - w.ls1=f32(N("ls1")); w.ls2=f32(N("ls2")); w.Wqkv=mm(N("qkv.weight")); w.bqkv=f32(N("qkv.bias")); - w.Wproj=mm(N("proj.weight")); w.bproj=f32(N("proj.bias")); w.Wfc1=mm(N("fc1.weight")); w.bfc1=f32(N("fc1.bias")); - w.Wfc2=mm(N("fc2.weight")); w.bfc2=f32(N("fc2.bias")); } - - m->s_patch_w=mk("vis.s.patch.weight",GGML_TYPE_F32); m->s_patch_b=f32("vis.s.patch.bias"); m->s_pos=f32("vis.s.pos"); - m->svit.resize(m->s_layers); - for(int i=0;is_layers;++i){ auto&w=m->svit[i]; char b[64]; - auto N=[&](const char*s){ std::snprintf(b,sizeof(b),"vis.s.blk.%d.%s",i,s); return (const char*)b; }; - w.n1w=f32(N("ln1.weight")); w.n1b=f32(N("ln1.bias")); w.n2w=f32(N("ln2.weight")); w.n2b=f32(N("ln2.bias")); - w.ls1=nullptr; w.ls2=nullptr; w.Wqkv=mm(N("qkv.weight")); w.bqkv=f32(N("qkv.bias")); - w.Wproj=mm(N("proj.weight")); w.bproj=f32(N("proj.bias")); w.Wfc1=mm(N("fc1.weight")); w.bfc1=f32(N("fc1.bias")); - w.Wfc2=mm(N("fc2.weight")); w.bfc2=f32(N("fc2.bias")); } - - m->pj_fc1w=mm("vis.proj.fc1.weight"); m->pj_fc1b=f32("vis.proj.fc1.bias"); - m->pj_fc2w=mm("vis.proj.fc2.weight"); m->pj_fc2b=f32("vis.proj.fc2.bias"); - m->pj_fc3w=mm("vis.proj.fc3.weight"); m->pj_fc3b=f32("vis.proj.fc3.bias"); + m->vis.declare(L, m->d_layers, m->s_layers); m->token_embd=mm("token_embd.weight"); m->action_queries=mm("action_queries.weight"); m->lm_out_norm=f32("lm.output_norm.weight"); m->lm.resize(m->lm_layers); @@ -263,17 +275,16 @@ std::unique_ptr vla_adapter_create(const std::string& mmproj_path w.flnw=f32(N("ffn_ln.weight")); w.flnb=f32(N("ffn_ln.bias")); w.flw=mm(N("ffn_lin.weight")); w.flb=f32(N("ffn_lin.bias")); std::vector gv=g.read_f32(N("gating")); w.rg = gv.empty()?0.0f:std::tanh(gv[0]); } (void)P; - if(!ok){ std::fprintf(stderr,"vla(vla_adapter): weight setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if(!m->weight_buf){ std::fprintf(stderr,"vla(vla_adapter): alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for(ggml_tensor*t=ggml_get_first_tensor(W); t; t=ggml_get_next_tensor(W,t)){ - std::vector bytes=g.read_convert(ggml_get_name(t),t->type); - if(bytes.empty()||bytes.size()!=ggml_nbytes(t)){ std::fprintf(stderr,"vla(vla_adapter): load %s (%zu vs %zu)\n",ggml_get_name(t),bytes.size(),ggml_nbytes(t)); return nullptr; } - ggml_backend_tensor_set(t,bytes.data(),0,bytes.size()); + if(!ok){ + std::fprintf(stderr,"vla(vla_adapter): weight setup failed\n"); + return nullptr; } + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + std::printf("vla(vla_adapter): weights resident %.2f GiB (%s) - DINOv2+SigLIP towers + Qwen2.5-0.5B + Bridge head\n", - ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), m->mt==GGML_TYPE_F32?"F32":"BF16"); + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), dtype_name(m->mt)); m->cfg.n_suffix = m->chunk; m->cfg.max_action_dim = m->action_dim; m->cfg.real_action_dim = m->action_dim; m->cfg.real_state_dim = m->proprio_dim; @@ -296,7 +307,7 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { if (n_views < 1) { std::fprintf(stderr, "vla(vla_adapter): need >=1 image view\n"); return {}; } if (!in.images) { std::fprintf(stderr, "vla(vla_adapter): n_images=%d but the images pointer is null\n", in.n_images); return {}; } // towers read S*S*3 per view; reject any view that is not exactly SxS. - for (int64_t v = 0; v < n_views; ++v) { + for (int64_t v=0; v VlaAdapterModelArch::predict(const Inputs& in) { for(int v=0; v dbuf, sbuf; @@ -338,7 +349,7 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { const int64_t NPROMPT = in.n_lang; // ggml_get_rows does not bound-check, so reject out-of-range tokens here. - for (int64_t i = 0; i < NPROMPT; ++i) + for (int64_t i=0; i= vocab) { std::fprintf(stderr, "vla(vla_adapter): token %d out of vocab\n", in.lang_tokens[i]); return {}; @@ -347,9 +358,9 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { std::fprintf(stderr, "vla(vla_adapter): stop_id %lld out of vocab\n", (long long) stop_id); return {}; } - const int64_t NUM_PROMPT_TOKENS = NPROMPT - 1; + const int64_t NUM_PROMPT_TOKENS = NPROMPT-1; const int64_t NPATCH = NP * n_views; - const int64_t SEQ = 1 + NPATCH + (NPROMPT-1) + num_tokens + 1; + const int64_t SEQ = 1+NPATCH+(NPROMPT-1)+num_tokens+1; const auto ti=clock::now(); // LM + action head graph depends only on the sequence layout. const MainKey mkey{ SEQ, n_views, NPROMPT }; @@ -357,7 +368,8 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { [&](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); + if(emb->type!=GGML_TYPE_F32) + emb=ggml_cast(C,emb,GGML_TYPE_F32); ggml_tensor*aqf=action_queries->type==GGML_TYPE_F32?action_queries:ggml_cast(C,action_queries,GGML_TYPE_F32); ggml_tensor*pre=ggml_cont(C,ggml_view_2d(C,emb,HC,NPROMPT,emb->nb[1],0)); ggml_tensor*stop=ggml_cont(C,ggml_view_2d(C,emb,HC,1,emb->nb[1],(NPROMPT+num_tokens)*emb->nb[1])); @@ -390,7 +402,8 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { ggml_tensor*final_norm=ggml_mul(C,ggml_rms_norm(C,lout[lm_layers-1],lm_rms_eps),lm_out_norm); std::vector cond(head_blocks); - for(int i=0;i VlaAdapterModelArch::predict(const Inputs& in) { 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 pp(SEQ); for(int64_t i=0;i pp(SEQ); + for(int64_t i=0;i 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 zx((size_t)action_dim*HC*chunk,0.0f); ggml_backend_tensor_set(t_x0,zx.data(),0,ggml_nbytes(t_x0)); } + { + std::vector zx((size_t)action_dim*HC*chunk,0.0f); + ggml_backend_tensor_set(t_x0,zx.data(),0,ggml_nbytes(t_x0)); + } auto fill_cs=[&](ggml_tensor*cc,ggml_tensor*ss,int64_t Lh){ std::vector cb(HD*Lh),sb(HD*Lh); const int64_t half=HD/2; - for(int64_t t=0;t #include @@ -42,9 +48,6 @@ namespace vla { namespace { -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; }; - } struct VlaJepaModelArch : public ModelArchBase { @@ -58,7 +61,9 @@ struct VlaJepaModelArch : public ModelArchBase { 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; } + 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; @@ -67,7 +72,9 @@ struct VlaJepaModelArch : public ModelArchBase { }; struct HeadKey { int64_t nsteps=-1; - bool operator==(const HeadKey & o) const { return nsteps==o.nsteps; } + 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; @@ -90,20 +97,14 @@ struct VlaJepaModelArch : public ModelArchBase { float vit_ln_eps=1e-6f, vit_rope_base=10000.0f, lm_rms_eps=1e-6f, lm_rope_base=5000000.0f, connector_ln_eps=1e-6f; float dit_ln_eps=1e-5f, dit_norm_out_eps=1e-6f; - ggml_tensor *vit_patch_w=nullptr,*vit_patch_b=nullptr,*vit_pos=nullptr; - std::vector vit; - MergerW deepstack[3]; - MergerW merger; - ggml_tensor *lm_output_norm=nullptr; - std::vector lm; + Qwen3VLTower vit; + Qwen3LM lm; + DitHead dit; ggml_tensor *ae_l1W=nullptr,*ae_l1b=nullptr,*ae_l2W=nullptr,*ae_l2b=nullptr,*ae_l3W=nullptr,*ae_l3b=nullptr; ggml_tensor *se_l1W=nullptr,*se_l1b=nullptr,*se_l2W=nullptr,*se_l2b=nullptr; ggml_tensor *ad_l1W=nullptr,*ad_l1b=nullptr,*ad_l2W=nullptr,*ad_l2b=nullptr; ggml_tensor *future_tokens=nullptr,*pos_embd=nullptr; - ggml_tensor *te_l1W=nullptr,*te_l1b=nullptr,*te_l2W=nullptr,*te_l2b=nullptr; - std::vector dit; - ggml_tensor *po1W=nullptr,*po1b=nullptr,*po2W=nullptr,*po2b=nullptr; bool caches_ready = false; std::vector c_grow, c_gcol; @@ -118,56 +119,7 @@ struct VlaJepaModelArch : public ModelArchBase { namespace { -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; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * hn = ggml_mul(C, ggml_rms_norm(C, h, m.lm_rms_eps), w.attn_norm); - ggml_tensor * qp = ggml_mul_mat(C, w.Wq, hn); - ggml_tensor * kp = ggml_mul_mat(C, w.Wk, hn); - ggml_tensor * vp = ggml_mul_mat(C, w.Wv, hn); - ggml_tensor * qh = ggml_reshape_3d(C, qp, hd, n_q, seq); - ggml_tensor * kh = ggml_reshape_3d(C, kp, hd, n_kv, seq); - ggml_tensor * vh = ggml_reshape_3d(C, vp, hd, n_kv, seq); - ggml_tensor * qn = ggml_mul(C, ggml_rms_norm(C, qh, m.lm_rms_eps), w.q_norm); - ggml_tensor * kn = ggml_mul(C, ggml_rms_norm(C, kh, m.lm_rms_eps), w.k_norm); - int sections[4] = { 24, 20, 20, 0 }; - ggml_tensor * qr = ggml_rope_multi(C, qn, positions, nullptr, (int) hd, sections, GGML_ROPE_TYPE_IMROPE, 0, m.lm_rope_base, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - ggml_tensor * kr = ggml_rope_multi(C, kn, positions, nullptr, (int) hd, sections, GGML_ROPE_TYPE_IMROPE, 0, m.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)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, kr, 0, 2, 1, 3)); - 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); - ggml_tensor * aw = ggml_soft_max_ext(C, kq, mask, 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)), hq, seq); - ggml_tensor * h_attn = ggml_add(C, h, ggml_mul_mat(C, w.Wo, att)); - ggml_tensor * hn2 = ggml_mul(C, ggml_rms_norm(C, h_attn, m.lm_rms_eps), w.ffn_norm); - ggml_tensor * gate = ggml_silu(C, ggml_mul_mat(C, w.Wgate, hn2)); - ggml_tensor * up = ggml_mul_mat(C, w.Wup, hn2); - return ggml_add(C, h_attn, ggml_mul_mat(C, w.Wdown, ggml_mul(C, gate, up))); -} -ggml_tensor * build_dit_block(ggml_context * C, const VlaJepaModelArch & m, const DitLayerW & w, - ggml_tensor * h, ggml_tensor * temb, ggml_tensor * enc) { - const int64_t hd = m.dit_head_dim, heads = m.dit_heads, dim = m.dit_hidden, Tk = h->ne[1]; - const float scale = 1.0f / std::sqrt((float) hd); - ggml_tensor * n = adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, m.dit_ln_eps); - ggml_tensor * kv = enc ? enc : n; - const int64_t Tkv = kv->ne[1]; - ggml_tensor * q = ggml_add(C, ggml_mul_mat(C, w.Wq, n), w.bq); - ggml_tensor * k = ggml_add(C, ggml_mul_mat(C, w.Wk, kv), w.bk); - ggml_tensor * v = ggml_add(C, ggml_mul_mat(C, w.Wv, kv), w.bv); - ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, q, hd, heads, Tk), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, k, hd, heads, Tkv), 0, 2, 1, 3)); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, v, hd, heads, Tkv), 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)), dim, Tk); - ggml_tensor * h1 = ggml_add(C, h, ggml_add(C, ggml_mul_mat(C, w.Wo, att), w.bo)); - ggml_tensor * n3 = ggml_norm(C, h1, m.dit_ln_eps); - ggml_tensor * ff = ggml_add(C, ggml_mul_mat(C, w.Wff2, ggml_gelu(C, ggml_add(C, ggml_mul_mat(C, w.Wff0, n3), w.bff0))), w.bff2); - return ggml_add(C, h1, ff); -} 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); }; @@ -187,16 +139,20 @@ bool load_config(const gguf_reader & g, VlaJepaModelArch & m, Config & cfg) { U(fk("num_future_tokens"), m.num_future); U(fk("num_inference_timesteps"), m.num_steps); U(fk("num_timestep_buckets"), m.num_buckets); 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(vla_jepa): VLA_NUM_STEPS override → num_steps=%lld\n", (long long) v); } + if (end && *end == '\0' && v >= 1) { + m.num_steps = (int64_t) v; + std::fprintf(stderr, "vla(vla_jepa): VLA_NUM_STEPS override → num_steps=%lld\n", (long long) v); + } } F(fk("vit_ln_eps"), m.vit_ln_eps); F(fk("lm_rms_eps"), m.lm_rms_eps); F(fk("connector_ln_eps"), m.connector_ln_eps); 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")); + 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) { + 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; @@ -207,8 +163,31 @@ bool load_config(const gguf_reader & g, VlaJepaModelArch & m, Config & cfg) { return false; } + m.lm.cfg.hidden = m.lm_hidden; + m.lm.cfg.layers = m.lm_layers; + m.lm.cfg.n_q = m.n_q; + m.lm.cfg.n_kv = m.n_kv; + m.lm.cfg.head_dim = m.lm_head_dim; + m.lm.cfg.inter = m.lm_inter; + m.lm.cfg.rms_eps = m.lm_rms_eps; + m.lm.cfg.flash_attn = flash_attn_enabled(); + m.lm.cfg.rope.type = GGML_ROPE_TYPE_IMROPE; + m.lm.cfg.rope.n_dims = (int) m.lm_head_dim; + m.lm.cfg.rope.freq_base = m.lm_rope_base; + m.lm.cfg.rope.sections[0] = 24; + m.lm.cfg.rope.sections[1] = 20; + m.lm.cfg.rope.sections[2] = 20; + m.lm.cfg.rope.sections[3] = 0; + + m.dit.cfg.hidden = m.dit_hidden; + m.dit.cfg.heads = m.dit_heads; + m.dit.cfg.head_dim = m.dit_head_dim; + m.dit.cfg.layers = m.dit_layers; + m.dit.cfg.ln_eps = m.dit_ln_eps; + m.dit.cfg.norm_out_eps = m.dit_norm_out_eps; + 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_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; cfg.n_suffix = m.action_horizon; cfg.max_state_dim = m.state_dim; cfg.max_action_dim = m.action_dim; cfg.real_state_dim = m.state_dim; cfg.real_action_dim = m.action_dim; @@ -225,25 +204,34 @@ bool load_config(const gguf_reader & g, VlaJepaModelArch & m, Config & cfg) { } VlaJepaModelArch::~VlaJepaModelArch() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + if (weight_buf) + ggml_backend_buffer_free(weight_buf); + if (ctx_weights) + ggml_free(ctx_weights); + if (backend) + ggml_backend_free(backend); } std::unique_ptr vla_jepa_create(const std::string& mmproj_path, const std::string& ckpt_path, - const std::string& ) { + const std::string&, + const Options& opts) { if (!mmproj_path.empty()) std::printf("vla(vla_jepa): note - mmproj '%s' is ignored (the vision tower is bundled in the combined GGUF)\n", mmproj_path.c_str()); auto m = std::make_unique(); m->gguf_path = ckpt_path; - m->matmul_type = std::getenv("VLA_JEPA_BF16_WEIGHTS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; + m->matmul_type = opts.weight_dtype.value_or(GGML_TYPE_BF16); gguf_reader g("vla_jepa"); - if (!g.open(ckpt_path)) return nullptr; - if (!g.has("vla_jepa.architecture")) { std::fprintf(stderr, "vla(vla_jepa): %s is not a vla_jepa GGUF\n", ckpt_path.c_str()); return nullptr; } - if (!load_config(g, *m, m->cfg)) return nullptr; + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("vla_jepa.architecture")) { + std::fprintf(stderr, "vla(vla_jepa): %s is not a vla_jepa GGUF\n", ckpt_path.c_str()); + return nullptr; + } + if (!load_config(g, *m, m->cfg)) + return nullptr; std::printf("vla(vla_jepa): vit=Qwen3-VL %lldd×%lldL (deepstack@{%lld,%lld,%lld}, merge÷%lld) lm=Qwen3-VL %lldd×%lldL (%lldq/%lldkv×%lld, θ=%g) " "dit-B %lldL×%lldh×%lld(inner %lld, cross %lld, out %lld) horizon=%lld action_dim=%lld state_dim=%lld future=%lld N_steps=%lld resident=%s\n", (long long) m->vit_hidden, (long long) m->vit_layers, (long long) m->deepstack_idx[0], (long long) m->deepstack_idx[1], (long long) m->deepstack_idx[2], (long long) m->spatial_merge, @@ -254,108 +242,63 @@ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, { const Backend b = backend_init("vla(vla_jepa)", m->n_threads); - if (!b.handle) { return nullptr; } + if (!b.handle) { + return nullptr; + } m->backend = b.handle; } - ggml_init_params wp = { (size_t) 32 * 1024 * 1024, nullptr, true }; + ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - if (!m->ctx_weights) { std::fprintf(stderr, "vla(vla_jepa): ggml_init(ctx_weights) failed\n"); return nullptr; } - ggml_context * W = m->ctx_weights; - auto mk = [&](const char * name, ggml_type type) -> ggml_tensor * { - const ggml_tensor * gt = g.meta(name); - if (!gt) { std::fprintf(stderr, "vla(vla_jepa): missing tensor %s\n", name); return nullptr; } - ggml_tensor * t = ggml_new_tensor(W, g.resident_type(gt, type), ggml_n_dims(gt), gt->ne); - ggml_set_name(t, name); return t; - }; - auto mk_mm = [&](const char * name) { return mk(name, m->matmul_type); }; - auto mk_f32 = [&](const char * name) { return mk(name, GGML_TYPE_F32); }; - - bool ok = true; - - m->vit_patch_w = mk_mm("vit.patch_embd.weight"); m->vit_patch_b = mk_f32("vit.patch_embd.bias"); m->vit_pos = mk_f32("vit.pos_embd"); - m->vit.resize(m->vit_layers); - for (int64_t i = 0; i < m->vit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->vit[i]; - w.ln1w=mk_f32(N("ln1.weight")); w.ln1b=mk_f32(N("ln1.bias")); w.ln2w=mk_f32(N("ln2.weight")); w.ln2b=mk_f32(N("ln2.bias")); - w.Wqkv=mk_mm(N("attn_qkv.weight")); w.bqkv=mk_f32(N("attn_qkv.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wfc1=mk_mm(N("fc1.weight")); w.bfc1=mk_f32(N("fc1.bias")); w.Wfc2=mk_mm(N("fc2.weight")); w.bfc2=mk_f32(N("fc2.bias")); - ok &= w.ln1w&&w.ln1b&&w.ln2w&&w.ln2b&&w.Wqkv&&w.bqkv&&w.Wo&&w.bo&&w.Wfc1&&w.bfc1&&w.Wfc2&&w.bfc2; - } - for (int j = 0; j < 3; ++j) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vit.deepstack.%d.%s", j, s); return p; }; - auto & w = m->deepstack[j]; - w.nw=mk_f32(N("norm.weight")); w.nb=mk_f32(N("norm.bias")); w.fc1w=mk_mm(N("fc1.weight")); w.fc1b=mk_f32(N("fc1.bias")); w.fc2w=mk_mm(N("fc2.weight")); w.fc2b=mk_f32(N("fc2.bias")); - ok &= w.nw&&w.nb&&w.fc1w&&w.fc1b&&w.fc2w&&w.fc2b; - } - { auto & w = m->merger; - w.nw=mk_f32("vit.merger.norm.weight"); w.nb=mk_f32("vit.merger.norm.bias"); w.fc1w=mk_mm("vit.merger.fc1.weight"); w.fc1b=mk_f32("vit.merger.fc1.bias"); w.fc2w=mk_mm("vit.merger.fc2.weight"); w.fc2b=mk_f32("vit.merger.fc2.bias"); - ok &= w.nw&&w.nb&&w.fc1w&&w.fc1b&&w.fc2w&&w.fc2b; } - - m->lm_output_norm = mk_f32("vlm.output_norm.weight"); - m->lm.resize(m->lm_layers); - for (int64_t i = 0; i < m->lm_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "vlm.blk.%lld.%s", (long long) i, s); return p; }; - auto & w = m->lm[i]; - w.attn_norm=mk_f32(N("attn_norm.weight")); - w.Wq=mk_mm(N("attn_q.weight")); w.Wk=mk_mm(N("attn_k.weight")); w.Wv=mk_mm(N("attn_v.weight")); w.Wo=mk_mm(N("attn_o.weight")); - w.q_norm=mk_f32(N("attn_q_norm.weight")); w.k_norm=mk_f32(N("attn_k_norm.weight")); w.ffn_norm=mk_f32(N("ffn_norm.weight")); - w.Wgate=mk_mm(N("ffn_gate.weight")); w.Wup=mk_mm(N("ffn_up.weight")); w.Wdown=mk_mm(N("ffn_down.weight")); - ok &= w.attn_norm&&w.Wq&&w.Wk&&w.Wv&&w.Wo&&w.q_norm&&w.k_norm&&w.ffn_norm&&w.Wgate&&w.Wup&&w.Wdown; + if (!m->ctx_weights) { + std::fprintf(stderr, "vla(vla_jepa): ggml_init(ctx_weights) failed\n"); + return nullptr; } - m->ae_l1W=mk_f32("ah.act_enc.l1.weight"); m->ae_l1b=mk_f32("ah.act_enc.l1.bias"); - m->ae_l2W=mk_f32("ah.act_enc.l2.weight"); m->ae_l2b=mk_f32("ah.act_enc.l2.bias"); - m->ae_l3W=mk_f32("ah.act_enc.l3.weight"); m->ae_l3b=mk_f32("ah.act_enc.l3.bias"); - m->se_l1W=mk_f32("ah.state_enc.l1.weight"); m->se_l1b=mk_f32("ah.state_enc.l1.bias"); - m->se_l2W=mk_f32("ah.state_enc.l2.weight"); m->se_l2b=mk_f32("ah.state_enc.l2.bias"); - m->ad_l1W=mk_f32("ah.act_dec.l1.weight"); m->ad_l1b=mk_f32("ah.act_dec.l1.bias"); - m->ad_l2W=mk_f32("ah.act_dec.l2.weight"); m->ad_l2b=mk_f32("ah.act_dec.l2.bias"); - m->future_tokens=mk_f32("ah.future_tokens"); m->pos_embd=mk_f32("ah.pos_embd"); - m->te_l1W=mk_mm("ah.time_emb.l1.weight"); m->te_l1b=mk_f32("ah.time_emb.l1.bias"); - m->te_l2W=mk_mm("ah.time_emb.l2.weight"); m->te_l2b=mk_f32("ah.time_emb.l2.bias"); - m->dit.resize(m->dit_layers); - for (int64_t i = 0; i < m->dit_layers && ok; ++i) { - char p[64]; auto N = [&](const char * s) { std::snprintf(p, sizeof(p), "ah.dit.%lld.%s", (long long) i, s); return p; }; - auto & w = m->dit[i]; - w.adaln_w=mk_mm(N("adaln.weight")); w.adaln_b=mk_f32(N("adaln.bias")); - w.Wq=mk_mm(N("attn_q.weight")); w.bq=mk_f32(N("attn_q.bias")); w.Wk=mk_mm(N("attn_k.weight")); w.bk=mk_f32(N("attn_k.bias")); - w.Wv=mk_mm(N("attn_v.weight")); w.bv=mk_f32(N("attn_v.bias")); w.Wo=mk_mm(N("attn_o.weight")); w.bo=mk_f32(N("attn_o.bias")); - w.Wff0=mk_mm(N("ff0.weight")); w.bff0=mk_f32(N("ff0.bias")); w.Wff2=mk_mm(N("ff2.weight")); w.bff2=mk_f32(N("ff2.bias")); - ok &= w.adaln_w&&w.adaln_b&&w.Wq&&w.bq&&w.Wk&&w.bk&&w.Wv&&w.bv&&w.Wo&&w.bo&&w.Wff0&&w.bff0&&w.Wff2&&w.bff2; - } - m->po1W=mk_mm("ah.proj_out1.weight"); m->po1b=mk_f32("ah.proj_out1.bias"); m->po2W=mk_mm("ah.proj_out2.weight"); m->po2b=mk_f32("ah.proj_out2.bias"); - ok &= m->vit_patch_w&&m->vit_patch_b&&m->vit_pos&&m->lm_output_norm&&m->ae_l1W&&m->ae_l2W&&m->ae_l3W&&m->se_l1W&&m->se_l2W&&m->ad_l1W&&m->ad_l2W&& - m->future_tokens&&m->pos_embd&&m->te_l1W&&m->te_l1b&&m->te_l2W&&m->te_l2b&&m->po1W&&m->po1b&&m->po2W&&m->po2b; - if (!ok) { std::fprintf(stderr, "vla(vla_jepa): weight tensor setup failed\n"); return nullptr; } - - m->weight_buf = ggml_backend_alloc_ctx_tensors(m->ctx_weights, m->backend); - if (!m->weight_buf) { std::fprintf(stderr, "vla(vla_jepa): ggml_backend_alloc_ctx_tensors failed (OOM?)\n"); return nullptr; } - for (ggml_tensor * t = ggml_get_first_tensor(W); t; t = ggml_get_next_tensor(W, t)) { - std::vector bytes = g.read_convert(ggml_get_name(t), t->type); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { - std::fprintf(stderr, "vla(vla_jepa): failed to load %s (%zu vs %zu bytes)\n", ggml_get_name(t), bytes.size(), ggml_nbytes(t)); return nullptr; - } - ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); + WeightLoader L("vla_jepa", g, m->ctx_weights, m->matmul_type); + + m->vit.declare(L, "vit", m->vit_layers); + m->lm.declare(L, "vlm"); + + m->ae_l1W = L.f32("ah.act_enc.l1.weight"); m->ae_l1b = L.f32("ah.act_enc.l1.bias"); + m->ae_l2W = L.f32("ah.act_enc.l2.weight"); m->ae_l2b = L.f32("ah.act_enc.l2.bias"); + m->ae_l3W = L.f32("ah.act_enc.l3.weight"); m->ae_l3b = L.f32("ah.act_enc.l3.bias"); + m->se_l1W = L.f32("ah.state_enc.l1.weight"); m->se_l1b = L.f32("ah.state_enc.l1.bias"); + m->se_l2W = L.f32("ah.state_enc.l2.weight"); m->se_l2b = L.f32("ah.state_enc.l2.bias"); + m->ad_l1W = L.f32("ah.act_dec.l1.weight"); m->ad_l1b = L.f32("ah.act_dec.l1.bias"); + m->ad_l2W = L.f32("ah.act_dec.l2.weight"); m->ad_l2b = L.f32("ah.act_dec.l2.bias"); + m->future_tokens = L.f32("ah.future_tokens"); + m->pos_embd = L.f32("ah.pos_embd"); + + m->dit.declare(L, "ah.dit", false, false, "ah"); + + if (!L.upload(m->backend, &m->weight_buf)) + return nullptr; + + std::printf("vla(vla_jepa): weights resident in %.2f GiB (%s)\n", + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), dtype_name(m->matmul_type)); + if (!m->build_caches()) { + std::fprintf(stderr, "vla(vla_jepa): build_caches failed\n"); + return nullptr; } - std::printf("vla(vla_jepa): weights resident in %.2f GiB (%s) - Qwen3-VL backbone + deepstack + DiT-B head\n", - ggml_backend_buffer_get_size(m->weight_buf) / (1024.0 * 1024.0 * 1024.0), m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); - if (!m->build_caches()) { std::fprintf(stderr, "vla(vla_jepa): build_caches failed\n"); return nullptr; } return m; } bool VlaJepaModelArch::build_caches() { - if (caches_ready) return true; + if (caches_ready) + return true; const int64_t side = image_target_size, ps = patch_size, m2 = spatial_merge; - const int64_t grid = side / ps; - const int64_t hd_vit = vit_hidden / vit_heads; + const int64_t grid = side/ps; + const int64_t hd_vit = vit_hidden/vit_heads; const int64_t num_side = (int64_t) std::lround(std::sqrt((double) vit_num_pos)); merge_block_coords(grid, grid, m2, c_grow, c_gcol); vit_rope_tables(c_grow, c_gcol, hd_vit, (double) vit_rope_base, c_rope_cos, c_rope_sin); - if (!io.open(gguf_path)) { std::fprintf(stderr, "vla(vla_jepa): build_caches: io.open(%s) failed\n", gguf_path.c_str()); return false; } + if (!io.open(gguf_path)) { + std::fprintf(stderr, "vla(vla_jepa): build_caches: io.open(%s) failed\n", gguf_path.c_str()); + return false; + } std::vector pos_table = io.read_f32("vit.pos_embd"); if (pos_table.empty() || (int64_t) pos_table.size() != vit_num_pos * vit_hidden) { std::fprintf(stderr, "vla(vla_jepa): build_caches: vit.pos_embd unreadable\n"); return false; @@ -363,8 +306,8 @@ bool VlaJepaModelArch::build_caches() { interp_pos_embed(pos_table, num_side, vit_hidden, c_grow, c_gcol, grid, grid, c_pos_interp); c_tau.assign((size_t) num_steps, {}); c_tproj.assign((size_t) num_steps, {}); - for (int64_t s = 0; s < num_steps; ++s) { - const int64_t bucket = (int64_t) ((double) s / (double) num_steps * (double) num_buckets); + for (int64_t s=0; s VlaJepaModelArch::predict(const Inputs& in) { const int64_t H = lm_hidden, E = dit_hidden, AD = action_dim, AH = action_horizon, OUTD = output_dim; const int64_t side = image_target_size, ps = patch_size, m2 = spatial_merge; - const int64_t grid = side / ps, n_patches = grid * grid, K = (grid / m2) * (grid / m2); - const int64_t hd_vit = vit_hidden / vit_heads; - const int64_t Nseq = 1 + num_future + AH; + const int64_t grid = side/ps, n_patches = grid * grid, K = (grid/m2)*(grid/m2); + const int64_t hd_vit = vit_hidden/vit_heads; + const int64_t Nseq = 1+num_future+AH; const char * dump_prefix = std::getenv("VLA_JEPA_DUMP"); if (!caches_ready) { std::fprintf(stderr, "vla(vla_jepa): caches not ready\n"); return {}; } auto dump_t = [&](const char * name, ggml_tensor * t) { - if (!dump_prefix) return; + if (!dump_prefix) + return; const int64_t n0 = t->ne[0], n1 = t->ne[1]; - std::vector buf((size_t) n0 * std::max(1, n1)); - ggml_backend_tensor_get(t, buf.data(), 0, buf.size() * sizeof(float)); + std::vector buf((size_t) n0*std::max(1, n1)); + ggml_backend_tensor_get(t, buf.data(), 0, buf.size()*sizeof(float)); char path[1024]; std::snprintf(path, sizeof(path), "%s_%s_%lldx%lld.f32", dump_prefix, name, (long long) n0, (long long) n1); - FILE * fp = std::fopen(path, "wb"); if (fp) { std::fwrite(buf.data(), sizeof(float), buf.size(), fp); std::fclose(fp); } + FILE * fp = std::fopen(path, "wb"); if (fp) { + std::fwrite(buf.data(), sizeof(float), buf.size(), fp); + std::fclose(fp); + } }; std::vector x_init((size_t) AH * AD); - 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); } + 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); + } std::vector cond_host((size_t) H * num_future, 0.0f); const char * cond_file = std::getenv("VLA_JEPA_COND"); @@ -415,7 +368,8 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { int64_t n_views = in.n_images; if (n_views <= 0) { std::fprintf(stderr, "vla(vla_jepa): no images in the request\n"); return {}; } std::vector img_emb_host((size_t) n_views * K * H), ds_host[3]; - for (int j = 0; j < 3; ++j) ds_host[j].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); std::vector inj_patches; const char * patches_file = std::getenv("VLA_JEPA_PATCHES"); if (patches_file) { @@ -428,59 +382,77 @@ 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_context * VC = vision_scratch.reset((size_t) 512 * 1024 * 1024); + 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); ggml_tensor * t_cos = ggml_new_tensor_2d(VC, GGML_TYPE_F32, hd_vit, n_patches); ggml_set_input(t_cos); ggml_tensor * t_sin = ggml_new_tensor_2d(VC, GGML_TYPE_F32, hd_vit, n_patches); ggml_set_input(t_sin); - ggml_tensor * h = ggml_add(VC, ggml_add(VC, ggml_mul_mat(VC, vit_patch_w, t_patches), vit_patch_b), t_pos); + ggml_tensor * h = ggml_add(VC, ggml_add(VC, ggml_mul_mat(VC, vit.patch_w, t_patches), vit.patch_b), t_pos); ggml_set_output(h); ggml_tensor * stash[3] = {nullptr, nullptr, nullptr}; - for (int64_t i = 0; i < vit_layers; ++i) { - h = build_vit_layer(VC, vit[i], h, t_cos, t_sin, n_patches, vit_heads, hd_vit, vit_hidden, vit_ln_eps); + for (int64_t i=0; i patches; bool vok = true; - for (int64_t v = 0; v < n_views && vok; ++v) { + for (int64_t v=0; v(std::chrono::steady_clock::now() - tv0).count(); + stats.ms_vision = std::chrono::duration(std::chrono::steady_clock::now()-tv0).count(); if (!vok) return {}; const int64_t n_img = n_views * K; std::vector input_ids; int64_t n_img_slots = 0; - for (int j = 0; j < in.n_lang; ++j) if (in.lang_tokens[j] == (int32_t) image_token_index) ++n_img_slots; + for (int j=0; j VlaJepaModelArch::predict(const Inputs& in) { std::vector inputs_embeds((size_t) SEQ * H); 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) { std::memcpy(inputs_embeds.data() + p * H, img_emb_host.data() + k * H, H * sizeof(float)); ++k; } } + { int64_t k = 0; for (int64_t p=0; p image_pos_idx, emb_pos_idx; - for (int64_t p = 0; p < SEQ; ++p) { - if (input_ids[p] == (int32_t) image_token_index) image_pos_idx.push_back((int32_t) p); - if (input_ids[p] == (int32_t) embodied_token_id) emb_pos_idx.push_back((int32_t) p); + for (int64_t p=0; p> 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)); } + for (int j=0; j<3; ++j) { + ds_pad[j].assign((size_t) SEQ * H, 0.0f); + for (int64_t k=0; k 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_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); ggml_tensor * t_emb_idx= ggml_new_tensor_1d(C, GGML_TYPE_I32, num_future); ggml_set_input(t_emb_idx); ggml_tensor * t_ds[3]; - for (int j = 0; j < 3; ++j) { t_ds[j] = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); ggml_set_input(t_ds[j]); } + for (int j=0; j<3; ++j) { + t_ds[j] = ggml_new_tensor_2d(C, GGML_TYPE_F32, H, SEQ); + ggml_set_input(t_ds[j]); + } ggml_tensor * hh = t_embeds; - for (int64_t i = 0; i < lm_layers; ++i) { - hh = build_qwen3_layer(C, *this, lm[i], hh, t_pos2, t_lmmask, SEQ); - if (i < 3) hh = ggml_add(C, hh, t_ds[i]); + for (int64_t i=0; i VlaJepaModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(t_embeds, inputs_embeds.data(), 0, ggml_nbytes(t_embeds)); { - const int64_t llm_grid = side / ps / m2; - std::vector pp((size_t) 4 * SEQ, 0); + const int64_t llm_grid = side/ps/m2; + std::vector pp((size_t) 4*SEQ, 0); int64_t st = 0, st_idx = 0; while (st < SEQ) { int64_t img_start = -1; - for (int64_t i = st; i < SEQ; ++i) if (input_ids[i] == (int32_t) image_token_index) { img_start = i; break; } + for (int64_t i=st; i max_image_pos) max_image_pos = llm_grid - 1; - st_idx = image_offset + max_image_pos + 1; st = img_end; + int64_t max_image_pos = this_t-1; if (llm_grid-1 > max_image_pos) max_image_pos = llm_grid-1; + st_idx = image_offset+max_image_pos+1; st = img_end; } - std::memcpy(pp.data() + (size_t) 3 * SEQ, pp.data(), (size_t) SEQ * sizeof(int32_t)); + 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) { build_causal_mask(SEQ, c_mask); 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])); + 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"); 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)); + 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)); } // Dumping adds graph outputs, so it always rebuilds. std::vector step_seq, step_pred, step_vel, step_act; - if (dump_prefix) head_graph.release(); + 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, + 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); std::vector t_tau(num_steps), t_tproj(num_steps); - for (int64_t s = 0; s < num_steps; ++s) { t_tau[s] = ggml_new_tensor_2d(C, GGML_TYPE_F32, E, AH); ggml_set_input(t_tau[s]); t_tproj[s] = ggml_new_tensor_1d(C, GGML_TYPE_F32, time_proj_dim); ggml_set_input(t_tproj[s]); } + for (int64_t s=0; s VlaJepaModelArch::predict(const Inputs& in) { ggml_tensor * seq = ggml_concat(C, ggml_concat(C, state_features, future, 1), af, 1); step_seq[s] = seq; ggml_tensor * x = seq; - for (int64_t i = 0; i < dit_layers; ++i) { - ggml_tensor * enc = (i % 2 == 0) ? t_cond : nullptr; - x = build_dit_block(C, *this, dit[i], x, temb, enc); + for (int64_t i=0; inb[1], (size_t) (Nseq - AH) * model_output->nb[1])); + ggml_tensor * last = ggml_cont(C, ggml_view_2d(C, model_output, OUTD, AH, model_output->nb[1], (size_t) (Nseq-AH)*model_output->nb[1])); ggml_tensor * vel = ggml_add(C, ggml_mul_mat(C, ad_l2W, ggml_relu(C, ggml_add(C, ggml_mul_mat(C, ad_l1W, last), ad_l1b))), ad_l2b); step_vel[s] = vel; actions = ggml_add(C, actions, ggml_scale(C, vel, dt)); step_act[s] = actions; - 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]); } + 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; @@ -630,7 +641,12 @@ std::vector VlaJepaModelArch::predict(const Inputs& in) { 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]); } + if (dump_prefix) for (int64_t s=0; s VlaJepaModelArch::predict(const Inputs& in) { 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)); } + { + std::vector st(state_dim, 0.0f); + for (int64_t i=0; i(std::chrono::steady_clock::now() - td0).count(); - stats.ms_inference = stats.ms_prefill + stats.ms_denoise; + stats.ms_denoise = std::chrono::duration(std::chrono::steady_clock::now()-td0).count(); + stats.ms_inference = stats.ms_prefill+stats.ms_denoise; - if (dump_prefix) for (int64_t s = 0; s < num_steps; ++s) { + if (dump_prefix) for (int64_t s=0; s VlaJepaModelArch::predict(const Inputs& in) { } std::vector out((size_t) AH * AD); - ggml_backend_tensor_get(actions, out.data(), 0, out.size() * sizeof(float)); - stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + ggml_backend_tensor_get(actions, out.data(), 0, out.size()*sizeof(float)); + stats.ms_total = std::chrono::duration(std::chrono::steady_clock::now()-t0).count(); return out; } diff --git a/src/modules/action_expert.cpp b/src/modules/action_expert.cpp new file mode 100644 index 0000000..0f21fca --- /dev/null +++ b/src/modules/action_expert.cpp @@ -0,0 +1,60 @@ +// 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. + +#include "modules/action_expert.h" + +#include "layers/linear.h" + +namespace vla { + +void ActionExpert::declare(WeightLoader & L, const char * prefix) { + se_l1W = L.f32("%s.state_enc.l1.W", prefix); + se_l1b = L.f32("%s.state_enc.l1.b", prefix); + se_l2W = L.f32("%s.state_enc.l2.W", prefix); + se_l2b = L.f32("%s.state_enc.l2.b", prefix); + + ae_W1W = L.f32("%s.act_enc.W1.W", prefix); + ae_W1b = L.f32("%s.act_enc.W1.b", prefix); + ae_W2W = L.f32("%s.act_enc.W2.W", prefix); + ae_W2b = L.f32("%s.act_enc.W2.b", prefix); + ae_W3W = L.f32("%s.act_enc.W3.W", prefix); + ae_W3b = L.f32("%s.act_enc.W3.b", prefix); + + ad_l1W = L.f32("%s.act_dec.l1.W", prefix); + ad_l1b = L.f32("%s.act_dec.l1.b", prefix); + ad_l2W = L.f32("%s.act_dec.l2.W", prefix); + ad_l2b = L.f32("%s.act_dec.l2.b", prefix); + + pos_embd = L.f32("%s.pos_embd", prefix); +} + +ggml_tensor * ActionExpert::encode_state(ggml_context * C, ggml_tensor * state) const { + ggml_tensor * h = ggml_relu(C, cat_linear(C, se_l1W, se_l1b, embodiment_id, state)); + return cat_linear(C, se_l2W, se_l2b, embodiment_id, h); +} + +ggml_tensor * ActionExpert::encode_action(ggml_context * C, ggml_tensor * actions, ggml_tensor * tau, + int64_t embed_dim, int64_t horizon) const { + ggml_tensor * a_emb = cat_linear(C, ae_W1W, ae_W1b, embodiment_id, actions); + ggml_tensor * x_w2 = ggml_silu(C, cat_linear(C, ae_W2W, ae_W2b, embodiment_id, ggml_concat(C, a_emb, tau, 0))); + ggml_tensor * pos = ggml_view_2d(C, pos_embd, embed_dim, horizon, pos_embd->nb[1], 0); + return ggml_add(C, cat_linear(C, ae_W3W, ae_W3b, embodiment_id, x_w2), pos); +} + +ggml_tensor * ActionExpert::decode(ggml_context * C, ggml_tensor * model_out) const { + ggml_tensor * h = ggml_relu(C, cat_linear(C, ad_l1W, ad_l1b, embodiment_id, model_out)); + return cat_linear(C, ad_l2W, ad_l2b, embodiment_id, h); +} + +} diff --git a/src/modules/action_expert.h b/src/modules/action_expert.h new file mode 100644 index 0000000..0804a74 --- /dev/null +++ b/src/modules/action_expert.h @@ -0,0 +1,46 @@ +// 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. + +// Every projection is a cat_linear row selected by embodiment_id. + +#pragma once + +#include "loader.h" + +#include "ggml.h" + +#include + +namespace vla { + +struct ActionExpert { + ggml_tensor *se_l1W = nullptr, *se_l1b = nullptr, *se_l2W = nullptr, *se_l2b = nullptr; + ggml_tensor *ae_W1W = nullptr, *ae_W1b = nullptr, *ae_W2W = nullptr, *ae_W2b = nullptr; + ggml_tensor *ae_W3W = nullptr, *ae_W3b = nullptr; + ggml_tensor *ad_l1W = nullptr, *ad_l1b = nullptr, *ad_l2W = nullptr, *ad_l2b = nullptr; + ggml_tensor *pos_embd = nullptr; + + int64_t embodiment_id = 0; + + void declare(WeightLoader & L, const char * prefix); + + ggml_tensor * encode_state(ggml_context * C, ggml_tensor * state) const; + + ggml_tensor * encode_action(ggml_context * C, ggml_tensor * actions, ggml_tensor * tau, + int64_t embed_dim, int64_t horizon) const; + + ggml_tensor * decode(ggml_context * C, ggml_tensor * model_out) const; +}; + +} diff --git a/src/modules/dit_head.cpp b/src/modules/dit_head.cpp new file mode 100644 index 0000000..f081288 --- /dev/null +++ b/src/modules/dit_head.cpp @@ -0,0 +1,149 @@ +// 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. + +#include "modules/dit_head.h" + +#include "layers/attn.h" +#include "layers/ffn.h" +#include "layers/linear.h" +#include "layers/norm.h" + +#include +#include + +namespace vla { + +void DitHead::declare(WeightLoader & L, const char * prefix, bool fuse_qkv, bool interleave, const char * outer) { + if (!outer) + outer = prefix; + + te_l1W = L.gemm("%s.time_emb.l1.weight", outer); + te_l1b = L.f32 ("%s.time_emb.l1.bias", outer); + te_l2W = L.gemm("%s.time_emb.l2.weight", outer); + te_l2b = L.f32 ("%s.time_emb.l2.bias", outer); + + blk.resize(cfg.layers); + for (int64_t i=0; ine[1]; + + if (w.Wkv) { + ggml_tensor * kvp = linear(C, w.Wkv, w.bkv, src); + *K_out = ggml_cont(C, ggml_permute(C, head_view(C, kvp, hd, heads, Tkv, cfg.hidden, 2, 0), 0, 2, 1, 3)); + *V_out = ggml_cont(C, ggml_permute(C, head_view(C, kvp, hd, heads, Tkv, cfg.hidden, 2, 1), 1, 2, 0, 3)); + return; + } + *K_out = to_heads (C, linear(C, w.Wk, w.bk, src), hd, heads, Tkv); + *V_out = to_heads_v(C, linear(C, w.Wv, w.bv, src), hd, heads, Tkv); +} + +ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * temb, + ggml_tensor * enc, ggml_tensor * K_pre, ggml_tensor * V_pre) const { + const int64_t hd = cfg.head_dim; + const int64_t heads = cfg.heads; + const int64_t dim = cfg.hidden; + const int64_t Tk = h->ne[1]; + const float scale = 1.0f/std::sqrt((float)hd); + + ggml_tensor * n = adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, cfg.ln_eps); + ggml_tensor *Q, *K, *V; + if (!enc && w.Wqkv) { + ggml_tensor * qkv = linear(C, w.Wqkv, w.bqkv, n); + Q = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 0), 0, 2, 1, 3)); + K = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 1), 0, 2, 1, 3)); + V = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 2), 1, 2, 0, 3)); + } else { + Q = to_heads(C, linear(C, w.Wq, w.bq, n), hd, heads, Tk); + if (K_pre) { + K = K_pre; + V = V_pre; + } + else { + kv(C, w, enc ? enc : n, &K, &V); + } + } + + ggml_tensor * att = attention(C, Q, K, V, nullptr, scale, dim, Tk); + ggml_tensor * h1 = ggml_add(C, h, linear(C, w.Wo, w.bo, att)); + ggml_tensor * n3 = ggml_norm(C, h1, cfg.ln_eps); + return ggml_add(C, h1, ffn_gelu(C, w.Wff0, w.bff0, w.Wff2, w.bff2, n3)); +} + +ggml_tensor * DitHead::time_emb(ggml_context * C, ggml_tensor * tproj) const { + return linear(C, te_l2W, te_l2b, ggml_silu(C, linear(C, te_l1W, te_l1b, tproj))); +} + +ggml_tensor * DitHead::proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * temb) const { + ggml_tensor * po = linear(C, po1W, po1b, ggml_silu(C, temb)); + ggml_tensor * sh = ggml_view_1d(C, po, cfg.hidden, 0); + ggml_tensor * sc = ggml_view_1d(C, po, cfg.hidden, (size_t)cfg.hidden*sizeof(float)); + + ggml_tensor * hn = ggml_norm(C, h, cfg.norm_out_eps); + ggml_tensor * h_mod = ggml_add(C, ggml_add(C, hn, ggml_mul(C, hn, sc)), sh); + return linear(C, po2W, po2b, h_mod); +} + +} diff --git a/src/modules/dit_head.h b/src/modules/dit_head.h new file mode 100644 index 0000000..6c00b2c --- /dev/null +++ b/src/modules/dit_head.h @@ -0,0 +1,70 @@ +// 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. + +// A block cross-attends when `enc` is non-null and self-attends otherwise; the +// per-layer choice is the model's. Cross-attention K/V depend only on `enc`, so +// a model hoists them out of the solver loop and passes them via K_pre/V_pre. + +#pragma once + +#include "loader.h" + +#include "ggml.h" + +#include +#include + +namespace vla { + +struct DitLayerW { + ggml_tensor *adaln_w, *adaln_b; + ggml_tensor *Wq, *bq, *Wk, *bk, *Wv, *bv, *Wo, *bo; + ggml_tensor *Wff0, *bff0, *Wff2, *bff2; + + ggml_tensor *Wqkv = nullptr, *bqkv = nullptr, *Wkv = nullptr, *bkv = nullptr; +}; + +struct DitCfg { + int64_t hidden = 1536; + int64_t heads = 32; + int64_t head_dim = 48; + int64_t layers = 16; + float ln_eps = 1e-5f; + float norm_out_eps = 1e-6f; +}; + +struct DitHead { + DitCfg cfg; + std::vector blk; + ggml_tensor *te_l1W = nullptr, *te_l1b = nullptr, *te_l2W = nullptr, *te_l2b = nullptr; + ggml_tensor *po1W = nullptr, *po1b = nullptr, *po2W = nullptr, *po2b = nullptr; + + // outer names time_emb and proj_out when they do not sit under the block + // prefix; null means they do. + void declare(WeightLoader & L, const char * prefix, bool fuse_qkv = false, bool interleave = false, + const char * outer = nullptr); + + void kv(ggml_context * C, const DitLayerW & w, ggml_tensor * src, + ggml_tensor ** K_out, ggml_tensor ** V_out) const; + + ggml_tensor * block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * temb, + ggml_tensor * enc, ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) const; + + ggml_tensor * time_emb(ggml_context * C, ggml_tensor * tproj) const; + + // (shift, scale) adaLN, opposite to layers/norm.h adaln. + ggml_tensor * proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * temb) const; +}; + +} diff --git a/src/models/dual_tower.h b/src/modules/dual_tower.h similarity index 56% rename from src/models/dual_tower.h rename to src/modules/dual_tower.h index 20bc40a..ffee524 100644 --- a/src/models/dual_tower.h +++ b/src/modules/dual_tower.h @@ -19,6 +19,7 @@ #pragma once #include "ggml.h" +#include "loader.h" #include "model.h" #include @@ -29,7 +30,57 @@ 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); } +// DINOv2 carries CLS + 4 register tokens and LayerScale; SigLIP carries +// neither, so its ls1/ls2 stay null and vit_block is told to skip them. +struct DualTower { + ggml_tensor *d_patch_w=nullptr,*d_patch_b=nullptr,*d_cls=nullptr,*d_reg=nullptr,*d_pos=nullptr; + ggml_tensor *s_patch_w=nullptr,*s_patch_b=nullptr,*s_pos=nullptr; + std::vector dvit, svit; + ggml_tensor *pj_fc1w=nullptr,*pj_fc1b=nullptr,*pj_fc2w=nullptr,*pj_fc2b=nullptr,*pj_fc3w=nullptr,*pj_fc3b=nullptr; + + void declare(WeightLoader & L, int64_t d_layers, int64_t s_layers) { + auto blocks = [&](std::vector & v, const char * pre, int64_t n, bool layer_scale) { + v.resize(n); + for (int64_t i=0; inb[1],nprefix*x->nb[1])); + 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; } diff --git a/src/modules/encoder.cpp b/src/modules/encoder.cpp new file mode 100644 index 0000000..f9dc636 --- /dev/null +++ b/src/modules/encoder.cpp @@ -0,0 +1,81 @@ +// 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. + +#include "modules/encoder.h" + +#include "layers/attn.h" +#include "layers/ffn.h" +#include "layers/linear.h" +#include "layers/norm.h" + +#include + +namespace vla { + +void EncStack::declare(WeightLoader & L, const char * prefix, int64_t layers, const EncNames & n) { + blk.resize(layers); + + for (int64_t i=0; i +#include + +namespace vla { + +struct EncBlockW { + ggml_tensor *ln1w, *ln1b, *ln2w, *ln2b; + ggml_tensor *Wq, *bq, *Wk, *bk, *Wv, *bv, *Wo, *bo; + ggml_tensor *Wfc1, *bfc1, *Wfc2, *bfc2; +}; + +struct EncNames { + const char * ln1 = "ln1"; + const char * ln2 = "ln2"; + const char * fc1 = "fc1"; + const char * fc2 = "fc2"; +}; + +struct EncCfg { + int64_t hidden = 0; + int64_t heads = 0; + int64_t head_dim = 0; + float ln_eps = 1e-6f; + bool flash_attn = false; +}; + +struct EncStack { + EncCfg cfg; + std::vector blk; + + void declare(WeightLoader & L, const char * prefix, int64_t layers, const EncNames & n = EncNames{}); + + ggml_tensor * block(ggml_context * C, const EncBlockW & w, ggml_tensor * x, + int64_t seq, int64_t nv = 1) const; + + ggml_tensor * build(ggml_context * C, ggml_tensor * x, int64_t seq, int64_t nv = 1) const; +}; + +} diff --git a/src/modules/gemma_expert.h b/src/modules/gemma_expert.h new file mode 100644 index 0000000..33e1001 --- /dev/null +++ b/src/modules/gemma_expert.h @@ -0,0 +1,65 @@ +// 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. + +// Gemma decoder stack. pi0, pi0.5 and SmolVLA each run two of these with a +// shared attention: a prefix tower over the image and language tokens, and an +// action expert over the state and noisy-action tokens. + +#pragma once + +#include "loader.h" + +#include "ggml.h" + +#include +#include + +namespace vla { + +struct GemmaLayerW { + ggml_tensor * ln_in = nullptr; + ggml_tensor * Wq = nullptr; + ggml_tensor * Wk = nullptr; + ggml_tensor * Wv = nullptr; + ggml_tensor * Wo = nullptr; + ggml_tensor * ln_post = nullptr; + ggml_tensor * Wgate = nullptr; + ggml_tensor * Wup = nullptr; + ggml_tensor * Wdown = nullptr; +}; + +struct GemmaStack { + std::vector blk; + ggml_tensor * output_norm = nullptr; + + void declare(WeightLoader & L, const char * prefix, int64_t layers, bool with_output_norm) { + blk.resize(layers); + for (int64_t i=0; i +#include +#include +#include + +namespace vla { + +// A tower that reads side*side*3 from a view needs the view to be exactly that +// size with real data, else it runs past the buffer. +inline bool view_is_side(const void * data, int w, int h, int64_t side) { + return data != nullptr && (int64_t) w == side && (int64_t) h == side; +} + +// IDEFICS3/SmolVLM pixel-shuffle (space-to-depth), c-innermost channel order. +// src [embed, n_patches] row-major (patch p, channel e) -> dst [embed*s^2, (grid/s)^2]. +inline void pixel_shuffle_hf(const float * src, float * dst, + int64_t embed, int64_t grid, int64_t s) { + const int64_t g2 = grid/s, c4 = embed * s * s; + for (int64_t h2=0; h2 & 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 & 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 grid = side/ps, pd = 3*ps*ps, np = grid*grid; + out.assign((size_t) pd*np, 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 row=0; row +#include +#include +#include + +namespace vla { + +bool build_prompt(const char * arch, const Inputs & in, int64_t n_img, + int32_t image_token, int64_t max_seq, Prompt & out) { + out = Prompt{}; + + int64_t slots = 0; + for (int j=0; j max_seq) { + std::fprintf(stderr, "vla(%s): prompt too long (%lld > %lld)\n", arch, (long long) seq, (long long) max_seq); + return false; + } + + out.image_pos.reserve((size_t) n_img); + out.text_pos.reserve((size_t)(seq-n_img)); + for (int64_t p=0; p & out) { + const int64_t seq = p.len(); + out.assign((size_t) seq*hidden, 0.0f); + if (!io.fetch_rows_f32("token_embd.weight", p.ids, out.data(), hidden)) + return false; + + for (size_t k=0; k & out) { + out.assign(n, 0.0f); + if (in.noise) { + std::memcpy(out.data(), in.noise, n*sizeof(float)); + return; + } + + 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 : out) + v = nd(rng); +} + +} diff --git a/src/modules/prompt.h b/src/modules/prompt.h new file mode 100644 index 0000000..7808190 --- /dev/null +++ b/src/modules/prompt.h @@ -0,0 +1,51 @@ +// 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. + +// Token-sequence assembly for backbones taking one interleaved image/text +// stream. + +#pragma once + +#include "gguf_reader.h" +#include "model.h" + +#include +#include + +namespace vla { + +struct Prompt { + std::vector ids; + std::vector image_pos; + std::vector text_pos; + + int64_t len() const { + return (int64_t) ids.size(); + } + int64_t n_text() const { + return (int64_t) text_pos.size(); + } +}; + +// Accepts a stream carrying exactly n_img placeholders, or none, in which case +// they are prepended. +bool build_prompt(const char * arch, const Inputs & in, int64_t n_img, + int32_t image_token, int64_t max_seq, Prompt & out); + +bool fetch_embeds(const char * arch, gguf_reader & io, const Prompt & p, + const float * img_emb, int64_t hidden, std::vector & out); + +void init_noise(const Inputs & in, size_t n, std::vector & out); + +} diff --git a/src/modules/qwen3_lm.cpp b/src/modules/qwen3_lm.cpp new file mode 100644 index 0000000..ef66f6e --- /dev/null +++ b/src/modules/qwen3_lm.cpp @@ -0,0 +1,85 @@ +// 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. + +#include "modules/qwen3_lm.h" + +#include "layers/attn.h" +#include "layers/ffn.h" +#include "layers/norm.h" + +#include + +namespace vla { + +void Qwen3LM::declare(WeightLoader & L, const char * prefix) { + output_norm = L.f32("%s.output_norm.weight", prefix); + blk.resize(cfg.layers); + + for (int64_t i=0; i +#include + +namespace vla { + +struct Qwen3LayerW { + ggml_tensor *attn_norm, *Wq, *Wk, *Wv, *Wo, *q_norm, *k_norm, *ffn_norm, *Wgate, *Wup, *Wdown; +}; + +struct Qwen3Cfg { + int64_t hidden = 2048; + int64_t layers = 16; + int64_t n_q = 16; + int64_t n_kv = 8; + int64_t head_dim = 128; + int64_t inter = 6144; + float rms_eps = 1e-6f; + RopeSpec rope; + bool flash_attn = false; +}; + +struct Qwen3LM { + Qwen3Cfg cfg; + std::vector blk; + ggml_tensor * output_norm = nullptr; + + void declare(WeightLoader & L, const char * prefix); + + ggml_tensor * block(ggml_context * C, const Qwen3LayerW & w, ggml_tensor * h, + ggml_tensor * pos, ggml_tensor * mask, int64_t seq) const; + + ggml_tensor * build(ggml_context * C, ggml_tensor * h, + ggml_tensor * pos, ggml_tensor * mask, int64_t seq) const; +}; + +} diff --git a/src/models/qwen3vl_vit.h b/src/modules/qwen3vl_vit.h similarity index 50% rename from src/models/qwen3vl_vit.h rename to src/modules/qwen3vl_vit.h index 9abc09a..851f471 100644 --- a/src/models/qwen3vl_vit.h +++ b/src/modules/qwen3vl_vit.h @@ -16,9 +16,13 @@ #pragma once +#include "layers/attn.h" +#include "loader.h" +#include "layers/rope.h" #include "model.h" #include "ggml.h" +#include "options.h" #include #include @@ -35,47 +39,71 @@ 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]); -} +struct Qwen3VLTower { + std::vector blk; + MergerW deepstack[3]; + MergerW merger; + ggml_tensor * patch_w = nullptr; + ggml_tensor * patch_b = nullptr; + ggml_tensor * pos = nullptr; + + void declare(WeightLoader & L, const char * prefix, int64_t layers) { + patch_w = L.gemm("%s.patch_embd.weight", prefix); + patch_b = L.f32 ("%s.patch_embd.bias", prefix); + pos = L.f32 ("%s.pos_embd", prefix); + + blk.resize(layers); + for (int64_t i=0; inb[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_view_2d(C, qkv, hidden, seq, qkv->nb[1], (size_t) 2*hidden * qkv->nb[0])); + ggml_tensor * Q = rope_2d(C, to_heads(C, q, hd, heads, seq), cos_t, sin_t); + ggml_tensor * K = rope_2d(C, to_heads(C, k, hd, heads, seq), 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); - } + if (vla::flash_attn_enabled()) + att = flash_attention(C, Q, K, to_heads (C, v, hd, heads, seq), nullptr, scale); + else + att = attention (C, Q, K, to_heads_v(C, v, hd, heads, seq), nullptr, scale, 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); @@ -85,7 +113,7 @@ inline ggml_tensor * build_vit_layer(ggml_context * C, const VitLayerW & w, ggml // 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); + 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); @@ -101,24 +129,32 @@ inline ggml_tensor * build_merger(ggml_context * C, const MergerW & w, ggml_tens // 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; + for (int64_t s=0; s & 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; + 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)); + for (int64_t i=0; i 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]); } + for (int64_t i=0; i & table, int64_t num_side, 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) { + 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