From 74f6525c6ea482719218527d9e7f0a7053387af2 Mon Sep 17 00:00:00 2001 From: Minh Hieu Date: Thu, 24 Sep 2026 03:55:29 +0700 Subject: [PATCH 1/5] feat(turbovla): add and clean TurboVLA support --- CMakeLists.txt | 1 + eval/client/run_sim_client_direct.py | 3 +- eval/client/vla_cpp_client.py | 131 ++- scripts/convert_turbovla_to_gguf.py | 845 ++++++++++++++ scripts/tokenize_prompt.py | 1 + src/arch.h | 10 + src/layers/ffn.h | 5 + src/layers/rope.h | 19 + src/loader.cpp | 5 +- src/model.cpp | 11 +- src/models/turbovla.cpp | 1579 ++++++++++++++++++++++++++ src/serving/vla-cli.cpp | 1 + 12 files changed, 2606 insertions(+), 5 deletions(-) create mode 100644 scripts/convert_turbovla_to_gguf.py create mode 100644 src/models/turbovla.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 708004d..11fece1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -144,6 +144,7 @@ add_library(vla_core src/models/vla_adapter.cpp src/models/openvla_oft.cpp src/models/vla_jepa.cpp + src/models/turbovla.cpp ) target_include_directories(vla_core PUBLIC diff --git a/eval/client/run_sim_client_direct.py b/eval/client/run_sim_client_direct.py index 6d65a2c..0bb62f5 100644 --- a/eval/client/run_sim_client_direct.py +++ b/eval/client/run_sim_client_direct.py @@ -88,7 +88,8 @@ parser.add_argument( "--stats-json", type=str, default=None, - help="[bitvla/gr00t_n1_6/gr00t_n1_7] path to dataset_statistics.json. Default for bitvla: " + help="[turbovla/bitvla/gr00t_n1_6/gr00t_n1_7] path to dataset statistics. TurboVLA defaults " + "to the H-EmbodVis/TurboVLA release stats. Default for bitvla: " "/dataset_statistics.json. For gr00t_n1_{6,7}, pass the canonical " "/experiment_cfg/dataset_statistics.json explicitly. Without " "it, gr00t_n1_7 returns the raw normalized [40, 132] chunk and gr00t_n1_6 " diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index 6d81559..75cf093 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -22,8 +22,9 @@ import sys import tempfile from collections import deque +from collections.abc import Sequence from pathlib import Path -from typing import Any, Sequence +from typing import Any import numpy as np import torch @@ -49,6 +50,12 @@ "max_state_dim": 8}, "vla_jepa": {"image_size": 256, "tokenizer": "Qwen/Qwen3-VL-2B-Instruct", "max_state_dim": 8, "use_processor": True}, + "turbovla": { + "image_size": 256, + "tokenizer": "bert-base-uncased", + "max_state_dim": 8, + "max_length": 21, + }, "gr00t_n1_7": {"image_size": 256, "tokenizer": "nvidia/Cosmos-Reason2-2B", "max_state_dim": 132}, "gr00t_n1_5": {"image_size": 224, "tokenizer": "lerobot/eagle2hg-processor-groot-n1p5", @@ -285,6 +292,64 @@ def _oft_norm(x, q01=q01, q99=q99, mask=mask): f"{stats_path}::observation.state ({self._pi05_state_q01.shape[0]}-D)", flush=True) + self._turbovla_state_norm = None + self._turbovla_action_unnorm = None + if arch == "turbovla": + if stats_json: + stats_path = Path(stats_json) + else: + from huggingface_hub import hf_hub_download + stats_path = Path(hf_hub_download( + repo_id="H-EmbodVis/TurboVLA", filename="libero_all4_stats.json")) + if not stats_path.exists(): + raise FileNotFoundError( + f"TurboVLA stats not found at {stats_path}. Pass --stats-json " + ".") + blob = json.loads(stats_path.read_text()) + try: + if "libero_all4_no_noops" in blob: + blob = blob["libero_all4_no_noops"] + state = blob["proprio"] if "proprio" in blob else blob["observation.state"] + action = blob["action"] + state_mean = np.asarray(state["mean"], dtype=np.float32) + state_std = np.asarray(state["std"], dtype=np.float32) + action_min = np.asarray(action["min"], dtype=np.float32) + action_max = np.asarray(action["max"], dtype=np.float32) + except (KeyError, TypeError) as exc: + raise ValueError( + f"{stats_path} must contain observation.state mean/std and action min/max") from exc + if ( + state_mean.shape != (8,) + or state_std.shape != (8,) + or action_min.shape != (7,) + or action_max.shape != (7,) + ): + raise ValueError( + f"{stats_path} does not contain TurboVLA's 8-D state and 7-D action stats" + ) + if np.any(state_std <= 0) or np.any(action_max <= action_min): + raise ValueError(f"{stats_path} contains invalid TurboVLA normalization ranges") + + def _state_norm(state_8d, mean=state_mean, std=state_std): + return ((state_8d - mean) / (std + 1e-8)).astype(np.float32) + + def _action_unnorm(chunk, mn=action_min, mx=action_max): + normalized = np.clip(chunk.astype(np.float32), -1.0, 1.0) + action = normalized.copy() + action[..., :6] = ((normalized[..., :6] + 1.0) * 0.5 + * (mx[:6] - mn[:6]) + mn[:6]) + # The released LIBERO policy thresholds its normalized gripper + # output instead of min/max denormalizing it as a continuous axis. + action[..., 6] = np.where(normalized[..., 6] > 0.0, 1.0, -1.0) + return action.astype(np.float32) + + self._turbovla_state_norm = _state_norm + self._turbovla_action_unnorm = _action_unnorm + print( + "vla-cpp-direct[arch=turbovla]: release MEAN_STD state / MIN_MAX " + f"arm stats via {stats_path}", + flush=True, + ) self._gr00t_action_unnorm = None self._gr00t_state_norm = None @@ -581,6 +646,8 @@ def get_action(self, observations: dict[str, Any]) -> np.ndarray: def _predict_chunk(self, observations: dict[str, Any]) -> np.ndarray: + if self.arch == "turbovla": + return self._predict_chunk_turbovla(observations) if self.arch == "evo1": return self._predict_chunk_evo1(observations) if self.arch == "gr00t_n1_7": @@ -650,6 +717,68 @@ def _predict_chunk(self, observations: dict[str, Any]) -> np.ndarray: return (np.array(resp.action_chunk, dtype=np.float32) .reshape(resp.chunk_size, resp.action_dim)) + def _predict_chunk_turbovla(self, observations: dict[str, Any]) -> np.ndarray: + if self._turbovla_state_norm is None or self._turbovla_action_unnorm is None: + raise RuntimeError("TurboVLA statistics were not initialized") + images_f32: list[np.ndarray] = [] + for key in self.image_keys: + if key not in observations: + raise KeyError( + f"TurboVLA image key {key!r} missing; got {list(observations)}" + ) + image = observations[key] + if isinstance(image, torch.Tensor): + image = image.numpy() + image = np.asarray(image, dtype=np.float32) + if image.ndim != 3 or image.shape[0] != 3: + raise ValueError(f"{key}: expected CHW float32 [3,H,W], got {image.shape}") + image = _resize_with_pad(image, self.image_size, self.image_size) + images_f32.append( + np.ascontiguousarray(np.transpose(image, (1, 2, 0)), dtype=np.float32) + ) + + state = observations["observation.state"] + if isinstance(state, torch.Tensor): + state = state.numpy() + state = np.asarray(state, dtype=np.float32).reshape(-1) + if state.shape != (8,): + raise ValueError(f"TurboVLA requires an 8-D state, got {state.shape}") + + task = observations.get("task", "") + if isinstance(task, bytes): + task = task.decode() + if not isinstance(task, str) or not task.strip(): + raise ValueError("TurboVLA requires a non-empty task instruction") + lang = self.tok( + task, + padding=False, + truncation=True, + max_length=self.max_length, + return_tensors="np", + )["input_ids"][0].astype(np.int32) + + req = self.pb.PredictRequest() + req.request_id = self._step + self._step += 1 + for image in images_f32: + proto_image = req.images.add() + proto_image.encoding = self.pb.Image.F32_RGB_01 + proto_image.height, proto_image.width = image.shape[:2] + proto_image.data = image.tobytes() + req.lang_tokens.extend(lang.tolist()) + req.state.extend(self._turbovla_state_norm(state).tolist()) + + self.sock.send(req.SerializeToString()) + resp = self.pb.PredictResponse() + resp.ParseFromString(self.sock.recv()) + if resp.error: + raise RuntimeError(f"vla-server error: {resp.error}") + self._last_response = resp + chunk = np.array(resp.action_chunk, dtype=np.float32).reshape( + resp.chunk_size, resp.action_dim + ) + return self._turbovla_action_unnorm(chunk) + _VJ_PS, _VJ_TPS, _VJ_MERGE, _VJ_SIDE = 16, 2, 2, 256 @classmethod diff --git a/scripts/convert_turbovla_to_gguf.py b/scripts/convert_turbovla_to_gguf.py new file mode 100644 index 0000000..9781c48 --- /dev/null +++ b/scripts/convert_turbovla_to_gguf.py @@ -0,0 +1,845 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Convert TurboVLA (LeRobot format) checkpoint to GGUF. + +Supports any TurboVLA variant with auto-detection of dimensions from checkpoint. +Architecture must be compatible with the LeRobot TurboVLA format. + +Usage: + python convert_turbovla_to_gguf.py /path/to/checkpoint [--out output.gguf] [--verify] +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from safetensors import safe_open + +import gguf + +F32 = gguf.GGMLQuantizationType.F32 +BF16 = gguf.GGMLQuantizationType.BF16 + +ARCH = "turbovla" + + +# ============================================================================= +# TENSOR KEY PREFIXES +# ============================================================================= + +# Prefixes for official TurboVLA checkpoint (no "model." prefix) +PREFIX_VIT = "vision_encoder.backbone" +PREFIX_TEXT = "text_encoder.bert" +PREFIX_VIT_PROJ = "vision_projection" +PREFIX_FUSION = "vision_language_interaction.fusion_layers" +PREFIX_VL_TEXT = "vision_language_interaction.text_layers" +PREFIX_ACT_DEC = "action_head.decoder" +PREFIX_STATE_PROJ = "action_head.state_projection" +KEY_VIEW_EMB = "view_embedding" +KEY_TEXT_PROJ = "text_encoder.text_projection" + + +# ============================================================================= +# UTILITY FUNCTIONS +# ============================================================================= + +def kv_prefix(name: str) -> str: + return f"{ARCH}.{name}" + + +def bf16_u16(t: torch.Tensor) -> np.ndarray: + return t.contiguous().view(torch.uint16).cpu().numpy() + + +def add_tensor(writer: gguf.GGUFWriter, name: str, t: torch.Tensor) -> None: + """Add tensor with preserved dtype.""" + if t.dtype == torch.float32: + writer.add_tensor(name, t.contiguous().cpu().numpy(), raw_dtype=F32) + elif t.dtype == torch.bfloat16: + writer.add_tensor(name, bf16_u16(t), raw_shape=list(t.shape), raw_dtype=BF16) + elif t.dtype == torch.float16: + writer.add_tensor(name, t.contiguous().cpu().numpy().astype(np.float32), raw_dtype=F32) + else: + raise NotImplementedError(f"unsupported dtype {t.dtype} for {name}") + + +def max_layer(keys: set[str], pfx: str) -> int: + """Count number of layers with given prefix.""" + m = -1 + for k in keys: + if k.startswith(pfx): + try: + m = max(m, int(k[len(pfx):].split(".", 1)[0])) + except ValueError: + pass + return m + 1 + + +def load_safetensors(ckpt: Path) -> dict[str, torch.Tensor]: + """Load all tensors from safetensors file.""" + tensors = {} + safetensors_path = ckpt / "model.safetensors" + if not safetensors_path.exists(): + raise SystemExit(f"model.safetensors not found in {ckpt}") + + with safe_open(str(safetensors_path), framework="pt", device="cpu") as f: + for key in f: + tensors[key] = f.get_tensor(key) + return tensors + + +# ============================================================================= +# DIMENSION INFERENCE +# ============================================================================= + +class TurboVLADims: + """Auto-detect all TurboVLA dimensions from checkpoint tensor shapes.""" + + def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: dict | None = None, + dinov3_cfg: dict | None = None): + self.tensors = tensors + self.keys = keys + self.cfg = cfg_json or {} + self.dinov3_cfg = dinov3_cfg or {} + + # Vision encoder (DINOv3 ViT) + q0 = self._get(f"{PREFIX_VIT}.layer.0.attention.q_proj.weight") + self.vit_dim = int(q0.shape[0]) + self.vit_layers = max_layer(keys, f"{PREFIX_VIT}.layer.") + + # Text encoder (BERT) + q0 = self._get(f"{PREFIX_TEXT}.encoder.layer.0.attention.self.query.weight") + self.text_dim = int(q0.shape[0]) + self.text_layers = max_layer(keys, f"{PREFIX_TEXT}.encoder.layer.") + + # VL Fusion layers + self.num_fusion_layers = max_layer(keys, f"{PREFIX_FUSION}.") + self.num_text_layers = max_layer(keys, f"{PREFIX_VL_TEXT}.") + + # BiMultiHeadAttention projects to interaction.enhancer_inner_dim, whereas + # the text enhancer and ACT decoder operate at interaction.hidden_dim. + interaction_cfg = self.cfg.get("interaction", {}) + self.hidden_dim = int(interaction_cfg.get("hidden_dim", 256)) + raw_nheads = int(interaction_cfg.get("nheads", 8)) + self.fusion_heads = max(1, raw_nheads // 2) + self.fusion_dim = int(interaction_cfg.get("enhancer_inner_dim", 0)) + if self.fusion_dim <= 0 or self.fusion_dim % self.fusion_heads: + raise SystemExit("interaction.enhancer_inner_dim must be divisible by fusion_heads") + self.fusion_head_dim = self.fusion_dim // self.fusion_heads + + # Text enhancer heads: same split as fusion + self.text_enhancer_heads = max(1, raw_nheads // 2) # = 4 + self.text_enhancer_head_dim = self.hidden_dim // self.text_enhancer_heads # = 64 + + # TurboVLA constructs the ACT decoder with interaction.nheads. + self.action_heads = raw_nheads + self.action_head_dim = self.hidden_dim // self.action_heads # = 32 + + # State projection config + self.num_state_tokens = int(self.cfg.get("action", {}).get("num_state_tokens", 2)) + + # Action decoder + self.num_action_decoder_layers = max_layer(keys, f"{PREFIX_ACT_DEC}.decoder.layers.") + action_q = self._get(f"{PREFIX_ACT_DEC}.action_queries.weight") + self.action_horizon = int(action_q.shape[0]) + act_proj_2 = self._get(f"{PREFIX_ACT_DEC}.action_projection.layers.2.weight") + self.action_dim = int(act_proj_2.shape[0]) + + # State projection + state_weight = self._get(f"{PREFIX_STATE_PROJ}.net.1.weight") + self.state_dim = int(state_weight.shape[1]) + + # Image specs - infer from config or tensor shapes + view_emb = self._get(KEY_VIEW_EMB) + if view_emb.ndim == 3: + if view_emb.shape[0] != 1: + raise SystemExit(f"Unexpected view_embedding shape: {tuple(view_emb.shape)}") + self.num_views = int(view_emb.shape[1]) + elif view_emb.ndim == 2: + self.num_views = int(view_emb.shape[0]) + else: + raise SystemExit(f"Unexpected view_embedding shape: {tuple(view_emb.shape)}") + + # image_size: config first, then infer + cfg_image_size = self.cfg.get("vision", {}).get("image_size", self.cfg.get("image_size")) + if cfg_image_size is not None: + if isinstance(cfg_image_size, (list, tuple)): + if not cfg_image_size: + raise SystemExit("empty image_size in config") + self.image_size = int(cfg_image_size[0]) + else: + self.image_size = int(cfg_image_size) + else: + inferred = self._infer_image_size() + if inferred is None: + raise SystemExit( + "Cannot determine image_size. " + "Checkpoint does not encode enough information " + "and config.json has no image_size." + ) + self.image_size = inferred + + # Patch size from vision encoder + patch_weight = self._get(f"{PREFIX_VIT}.embeddings.patch_embeddings.weight") + if len(patch_weight.shape) == 4: + self.patch_size = patch_weight.shape[2] + else: + raise SystemExit("Cannot determine patch_size from patch_embeddings") + + # Register tokens (optional) + reg_key = f"{PREFIX_VIT}.embeddings.register_tokens" + if reg_key in tensors: + reg = tensors[reg_key] + self.num_register_tokens = int(reg.shape[1]) if reg.ndim == 3 else 0 + else: + self.num_register_tokens = 0 + + # DINOv3 RoPE is architectural data, not a shape-derived default. The + # checkpoint's TurboVLA config names the backbone, while the matching + # DINO config supplies rope_theta and coordinate normalization. + if "rope_theta" not in self.dinov3_cfg: + raise SystemExit("DINOv3 config with rope_theta is required; pass --dinov3-config") + self.rope_theta = float(self.dinov3_cfg["rope_theta"]) + self.rope_normalize_coords = str(self.dinov3_cfg.get("rope_normalize_coords", "separate")) + if self.rope_normalize_coords != "separate": + raise SystemExit( + f"unsupported DINOv3 rope coordinate normalization: {self.rope_normalize_coords!r}" + ) + if int(self.dinov3_cfg.get("num_register_tokens", self.num_register_tokens)) != self.num_register_tokens: + raise SystemExit("DINOv3 config num_register_tokens disagrees with checkpoint weights") + + # Infer vocab_size from word embeddings + word_emb = self._get(f"{PREFIX_TEXT}.embeddings.word_embeddings.weight") + self.vocab_size = int(word_emb.shape[0]) + + # dropout: only from config, no fallback + self.dropout = float(self.cfg.get("vision", {}).get("dropout", 0.0)) + + def _get(self, key: str) -> torch.Tensor: + if key not in self.tensors: + raise KeyError(f"Required tensor not found: {key}") + return self.tensors[key] + + def _infer_image_size(self) -> int | None: + """Infer image size from tensor shapes if config not available. + + Note: DINOv3 uses RoPE (rotary position embeddings) instead of learned + position embeddings, so position_embeddings.weight may not exist or may + have different semantics. Use patch_size * grid_size from config instead. + """ + # Try the validated DINOv3 config first. + if self.dinov3_cfg: + image_size = self.dinov3_cfg.get("image_size") + if image_size: + return int(image_size) + + # Fallback: try to infer from position embeddings + # DINOv2 with position embeddings + pos_key = f"{PREFIX_VIT}.embeddings.position_embeddings.weight" + if pos_key in self.tensors: + pos = self.tensors[pos_key] + # Position embeddings format: [seq_len, dim] or [1, seq_len, dim] + seq_len = pos.shape[0] if pos.ndim == 2 else pos.shape[1] + num_patches = seq_len - 1 # exclude CLS token + patch = self.tensors.get(f"{PREFIX_VIT}.embeddings.patch_embeddings.weight") + if patch is not None and len(patch.shape) == 4: + patch_size = patch.shape[2] + grid_size = int(num_patches ** 0.5) + if grid_size * grid_size == num_patches: + return grid_size * patch_size + + return None + + def __repr__(self): + return ( + f"TurboVLADims(vit={self.vit_dim}d×{self.vit_layers}L, " + f"text={self.text_dim}d×{self.text_layers}L, " + f"hidden={self.hidden_dim}, " + f"fusion={self.num_fusion_layers}L, " + f"decoder={self.num_action_decoder_layers}L, " + f"action={self.action_dim}d×{self.action_horizon}horizon, " + f"state={self.state_dim}d, views={self.num_views})" + ) + + +# ============================================================================= +# TENSOR WRITERS +# ============================================================================= + +def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write DINOv3 ViT vision encoder. + + Tensor naming follows TurboVLA naming convention with attn_q/k/v/o. + LayerScale is baked into weights. + """ + root = PREFIX_VIT + + # embeddings + # PyTorch patch_embedding weight: [H, IC, KH, KW] = [768, 3, 16, 16] + # For matmul-based patch embedding: flatten weight to [H, IC*KH*KW] = [768, 768] + patch_weight = tensors[f"{root}.embeddings.patch_embeddings.weight"] + patch_weight = patch_weight.reshape(dims.vit_dim, -1) # [768, 768] + add_tensor(writer, "vit.cls_token", tensors[f"{root}.embeddings.cls_token"].squeeze(0)) + add_tensor(writer, "vit.patch_embed.weight", patch_weight) + add_tensor(writer, "vit.patch_embed.bias", tensors[f"{root}.embeddings.patch_embeddings.bias"]) + + # register tokens (optional - DINOv3 uses them) + reg_key = f"{root}.embeddings.register_tokens" + if reg_key in tensors: + add_tensor(writer, "vit.register_tokens", tensors[reg_key].squeeze(0)) + + # transformer layers + for i in range(dims.vit_layers): + lr = f"{root}.layer.{i}" + + # Get LayerScale values for baking into weights + ls1 = tensors[f"{lr}.layer_scale1.lambda1"].clone() + ls2 = tensors[f"{lr}.layer_scale2.lambda1"].clone() + + # attention projections - bake LayerScale into o_proj + w_q = tensors[f"{lr}.attention.q_proj.weight"] + b_q = tensors[f"{lr}.attention.q_proj.bias"] + w_k = tensors[f"{lr}.attention.k_proj.weight"] + k_bias_key = f"{lr}.attention.k_proj.bias" + b_k = tensors[k_bias_key] if k_bias_key in tensors else torch.zeros(dims.vit_dim, dtype=torch.float32) + w_v = tensors[f"{lr}.attention.v_proj.weight"] + b_v = tensors[f"{lr}.attention.v_proj.bias"] + # Bake ls1 into o_proj + w_o = tensors[f"{lr}.attention.o_proj.weight"] * ls1.view(-1, 1) + b_o = tensors[f"{lr}.attention.o_proj.bias"] * ls1 + + add_tensor(writer, f"vit.blk.{i}.attn_q.weight", w_q) + add_tensor(writer, f"vit.blk.{i}.attn_q.bias", b_q) + add_tensor(writer, f"vit.blk.{i}.attn_k.weight", w_k) + add_tensor(writer, f"vit.blk.{i}.attn_k.bias", b_k) + add_tensor(writer, f"vit.blk.{i}.attn_v.weight", w_v) + add_tensor(writer, f"vit.blk.{i}.attn_v.bias", b_v) + add_tensor(writer, f"vit.blk.{i}.attn_o.weight", w_o) + add_tensor(writer, f"vit.blk.{i}.attn_o.bias", b_o) + + # layer norms + add_tensor(writer, f"vit.blk.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) + add_tensor(writer, f"vit.blk.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) + add_tensor(writer, f"vit.blk.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) + add_tensor(writer, f"vit.blk.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) + + # MLP - bake ls2 into down_proj + w_fc1 = tensors[f"{lr}.mlp.up_proj.weight"] + b_fc1 = tensors[f"{lr}.mlp.up_proj.bias"] + w_fc2 = tensors[f"{lr}.mlp.down_proj.weight"] * ls2.view(-1, 1) + b_fc2 = tensors[f"{lr}.mlp.down_proj.bias"] * ls2 + + add_tensor(writer, f"vit.blk.{i}.fc1.weight", w_fc1) + add_tensor(writer, f"vit.blk.{i}.fc1.bias", b_fc1) + add_tensor(writer, f"vit.blk.{i}.fc2.weight", w_fc2) + add_tensor(writer, f"vit.blk.{i}.fc2.bias", b_fc2) + + # final norm + add_tensor(writer, "vit.final_norm.weight", tensors[f"{root}.norm.weight"]) + add_tensor(writer, "vit.final_norm.bias", tensors[f"{root}.norm.bias"]) + + +def write_text_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write BERT text encoder.""" + root = PREFIX_TEXT + + # embeddings + add_tensor(writer, "text.embed.word_embeddings", tensors[f"{root}.embeddings.word_embeddings.weight"]) + add_tensor(writer, "text.embed.position_embeddings", tensors[f"{root}.embeddings.position_embeddings.weight"]) + add_tensor(writer, "text.embed.token_type_embeddings", tensors[f"{root}.embeddings.token_type_embeddings.weight"]) + add_tensor(writer, "text.embed.LayerNorm.weight", tensors[f"{root}.embeddings.LayerNorm.weight"]) + add_tensor(writer, "text.embed.LayerNorm.bias", tensors[f"{root}.embeddings.LayerNorm.bias"]) + + # encoder layers + for i in range(dims.text_layers): + lr = f"{root}.encoder.layer.{i}" + + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.query.weight", tensors[f"{lr}.attention.self.query.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.query.bias", tensors[f"{lr}.attention.self.query.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.key.weight", tensors[f"{lr}.attention.self.key.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.key.bias", tensors[f"{lr}.attention.self.key.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.value.weight", tensors[f"{lr}.attention.self.value.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.self.value.bias", tensors[f"{lr}.attention.self.value.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.dense.weight", tensors[f"{lr}.attention.output.dense.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.dense.bias", tensors[f"{lr}.attention.output.dense.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.LayerNorm.weight", tensors[f"{lr}.attention.output.LayerNorm.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.attention.output.LayerNorm.bias", tensors[f"{lr}.attention.output.LayerNorm.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.intermediate.dense.weight", tensors[f"{lr}.intermediate.dense.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.intermediate.dense.bias", tensors[f"{lr}.intermediate.dense.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.dense.weight", tensors[f"{lr}.output.dense.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.dense.bias", tensors[f"{lr}.output.dense.bias"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.weight", tensors[f"{lr}.output.LayerNorm.weight"]) + add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.bias", tensors[f"{lr}.output.LayerNorm.bias"]) + + # pooler + add_tensor(writer, "text.pooler.dense.weight", tensors[f"{root}.pooler.dense.weight"]) + add_tensor(writer, "text.pooler.dense.bias", tensors[f"{root}.pooler.dense.bias"]) + + # text projection + add_tensor(writer, "text_proj.weight", tensors[KEY_TEXT_PROJ + ".weight"]) + add_tensor(writer, "text_proj.bias", tensors[KEY_TEXT_PROJ + ".bias"]) + + +def write_vision_projection(writer: gguf.GGUFWriter, tensors: dict) -> None: + """Write vision projection MLP.""" + root = PREFIX_VIT_PROJ + + add_tensor(writer, "vit_proj.input_norm.weight", tensors[f"{root}.input_norm.weight"]) + add_tensor(writer, "vit_proj.input_norm.bias", tensors[f"{root}.input_norm.bias"]) + add_tensor(writer, "vit_proj.mlp.0.weight", tensors[f"{root}.mlp.0.weight"]) + add_tensor(writer, "vit_proj.mlp.0.bias", tensors[f"{root}.mlp.0.bias"]) + add_tensor(writer, "vit_proj.mlp.3.weight", tensors[f"{root}.mlp.3.weight"]) + add_tensor(writer, "vit_proj.mlp.3.bias", tensors[f"{root}.mlp.3.bias"]) + add_tensor(writer, "vit_proj.skip.weight", tensors[f"{root}.skip.weight"]) + add_tensor(writer, "vit_proj.output_norm.weight", tensors[f"{root}.output_norm.weight"]) + add_tensor(writer, "vit_proj.output_norm.bias", tensors[f"{root}.output_norm.bias"]) + + +def write_vision_language_interaction(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write VL fusion and text self-attention layers.""" + fusion_root = PREFIX_FUSION + text_root = PREFIX_VL_TEXT + + # VL Fusion layers + for i in range(dims.num_fusion_layers): + lr = f"{fusion_root}.{i}" + + add_tensor(writer, f"vl_fusion.{i}.v_proj.weight", tensors[f"{lr}.attn.v_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.v_proj.bias", tensors[f"{lr}.attn.v_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.l_proj.weight", tensors[f"{lr}.attn.l_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.l_proj.bias", tensors[f"{lr}.attn.l_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.values_v.weight", tensors[f"{lr}.attn.values_v_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.values_v.bias", tensors[f"{lr}.attn.values_v_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.values_l.weight", tensors[f"{lr}.attn.values_l_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.values_l.bias", tensors[f"{lr}.attn.values_l_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.out_v.weight", tensors[f"{lr}.attn.out_v_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.out_v.bias", tensors[f"{lr}.attn.out_v_proj.bias"]) + add_tensor(writer, f"vl_fusion.{i}.out_l.weight", tensors[f"{lr}.attn.out_l_proj.weight"]) + add_tensor(writer, f"vl_fusion.{i}.out_l.bias", tensors[f"{lr}.attn.out_l_proj.bias"]) + + add_tensor(writer, f"vl_fusion.{i}.norm_v.weight", tensors[f"{lr}.layer_norm_v.weight"]) + add_tensor(writer, f"vl_fusion.{i}.norm_v.bias", tensors[f"{lr}.layer_norm_v.bias"]) + add_tensor(writer, f"vl_fusion.{i}.norm_l.weight", tensors[f"{lr}.layer_norm_l.weight"]) + add_tensor(writer, f"vl_fusion.{i}.norm_l.bias", tensors[f"{lr}.layer_norm_l.bias"]) + + add_tensor(writer, f"vl_fusion.{i}.gamma_v", tensors[f"{lr}.gamma_v"]) + add_tensor(writer, f"vl_fusion.{i}.gamma_l", tensors[f"{lr}.gamma_l"]) + + # VL Text self-attention layers + for i in range(dims.num_text_layers): + lr = f"{text_root}.{i}" + + add_tensor(writer, f"vl_text.{i}.attn_qkv.weight", tensors[f"{lr}.self_attn.in_proj_weight"]) + add_tensor(writer, f"vl_text.{i}.attn_qkv.bias", tensors[f"{lr}.self_attn.in_proj_bias"]) + add_tensor(writer, f"vl_text.{i}.attn_o.weight", tensors[f"{lr}.self_attn.out_proj.weight"]) + add_tensor(writer, f"vl_text.{i}.attn_o.bias", tensors[f"{lr}.self_attn.out_proj.bias"]) + + add_tensor(writer, f"vl_text.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) + add_tensor(writer, f"vl_text.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) + + add_tensor(writer, f"vl_text.{i}.fc1.weight", tensors[f"{lr}.linear1.weight"]) + add_tensor(writer, f"vl_text.{i}.fc1.bias", tensors[f"{lr}.linear1.bias"]) + add_tensor(writer, f"vl_text.{i}.fc2.weight", tensors[f"{lr}.linear2.weight"]) + add_tensor(writer, f"vl_text.{i}.fc2.bias", tensors[f"{lr}.linear2.bias"]) + add_tensor(writer, f"vl_text.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) + add_tensor(writer, f"vl_text.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) + + +def write_action_decoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: + """Write ACT-style action decoder.""" + root = PREFIX_ACT_DEC + + add_tensor(writer, "act.q.weight", tensors[f"{root}.action_queries.weight"]) + + for i in range(dims.num_action_decoder_layers): + lr = f"{root}.decoder.layers.{i}" + + add_tensor(writer, f"act.dec.{i}.self_qkv.weight", tensors[f"{lr}.self_attn.in_proj_weight"]) + add_tensor(writer, f"act.dec.{i}.self_qkv.bias", tensors[f"{lr}.self_attn.in_proj_bias"]) + add_tensor(writer, f"act.dec.{i}.self_out.weight", tensors[f"{lr}.self_attn.out_proj.weight"]) + add_tensor(writer, f"act.dec.{i}.self_out.bias", tensors[f"{lr}.self_attn.out_proj.bias"]) + + add_tensor(writer, f"act.dec.{i}.cross_qkv.weight", tensors[f"{lr}.multihead_attn.in_proj_weight"]) + add_tensor(writer, f"act.dec.{i}.cross_qkv.bias", tensors[f"{lr}.multihead_attn.in_proj_bias"]) + add_tensor(writer, f"act.dec.{i}.cross_out.weight", tensors[f"{lr}.multihead_attn.out_proj.weight"]) + add_tensor(writer, f"act.dec.{i}.cross_out.bias", tensors[f"{lr}.multihead_attn.out_proj.bias"]) + + add_tensor(writer, f"act.dec.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) + add_tensor(writer, f"act.dec.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) + add_tensor(writer, f"act.dec.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) + add_tensor(writer, f"act.dec.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) + add_tensor(writer, f"act.dec.{i}.ln3.weight", tensors[f"{lr}.norm3.weight"]) + add_tensor(writer, f"act.dec.{i}.ln3.bias", tensors[f"{lr}.norm3.bias"]) + + add_tensor(writer, f"act.dec.{i}.fc1.weight", tensors[f"{lr}.linear1.weight"]) + add_tensor(writer, f"act.dec.{i}.fc1.bias", tensors[f"{lr}.linear1.bias"]) + add_tensor(writer, f"act.dec.{i}.fc2.weight", tensors[f"{lr}.linear2.weight"]) + add_tensor(writer, f"act.dec.{i}.fc2.bias", tensors[f"{lr}.linear2.bias"]) + + act_proj_root = f"{root}.action_projection" + add_tensor(writer, "act.proj.0.weight", tensors[f"{act_proj_root}.layers.0.weight"]) + add_tensor(writer, "act.proj.0.bias", tensors[f"{act_proj_root}.layers.0.bias"]) + add_tensor(writer, "act.proj.1.weight", tensors[f"{act_proj_root}.layers.1.weight"]) + add_tensor(writer, "act.proj.1.bias", tensors[f"{act_proj_root}.layers.1.bias"]) + add_tensor(writer, "act.proj.2.weight", tensors[f"{act_proj_root}.layers.2.weight"]) + add_tensor(writer, "act.proj.2.bias", tensors[f"{act_proj_root}.layers.2.bias"]) + + +def write_state_projection(writer: gguf.GGUFWriter, tensors: dict) -> None: + """Write state projection MLP.""" + root = PREFIX_STATE_PROJ + + add_tensor(writer, "state.proj.0.weight", tensors[f"{root}.net.0.weight"]) + add_tensor(writer, "state.proj.0.bias", tensors[f"{root}.net.0.bias"]) + add_tensor(writer, "state.proj.1.weight", tensors[f"{root}.net.1.weight"]) + add_tensor(writer, "state.proj.1.bias", tensors[f"{root}.net.1.bias"]) + add_tensor(writer, "state.proj.4.weight", tensors[f"{root}.net.4.weight"]) + add_tensor(writer, "state.proj.4.bias", tensors[f"{root}.net.4.bias"]) + + add_tensor(writer, "state.proj.output_norm.weight", tensors[f"{root}.output_norm.weight"]) + add_tensor(writer, "state.proj.output_norm.bias", tensors[f"{root}.output_norm.bias"]) + add_tensor(writer, "state.proj.position", tensors[f"{root}.position"]) + + +def write_view_embeddings(writer: gguf.GGUFWriter, tensors: dict) -> None: + """Write view embeddings.""" + add_tensor(writer, "view_emb", tensors[KEY_VIEW_EMB]) + + +# ============================================================================= +# VERIFICATION +# ============================================================================= + +def verify_consumed_tensors(source_keys: set[str], dims: TurboVLADims) -> None: + """Verify all source tensors are consumed by the converter.""" + consumed = set() + + root_vit = PREFIX_VIT + consumed.add(f"{root_vit}.embeddings.cls_token") + consumed.add(f"{root_vit}.embeddings.patch_embeddings.weight") + consumed.add(f"{root_vit}.embeddings.patch_embeddings.bias") + reg_key = f"{root_vit}.embeddings.register_tokens" + if reg_key in source_keys: + consumed.add(reg_key) + + for i in range(dims.vit_layers): + lr = f"{root_vit}.layer.{i}" + consumed.add(f"{lr}.attention.q_proj.weight") + consumed.add(f"{lr}.attention.q_proj.bias") + consumed.add(f"{lr}.attention.k_proj.weight") + kb = f"{lr}.attention.k_proj.bias" + if kb in source_keys: + consumed.add(kb) + consumed.add(f"{lr}.attention.v_proj.weight") + consumed.add(f"{lr}.attention.v_proj.bias") + consumed.add(f"{lr}.attention.o_proj.weight") + consumed.add(f"{lr}.attention.o_proj.bias") + consumed.add(f"{lr}.norm1.weight") + consumed.add(f"{lr}.norm1.bias") + consumed.add(f"{lr}.norm2.weight") + consumed.add(f"{lr}.norm2.bias") + consumed.add(f"{lr}.mlp.up_proj.weight") + consumed.add(f"{lr}.mlp.up_proj.bias") + consumed.add(f"{lr}.mlp.down_proj.weight") + consumed.add(f"{lr}.mlp.down_proj.bias") + consumed.add(f"{lr}.layer_scale1.lambda1") + consumed.add(f"{lr}.layer_scale2.lambda1") + + consumed.add(f"{root_vit}.norm.weight") + consumed.add(f"{root_vit}.norm.bias") + + root_text = PREFIX_TEXT + consumed.add(f"{root_text}.embeddings.word_embeddings.weight") + consumed.add(f"{root_text}.embeddings.position_embeddings.weight") + consumed.add(f"{root_text}.embeddings.token_type_embeddings.weight") + consumed.add(f"{root_text}.embeddings.LayerNorm.weight") + consumed.add(f"{root_text}.embeddings.LayerNorm.bias") + consumed.add(f"{root_text}.pooler.dense.weight") + consumed.add(f"{root_text}.pooler.dense.bias") + + for i in range(dims.text_layers): + lr = f"{root_text}.encoder.layer.{i}" + consumed.add(f"{lr}.attention.self.query.weight") + consumed.add(f"{lr}.attention.self.query.bias") + consumed.add(f"{lr}.attention.self.key.weight") + consumed.add(f"{lr}.attention.self.key.bias") + consumed.add(f"{lr}.attention.self.value.weight") + consumed.add(f"{lr}.attention.self.value.bias") + consumed.add(f"{lr}.attention.output.dense.weight") + consumed.add(f"{lr}.attention.output.dense.bias") + consumed.add(f"{lr}.attention.output.LayerNorm.weight") + consumed.add(f"{lr}.attention.output.LayerNorm.bias") + consumed.add(f"{lr}.intermediate.dense.weight") + consumed.add(f"{lr}.intermediate.dense.bias") + consumed.add(f"{lr}.output.dense.weight") + consumed.add(f"{lr}.output.dense.bias") + consumed.add(f"{lr}.output.LayerNorm.weight") + consumed.add(f"{lr}.output.LayerNorm.bias") + + consumed.add(KEY_TEXT_PROJ + ".weight") + consumed.add(KEY_TEXT_PROJ + ".bias") + + vp = PREFIX_VIT_PROJ + consumed.add(f"{vp}.input_norm.weight") + consumed.add(f"{vp}.input_norm.bias") + consumed.add(f"{vp}.mlp.0.weight") + consumed.add(f"{vp}.mlp.0.bias") + consumed.add(f"{vp}.mlp.3.weight") + consumed.add(f"{vp}.mlp.3.bias") + consumed.add(f"{vp}.skip.weight") + consumed.add(f"{vp}.output_norm.weight") + consumed.add(f"{vp}.output_norm.bias") + + for i in range(dims.num_fusion_layers): + lr = f"{PREFIX_FUSION}.{i}" + consumed.add(f"{lr}.attn.v_proj.weight") + consumed.add(f"{lr}.attn.v_proj.bias") + consumed.add(f"{lr}.attn.l_proj.weight") + consumed.add(f"{lr}.attn.l_proj.bias") + consumed.add(f"{lr}.attn.values_v_proj.weight") + consumed.add(f"{lr}.attn.values_v_proj.bias") + consumed.add(f"{lr}.attn.values_l_proj.weight") + consumed.add(f"{lr}.attn.values_l_proj.bias") + consumed.add(f"{lr}.attn.out_v_proj.weight") + consumed.add(f"{lr}.attn.out_v_proj.bias") + consumed.add(f"{lr}.attn.out_l_proj.weight") + consumed.add(f"{lr}.attn.out_l_proj.bias") + consumed.add(f"{lr}.layer_norm_v.weight") + consumed.add(f"{lr}.layer_norm_v.bias") + consumed.add(f"{lr}.layer_norm_l.weight") + consumed.add(f"{lr}.layer_norm_l.bias") + consumed.add(f"{lr}.gamma_v") + consumed.add(f"{lr}.gamma_l") + + for i in range(dims.num_text_layers): + lr = f"{PREFIX_VL_TEXT}.{i}" + consumed.add(f"{lr}.self_attn.in_proj_weight") + consumed.add(f"{lr}.self_attn.in_proj_bias") + consumed.add(f"{lr}.self_attn.out_proj.weight") + consumed.add(f"{lr}.self_attn.out_proj.bias") + consumed.add(f"{lr}.norm1.weight") + consumed.add(f"{lr}.norm1.bias") + consumed.add(f"{lr}.linear1.weight") + consumed.add(f"{lr}.linear1.bias") + consumed.add(f"{lr}.linear2.weight") + consumed.add(f"{lr}.linear2.bias") + consumed.add(f"{lr}.norm2.weight") + consumed.add(f"{lr}.norm2.bias") + + consumed.add(f"{PREFIX_ACT_DEC}.action_queries.weight") + for i in range(dims.num_action_decoder_layers): + lr = f"{PREFIX_ACT_DEC}.decoder.layers.{i}" + consumed.add(f"{lr}.self_attn.in_proj_weight") + consumed.add(f"{lr}.self_attn.in_proj_bias") + consumed.add(f"{lr}.self_attn.out_proj.weight") + consumed.add(f"{lr}.self_attn.out_proj.bias") + consumed.add(f"{lr}.multihead_attn.in_proj_weight") + consumed.add(f"{lr}.multihead_attn.in_proj_bias") + consumed.add(f"{lr}.multihead_attn.out_proj.weight") + consumed.add(f"{lr}.multihead_attn.out_proj.bias") + consumed.add(f"{lr}.norm1.weight") + consumed.add(f"{lr}.norm1.bias") + consumed.add(f"{lr}.norm2.weight") + consumed.add(f"{lr}.norm2.bias") + consumed.add(f"{lr}.norm3.weight") + consumed.add(f"{lr}.norm3.bias") + consumed.add(f"{lr}.linear1.weight") + consumed.add(f"{lr}.linear1.bias") + consumed.add(f"{lr}.linear2.weight") + consumed.add(f"{lr}.linear2.bias") + + ap = f"{PREFIX_ACT_DEC}.action_projection" + consumed.add(f"{ap}.layers.0.weight") + consumed.add(f"{ap}.layers.0.bias") + consumed.add(f"{ap}.layers.1.weight") + consumed.add(f"{ap}.layers.1.bias") + consumed.add(f"{ap}.layers.2.weight") + consumed.add(f"{ap}.layers.2.bias") + + sp = PREFIX_STATE_PROJ + consumed.add(f"{sp}.net.0.weight") + consumed.add(f"{sp}.net.0.bias") + consumed.add(f"{sp}.net.1.weight") + consumed.add(f"{sp}.net.1.bias") + consumed.add(f"{sp}.net.4.weight") + consumed.add(f"{sp}.net.4.bias") + consumed.add(f"{sp}.output_norm.weight") + consumed.add(f"{sp}.output_norm.bias") + consumed.add(f"{sp}.position") + + consumed.add(KEY_VIEW_EMB) + + missing = source_keys - consumed + unexpected_missing = [ + k for k in missing + if not any(x in k for x in ["_optim", "_avg", "mask", "running_", ".nbytes"]) + and not k.endswith(".idx") + ] + + if unexpected_missing: + print(f"\n WARNING: {len(unexpected_missing)} keys not consumed:") + for k in sorted(unexpected_missing)[:20]: + print(f" {k}") + if len(unexpected_missing) > 20: + print(f" ... and {len(unexpected_missing) - 20} more") + print(f"\n Consumed: {len(consumed)} / {len(source_keys)} keys (missing {len(unexpected_missing)})") + else: + print(f"\n All {len(consumed)} keys consumed successfully!") + + +# ============================================================================= +# MAIN +# ============================================================================= + +def main() -> int: + import argparse + parser = argparse.ArgumentParser( + description=f"Convert TurboVLA checkpoint to GGUF ({ARCH})" + ) + parser.add_argument("ckpt", type=Path, help="Path to TurboVLA checkpoint directory") + parser.add_argument("-o", "--out", type=Path, default=None, help="Output GGUF path") + parser.add_argument("--dinov3-config", type=Path, required=True, + help="Validated config.json for the exact DINOv3 backbone") + parser.add_argument("--verify", action="store_true", help="Verify all tensor keys are consumed") + args = parser.parse_args() + + ckpt = args.ckpt.resolve() + if not (ckpt / "model.safetensors").exists(): + raise SystemExit(f"model.safetensors not found in {ckpt}") + + out = args.out or (ckpt / f"{ARCH}.gguf") + + print(f"Loading checkpoint from {ckpt}...") + tensors = load_safetensors(ckpt) + keys = set(tensors.keys()) + print(f" Loaded {len(tensors)} tensors") + + cfg_json = {} + cfg_path = ckpt / "config.json" + if cfg_path.exists(): + cfg_json = json.loads(cfg_path.read_text()) + print(f" Config: {cfg_path}") + + dinov3_cfg_path = args.dinov3_config.resolve() + if not dinov3_cfg_path.is_file(): + raise SystemExit(f"DINOv3 config not found: {dinov3_cfg_path}") + dinov3_cfg = json.loads(dinov3_cfg_path.read_text()) + print(f" DINOv3 config: {dinov3_cfg_path}") + + print(" Inferring dimensions from checkpoint...") + dims = TurboVLADims(tensors, keys, cfg_json, dinov3_cfg) + print(f" Detected: {dims}") + + print(f"Writing GGUF to {out}...") + out.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter(str(out), ARCH) + + kv = kv_prefix + writer.add_string(kv("architecture"), ARCH) + writer.add_uint32(kv("hidden"), dims.hidden_dim) + writer.add_uint32(kv("vit_dim"), dims.vit_dim) + writer.add_uint32(kv("vit_layers"), dims.vit_layers) + writer.add_uint32(kv("vit_head_dim"), 64) + writer.add_uint32(kv("vit_heads"), 12) + writer.add_uint32(kv("text_dim"), dims.text_dim) + writer.add_uint32(kv("text_layers"), dims.text_layers) + writer.add_uint32(kv("text_head_dim"), 64) + writer.add_uint32(kv("text_heads"), 12) + writer.add_uint32(kv("num_fusion_layers"), dims.num_fusion_layers) + # CORRECT: fusion_heads = nheads // 2, NOT nheads + writer.add_uint32(kv("fusion_heads"), dims.fusion_heads) + # CORRECT: fusion_head_dim = hidden / fusion_heads + writer.add_uint32(kv("fusion_head_dim"), dims.fusion_head_dim) + writer.add_uint32(kv("num_text_layers"), dims.num_text_layers) + # Text enhancer: same split as fusion + writer.add_uint32(kv("text_enhancer_heads"), dims.text_enhancer_heads) + writer.add_uint32(kv("text_enhancer_head_dim"), dims.text_enhancer_head_dim) + # Action transformer config + writer.add_uint32(kv("action_heads"), dims.action_heads) + writer.add_uint32(kv("action_head_dim"), dims.action_head_dim) + writer.add_uint32(kv("num_action_decoder_layers"), dims.num_action_decoder_layers) + writer.add_uint32(kv("action_dim"), dims.action_dim) + writer.add_uint32(kv("state_dim"), dims.state_dim) + writer.add_uint32(kv("num_state_tokens"), dims.num_state_tokens) + writer.add_uint32(kv("action_horizon"), dims.action_horizon) + writer.add_uint32(kv("image_size"), dims.image_size) + writer.add_uint32(kv("num_views"), dims.num_views) + writer.add_uint32(kv("patch_size"), dims.patch_size) + writer.add_uint32(kv("num_register_tokens"), dims.num_register_tokens) + writer.add_float32(kv("rope_theta"), float(dims.rope_theta)) + writer.add_uint32(kv("vocab_size"), dims.vocab_size) + writer.add_uint32(kv("max_text_length"), int(cfg_json.get("text", {}).get("padding_length", 256))) + # TurboVLA's BERT wrapper uses these exact punctuation IDs when it creates + # sub-sentence attention masks. Persist them so the GGUF runtime does not + # silently depend on a tokenizer installation. + text_cfg = cfg_json.get("text", {}) + model_name = text_cfg.get("model_name_or_path", "bert-base-uncased") + if model_name not in ("bert-base-uncased", "google-bert/bert-base-uncased"): + raise SystemExit( + "unsupported TurboVLA text tokenizer for deterministic mask conversion: " + f"{model_name!r}" + ) + writer.add_uint32(kv("pad_token_id"), 0) + writer.add_uint32(kv("cls_token_id"), 101) + writer.add_uint32(kv("sep_token_id"), 102) + writer.add_uint32(kv("period_token_id"), 1012) + writer.add_uint32(kv("question_token_id"), 1029) + writer.add_float32(kv("dropout"), dims.dropout) + + print(" Writing vision encoder...") + write_vision_encoder(writer, tensors, dims) + + print(" Writing text encoder...") + write_text_encoder(writer, tensors, dims) + + print(" Writing vision projection...") + write_vision_projection(writer, tensors) + + print(" Writing VL interaction...") + write_vision_language_interaction(writer, tensors, dims) + + print(" Writing action decoder...") + write_action_decoder(writer, tensors, dims) + + print(" Writing state projection...") + write_state_projection(writer, tensors) + + print(" Writing view embeddings...") + write_view_embeddings(writer, tensors) + + if cfg_json: + writer.add_string(kv("config_json"), json.dumps(cfg_json)) + + if args.verify: + print(" Verifying tensor consumption...") + verify_consumed_tensors(keys, dims) + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + size_mb = out.stat().st_size / (1024 * 1024) + print(f" Done! Output: {out} ({size_mb:.1f} MiB)") + return 0 + + +if __name__ == "__main__": + import sys + sys.exit(main()) diff --git a/scripts/tokenize_prompt.py b/scripts/tokenize_prompt.py index 4968224..aa604d6 100644 --- a/scripts/tokenize_prompt.py +++ b/scripts/tokenize_prompt.py @@ -37,6 +37,7 @@ "vla_jepa": "Qwen/Qwen3-VL-2B-Instruct", "gr00t_n1_5": "lerobot/eagle2hg-processor-groot-n1p5", "gr00t_n1_7": "nvidia/Cosmos-Reason2-2B", + "turbovla": "bert-base-uncased", } TRUST_REMOTE_CODE = {"evo1", "gr00t_n1_5"} diff --git a/src/arch.h b/src/arch.h index 517f4cb..47be355 100644 --- a/src/arch.h +++ b/src/arch.h @@ -71,6 +71,7 @@ enum class Arch { VLA_ADAPTER,// OpenHelix VLA-Adapter DINOv2 + SigLIP + Bridge-Attention. OPENVLA_OFT,// DINOv2-L/14-reg4 + SigLIP-so400m/14 +Llama-2-7B + MLPResNet. VLA_JEPA, // LeRobot Qwen3-VL-2B-Instruct+V-JEPÀ+DiT-B FM. + TURBOVLA, // TurboVLA (DINOv3 + BERT + VL Fusion + ACT decoder). }; /** @@ -217,6 +218,15 @@ std::unique_ptr vla_jepa_create(const std::string& mmproj_path, const std::string& config_path, const Options& opts); +/** + * @brief Build a TurboVLA model. Vision is baked into @p ckpt_path. + * @copydetails smolvla_create + */ +std::unique_ptr turbovla_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string& config_path, + const Options& opts); + /** * @brief Inspect a GGUF and identify the architecture tag. * diff --git a/src/layers/ffn.h b/src/layers/ffn.h index f514d93..2e26690 100644 --- a/src/layers/ffn.h +++ b/src/layers/ffn.h @@ -26,6 +26,11 @@ inline ggml_tensor * ffn_gelu(ggml_context * C, ggml_tensor * W1, ggml_tensor * return linear(C, W2, b2, ggml_gelu(C, linear(C, W1, b1, x))); } +inline ggml_tensor * ffn_relu(ggml_context * C, ggml_tensor * W1, ggml_tensor * b1, + ggml_tensor * W2, ggml_tensor * b2, ggml_tensor * x) { + return linear(C, W2, b2, ggml_relu(C, linear(C, W1, b1, x))); +} + // DINOv2 and SigLIP-so400m were trained with the exact erf form. inline ggml_tensor * ffn_gelu_erf(ggml_context * C, ggml_tensor * W1, ggml_tensor * b1, ggml_tensor * W2, ggml_tensor * b2, ggml_tensor * x) { diff --git a/src/layers/rope.h b/src/layers/rope.h index ca06156..19d5a2b 100644 --- a/src/layers/rope.h +++ b/src/layers/rope.h @@ -59,6 +59,25 @@ inline ggml_tensor * rope_2d(ggml_context * C, ggml_tensor * x, ggml_tensor * co return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, rot, sin_t)); } +// DINOv3 uses axial RoPE. x is [head_dim, patches, heads]; cos/sin are +// [head_dim, patches] generated from the official patch-centre coordinates. +inline ggml_tensor * rope_dinov3_axial(ggml_context * C, ggml_tensor * x, + ggml_tensor * cos_t, ggml_tensor * sin_t) { + GGML_ASSERT(x->ne[0] == cos_t->ne[0]); + GGML_ASSERT(x->ne[1] == cos_t->ne[1]); + ggml_tensor * cos = ggml_reshape_3d(C, cos_t, cos_t->ne[0], cos_t->ne[1], 1); + ggml_tensor * sin = ggml_reshape_3d(C, sin_t, sin_t->ne[0], sin_t->ne[1], 1); + cos = ggml_repeat(C, cos, x); + sin = ggml_repeat(C, sin, x); + const int64_t half = x->ne[0] / 2; + ggml_tensor * x0 = ggml_cont(C, ggml_view_3d(C, x, half, x->ne[1], x->ne[2], + x->nb[1], x->nb[2], 0)); + ggml_tensor * x1 = ggml_cont(C, ggml_view_3d(C, x, half, x->ne[1], x->ne[2], + x->nb[1], x->nb[2], half * x->nb[0])); + ggml_tensor * rotated = ggml_concat(C, ggml_neg(C, x1), x0, 0); + return ggml_add(C, ggml_mul(C, x, cos), ggml_mul(C, rotated, sin)); +} + 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]; diff --git a/src/loader.cpp b/src/loader.cpp index 8e231c0..0137731 100644 --- a/src/loader.cpp +++ b/src/loader.cpp @@ -154,8 +154,9 @@ bool WeightLoader::upload(ggml_backend_t backend, ggml_backend_buffer_t * out_bu 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)); + std::fprintf(stderr, + "vla(%s): failed to load %s (got %zu bytes, expected %zu, type=%d)\n", + arch_, name, bytes.size(), ggml_nbytes(t), (int) t->type); return false; } ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); diff --git a/src/model.cpp b/src/model.cpp index 5242efb..1b14754 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -76,7 +76,8 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { try_str("bitvla.architecture", arch_str) || try_str("openvla_oft.architecture", arch_str) || try_str("vla_jepa.architecture", arch_str) || - try_str("vla_adapter.architecture", arch_str)) { + try_str("vla_adapter.architecture", arch_str) || + try_str("turbovla.architecture", arch_str)) { if (arch_str == "smolvla") { *out = Arch::SMOLVLA; ok = true; @@ -127,6 +128,10 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { *out = Arch::VLA_JEPA; ok = true; } + else if (arch_str == "turbovla") { + *out = Arch::TURBOVLA; + ok = true; + } } gguf_free(gctx); @@ -280,6 +285,10 @@ Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, std::printf("vla: arch = vla_jepa\n"); impl = vla_jepa_create(mmproj_path, ckpt_path, config_path, opts); break; + case Arch::TURBOVLA: + std::printf("vla: arch = turbovla\n"); + impl = turbovla_create(mmproj_path, ckpt_path, config_path, opts); + break; } if (!impl) return nullptr; diff --git a/src/models/turbovla.cpp b/src/models/turbovla.cpp new file mode 100644 index 0000000..d23720e --- /dev/null +++ b/src/models/turbovla.cpp @@ -0,0 +1,1579 @@ + // 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 turbovla.cpp + * @brief TurboVLA implementation: DINOv3 ViT + BERT + VL Fusion + ACT decoder. + */ + + #include "arch.h" + #include "loader.h" + #include "model.h" + #include "options.h" + + #include "layers/linear.h" + #include "layers/norm.h" + #include "layers/ffn.h" + #include "layers/attn.h" + #include "layers/rope.h" + #include "modules/preprocess.h" + #include "gguf_reader.h" + #include "scratch_ctx.h" + #include "backend.h" + + #include "ggml.h" + #include "ggml-cpu.h" + #include "ggml-backend.h" + + #include + #include + #include + #include + #include + #include + #include + #include + #include + + namespace vla { + + // ============================================================================= + // CONFIG + // ============================================================================= + + struct TurboVLAConfig { + int64_t hidden = 256; + int64_t vit_dim = 768; + int64_t vit_layers = 12; + int64_t vit_head_dim = 64; + int64_t vit_heads = 12; + int64_t text_dim = 768; + int64_t text_layers = 12; + int64_t text_head_dim = 64; + int64_t text_heads = 12; + int64_t vocab_size = 30522; + int64_t max_text_length = 256; + int64_t num_fusion_layers = 6; + int64_t fusion_heads = 4; // TurboVLA uses 4 fusion heads + int64_t fusion_head_dim = 256; // 256 dim per fusion head + int64_t text_enhancer_heads = 4; + int64_t text_enhancer_head_dim = 64; + int64_t action_heads = 8; + int64_t action_head_dim = 32; + int64_t num_text_layers = 6; + int64_t num_action_decoder_layers = 3; + int64_t action_dim = 7; + int64_t state_dim = 8; + int64_t num_state_tokens = 2; + int64_t action_horizon = 12; + int64_t image_size = 256; + int64_t patch_size = 16; + int64_t num_views = 2; + int64_t num_register_tokens = 0; + float rope_theta = 10000.0f; // DINOv3 2D RoPE base frequency + float dropout = 0.1f; + int64_t pad_token_id = 0; + int64_t cls_token_id = 101; + int64_t sep_token_id = 102; + int64_t period_token_id = 1012; + int64_t question_token_id = 1029; + }; + + // ============================================================================= + // WEIGHT STRUCTURES + // ============================================================================= + + namespace { + + struct VitBlockW { + ggml_tensor* ln1_w, *ln1_b; + ggml_tensor* ln2_w, *ln2_b; + ggml_tensor* attn_q_w, *attn_q_b; + ggml_tensor* attn_k_w, *attn_k_b; + ggml_tensor* attn_v_w, *attn_v_b; + ggml_tensor* attn_o_w, *attn_o_b; + ggml_tensor* fc1_w, *fc1_b; + ggml_tensor* fc2_w, *fc2_b; + }; + + struct VisionEncoderW { + ggml_tensor* cls_token; + ggml_tensor* patch_embed_w; + ggml_tensor* patch_embed_b; + ggml_tensor* register_tokens; ///< DINOv3 register tokens (num_register_tokens x vit_dim) + std::vector layers; + ggml_tensor* final_ln_w; + ggml_tensor* final_ln_b; + // Precomputed 2D RoPE frequencies for patch tokens: [n_patches, head_dim/2] + std::vector rope_freqs; + }; + + struct VisionProjectionW { + ggml_tensor* input_norm_w, *input_norm_b; + ggml_tensor* mlp_0_w, *mlp_0_b; + ggml_tensor* mlp_3_w, *mlp_3_b; + ggml_tensor* skip_w; + ggml_tensor* output_norm_w, *output_norm_b; + }; + + struct BertLayerW { + ggml_tensor* attn_q_w, *attn_q_b; + ggml_tensor* attn_k_w, *attn_k_b; + ggml_tensor* attn_v_w, *attn_v_b; + ggml_tensor* attn_o_w, *attn_o_b; + ggml_tensor* ln1_w, *ln1_b; + ggml_tensor* fc1_w, *fc1_b; + ggml_tensor* fc2_w, *fc2_b; + ggml_tensor* ln2_w, *ln2_b; + }; + + struct TextEncoderW { + ggml_tensor* word_embed; + ggml_tensor* pos_embed; + ggml_tensor* tok_type_embed; + ggml_tensor* embed_ln_w, *embed_ln_b; + std::vector layers; + ggml_tensor* pooler_w, *pooler_b; + ggml_tensor* text_proj_w, *text_proj_b; + }; + + struct VlFusionLayerW { + ggml_tensor* v_proj_w, *v_proj_b; + ggml_tensor* l_proj_w, *l_proj_b; + ggml_tensor* values_v_w, *values_v_b; + ggml_tensor* values_l_w, *values_l_b; + ggml_tensor* out_v_w, *out_v_b; + ggml_tensor* out_l_w, *out_l_b; + ggml_tensor* norm_v_w, *norm_v_b; + ggml_tensor* norm_l_w, *norm_l_b; + ggml_tensor* gamma_v; + ggml_tensor* gamma_l; + }; + + struct VlTextLayerW { + ggml_tensor* attn_qkv_w, *attn_qkv_b; + ggml_tensor* attn_o_w, *attn_o_b; + ggml_tensor* ln1_w, *ln1_b; + ggml_tensor* fc1_w, *fc1_b; + ggml_tensor* fc2_w, *fc2_b; + ggml_tensor* ln2_w, *ln2_b; + }; + + struct ActDecoderLayerW { + ggml_tensor* self_qkv_w, *self_qkv_b; + ggml_tensor* self_out_w, *self_out_b; + ggml_tensor* cross_qkv_w, *cross_qkv_b; + ggml_tensor* cross_out_w, *cross_out_b; + ggml_tensor* ln1_w, *ln1_b; + ggml_tensor* ln2_w, *ln2_b; + ggml_tensor* ln3_w, *ln3_b; + ggml_tensor* fc1_w, *fc1_b; + ggml_tensor* fc2_w, *fc2_b; + }; + + struct ActionDecoderW { + ggml_tensor* action_q; + std::vector layers; + ggml_tensor* proj_0_w, *proj_0_b; + ggml_tensor* proj_1_w, *proj_1_b; + ggml_tensor* proj_2_w, *proj_2_b; + }; + + struct StateProjectionW { + ggml_tensor* net_0_w, *net_0_b; + ggml_tensor* net_1_w, *net_1_b; + ggml_tensor* net_4_w, *net_4_b; + ggml_tensor* output_norm_w, *output_norm_b; + ggml_tensor* position; + }; + + static void check_heads(const char* tag, ggml_tensor* t, int64_t hd, int64_t heads, int64_t T, int64_t nv = 1) { + (void) tag; + const int64_t got = ggml_nelements(t); + const int64_t expected = hd * heads * T * nv; + GGML_ASSERT(got == expected); + } + + } // anonymous namespace + + // ============================================================================= + // MODEL + // ============================================================================= + + struct TurboVLA : public ModelArchBase { + TurboVLA() : ModelArchBase(Arch::TURBOVLA) {} + ~TurboVLA() override; + + ggml_backend_t backend = nullptr; + int n_threads = default_cpu_threads(); + ggml_context* ctx_weights = nullptr; + ggml_backend_buffer_t weight_buf = nullptr; + scratch_ctx scratch; + + TurboVLAConfig cfg_; + + VisionEncoderW vision_encoder; + VisionProjectionW vision_proj; + TextEncoderW text_encoder; + std::vector fusion_layers; + std::vector text_layers; + ActionDecoderW action_decoder; + StateProjectionW state_proj; + ggml_tensor* view_emb = nullptr; + + std::vector predict(const Inputs& in) override; + + private: + ggml_tensor* vit_layer(ggml_context* C, const VitBlockW& w, ggml_tensor* x, + ggml_tensor* rope_cos, ggml_tensor* rope_sin, + int64_t internal_seq, int64_t n_heads, int64_t head_dim) const; + ggml_tensor* bert_layer(ggml_context* C, const BertLayerW& w, ggml_tensor* x, + ggml_tensor* self_mask, + int64_t seq, int64_t n_heads, int64_t head_dim) const; + ggml_tensor* vl_fusion_layer(ggml_context* C, const VlFusionLayerW& w, + ggml_tensor* v, ggml_tensor*& l, + ggml_tensor* language_key_mask, + int64_t v_seq, int64_t l_seq, + int64_t n_heads, int64_t head_dim) const; + ggml_tensor* vl_text_layer(ggml_context* C, const VlTextLayerW& w, ggml_tensor* x, + ggml_tensor* self_mask, + int64_t seq, int64_t n_heads, int64_t head_dim) const; + ggml_tensor* act_decoder_layer(ggml_context* C, const ActDecoderLayerW& w, + ggml_tensor* queries, ggml_tensor* memory, + int64_t q_seq, int64_t m_seq, + int64_t n_heads, int64_t head_dim) const; + }; + + TurboVLA::~TurboVLA() { + if (weight_buf) ggml_backend_buffer_free(weight_buf); + if (ctx_weights) ggml_free(ctx_weights); + if (backend) ggml_backend_free(backend); + } + + // ============================================================================= + // DINOv3 ViT LAYER with 2D RoPE + // ============================================================================= + // DINOv3 2D RoPE is applied in vit_layer() - see implementation below. + // The 2D positions are handled via ggml_rope which computes: + // x'[2i] = x[2i] * cos(pos * freq_i) - x[2i+1] * sin(pos * freq_i) + // x'[2i+1] = x[2i] * sin(pos * freq_i) + x[2i+1] * cos(pos * freq_i) + // where freq_i = rope_theta^(-2i/head_dim) + // Position IDs are patch indices (0, 1, 2, ..., n_patches-1) + + // ============================================================================= + // FORWARD FUNCTIONS + // ============================================================================= + + ggml_tensor* TurboVLA::vit_layer(ggml_context* C, const VitBlockW& w, ggml_tensor* x, + ggml_tensor* rope_cos, ggml_tensor* rope_sin, + int64_t internal_seq, int64_t n_heads, int64_t head_dim) const { + const float scale = 1.0f / std::sqrt((float)head_dim); + const int64_t hidden = n_heads * head_dim; + const int64_t patch_start = 1 + cfg_.num_register_tokens; + const int64_t n_patches = internal_seq - patch_start; + GGML_ASSERT(n_patches > 0 && rope_cos->ne[1] == n_patches && rope_sin->ne[1] == n_patches); + + ggml_tensor* normed = layer_norm(C, x, w.ln1_w, w.ln1_b, 1e-5f); + ggml_tensor* q = linear(C, w.attn_q_w, w.attn_q_b, normed); + ggml_tensor* k = linear(C, w.attn_k_w, w.attn_k_b, normed); + ggml_tensor* v = linear(C, w.attn_v_w, w.attn_v_b, normed); + + ggml_tensor* Q = to_heads(C, q, head_dim, n_heads, internal_seq); + ggml_tensor* K = to_heads(C, k, head_dim, n_heads, internal_seq); + ggml_tensor* V = to_heads_v(C, v, head_dim, n_heads, internal_seq); + const size_t q_offset = (size_t) patch_start * Q->nb[1]; + ggml_tensor* Q_prefix = ggml_cont(C, ggml_view_3d(C, Q, head_dim, patch_start, n_heads, + Q->nb[1], Q->nb[2], 0)); + ggml_tensor* K_prefix = ggml_cont(C, ggml_view_3d(C, K, head_dim, patch_start, n_heads, + K->nb[1], K->nb[2], 0)); + ggml_tensor* Q_patch = ggml_cont(C, ggml_view_3d(C, Q, head_dim, n_patches, n_heads, + Q->nb[1], Q->nb[2], q_offset)); + ggml_tensor* K_patch = ggml_cont(C, ggml_view_3d(C, K, head_dim, n_patches, n_heads, + K->nb[1], K->nb[2], (size_t) patch_start*K->nb[1])); + Q_patch = rope_dinov3_axial(C, Q_patch, rope_cos, rope_sin); + K_patch = rope_dinov3_axial(C, K_patch, rope_cos, rope_sin); + Q = ggml_concat(C, Q_prefix, Q_patch, 1); + K = ggml_concat(C, K_prefix, K_patch, 1); + ggml_tensor* att = attention(C, Q, K, V, nullptr, scale, hidden, internal_seq); + ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); + ggml_tensor* h = ggml_add(C, x, o); + ggml_tensor* mlp = ffn_gelu_erf(C, w.fc1_w, w.fc1_b, w.fc2_w, w.fc2_b, + layer_norm(C, h, w.ln2_w, w.ln2_b, 1e-5f)); + return ggml_add(C, h, mlp); + } + + ggml_tensor* TurboVLA::bert_layer(ggml_context* C, const BertLayerW& w, ggml_tensor* x, + ggml_tensor* self_mask, + int64_t seq, int64_t n_heads, int64_t head_dim) const { + const float scale = 1.0f / std::sqrt((float)head_dim); + const int64_t hidden = n_heads * head_dim; + + // BERT uses post-norm blocks: attention on x, residual + LayerNorm, + // then FFN, residual + LayerNorm. + ggml_tensor* q = linear(C, w.attn_q_w, w.attn_q_b, x); + ggml_tensor* k = linear(C, w.attn_k_w, w.attn_k_b, x); + ggml_tensor* v = linear(C, w.attn_v_w, w.attn_v_b, x); + + ggml_tensor* Q = to_heads(C, q, head_dim, n_heads, seq); + ggml_tensor* K = to_heads(C, k, head_dim, n_heads, seq); + ggml_tensor* V = to_heads_v(C, v, head_dim, n_heads, seq); + + ggml_tensor* att = attention(C, Q, K, V, self_mask, scale, hidden, seq); + ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); + ggml_tensor* h = layer_norm(C, + ggml_add(C, x, o), + w.ln1_w, + w.ln1_b, + 1e-12f); + + ggml_tensor* ffn_out = ffn_gelu_erf(C, + w.fc1_w, + w.fc1_b, + w.fc2_w, + w.fc2_b, + h); + return layer_norm(C, + ggml_add(C, h, ffn_out), + w.ln2_w, + w.ln2_b, + 1e-12f); + } + + ggml_tensor* TurboVLA::vl_fusion_layer(ggml_context* C, const VlFusionLayerW& w, + ggml_tensor* v, ggml_tensor*& l, + ggml_tensor* language_key_mask, + int64_t v_seq, int64_t l_seq, + int64_t n_heads, int64_t head_dim) const { + const float scale = 1.0f / std::sqrt((float)head_dim); + + // TurboVLA fusion: 4 heads × 64 dim = 256 att_dim + const int64_t att_dim = n_heads * head_dim; + + // FIXED (Point 5): Official uses residual_style="normalized" + // v_norm = LayerNorm(v), l_norm = LayerNorm(l) + // delta_v, delta_l = bi_attention(v_norm, l_norm) + // v_new = v_norm + gamma_v * delta_v + // l_new = l_norm + gamma_l * delta_l + ggml_tensor* v_norm = layer_norm(C, v, w.norm_v_w, w.norm_v_b, 1e-5f); + ggml_tensor* l_norm = layer_norm(C, l, w.norm_l_w, w.norm_l_b, 1e-5f); + + // V attends to L. Keep this as explicit QK^T + masked softmax rather than + // flash attention: the official implementation masks padded language keys + // on this branch only. + ggml_tensor* q_v = linear(C, w.v_proj_w, w.v_proj_b, v_norm); + ggml_tensor* k_l = linear(C, w.l_proj_w, w.l_proj_b, l_norm); + ggml_tensor* v_l = linear(C, w.values_l_w, w.values_l_b, l_norm); + + ggml_tensor* Q = to_heads(C, q_v, head_dim, n_heads, v_seq); + ggml_tensor* K = to_heads(C, k_l, head_dim, n_heads, l_seq); + ggml_tensor* V = to_heads_v(C, v_l, head_dim, n_heads, l_seq); + ggml_tensor* o = attention(C, Q, K, V, language_key_mask, scale, att_dim, v_seq); + ggml_tensor* o_proj = linear(C, w.out_v_w, w.out_v_b, o); + + // FIXED (Point 5): Apply gamma_v LayerScale with proper broadcasting + // gamma_v is [hidden], delta is [hidden, v_seq] + // Use ggml_repeat to broadcast gamma_v across sequence dimension + ggml_tensor* gamma_v_scaled = ggml_repeat(C, w.gamma_v, o_proj); + ggml_tensor* delta_v = ggml_mul(C, gamma_v_scaled, o_proj); + + // FIXED (Point 5): v_new = v_norm + gamma_v * delta_v + ggml_tensor* v_out = ggml_add(C, v_norm, delta_v); + + // L attends to V. There is no visual padding, hence no mask on this branch. + ggml_tensor* q_l = linear(C, w.l_proj_w, w.l_proj_b, l_norm); + ggml_tensor* k_v = linear(C, w.v_proj_w, w.v_proj_b, v_norm); + ggml_tensor* v_v = linear(C, w.values_v_w, w.values_v_b, v_norm); + + ggml_tensor* Q_l = to_heads(C, q_l, head_dim, n_heads, l_seq); + ggml_tensor* K_v = to_heads(C, k_v, head_dim, n_heads, v_seq); + ggml_tensor* V_v = to_heads_v(C, v_v, head_dim, n_heads, v_seq); + ggml_tensor* o_l = attention(C, Q_l, K_v, V_v, nullptr, scale, att_dim, l_seq); + ggml_tensor* o_proj_l = linear(C, w.out_l_w, w.out_l_b, o_l); + + // FIXED (Point 5): Apply gamma_l LayerScale with proper broadcasting + ggml_tensor* gamma_l_scaled = ggml_repeat(C, w.gamma_l, o_proj_l); + ggml_tensor* delta_l = ggml_mul(C, gamma_l_scaled, o_proj_l); + + // FIXED (Point 5): l_new = l_norm + gamma_l * delta_l + l = ggml_add(C, l_norm, delta_l); + + return v_out; + } + + ggml_tensor* TurboVLA::vl_text_layer(ggml_context* C, const VlTextLayerW& w, ggml_tensor* x, + ggml_tensor* self_mask, + int64_t seq, int64_t n_heads, int64_t head_dim) const { + // FIXED (Point 7): Official TurboVLA VL text enhancer uses POST-NORM and ReLU + // Structure: + // qkv = self_attn(x) + // h = x + attention_output + // h = LayerNorm(h) <- ln1 + // ff = Linear1(h) + // ff = ReLU(ff) + // ff = Linear2(ff) + // out = h + ff + // out = LayerNorm(out) <- ln2 + const float scale = 1.0f / std::sqrt((float)head_dim); + const int64_t att_dim = n_heads * head_dim; + const int64_t n = n_heads; + + ggml_tensor* qkv = linear(C, w.attn_qkv_w, w.attn_qkv_b, x); + + // Most checkpoints store fused [Q; K; V]. Keep compatibility with a + // checkpoint that stores a single shared projection as well. + ggml_tensor* q = nullptr; + ggml_tensor* k = nullptr; + ggml_tensor* v = nullptr; + if (qkv->ne[0] == 3 * att_dim) { + const size_t es = ggml_element_size(qkv); + q = ggml_cont(C, ggml_view_2d(C, qkv, att_dim, seq, qkv->nb[1], 0)); + k = ggml_cont(C, ggml_view_2d(C, qkv, att_dim, seq, qkv->nb[1], (size_t)att_dim * es)); + v = ggml_cont(C, ggml_view_2d(C, qkv, att_dim, seq, qkv->nb[1], (size_t)2 * att_dim * es)); + } else { + GGML_ASSERT(qkv->ne[0] == att_dim); + q = qkv; + k = qkv; + v = qkv; + } + + check_heads("VLText.Q", q, head_dim, n, seq); + check_heads("VLText.K", k, head_dim, n, seq); + check_heads("VLText.V", v, head_dim, n, seq); + ggml_tensor* Q = to_heads(C, q, head_dim, n, seq); + ggml_tensor* K = to_heads(C, k, head_dim, n, seq); + ggml_tensor* V = to_heads_v(C, v, head_dim, n, seq); + + ggml_tensor* att = attention(C, Q, K, V, self_mask, scale, att_dim, seq); + ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); + ggml_tensor* h = ggml_add(C, x, o); + + // POST-NORM after attention: LayerNorm(h) + h = layer_norm(C, h, w.ln1_w, w.ln1_b, 1e-5f); + + // FFN with ReLU (not GELU!) + ggml_tensor* ff = linear(C, w.fc1_w, w.fc1_b, h); + ff = ggml_relu(C, ff); + ggml_tensor* ff_out = linear(C, w.fc2_w, w.fc2_b, ff); + + // Add residual and final POST-NORM: LayerNorm(h + ff) + ggml_tensor* out = ggml_add(C, h, ff_out); + return layer_norm(C, out, w.ln2_w, w.ln2_b, 1e-5f); + } + + // ============================================================================= + // ACT DECODER LAYER + // ============================================================================= + + ggml_tensor* TurboVLA::act_decoder_layer(ggml_context* C, const ActDecoderLayerW& w, + ggml_tensor* queries, ggml_tensor* memory, + int64_t q_seq, int64_t m_seq, + int64_t n_heads, int64_t head_dim) const { + const float scale = 1.0f / std::sqrt((float)head_dim); + const int64_t hidden = n_heads * head_dim; + + // Self-attention on queries. + ggml_tensor* normed = layer_norm(C, queries, w.ln1_w, w.ln1_b, 1e-5f); + ggml_tensor* self_qkv = linear(C, w.self_qkv_w, w.self_qkv_b, normed); + + ggml_tensor* self_q = nullptr; + ggml_tensor* self_k = nullptr; + ggml_tensor* self_v = nullptr; + if (self_qkv->ne[0] == 3 * hidden) { + const size_t self_es = ggml_element_size(self_qkv); + self_q = ggml_cont(C, ggml_view_2d(C, self_qkv, hidden, q_seq, self_qkv->nb[1], 0)); + self_k = ggml_cont(C, ggml_view_2d(C, self_qkv, hidden, q_seq, self_qkv->nb[1], (size_t)hidden * self_es)); + self_v = ggml_cont(C, ggml_view_2d(C, self_qkv, hidden, q_seq, self_qkv->nb[1], (size_t)2 * hidden * self_es)); + } else { + GGML_ASSERT(self_qkv->ne[0] == hidden); + self_q = self_qkv; + self_k = self_qkv; + self_v = self_qkv; + } + + check_heads("ACT.self.Q", self_q, head_dim, n_heads, q_seq); + check_heads("ACT.self.K", self_k, head_dim, n_heads, q_seq); + check_heads("ACT.self.V", self_v, head_dim, n_heads, q_seq); + ggml_tensor* Q = to_heads(C, self_q, head_dim, n_heads, q_seq); + ggml_tensor* K = to_heads(C, self_k, head_dim, n_heads, q_seq); + ggml_tensor* V = to_heads_v(C, self_v, head_dim, n_heads, q_seq); + + ggml_tensor* att = attention(C, Q, K, V, nullptr, scale, hidden, q_seq); + ggml_tensor* o = linear(C, w.self_out_w, w.self_out_b, att); + ggml_tensor* h = ggml_add(C, queries, o); + + // Cross-attention. If cross_qkv is fused [Q;K;V], project both query and + // memory and select their respective blocks. If the checkpoint stores only + // a single query projection, preserve the original direct-memory K/V path. + normed = layer_norm(C, h, w.ln2_w, w.ln2_b, 1e-5f); + ggml_tensor* cross_qkv_q = linear(C, w.cross_qkv_w, w.cross_qkv_b, normed); + + ggml_tensor* cross_q = nullptr; + ggml_tensor* cross_k = nullptr; + ggml_tensor* cross_v = nullptr; + if (cross_qkv_q->ne[0] == 3 * hidden) { + ggml_tensor* cross_qkv_m = linear(C, w.cross_qkv_w, w.cross_qkv_b, memory); + GGML_ASSERT(cross_qkv_m->ne[0] == 3 * hidden); + const size_t cross_m_es = ggml_element_size(cross_qkv_m); + cross_q = ggml_cont(C, ggml_view_2d(C, cross_qkv_q, hidden, q_seq, cross_qkv_q->nb[1], 0)); + cross_k = ggml_cont(C, ggml_view_2d(C, cross_qkv_m, hidden, m_seq, cross_qkv_m->nb[1], (size_t)hidden * cross_m_es)); + cross_v = ggml_cont(C, ggml_view_2d(C, cross_qkv_m, hidden, m_seq, cross_qkv_m->nb[1], (size_t)2 * hidden * cross_m_es)); + } else { + GGML_ASSERT(cross_qkv_q->ne[0] == hidden); + cross_q = cross_qkv_q; + cross_k = ggml_is_contiguous(memory) ? memory : ggml_cont(C, memory); + cross_v = cross_k; + } + check_heads("ACT.cross.Q", cross_q, head_dim, n_heads, q_seq); + check_heads("ACT.cross.K", cross_k, head_dim, n_heads, m_seq); + check_heads("ACT.cross.V", cross_v, head_dim, n_heads, m_seq); + ggml_tensor* Q_c = to_heads(C, cross_q, head_dim, n_heads, q_seq); + ggml_tensor* K_c = to_heads(C, cross_k, head_dim, n_heads, m_seq); + ggml_tensor* V_c = to_heads_v(C, cross_v, head_dim, n_heads, m_seq); + + // Cross-attention output length is q_seq, not m_seq. + ggml_tensor* att_c = attention(C, Q_c, K_c, V_c, nullptr, scale, hidden, q_seq); + ggml_tensor* o_c = linear(C, w.cross_out_w, w.cross_out_b, att_c); + h = ggml_add(C, h, o_c); + + // FIXED (Point 10): Official PyTorch TransformerDecoderLayer uses ReLU, not GELU + normed = layer_norm(C, h, w.ln3_w, w.ln3_b, 1e-5f); + ggml_tensor* ffn_out = ffn_relu(C, w.fc1_w, w.fc1_b, w.fc2_w, w.fc2_b, normed); + + return ggml_add(C, h, ffn_out); + } + + // ============================================================================= + // INFERENCE + // ============================================================================= + + namespace { + + static const float IMAGENET_MEAN[3] = {0.485f, 0.456f, 0.406f}; + static const float IMAGENET_STD[3] = {0.229f, 0.224f, 0.225f}; + + // CHW planar layout matching ggml_conv_2d expectation + // Build the exact flattened Conv2D patch layout expected by a patch-embedding + // weight flattened from [out, C, patch_h, patch_w] to [out, C*patch_h*patch_w]. + // Output is GGML matrix layout [patch_dim, n_patches], i.e. each column is one + // normalized Cxpatchxpatch image patch. + void normalize_imagenet_patches(const ImageView& v, + int64_t side, + int64_t patch_size, + std::vector& out) { + const int64_t grid = side / patch_size; + const int64_t n_patches = grid * grid; + const int64_t patch_dim = 3 * patch_size * patch_size; + out.assign((size_t)patch_dim * n_patches, 0.0f); + + for (int64_t py = 0; py < grid; ++py) { + for (int64_t px_i = 0; px_i < grid; ++px_i) { + const int64_t patch_idx = py * grid + px_i; + for (int64_t c = 0; c < 3; ++c) { + for (int64_t ky = 0; ky < patch_size; ++ky) { + for (int64_t kx = 0; kx < patch_size; ++kx) { + const int64_t h = py * patch_size + ky; + const int64_t w = px_i * patch_size + kx; + float pixel; + if (v.format == PixelFormat::U8) { + pixel = ((const uint8_t*)v.data)[(h * side + w) * 3 + c] / 255.0f; + } else { + pixel = ((const float*)v.data)[(h * side + w) * 3 + c]; + } + const float norm = (pixel - IMAGENET_MEAN[c]) / IMAGENET_STD[c]; + const int64_t feature_idx = + (c * patch_size + ky) * patch_size + kx; + out[(size_t)feature_idx + (size_t)patch_idx * patch_dim] = norm; + } + } + } + } + } + } + + void build_dinov3_rope(int64_t grid, int64_t head_dim, float base, + std::vector& cos, std::vector& sin) { + const int64_t n_patches = grid * grid; + const int64_t axis_dim = head_dim / 4; + cos.resize((size_t)head_dim * n_patches); + sin.resize((size_t)head_dim * n_patches); + for (int64_t row = 0; row < grid; ++row) { + const float h = 2.0f * ((float(row) + 0.5f) / float(grid)) - 1.0f; + for (int64_t col = 0; col < grid; ++col) { + const float w = 2.0f * ((float(col) + 0.5f) / float(grid)) - 1.0f; + const size_t offset = (size_t)(row * grid + col) * head_dim; + for (int64_t i = 0; i < axis_dim; ++i) { + const float period = std::pow(base, 4.0f * float(i) / float(head_dim)); + constexpr float two_pi = 6.28318530717958647692f; + const float ah = two_pi * h / period; + const float aw = two_pi * w / period; + const int64_t dims[4] = {i, axis_dim + i, 2 * axis_dim + i, 3 * axis_dim + i}; + const float angles[4] = {ah, aw, ah, aw}; + for (int j = 0; j < 4; ++j) { + cos[offset + dims[j]] = std::cos(angles[j]); + sin[offset + dims[j]] = std::sin(angles[j]); + } + } + } + } + } + + struct TurboVLATextInputs { + std::vector token_ids; + std::vector position_ids; + std::vector bert_self_mask; + std::vector enhancer_self_mask; + std::vector fusion_language_key_mask; + }; + + TurboVLATextInputs build_turbovla_text_inputs(const Inputs& in, + const TurboVLAConfig& cfg, + int64_t vision_seq) { + const int64_t seq = cfg.max_text_length; + TurboVLATextInputs result; + result.token_ids.assign((size_t)seq, (int32_t)cfg.pad_token_id); + result.position_ids.assign((size_t)seq, 0); + + if (in.n_lang > seq) { + throw std::runtime_error("TurboVLA language input exceeds the checkpoint padding length"); + } + if (in.lang_tokens && in.n_lang > 0) { + std::memcpy(result.token_ids.data(), in.lang_tokens, + (size_t)in.n_lang * sizeof(int32_t)); + } + + std::vector valid((size_t)seq, 0); + for (int64_t i = 0; i < seq; ++i) { + const bool from_mask = in.attention_mask && + ((in.attention_mask_n == seq) || (i < in.attention_mask_n)); + valid[(size_t)i] = from_mask ? in.attention_mask[i] != 0 + : result.token_ids[(size_t)i] != cfg.pad_token_id; + } + + // Port of generate_masks_with_special_tokens from TurboVLA's BERT wrapper. + // Its block masks separate sub-sentences at [CLS], [SEP], '.' and '?'. + std::vector allowed((size_t)seq * seq, 0); + for (int64_t i = 0; i < seq; ++i) { + allowed[(size_t)i * seq + i] = 1; + } + int64_t previous_special = 0; + for (int64_t col = 0; col < seq; ++col) { + const int32_t id = result.token_ids[(size_t)col]; + const bool special = id == cfg.cls_token_id || id == cfg.sep_token_id || + id == cfg.period_token_id || id == cfg.question_token_id; + if (!special) { + continue; + } + if (col == 0 || col == seq - 1) { + result.position_ids[(size_t)col] = 0; + } else { + for (int64_t q = previous_special + 1; q <= col; ++q) { + result.position_ids[(size_t)q] = (int32_t)(q - previous_special - 1); + for (int64_t k = previous_special + 1; k <= col; ++k) { + allowed[(size_t)q * seq + k] = 1; + } + } + } + previous_special = col; + } + + result.bert_self_mask.resize((size_t)seq * seq); + result.enhancer_self_mask.resize((size_t)seq * seq); + for (int64_t q = 0; q < seq; ++q) { + for (int64_t k = 0; k < seq; ++k) { + const bool is_allowed = allowed[(size_t)q * seq + k] != 0; + result.bert_self_mask[(size_t)q * seq + k] = is_allowed ? 0.0f : -FLT_MAX; + // The post-BERT enhancer receives both the inverse source mask and + // BERT's key-padding mask. + result.enhancer_self_mask[(size_t)q * seq + k] = + is_allowed && valid[(size_t)k] ? 0.0f : -FLT_MAX; + } + } + + result.fusion_language_key_mask.resize((size_t)vision_seq * seq); + for (int64_t q = 0; q < vision_seq; ++q) { + for (int64_t k = 0; k < seq; ++k) { + result.fusion_language_key_mask[(size_t)q * seq + k] = + valid[(size_t)k] ? 0.0f : -FLT_MAX; + } + } + return result; + } + + } // anonymous namespace + + std::vector TurboVLA::predict(const Inputs& in) { + const auto& cfg = cfg_; + + const int64_t patch_size = cfg.patch_size; + const int64_t img_size = cfg.image_size; + const int64_t n_patches_per_side = img_size / patch_size; + const int64_t n_patches = n_patches_per_side * n_patches_per_side; + const int64_t n_reg = cfg.num_register_tokens; // DINOv3 register tokens + const int64_t n_cls = 1; // CLS token + const int64_t internal_vit_seq = n_cls + n_reg + n_patches; // Full internal ViT sequence + const int64_t n_views = cfg.num_views; + // vision_seq is n_patches per view (CLS + register tokens are sliced away) + const int64_t vision_seq = n_views * n_patches; + const int64_t hidden = cfg.hidden; // VLA hidden dim (256) + const int64_t vit_dim = cfg.vit_dim; // ViT hidden dim (768) + + // Basic shape sanity checks. These fail early with a useful message instead + // of failing later inside ggml_reshape_*. + if (patch_size <= 0 || img_size <= 0 || img_size % patch_size != 0) { + fprintf(stderr, "ERROR: invalid image/patch configuration\n"); + return {}; + } + if (cfg.vit_heads * cfg.vit_head_dim != vit_dim) { + fprintf(stderr, + "ERROR: vit_heads * vit_head_dim = %lld, expected vit_dim = %lld\n", + (long long)(cfg.vit_heads * cfg.vit_head_dim), + (long long)vit_dim); + return {}; + } + if (n_views <= 0 || !in.images || in.n_images < n_views) { + fprintf(stderr, + "ERROR: model expects %lld image view(s), but input has %lld\n", + (long long)n_views, + (long long)in.n_images); + return {}; + } + + // DINOv3 ViT processes internal sequence: [CLS, register_0..register_N-1, patch_0..patch_255] + // After final DINO norm, we slice away CLS and ALL register tokens + // Only n_patches per view go to vision_projection + // Output per view: [vit_dim, n_patches] (not internal_vit_seq) + const int64_t vision_out_size = n_patches * vit_dim; + std::vector vision_outputs((size_t)n_views * vision_out_size, 0.0f); + + for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { + ggml_context* VC = scratch.reset(256 * 1024 * 1024); // metadata arena + if (!VC) { + fprintf(stderr, "ERROR: vision scratch reset failed\n"); + return {}; + } + + // Patch-embedding input [patch_dim, n_patches]. Patch extraction is + // performed on the host so that each column is exactly one spatial + // Cxpatchxpatch block; a plain reshape of CHW image memory does not + // create spatial patches correctly. + const int64_t patch_dim = 3 * patch_size * patch_size; + ggml_tensor* patch_input = + ggml_new_tensor_2d(VC, GGML_TYPE_F32, patch_dim, n_patches); + ggml_set_input(patch_input); + + // patch_embed_w is a flattened Conv2D patch projection. GGML linear + // layout is [in=patch_dim, out=vit_dim], so mul_mat produces + // [vit_dim, n_patches]. + if (vision_encoder.patch_embed_w->ne[0] != patch_dim) { + fprintf(stderr, + "ERROR: patch_embed_w input dim=%lld, expected patch_dim=%lld\n", + (long long)vision_encoder.patch_embed_w->ne[0], + (long long)patch_dim); + return {}; + } + + ggml_tensor* patches_emb = + ggml_mul_mat(VC, vision_encoder.patch_embed_w, patch_input); + ggml_tensor* patches_2d = ggml_cont(VC, patches_emb); + + if (patches_2d->ne[0] != vit_dim || patches_2d->ne[1] != n_patches) { + fprintf(stderr, + "ERROR: patch embedding shape is [%lld,%lld], expected [%lld,%lld]\n", + (long long)patches_2d->ne[0], + (long long)patches_2d->ne[1], + (long long)vit_dim, + (long long)n_patches); + return {}; + } + + ggml_tensor* patches_with_bias = + ggml_add(VC, patches_2d, vision_encoder.patch_embed_b); + + ggml_tensor* rope_cos = ggml_new_tensor_2d(VC, GGML_TYPE_F32, cfg.vit_head_dim, n_patches); + ggml_tensor* rope_sin = ggml_new_tensor_2d(VC, GGML_TYPE_F32, cfg.vit_head_dim, n_patches); + ggml_set_input(rope_cos); + ggml_set_input(rope_sin); + std::vector rope_cos_host, rope_sin_host; + build_dinov3_rope(n_patches_per_side, cfg.vit_head_dim, cfg.rope_theta, + rope_cos_host, rope_sin_host); + + // ========================================================================= + // DINOv3 Token Flow: Build [CLS, register_tokens..., patches...] + // ========================================================================= + // Concatenate: [vit_dim, 1] + [vit_dim, n_reg] + [vit_dim, n_patches] + ggml_tensor* cls_2d = ggml_reshape_2d(VC, vision_encoder.cls_token, vit_dim, 1); + + ggml_tensor* vit_out = cls_2d; + if (n_reg > 0 && vision_encoder.register_tokens) { + // Add register tokens + ggml_tensor* reg_2d = ggml_reshape_2d(VC, vision_encoder.register_tokens, vit_dim, n_reg); + vit_out = ggml_concat(VC, vit_out, reg_2d, 1); + } + // Add patches + vit_out = ggml_concat(VC, vit_out, patches_with_bias, 1); + + // Run all DINO layers over FULL internal sequence [CLS, reg, patches] + for (int64_t i = 0; i < cfg.vit_layers; ++i) { + vit_out = vit_layer(VC, + vision_encoder.layers[i], + vit_out, + rope_cos, + rope_sin, + internal_vit_seq, + cfg.vit_heads, + cfg.vit_head_dim); + } + // ========================================================================= + // DINOv3: Slice away CLS and ALL register tokens + // Only patch tokens (n_patches) go to vision_projection + // ========================================================================= + // The reference TurboVLA encoder uses DINOv3's final hidden state + // (`outputs.hidden_states[-1]`), before the model's final LayerNorm. + // Its shape is [vit_dim, internal_vit_seq]. + // We need to extract only the patch tokens: positions [n_cls + n_reg, internal_vit_seq) + // For ViT-B with 256 patches and 4 register tokens: extract positions [5, 261) + const int64_t patch_token_start = n_cls + n_reg; // First patch position + + ggml_tensor* vit_patches_only = ggml_view_2d(VC, + vit_out, + vit_dim, + n_patches, + vit_out->nb[1], + patch_token_start * vit_out->nb[1]); + vit_patches_only = ggml_cont(VC, vit_patches_only); + + ggml_set_output(vit_patches_only); + + ggml_cgraph* vg = ggml_new_graph_custom(VC, 8192, false); + ggml_build_forward_expand(vg, vit_patches_only); + + if (!scratch.alloc(backend, vg)) { + fprintf(stderr, "ERROR: vision gallocr alloc failed\n"); + return {}; + } + + // Only now should patch_input have a backend buffer. + if (!patch_input->buffer) { + fprintf(stderr, + "ERROR: patch_input has no backend buffer after vision graph allocation; " + "vision graph is disconnected\n"); + return {}; + } + + std::vector patch_buf; + normalize_imagenet_patches(in.images[view_idx], + img_size, + patch_size, + patch_buf); + ggml_backend_tensor_set(patch_input, + patch_buf.data(), + 0, + patch_buf.size() * sizeof(float)); + ggml_backend_tensor_set(rope_cos, rope_cos_host.data(), 0, + rope_cos_host.size() * sizeof(float)); + ggml_backend_tensor_set(rope_sin, rope_sin_host.data(), 0, + rope_sin_host.size() * sizeof(float)); + + if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "ERROR: vision compute failed\n"); + return {}; + } + + float* out_ptr = vision_outputs.data() + view_idx * vision_out_size; + ggml_backend_tensor_get(vit_patches_only, + out_ptr, + 0, + vision_out_size * sizeof(float)); + } + + ggml_context* C = scratch.reset(512 * 1024 * 1024); // metadata arena + if (!C) { + fprintf(stderr, "ERROR: main scratch reset failed\n"); + return {}; + } + + // Create one graph input per view. Do NOT copy data yet: context is + // no_alloc=true and buffers are created by scratch.alloc(gf) below. + // NOTE: vision_seq = n_views * n_patches (CLS and register tokens were sliced away) + std::vector view_tokens((size_t)n_views, nullptr); + for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { + view_tokens[view_idx] = + ggml_new_tensor_2d(C, GGML_TYPE_F32, vit_dim, n_patches); + ggml_set_input(view_tokens[view_idx]); + } + + ggml_tensor* vit_tokens = view_tokens[0]; + for (int64_t view_idx = 1; view_idx < n_views; ++view_idx) { + vit_tokens = ggml_concat(C, vit_tokens, view_tokens[view_idx], 1); + } + + // Vision projection: [vit_dim, vision_seq] -> [hidden, vision_seq]. + ggml_tensor* proj = + layer_norm(C, + vit_tokens, + vision_proj.input_norm_w, + vision_proj.input_norm_b, + 1e-5f); + ggml_tensor* mlp_out = + linear(C, vision_proj.mlp_0_w, vision_proj.mlp_0_b, proj); + mlp_out = ggml_gelu_erf(C, mlp_out); + mlp_out = + linear(C, vision_proj.mlp_3_w, vision_proj.mlp_3_b, mlp_out); + + ggml_tensor* skip_out = + ggml_mul_mat(C, vision_proj.skip_w, vit_tokens); + proj = ggml_add(C, mlp_out, skip_out); + proj = layer_norm(C, + proj, + vision_proj.output_norm_w, + vision_proj.output_norm_b, + 1e-5f); + + // The checkpoint stores one embedding per camera. Each embedding must form + // a contiguous patch block: [view0 x 256 patches, view1 x 256 patches]. + const int64_t view_emb_elems = ggml_nelements(view_emb); + ggml_tensor* view_emb_used = nullptr; + if (view_emb_elems == (int64_t)hidden * n_views) { + ggml_tensor* by_view = ggml_reshape_2d(C, view_emb, hidden, n_views); + for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { + ggml_tensor* one_view = ggml_cont(C, ggml_view_2d( + C, by_view, hidden, 1, by_view->nb[1], view_idx * by_view->nb[1])); + ggml_tensor* view_block = ggml_repeat( + C, one_view, ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden, n_patches)); + view_emb_used = view_emb_used ? ggml_concat(C, view_emb_used, view_block, 1) : view_block; + } + } else if (view_emb_elems == (int64_t)hidden) { + // Single shared [hidden] embedding: repeat to [hidden, vision_seq] + view_emb_used = ggml_repeat(C, view_emb, + ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden, vision_seq)); + } else if (view_emb_elems == (int64_t)hidden * vision_seq) { + view_emb_used = ggml_reshape_2d(C, view_emb, (int64_t)hidden, vision_seq); + } else { + fprintf(stderr, + "ERROR: view_emb has %lld values; expected %lld (n_views,hidden), " + "%lld (shared), or %lld (per-pos)\n", + (long long)view_emb_elems, + (long long)((int64_t)hidden * n_views), + (long long)(int64_t)hidden, + (long long)((int64_t)hidden * vision_seq)); + return {}; + } + proj = ggml_add(C, proj, view_emb_used); + + const int64_t text_seq = cfg.max_text_length; + if (text_seq <= 0) { + fprintf(stderr, "ERROR: text_seq must be positive\n"); + return {}; + } + TurboVLATextInputs text_inputs; + try { + text_inputs = build_turbovla_text_inputs(in, cfg, vision_seq); + } catch (const std::exception& e) { + fprintf(stderr, "ERROR: %s\n", e.what()); + return {}; + } + std::vector tok_type_ids_host((size_t)text_seq, 0); + + ggml_tensor* token_ids = + ggml_new_tensor_1d(C, GGML_TYPE_I32, text_seq); + ggml_set_input(token_ids); + + ggml_tensor* pos_ids = + ggml_new_tensor_1d(C, GGML_TYPE_I32, text_seq); + ggml_set_input(pos_ids); + + ggml_tensor* tok_type_ids = + ggml_new_tensor_1d(C, GGML_TYPE_I32, text_seq); + ggml_set_input(tok_type_ids); + + ggml_tensor* bert_self_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, text_seq, text_seq); + ggml_tensor* enhancer_self_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, text_seq, text_seq); + ggml_tensor* fusion_language_key_mask = + ggml_new_tensor_2d(C, GGML_TYPE_F32, text_seq, vision_seq); + ggml_set_input(bert_self_mask); + ggml_set_input(enhancer_self_mask); + ggml_set_input(fusion_language_key_mask); + + ggml_tensor* pos_emb = + ggml_get_rows(C, text_encoder.pos_embed, pos_ids); + ggml_tensor* text_emb = + ggml_get_rows(C, text_encoder.word_embed, token_ids); + ggml_tensor* tok_type_emb = + ggml_get_rows(C, text_encoder.tok_type_embed, tok_type_ids); + + ggml_tensor* text_tokens = + ggml_add(C, ggml_add(C, text_emb, pos_emb), tok_type_emb); + text_tokens = layer_norm(C, + text_tokens, + text_encoder.embed_ln_w, + text_encoder.embed_ln_b, + 1e-12f); + + for (int64_t i = 0; i < cfg.text_layers; ++i) { + text_tokens = bert_layer(C, + text_encoder.layers[i], + text_tokens, + bert_self_mask, + text_seq, + cfg.text_heads, + cfg.text_head_dim); + } + text_tokens = + linear(C, + text_encoder.text_proj_w, + text_encoder.text_proj_b, + text_tokens); + ggml_tensor* vl_features = proj; + ggml_tensor* text_features = text_tokens; + + // TurboVLA fusion: 4 heads × 256 dim (hidden=256, v_proj outputs 4*256=1024) + const int64_t vl_heads = cfg.fusion_heads; + const int64_t vl_head_dim = cfg.fusion_head_dim; + + // TurboVLA text enhancer: 4 heads × 64 dim (hidden=256) + const int64_t text_enh_heads = cfg.text_enhancer_heads; + const int64_t text_enh_head_dim = cfg.text_enhancer_head_dim; + + // FIXED (Point 4): VL interaction is interleaved + // Official forward: + // for i in 0..num_layers-1: + // visual_tokens, text_tokens = fusion_layer[i](visual_tokens, text_tokens) + // text_tokens = text_layer[i](text_tokens) + GGML_ASSERT(cfg.num_fusion_layers == cfg.num_text_layers && + "VL fusion and text layers must have same count for interleaved execution"); + + for (int64_t i = 0; i < cfg.num_fusion_layers; ++i) { + // Fusion layer + vl_features = vl_fusion_layer(C, + fusion_layers[i], + vl_features, + text_features, + fusion_language_key_mask, + vision_seq, + text_seq, + vl_heads, + vl_head_dim); + // Text enhancer layer (operates on text_features in-place) + text_features = vl_text_layer(C, + text_layers[i], + text_features, + enhancer_self_mask, + text_seq, + text_enh_heads, + text_enh_head_dim); + } + + ggml_tensor* vl_condition = + ggml_concat(C, vl_features, text_features, 1); + const int64_t vl_cond_seq = vision_seq + text_seq; + + const int64_t state_dim = cfg.state_dim; + + std::vector state_host((size_t)state_dim, 0.0f); + if (in.state) { + std::copy(in.state, in.state + state_dim, state_host.begin()); + } + + ggml_tensor* state_tensor = + ggml_new_tensor_2d(C, GGML_TYPE_F32, state_dim, 1); + ggml_set_input(state_tensor); + + // FIXED (Point 9): State projection order - output_norm AFTER position add + // Official order: + // s = LayerNorm(state) + // s = Linear(...) + // s = GELU(s) + // s = Linear(...) + // tokens = reshape(s, num_state_tokens, hidden) + // tokens = tokens + learned_position + // tokens = output_norm(tokens) <- MOVED HERE + ggml_tensor* s = + layer_norm(C, + state_tensor, + state_proj.net_0_w, + state_proj.net_0_b, + 1e-5f); + s = linear(C, state_proj.net_1_w, state_proj.net_1_b, s); + s = ggml_gelu_erf(C, s); + s = linear(C, state_proj.net_4_w, state_proj.net_4_b, s); + + // state projection is expected to produce hidden * num_state_tokens values. + if (ggml_nelements(s) != hidden * cfg.num_state_tokens) { + fprintf(stderr, + "ERROR: state projection produced %lld values; expected hidden*num_state_tokens=%lld\n", + (long long)ggml_nelements(s), + (long long)(hidden * cfg.num_state_tokens)); + return {}; + } + + ggml_tensor* state_tokens = + ggml_reshape_2d(C, s, hidden, cfg.num_state_tokens); + ggml_tensor* pos_emb_state = + ggml_reshape_2d(C, + state_proj.position, + hidden, + cfg.num_state_tokens); + state_tokens = ggml_add(C, state_tokens, pos_emb_state); + + // FIXED: output_norm AFTER position add + state_tokens = layer_norm(C, + state_tokens, + state_proj.output_norm_w, + state_proj.output_norm_b, + 1e-5f); + + // IMPORTANT: previous code computed state_tokens and then discarded them. + // That made state_tensor a dead graph input and also meant robot state never + // influenced the action decoder. Include state tokens in decoder memory. + ggml_tensor* condition = + ggml_concat(C, vl_condition, state_tokens, 1); + const int64_t cond_seq = vl_cond_seq + cfg.num_state_tokens; + + ggml_tensor* queries = + ggml_reshape_2d(C, + action_decoder.action_q, + hidden, + cfg.action_horizon); + + // TurboVLA ACT decoder: 8 heads × 32 dim + const int64_t act_heads = cfg.action_heads; + + if (hidden % act_heads != 0) { + fprintf(stderr, + "ERROR: hidden=%lld is not divisible by act_heads=%lld\n", + (long long)hidden, + (long long)act_heads); + return {}; + } + const int64_t act_head_dim = cfg.action_head_dim; + + for (int64_t i = 0; i < cfg.num_action_decoder_layers; ++i) { + queries = act_decoder_layer(C, + action_decoder.layers[i], + queries, + condition, + cfg.action_horizon, + cond_seq, + act_heads, + act_head_dim); + } + + ggml_tensor* actions = + linear(C, action_decoder.proj_0_w, action_decoder.proj_0_b, queries); + actions = ggml_relu(C, actions); + actions = + linear(C, action_decoder.proj_1_w, action_decoder.proj_1_b, actions); + actions = ggml_relu(C, actions); + actions = + linear(C, action_decoder.proj_2_w, action_decoder.proj_2_b, actions); + actions = ggml_tanh(C, actions); + ggml_set_output(actions); + + ggml_cgraph* gf = ggml_new_graph_custom(C, 16384, false); + ggml_build_forward_expand(gf, actions); + + if (!scratch.alloc(backend, gf)) { + fprintf(stderr, "ERROR: scratch alloc failed\n"); + return {}; + } + + // All graph inputs now have backend buffers. Copy them only at this point. + for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { + if (!view_tokens[view_idx]->buffer) { + fprintf(stderr, + "ERROR: view_tokens[%lld] has no backend buffer after main graph allocation\n", + (long long)view_idx); + return {}; + } + const float* out_ptr = + vision_outputs.data() + view_idx * vision_out_size; + ggml_backend_tensor_set(view_tokens[view_idx], + out_ptr, + 0, + vision_out_size * sizeof(float)); + } + + if (!token_ids->buffer || !pos_ids->buffer || !tok_type_ids->buffer || + !bert_self_mask->buffer || !enhancer_self_mask->buffer || + !fusion_language_key_mask->buffer || !state_tensor->buffer) { + fprintf(stderr, + "ERROR: one or more main graph inputs have no backend buffer " + "(token=%p pos=%p type=%p bert=%p enhancer=%p fusion=%p state=%p)\n", + (void*)token_ids->buffer, + (void*)pos_ids->buffer, + (void*)tok_type_ids->buffer, + (void*)bert_self_mask->buffer, + (void*)enhancer_self_mask->buffer, + (void*)fusion_language_key_mask->buffer, + (void*)state_tensor->buffer); + return {}; + } + + ggml_backend_tensor_set(token_ids, + text_inputs.token_ids.data(), + 0, + text_inputs.token_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(pos_ids, + text_inputs.position_ids.data(), + 0, + text_inputs.position_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(tok_type_ids, + tok_type_ids_host.data(), + 0, + tok_type_ids_host.size() * sizeof(int32_t)); + ggml_backend_tensor_set(bert_self_mask, + text_inputs.bert_self_mask.data(), + 0, + text_inputs.bert_self_mask.size() * sizeof(float)); + ggml_backend_tensor_set(enhancer_self_mask, + text_inputs.enhancer_self_mask.data(), + 0, + text_inputs.enhancer_self_mask.size() * sizeof(float)); + ggml_backend_tensor_set(fusion_language_key_mask, + text_inputs.fusion_language_key_mask.data(), + 0, + text_inputs.fusion_language_key_mask.size() * sizeof(float)); + ggml_backend_tensor_set(state_tensor, + state_host.data(), + 0, + state_host.size() * sizeof(float)); + + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "ERROR: compute failed\n"); + return {}; + } + + const int64_t expected_action_values = + cfg.action_horizon * cfg.action_dim; + if (ggml_nelements(actions) != expected_action_values) { + fprintf(stderr, + "ERROR: actions has %lld values, expected %lld (= horizon %lld * dim %lld)\n", + (long long)ggml_nelements(actions), + (long long)expected_action_values, + (long long)cfg.action_horizon, + (long long)cfg.action_dim); + return {}; + } + + std::vector output((size_t)expected_action_values, 0.0f); + ggml_backend_tensor_get(actions, + output.data(), + 0, + output.size() * sizeof(float)); + + return output; + } + + // ============================================================================= + // MODEL LOADING + // ============================================================================= + + std::unique_ptr turbovla_create(const std::string& /*mmproj_path*/, + const std::string& ckpt_path, + const std::string& /*config_path*/, + const Options& /*opts*/) { + // Open GGUF + gguf_reader g("turbovla"); + if (!g.open(ckpt_path)) { + std::fprintf(stderr, "vla(turbovla): failed to open %s\n", ckpt_path.c_str()); + return nullptr; + } + + auto* m = new vla::TurboVLA; + + // Load config from GGUF + m->cfg_.hidden = (int64_t)g.u32("turbovla.hidden"); + m->cfg_.vit_dim = (int64_t)g.u32("turbovla.vit_dim"); + m->cfg_.vit_layers = (int64_t)g.u32("turbovla.vit_layers"); + m->cfg_.vit_head_dim = (int64_t)g.u32("turbovla.vit_head_dim"); + m->cfg_.vit_heads = (int64_t)g.u32("turbovla.vit_heads"); + m->cfg_.text_dim = (int64_t)g.u32("turbovla.text_dim"); + m->cfg_.text_layers = (int64_t)g.u32("turbovla.text_layers"); + m->cfg_.text_head_dim = (int64_t)g.u32("turbovla.text_head_dim"); + m->cfg_.text_heads = (int64_t)g.u32("turbovla.text_heads"); + m->cfg_.num_fusion_layers = (int64_t)g.u32("turbovla.num_fusion_layers"); + m->cfg_.fusion_heads = (int64_t)g.u32("turbovla.fusion_heads"); + m->cfg_.fusion_head_dim = (int64_t)g.u32("turbovla.fusion_head_dim"); + m->cfg_.text_enhancer_heads = (int64_t)g.u32("turbovla.text_enhancer_heads"); + m->cfg_.text_enhancer_head_dim = (int64_t)g.u32("turbovla.text_enhancer_head_dim"); + m->cfg_.action_heads = (int64_t)g.u32("turbovla.action_heads"); + m->cfg_.action_head_dim = (int64_t)g.u32("turbovla.action_head_dim"); + m->cfg_.num_text_layers = (int64_t)g.u32("turbovla.num_text_layers"); + m->cfg_.num_action_decoder_layers = (int64_t)g.u32("turbovla.num_action_decoder_layers"); + m->cfg_.action_dim = (int64_t)g.u32("turbovla.action_dim"); + m->cfg_.state_dim = (int64_t)g.u32("turbovla.state_dim"); + m->cfg_.num_state_tokens = (int64_t)g.u32("turbovla.num_state_tokens"); + m->cfg_.action_horizon = (int64_t)g.u32("turbovla.action_horizon"); + m->cfg_.image_size = (int64_t)g.u32("turbovla.image_size"); + m->cfg_.patch_size = (int64_t)g.u32("turbovla.patch_size"); + m->cfg_.num_views = (int64_t)g.u32("turbovla.num_views"); + m->cfg_.num_register_tokens = (int64_t)g.u32("turbovla.num_register_tokens"); + if (!g.has("turbovla.rope_theta")) { + std::fprintf(stderr, "vla(turbovla): GGUF is missing required DINOv3 rope_theta metadata\n"); + return nullptr; + } + m->cfg_.rope_theta = g.f32("turbovla.rope_theta"); + m->cfg_.max_text_length = (int64_t)g.u32("turbovla.max_text_length"); + if (!g.has("turbovla.pad_token_id") || !g.has("turbovla.cls_token_id") || + !g.has("turbovla.sep_token_id") || !g.has("turbovla.period_token_id") || + !g.has("turbovla.question_token_id")) { + std::fprintf(stderr, "vla(turbovla): GGUF is missing required BERT special-token metadata\n"); + return nullptr; + } + m->cfg_.pad_token_id = (int64_t)g.u32("turbovla.pad_token_id"); + m->cfg_.cls_token_id = (int64_t)g.u32("turbovla.cls_token_id"); + m->cfg_.sep_token_id = (int64_t)g.u32("turbovla.sep_token_id"); + m->cfg_.period_token_id = (int64_t)g.u32("turbovla.period_token_id"); + m->cfg_.question_token_id = (int64_t)g.u32("turbovla.question_token_id"); + if (m->cfg_.fusion_heads * m->cfg_.fusion_head_dim != 1024 || + m->cfg_.text_enhancer_heads * m->cfg_.text_enhancer_head_dim != m->cfg_.hidden || + m->cfg_.action_heads * m->cfg_.action_head_dim != m->cfg_.hidden) { + std::fprintf(stderr, "vla(turbovla): incompatible attention metadata\n"); + return nullptr; + } + + // Publish the dimensions used by serving tools. Without this, callers + // allocate an empty state vector from ModelArchBase::cfg while predict() + // reads the checkpoint's state_dim values. + const int64_t patches_per_side = m->cfg_.image_size / m->cfg_.patch_size; + m->cfg.n_img = m->cfg_.num_views * patches_per_side * patches_per_side; + m->cfg.n_lang = m->cfg_.max_text_length; + m->cfg.n_state = m->cfg_.num_state_tokens; + m->cfg.n_prefix = m->cfg.n_img + m->cfg.n_lang + m->cfg.n_state; + m->cfg.n_suffix = m->cfg_.action_horizon; + m->cfg.n_full = m->cfg.n_prefix + m->cfg.n_suffix; + m->cfg.hidden = m->cfg_.hidden; + m->cfg.n_q_heads = m->cfg_.action_heads; + m->cfg.n_kv_heads = m->cfg_.action_heads; + m->cfg.head_dim = m->cfg_.action_head_dim; + m->cfg.q_full_dim = m->cfg_.hidden; + m->cfg.kv_full_dim = m->cfg_.hidden; + m->cfg.n_layers = m->cfg_.num_action_decoder_layers; + m->cfg.max_state_dim = m->cfg_.state_dim; + m->cfg.real_state_dim = m->cfg_.state_dim; + m->cfg.max_action_dim = m->cfg_.action_dim; + m->cfg.real_action_dim = m->cfg_.action_dim; + m->cfg.norm_eps = 1e-5f; + m->cfg.rms_eps = 1e-5f; + m->cfg.rope_n_dims = m->cfg_.vit_head_dim; + m->cfg.rope_freq_base = m->cfg_.rope_theta; + m->cfg.num_steps = m->cfg_.action_horizon; + + // Init backend + const Backend b = backend_init("vla(turbovla)", m->n_threads); + if (!b.handle) { + std::fprintf(stderr, "vla(turbovla): backend_init failed\n"); + return nullptr; + } + m->backend = b.handle; + + ggml_init_params wp = { (size_t)16*1024*1024, nullptr, true }; + m->ctx_weights = ggml_init(wp); + + // Load weights + WeightLoader L("turbovla", g, m->ctx_weights, GGML_TYPE_F32); + + // Vision encoder + m->vision_encoder.cls_token = L.f32("vit.cls_token"); + m->vision_encoder.patch_embed_w = L.gemm("vit.patch_embed.weight"); + m->vision_encoder.patch_embed_b = L.gemm("vit.patch_embed.bias"); + + // DINOv3 register tokens: load only when num_register_tokens > 0 + if (m->cfg_.num_register_tokens > 0) { + m->vision_encoder.register_tokens = L.f32("vit.register_tokens"); + if (!m->vision_encoder.register_tokens) { + std::fprintf(stderr, "vla(turbovla): WARNING: num_register_tokens=%lld but vit.register_tokens not found\n", + (long long)m->cfg_.num_register_tokens); + } + } + + m->vision_encoder.layers.resize(m->cfg_.vit_layers); + for (int64_t i = 0; i < m->cfg_.vit_layers; ++i) { + auto& w = m->vision_encoder.layers[i]; + w.ln1_w = L.f32("vit.blk.%lld.ln1.weight", (long long) i); + w.ln1_b = L.f32("vit.blk.%lld.ln1.bias", (long long) i); + w.attn_q_w = L.f32("vit.blk.%lld.attn_q.weight", (long long) i); + w.attn_q_b = L.f32("vit.blk.%lld.attn_q.bias", (long long) i); + w.attn_k_w = L.f32("vit.blk.%lld.attn_k.weight", (long long) i); + w.attn_k_b = L.f32("vit.blk.%lld.attn_k.bias", (long long) i); + w.attn_v_w = L.f32("vit.blk.%lld.attn_v.weight", (long long) i); + w.attn_v_b = L.f32("vit.blk.%lld.attn_v.bias", (long long) i); + w.attn_o_w = L.f32("vit.blk.%lld.attn_o.weight", (long long) i); + w.attn_o_b = L.f32("vit.blk.%lld.attn_o.bias", (long long) i); + w.ln2_w = L.f32("vit.blk.%lld.ln2.weight", (long long) i); + w.ln2_b = L.f32("vit.blk.%lld.ln2.bias", (long long) i); + w.fc1_w = L.f32("vit.blk.%lld.fc1.weight", (long long) i); + w.fc1_b = L.f32("vit.blk.%lld.fc1.bias", (long long) i); + w.fc2_w = L.f32("vit.blk.%lld.fc2.weight", (long long) i); + w.fc2_b = L.f32("vit.blk.%lld.fc2.bias", (long long) i); + } + m->vision_encoder.final_ln_w = L.f32("vit.final_norm.weight"); + m->vision_encoder.final_ln_b = L.f32("vit.final_norm.bias"); + + // Precompute DINOv3 2D RoPE frequencies + // This is done in predict() since we need n_patches_per_side which depends on image_size + + // Vision projection + m->vision_proj.input_norm_w = L.f32("vit_proj.input_norm.weight"); + m->vision_proj.input_norm_b = L.f32("vit_proj.input_norm.bias"); + m->vision_proj.mlp_0_w = L.f32("vit_proj.mlp.0.weight"); + m->vision_proj.mlp_0_b = L.f32("vit_proj.mlp.0.bias"); + m->vision_proj.mlp_3_w = L.f32("vit_proj.mlp.3.weight"); + m->vision_proj.mlp_3_b = L.f32("vit_proj.mlp.3.bias"); + m->vision_proj.skip_w = L.f32("vit_proj.skip.weight"); + m->vision_proj.output_norm_w = L.f32("vit_proj.output_norm.weight"); + m->vision_proj.output_norm_b = L.f32("vit_proj.output_norm.bias"); + + // Text encoder + m->text_encoder.word_embed = L.f32("text.embed.word_embeddings"); + m->text_encoder.pos_embed = L.f32("text.embed.position_embeddings"); + m->text_encoder.tok_type_embed = L.f32("text.embed.token_type_embeddings"); + m->text_encoder.embed_ln_w = L.f32("text.embed.LayerNorm.weight"); + m->text_encoder.embed_ln_b = L.f32("text.embed.LayerNorm.bias"); + m->text_encoder.pooler_w = L.f32("text.pooler.dense.weight"); + m->text_encoder.pooler_b = L.f32("text.pooler.dense.bias"); + m->text_encoder.text_proj_w = L.f32("text_proj.weight"); + m->text_encoder.text_proj_b = L.f32("text_proj.bias"); + + m->text_encoder.layers.resize(m->cfg_.text_layers); + for (int64_t i = 0; i < m->cfg_.text_layers; ++i) { + auto& w = m->text_encoder.layers[i]; + w.attn_q_w = L.f32("text.encoder.layer.%lld.attention.self.query.weight", (long long) i); + w.attn_q_b = L.f32("text.encoder.layer.%lld.attention.self.query.bias", (long long) i); + w.attn_k_w = L.f32("text.encoder.layer.%lld.attention.self.key.weight", (long long) i); + w.attn_k_b = L.f32("text.encoder.layer.%lld.attention.self.key.bias", (long long) i); + w.attn_v_w = L.f32("text.encoder.layer.%lld.attention.self.value.weight", (long long) i); + w.attn_v_b = L.f32("text.encoder.layer.%lld.attention.self.value.bias", (long long) i); + w.attn_o_w = L.f32("text.encoder.layer.%lld.attention.output.dense.weight", (long long) i); + w.attn_o_b = L.f32("text.encoder.layer.%lld.attention.output.dense.bias", (long long) i); + w.ln1_w = L.f32("text.encoder.layer.%lld.attention.output.LayerNorm.weight", (long long) i); + w.ln1_b = L.f32("text.encoder.layer.%lld.attention.output.LayerNorm.bias", (long long) i); + w.fc1_w = L.f32("text.encoder.layer.%lld.intermediate.dense.weight", (long long) i); + w.fc1_b = L.f32("text.encoder.layer.%lld.intermediate.dense.bias", (long long) i); + w.fc2_w = L.f32("text.encoder.layer.%lld.output.dense.weight", (long long) i); + w.fc2_b = L.f32("text.encoder.layer.%lld.output.dense.bias", (long long) i); + w.ln2_w = L.f32("text.encoder.layer.%lld.output.LayerNorm.weight", (long long) i); + w.ln2_b = L.f32("text.encoder.layer.%lld.output.LayerNorm.bias", (long long) i); + } + + // VL Fusion layers + m->fusion_layers.resize(m->cfg_.num_fusion_layers); + for (int64_t i = 0; i < m->cfg_.num_fusion_layers; ++i) { + auto& w = m->fusion_layers[i]; + w.v_proj_w = L.f32("vl_fusion.%lld.v_proj.weight", (long long) i); + w.v_proj_b = L.f32("vl_fusion.%lld.v_proj.bias", (long long) i); + w.l_proj_w = L.f32("vl_fusion.%lld.l_proj.weight", (long long) i); + w.l_proj_b = L.f32("vl_fusion.%lld.l_proj.bias", (long long) i); + w.values_v_w = L.f32("vl_fusion.%lld.values_v.weight", (long long) i); + w.values_v_b = L.f32("vl_fusion.%lld.values_v.bias", (long long) i); + w.values_l_w = L.f32("vl_fusion.%lld.values_l.weight", (long long) i); + w.values_l_b = L.f32("vl_fusion.%lld.values_l.bias", (long long) i); + w.out_v_w = L.f32("vl_fusion.%lld.out_v.weight", (long long) i); + w.out_v_b = L.f32("vl_fusion.%lld.out_v.bias", (long long) i); + w.out_l_w = L.f32("vl_fusion.%lld.out_l.weight", (long long) i); + w.out_l_b = L.f32("vl_fusion.%lld.out_l.bias", (long long) i); + w.norm_v_w = L.f32("vl_fusion.%lld.norm_v.weight", (long long) i); + w.norm_v_b = L.f32("vl_fusion.%lld.norm_v.bias", (long long) i); + w.norm_l_w = L.f32("vl_fusion.%lld.norm_l.weight", (long long) i); + w.norm_l_b = L.f32("vl_fusion.%lld.norm_l.bias", (long long) i); + w.gamma_v = L.f32("vl_fusion.%lld.gamma_v", (long long) i); + w.gamma_l = L.f32("vl_fusion.%lld.gamma_l", (long long) i); + } + + // VL Text layers + m->text_layers.resize(m->cfg_.num_text_layers); + for (int64_t i = 0; i < m->cfg_.num_text_layers; ++i) { + auto& w = m->text_layers[i]; + w.attn_qkv_w = L.f32("vl_text.%lld.attn_qkv.weight", (long long) i); + w.attn_qkv_b = L.f32("vl_text.%lld.attn_qkv.bias", (long long) i); + w.attn_o_w = L.f32("vl_text.%lld.attn_o.weight", (long long) i); + w.attn_o_b = L.f32("vl_text.%lld.attn_o.bias", (long long) i); + w.ln1_w = L.f32("vl_text.%lld.ln1.weight", (long long) i); + w.ln1_b = L.f32("vl_text.%lld.ln1.bias", (long long) i); + w.fc1_w = L.f32("vl_text.%lld.fc1.weight", (long long) i); + w.fc1_b = L.f32("vl_text.%lld.fc1.bias", (long long) i); + w.fc2_w = L.f32("vl_text.%lld.fc2.weight", (long long) i); + w.fc2_b = L.f32("vl_text.%lld.fc2.bias", (long long) i); + w.ln2_w = L.f32("vl_text.%lld.ln2.weight", (long long) i); + w.ln2_b = L.f32("vl_text.%lld.ln2.bias", (long long) i); + } + + // Action decoder + m->action_decoder.action_q = L.f32("act.q.weight"); + m->action_decoder.layers.resize(m->cfg_.num_action_decoder_layers); + for (int64_t i = 0; i < m->cfg_.num_action_decoder_layers; ++i) { + auto& w = m->action_decoder.layers[i]; + w.self_qkv_w = L.f32("act.dec.%lld.self_qkv.weight", (long long) i); + w.self_qkv_b = L.f32("act.dec.%lld.self_qkv.bias", (long long) i); + w.self_out_w = L.f32("act.dec.%lld.self_out.weight", (long long) i); + w.self_out_b = L.f32("act.dec.%lld.self_out.bias", (long long) i); + w.cross_qkv_w = L.f32("act.dec.%lld.cross_qkv.weight", (long long) i); + w.cross_qkv_b = L.f32("act.dec.%lld.cross_qkv.bias", (long long) i); + w.cross_out_w = L.f32("act.dec.%lld.cross_out.weight", (long long) i); + w.cross_out_b = L.f32("act.dec.%lld.cross_out.bias", (long long) i); + w.ln1_w = L.f32("act.dec.%lld.ln1.weight", (long long) i); + w.ln1_b = L.f32("act.dec.%lld.ln1.bias", (long long) i); + w.ln2_w = L.f32("act.dec.%lld.ln2.weight", (long long) i); + w.ln2_b = L.f32("act.dec.%lld.ln2.bias", (long long) i); + w.ln3_w = L.f32("act.dec.%lld.ln3.weight", (long long) i); + w.ln3_b = L.f32("act.dec.%lld.ln3.bias", (long long) i); + w.fc1_w = L.f32("act.dec.%lld.fc1.weight", (long long) i); + w.fc1_b = L.f32("act.dec.%lld.fc1.bias", (long long) i); + w.fc2_w = L.f32("act.dec.%lld.fc2.weight", (long long) i); + w.fc2_b = L.f32("act.dec.%lld.fc2.bias", (long long) i); + } + m->action_decoder.proj_0_w = L.f32("act.proj.0.weight"); + m->action_decoder.proj_0_b = L.f32("act.proj.0.bias"); + m->action_decoder.proj_1_w = L.f32("act.proj.1.weight"); + m->action_decoder.proj_1_b = L.f32("act.proj.1.bias"); + m->action_decoder.proj_2_w = L.f32("act.proj.2.weight"); + m->action_decoder.proj_2_b = L.f32("act.proj.2.bias"); + + // State projection + m->state_proj.net_0_w = L.f32("state.proj.0.weight"); + m->state_proj.net_0_b = L.f32("state.proj.0.bias"); + m->state_proj.net_1_w = L.f32("state.proj.1.weight"); + m->state_proj.net_1_b = L.f32("state.proj.1.bias"); + m->state_proj.net_4_w = L.f32("state.proj.4.weight"); + m->state_proj.net_4_b = L.f32("state.proj.4.bias"); + m->state_proj.output_norm_w = L.f32("state.proj.output_norm.weight"); + m->state_proj.output_norm_b = L.f32("state.proj.output_norm.bias"); + m->state_proj.position = L.f32("state.proj.position"); + + // View embeddings + m->view_emb = L.f32("view_emb"); + + if (!L.upload(m->backend, &m->weight_buf)) { + std::fprintf(stderr, "vla(turbovla): weight upload failed\n"); + return nullptr; + } + + // Print DINOv3 config + std::printf("vla(turbovla): hidden=%lld vit_layers=%lld text_layers=%lld fusion=%lld action_horizon=%lld action_dim=%lld\n", + (long long)m->cfg_.hidden, + (long long)m->cfg_.vit_layers, + (long long)m->cfg_.text_layers, + (long long)m->cfg_.num_fusion_layers, + (long long)m->cfg_.action_horizon, + (long long)m->cfg_.action_dim); + std::printf("vla(turbovla): DINOv3 config - image_size=%lld patch_size=%lld num_register_tokens=%lld rope_theta=%.1f\n", + (long long)m->cfg_.image_size, + (long long)m->cfg_.patch_size, + (long long)m->cfg_.num_register_tokens, + (double)m->cfg_.rope_theta); + + return std::unique_ptr(m); + } + + } // namespace vla diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp index 7d31060..ec48e13 100644 --- a/src/serving/vla-cli.cpp +++ b/src/serving/vla-cli.cpp @@ -130,6 +130,7 @@ const char * arch_slug(Arch a) { case Arch::OPENVLA_OFT: return "openvla_oft"; case Arch::VLA_JEPA: return "vla_jepa"; case Arch::OCTO: return "octo"; + case Arch::TURBOVLA: return "turbovla"; } return ""; } From d959889265e7ec2716b1cc2ddc865658c995217d Mon Sep 17 00:00:00 2001 From: Minh Hieu Date: Thu, 24 Sep 2026 20:05:29 +0700 Subject: [PATCH 2/5] docs(turbovla): document support and converter mapping --- README.md | 2 + eval/client/vla_cpp_client.py | 5 +- scripts/convert_turbovla_to_gguf.py | 73 ++--------- src/models/turbovla.cpp | 184 ++++++---------------------- tests/py/test_converters.py | 90 ++++++++++++++ 5 files changed, 138 insertions(+), 216 deletions(-) diff --git a/README.md b/README.md index 7978201..d00a12c 100644 --- a/README.md +++ b/README.md @@ -326,6 +326,7 @@ supported (released and benchmarked), `~` = in progress, `-` = planned. | [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | - | Y | Y | - | | [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | - | Y | Y | - | | [Octo-Small](https://hf.co/vrfai/octo-small-libero-gguf) | Y | Y | Y | Y | - | - | +| [TurboVLA](https://github.com/H-EmbodVis/TurboVLA) | ~ | ~ | - | - | - | - | --- @@ -364,6 +365,7 @@ Supported VLA models: - [GR00T N1.x](https://github.com/NVIDIA/Isaac-GR00T) - NVIDIA Isaac. - [VLA-JEPA](https://github.com/ginwind/VLA-JEPA) - Jingwen Sun et al. - [Octo](https://github.com/octo-models/octo) - Octo Model Team, UC Berkeley RAIL. +- [TurboVLA](https://github.com/H-EmbodVis/TurboVLA) - Hengyi Xie et al. Built on: diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index 75cf093..d7f2946 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -338,8 +338,9 @@ def _action_unnorm(chunk, mn=action_min, mx=action_max): action = normalized.copy() action[..., :6] = ((normalized[..., :6] + 1.0) * 0.5 * (mx[:6] - mn[:6]) + mn[:6]) - # The released LIBERO policy thresholds its normalized gripper - # output instead of min/max denormalizing it as a continuous axis. + # TurboVLA thresholds the normalized gripper instead of applying + # continuous min/max denormalization. See H-EmbodVis/TurboVLA@ + # b29ab142, turbovla/evaluation/policy.py:213-221. action[..., 6] = np.where(normalized[..., 6] > 0.0, 1.0, -1.0) return action.astype(np.float32) diff --git a/scripts/convert_turbovla_to_gguf.py b/scripts/convert_turbovla_to_gguf.py index 9781c48..5bee8fd 100644 --- a/scripts/convert_turbovla_to_gguf.py +++ b/scripts/convert_turbovla_to_gguf.py @@ -39,11 +39,7 @@ ARCH = "turbovla" -# ============================================================================= -# TENSOR KEY PREFIXES -# ============================================================================= - -# Prefixes for official TurboVLA checkpoint (no "model." prefix) +# Prefixes used by official TurboVLA checkpoints (without a "model." prefix). PREFIX_VIT = "vision_encoder.backbone" PREFIX_TEXT = "text_encoder.bert" PREFIX_VIT_PROJ = "vision_projection" @@ -55,10 +51,6 @@ KEY_TEXT_PROJ = "text_encoder.text_projection" -# ============================================================================= -# UTILITY FUNCTIONS -# ============================================================================= - def kv_prefix(name: str) -> str: return f"{ARCH}.{name}" @@ -104,10 +96,6 @@ def load_safetensors(ckpt: Path) -> dict[str, torch.Tensor]: return tensors -# ============================================================================= -# DIMENSION INFERENCE -# ============================================================================= - class TurboVLADims: """Auto-detect all TurboVLA dimensions from checkpoint tensor shapes.""" @@ -118,17 +106,14 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d self.cfg = cfg_json or {} self.dinov3_cfg = dinov3_cfg or {} - # Vision encoder (DINOv3 ViT) q0 = self._get(f"{PREFIX_VIT}.layer.0.attention.q_proj.weight") self.vit_dim = int(q0.shape[0]) self.vit_layers = max_layer(keys, f"{PREFIX_VIT}.layer.") - # Text encoder (BERT) q0 = self._get(f"{PREFIX_TEXT}.encoder.layer.0.attention.self.query.weight") self.text_dim = int(q0.shape[0]) self.text_layers = max_layer(keys, f"{PREFIX_TEXT}.encoder.layer.") - # VL Fusion layers self.num_fusion_layers = max_layer(keys, f"{PREFIX_FUSION}.") self.num_text_layers = max_layer(keys, f"{PREFIX_VL_TEXT}.") @@ -143,29 +128,24 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d raise SystemExit("interaction.enhancer_inner_dim must be divisible by fusion_heads") self.fusion_head_dim = self.fusion_dim // self.fusion_heads - # Text enhancer heads: same split as fusion - self.text_enhancer_heads = max(1, raw_nheads // 2) # = 4 - self.text_enhancer_head_dim = self.hidden_dim // self.text_enhancer_heads # = 64 + self.text_enhancer_heads = max(1, raw_nheads // 2) + self.text_enhancer_head_dim = self.hidden_dim // self.text_enhancer_heads # TurboVLA constructs the ACT decoder with interaction.nheads. self.action_heads = raw_nheads - self.action_head_dim = self.hidden_dim // self.action_heads # = 32 + self.action_head_dim = self.hidden_dim // self.action_heads - # State projection config self.num_state_tokens = int(self.cfg.get("action", {}).get("num_state_tokens", 2)) - # Action decoder self.num_action_decoder_layers = max_layer(keys, f"{PREFIX_ACT_DEC}.decoder.layers.") action_q = self._get(f"{PREFIX_ACT_DEC}.action_queries.weight") self.action_horizon = int(action_q.shape[0]) act_proj_2 = self._get(f"{PREFIX_ACT_DEC}.action_projection.layers.2.weight") self.action_dim = int(act_proj_2.shape[0]) - # State projection state_weight = self._get(f"{PREFIX_STATE_PROJ}.net.1.weight") self.state_dim = int(state_weight.shape[1]) - # Image specs - infer from config or tensor shapes view_emb = self._get(KEY_VIEW_EMB) if view_emb.ndim == 3: if view_emb.shape[0] != 1: @@ -176,7 +156,6 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d else: raise SystemExit(f"Unexpected view_embedding shape: {tuple(view_emb.shape)}") - # image_size: config first, then infer cfg_image_size = self.cfg.get("vision", {}).get("image_size", self.cfg.get("image_size")) if cfg_image_size is not None: if isinstance(cfg_image_size, (list, tuple)): @@ -195,14 +174,12 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d ) self.image_size = inferred - # Patch size from vision encoder patch_weight = self._get(f"{PREFIX_VIT}.embeddings.patch_embeddings.weight") if len(patch_weight.shape) == 4: self.patch_size = patch_weight.shape[2] else: raise SystemExit("Cannot determine patch_size from patch_embeddings") - # Register tokens (optional) reg_key = f"{PREFIX_VIT}.embeddings.register_tokens" if reg_key in tensors: reg = tensors[reg_key] @@ -224,11 +201,9 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d if int(self.dinov3_cfg.get("num_register_tokens", self.num_register_tokens)) != self.num_register_tokens: raise SystemExit("DINOv3 config num_register_tokens disagrees with checkpoint weights") - # Infer vocab_size from word embeddings word_emb = self._get(f"{PREFIX_TEXT}.embeddings.word_embeddings.weight") self.vocab_size = int(word_emb.shape[0]) - # dropout: only from config, no fallback self.dropout = float(self.cfg.get("vision", {}).get("dropout", 0.0)) def _get(self, key: str) -> torch.Tensor: @@ -243,20 +218,16 @@ def _infer_image_size(self) -> int | None: position embeddings, so position_embeddings.weight may not exist or may have different semantics. Use patch_size * grid_size from config instead. """ - # Try the validated DINOv3 config first. if self.dinov3_cfg: image_size = self.dinov3_cfg.get("image_size") if image_size: return int(image_size) - # Fallback: try to infer from position embeddings - # DINOv2 with position embeddings pos_key = f"{PREFIX_VIT}.embeddings.position_embeddings.weight" if pos_key in self.tensors: pos = self.tensors[pos_key] - # Position embeddings format: [seq_len, dim] or [1, seq_len, dim] seq_len = pos.shape[0] if pos.ndim == 2 else pos.shape[1] - num_patches = seq_len - 1 # exclude CLS token + num_patches = seq_len - 1 patch = self.tensors.get(f"{PREFIX_VIT}.embeddings.patch_embeddings.weight") if patch is not None and len(patch.shape) == 4: patch_size = patch.shape[2] @@ -278,10 +249,6 @@ def __repr__(self): ) -# ============================================================================= -# TENSOR WRITERS -# ============================================================================= - def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: """Write DINOv3 ViT vision encoder. @@ -290,29 +257,24 @@ def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLAD """ root = PREFIX_VIT - # embeddings # PyTorch patch_embedding weight: [H, IC, KH, KW] = [768, 3, 16, 16] # For matmul-based patch embedding: flatten weight to [H, IC*KH*KW] = [768, 768] patch_weight = tensors[f"{root}.embeddings.patch_embeddings.weight"] - patch_weight = patch_weight.reshape(dims.vit_dim, -1) # [768, 768] + patch_weight = patch_weight.reshape(dims.vit_dim, -1) add_tensor(writer, "vit.cls_token", tensors[f"{root}.embeddings.cls_token"].squeeze(0)) add_tensor(writer, "vit.patch_embed.weight", patch_weight) add_tensor(writer, "vit.patch_embed.bias", tensors[f"{root}.embeddings.patch_embeddings.bias"]) - # register tokens (optional - DINOv3 uses them) reg_key = f"{root}.embeddings.register_tokens" if reg_key in tensors: add_tensor(writer, "vit.register_tokens", tensors[reg_key].squeeze(0)) - # transformer layers for i in range(dims.vit_layers): lr = f"{root}.layer.{i}" - # Get LayerScale values for baking into weights ls1 = tensors[f"{lr}.layer_scale1.lambda1"].clone() ls2 = tensors[f"{lr}.layer_scale2.lambda1"].clone() - # attention projections - bake LayerScale into o_proj w_q = tensors[f"{lr}.attention.q_proj.weight"] b_q = tensors[f"{lr}.attention.q_proj.bias"] w_k = tensors[f"{lr}.attention.k_proj.weight"] @@ -320,7 +282,6 @@ def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLAD b_k = tensors[k_bias_key] if k_bias_key in tensors else torch.zeros(dims.vit_dim, dtype=torch.float32) w_v = tensors[f"{lr}.attention.v_proj.weight"] b_v = tensors[f"{lr}.attention.v_proj.bias"] - # Bake ls1 into o_proj w_o = tensors[f"{lr}.attention.o_proj.weight"] * ls1.view(-1, 1) b_o = tensors[f"{lr}.attention.o_proj.bias"] * ls1 @@ -333,13 +294,11 @@ def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLAD add_tensor(writer, f"vit.blk.{i}.attn_o.weight", w_o) add_tensor(writer, f"vit.blk.{i}.attn_o.bias", b_o) - # layer norms add_tensor(writer, f"vit.blk.{i}.ln1.weight", tensors[f"{lr}.norm1.weight"]) add_tensor(writer, f"vit.blk.{i}.ln1.bias", tensors[f"{lr}.norm1.bias"]) add_tensor(writer, f"vit.blk.{i}.ln2.weight", tensors[f"{lr}.norm2.weight"]) add_tensor(writer, f"vit.blk.{i}.ln2.bias", tensors[f"{lr}.norm2.bias"]) - # MLP - bake ls2 into down_proj w_fc1 = tensors[f"{lr}.mlp.up_proj.weight"] b_fc1 = tensors[f"{lr}.mlp.up_proj.bias"] w_fc2 = tensors[f"{lr}.mlp.down_proj.weight"] * ls2.view(-1, 1) @@ -350,7 +309,6 @@ def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLAD add_tensor(writer, f"vit.blk.{i}.fc2.weight", w_fc2) add_tensor(writer, f"vit.blk.{i}.fc2.bias", b_fc2) - # final norm add_tensor(writer, "vit.final_norm.weight", tensors[f"{root}.norm.weight"]) add_tensor(writer, "vit.final_norm.bias", tensors[f"{root}.norm.bias"]) @@ -359,14 +317,12 @@ def write_text_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADim """Write BERT text encoder.""" root = PREFIX_TEXT - # embeddings add_tensor(writer, "text.embed.word_embeddings", tensors[f"{root}.embeddings.word_embeddings.weight"]) add_tensor(writer, "text.embed.position_embeddings", tensors[f"{root}.embeddings.position_embeddings.weight"]) add_tensor(writer, "text.embed.token_type_embeddings", tensors[f"{root}.embeddings.token_type_embeddings.weight"]) add_tensor(writer, "text.embed.LayerNorm.weight", tensors[f"{root}.embeddings.LayerNorm.weight"]) add_tensor(writer, "text.embed.LayerNorm.bias", tensors[f"{root}.embeddings.LayerNorm.bias"]) - # encoder layers for i in range(dims.text_layers): lr = f"{root}.encoder.layer.{i}" @@ -387,11 +343,9 @@ def write_text_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADim add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.weight", tensors[f"{lr}.output.LayerNorm.weight"]) add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.bias", tensors[f"{lr}.output.LayerNorm.bias"]) - # pooler add_tensor(writer, "text.pooler.dense.weight", tensors[f"{root}.pooler.dense.weight"]) add_tensor(writer, "text.pooler.dense.bias", tensors[f"{root}.pooler.dense.bias"]) - # text projection add_tensor(writer, "text_proj.weight", tensors[KEY_TEXT_PROJ + ".weight"]) add_tensor(writer, "text_proj.bias", tensors[KEY_TEXT_PROJ + ".bias"]) @@ -416,7 +370,6 @@ def write_vision_language_interaction(writer: gguf.GGUFWriter, tensors: dict, di fusion_root = PREFIX_FUSION text_root = PREFIX_VL_TEXT - # VL Fusion layers for i in range(dims.num_fusion_layers): lr = f"{fusion_root}.{i}" @@ -441,7 +394,6 @@ def write_vision_language_interaction(writer: gguf.GGUFWriter, tensors: dict, di add_tensor(writer, f"vl_fusion.{i}.gamma_v", tensors[f"{lr}.gamma_v"]) add_tensor(writer, f"vl_fusion.{i}.gamma_l", tensors[f"{lr}.gamma_l"]) - # VL Text self-attention layers for i in range(dims.num_text_layers): lr = f"{text_root}.{i}" @@ -522,10 +474,6 @@ def write_view_embeddings(writer: gguf.GGUFWriter, tensors: dict) -> None: add_tensor(writer, "view_emb", tensors[KEY_VIEW_EMB]) -# ============================================================================= -# VERIFICATION -# ============================================================================= - def verify_consumed_tensors(source_keys: set[str], dims: TurboVLADims) -> None: """Verify all source tensors are consumed by the converter.""" consumed = set() @@ -703,10 +651,6 @@ def verify_consumed_tensors(source_keys: set[str], dims: TurboVLADims) -> None: print(f"\n All {len(consumed)} keys consumed successfully!") -# ============================================================================= -# MAIN -# ============================================================================= - def main() -> int: import argparse parser = argparse.ArgumentParser( @@ -762,15 +706,12 @@ def main() -> int: writer.add_uint32(kv("text_head_dim"), 64) writer.add_uint32(kv("text_heads"), 12) writer.add_uint32(kv("num_fusion_layers"), dims.num_fusion_layers) - # CORRECT: fusion_heads = nheads // 2, NOT nheads + # Fusion attention splits the configured head count in half. writer.add_uint32(kv("fusion_heads"), dims.fusion_heads) - # CORRECT: fusion_head_dim = hidden / fusion_heads writer.add_uint32(kv("fusion_head_dim"), dims.fusion_head_dim) writer.add_uint32(kv("num_text_layers"), dims.num_text_layers) - # Text enhancer: same split as fusion writer.add_uint32(kv("text_enhancer_heads"), dims.text_enhancer_heads) writer.add_uint32(kv("text_enhancer_head_dim"), dims.text_enhancer_head_dim) - # Action transformer config writer.add_uint32(kv("action_heads"), dims.action_heads) writer.add_uint32(kv("action_head_dim"), dims.action_head_dim) writer.add_uint32(kv("num_action_decoder_layers"), dims.num_action_decoder_layers) diff --git a/src/models/turbovla.cpp b/src/models/turbovla.cpp index d23720e..8372e9a 100644 --- a/src/models/turbovla.cpp +++ b/src/models/turbovla.cpp @@ -48,10 +48,6 @@ namespace vla { - // ============================================================================= - // CONFIG - // ============================================================================= - struct TurboVLAConfig { int64_t hidden = 256; int64_t vit_dim = 768; @@ -65,8 +61,8 @@ int64_t vocab_size = 30522; int64_t max_text_length = 256; int64_t num_fusion_layers = 6; - int64_t fusion_heads = 4; // TurboVLA uses 4 fusion heads - int64_t fusion_head_dim = 256; // 256 dim per fusion head + int64_t fusion_heads = 4; + int64_t fusion_head_dim = 256; int64_t text_enhancer_heads = 4; int64_t text_enhancer_head_dim = 64; int64_t action_heads = 8; @@ -81,7 +77,7 @@ int64_t patch_size = 16; int64_t num_views = 2; int64_t num_register_tokens = 0; - float rope_theta = 10000.0f; // DINOv3 2D RoPE base frequency + float rope_theta = 10000.0f; float dropout = 0.1f; int64_t pad_token_id = 0; int64_t cls_token_id = 101; @@ -90,10 +86,6 @@ int64_t question_token_id = 1029; }; - // ============================================================================= - // WEIGHT STRUCTURES - // ============================================================================= - namespace { struct VitBlockW { @@ -111,11 +103,11 @@ ggml_tensor* cls_token; ggml_tensor* patch_embed_w; ggml_tensor* patch_embed_b; - ggml_tensor* register_tokens; ///< DINOv3 register tokens (num_register_tokens x vit_dim) + ggml_tensor* register_tokens; std::vector layers; ggml_tensor* final_ln_w; ggml_tensor* final_ln_b; - // Precomputed 2D RoPE frequencies for patch tokens: [n_patches, head_dim/2] + // One frequency tensor per patch-grid position. std::vector rope_freqs; }; @@ -207,10 +199,6 @@ } // anonymous namespace - // ============================================================================= - // MODEL - // ============================================================================= - struct TurboVLA : public ModelArchBase { TurboVLA() : ModelArchBase(Arch::TURBOVLA) {} ~TurboVLA() override; @@ -261,20 +249,6 @@ if (backend) ggml_backend_free(backend); } - // ============================================================================= - // DINOv3 ViT LAYER with 2D RoPE - // ============================================================================= - // DINOv3 2D RoPE is applied in vit_layer() - see implementation below. - // The 2D positions are handled via ggml_rope which computes: - // x'[2i] = x[2i] * cos(pos * freq_i) - x[2i+1] * sin(pos * freq_i) - // x'[2i+1] = x[2i] * sin(pos * freq_i) + x[2i+1] * cos(pos * freq_i) - // where freq_i = rope_theta^(-2i/head_dim) - // Position IDs are patch indices (0, 1, 2, ..., n_patches-1) - - // ============================================================================= - // FORWARD FUNCTIONS - // ============================================================================= - ggml_tensor* TurboVLA::vit_layer(ggml_context* C, const VitBlockW& w, ggml_tensor* x, ggml_tensor* rope_cos, ggml_tensor* rope_sin, int64_t internal_seq, int64_t n_heads, int64_t head_dim) const { @@ -357,20 +331,16 @@ int64_t n_heads, int64_t head_dim) const { const float scale = 1.0f / std::sqrt((float)head_dim); - // TurboVLA fusion: 4 heads × 64 dim = 256 att_dim const int64_t att_dim = n_heads * head_dim; - // FIXED (Point 5): Official uses residual_style="normalized" - // v_norm = LayerNorm(v), l_norm = LayerNorm(l) - // delta_v, delta_l = bi_attention(v_norm, l_norm) - // v_new = v_norm + gamma_v * delta_v - // l_new = l_norm + gamma_l * delta_l + // TurboVLA's normalized residual style adds the deltas to normalized + // tokens, not the pre-normalization residuals. See + // H-EmbodVis/TurboVLA@b29ab142, models/components/fusion.py:336-347. ggml_tensor* v_norm = layer_norm(C, v, w.norm_v_w, w.norm_v_b, 1e-5f); ggml_tensor* l_norm = layer_norm(C, l, w.norm_l_w, w.norm_l_b, 1e-5f); - // V attends to L. Keep this as explicit QK^T + masked softmax rather than - // flash attention: the official implementation masks padded language keys - // on this branch only. + // Only vision-to-language attention masks language keys. See + // H-EmbodVis/TurboVLA@b29ab142, models/components/fusion.py:246-267. ggml_tensor* q_v = linear(C, w.v_proj_w, w.v_proj_b, v_norm); ggml_tensor* k_l = linear(C, w.l_proj_w, w.l_proj_b, l_norm); ggml_tensor* v_l = linear(C, w.values_l_w, w.values_l_b, l_norm); @@ -381,13 +351,9 @@ ggml_tensor* o = attention(C, Q, K, V, language_key_mask, scale, att_dim, v_seq); ggml_tensor* o_proj = linear(C, w.out_v_w, w.out_v_b, o); - // FIXED (Point 5): Apply gamma_v LayerScale with proper broadcasting - // gamma_v is [hidden], delta is [hidden, v_seq] - // Use ggml_repeat to broadcast gamma_v across sequence dimension ggml_tensor* gamma_v_scaled = ggml_repeat(C, w.gamma_v, o_proj); ggml_tensor* delta_v = ggml_mul(C, gamma_v_scaled, o_proj); - // FIXED (Point 5): v_new = v_norm + gamma_v * delta_v ggml_tensor* v_out = ggml_add(C, v_norm, delta_v); // L attends to V. There is no visual padding, hence no mask on this branch. @@ -401,11 +367,9 @@ ggml_tensor* o_l = attention(C, Q_l, K_v, V_v, nullptr, scale, att_dim, l_seq); ggml_tensor* o_proj_l = linear(C, w.out_l_w, w.out_l_b, o_l); - // FIXED (Point 5): Apply gamma_l LayerScale with proper broadcasting ggml_tensor* gamma_l_scaled = ggml_repeat(C, w.gamma_l, o_proj_l); ggml_tensor* delta_l = ggml_mul(C, gamma_l_scaled, o_proj_l); - // FIXED (Point 5): l_new = l_norm + gamma_l * delta_l l = ggml_add(C, l_norm, delta_l); return v_out; @@ -414,16 +378,7 @@ ggml_tensor* TurboVLA::vl_text_layer(ggml_context* C, const VlTextLayerW& w, ggml_tensor* x, ggml_tensor* self_mask, int64_t seq, int64_t n_heads, int64_t head_dim) const { - // FIXED (Point 7): Official TurboVLA VL text enhancer uses POST-NORM and ReLU - // Structure: - // qkv = self_attn(x) - // h = x + attention_output - // h = LayerNorm(h) <- ln1 - // ff = Linear1(h) - // ff = ReLU(ff) - // ff = Linear2(ff) - // out = h + ff - // out = LayerNorm(out) <- ln2 + // TurboVLA's text enhancer uses post-norm residuals and ReLU. const float scale = 1.0f / std::sqrt((float)head_dim); const int64_t att_dim = n_heads * head_dim; const int64_t n = n_heads; @@ -458,23 +413,16 @@ ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); ggml_tensor* h = ggml_add(C, x, o); - // POST-NORM after attention: LayerNorm(h) h = layer_norm(C, h, w.ln1_w, w.ln1_b, 1e-5f); - // FFN with ReLU (not GELU!) ggml_tensor* ff = linear(C, w.fc1_w, w.fc1_b, h); ff = ggml_relu(C, ff); ggml_tensor* ff_out = linear(C, w.fc2_w, w.fc2_b, ff); - // Add residual and final POST-NORM: LayerNorm(h + ff) ggml_tensor* out = ggml_add(C, h, ff_out); return layer_norm(C, out, w.ln2_w, w.ln2_b, 1e-5f); } - // ============================================================================= - // ACT DECODER LAYER - // ============================================================================= - ggml_tensor* TurboVLA::act_decoder_layer(ggml_context* C, const ActDecoderLayerW& w, ggml_tensor* queries, ggml_tensor* memory, int64_t q_seq, int64_t m_seq, @@ -482,7 +430,6 @@ const float scale = 1.0f / std::sqrt((float)head_dim); const int64_t hidden = n_heads * head_dim; - // Self-attention on queries. ggml_tensor* normed = layer_norm(C, queries, w.ln1_w, w.ln1_b, 1e-5f); ggml_tensor* self_qkv = linear(C, w.self_qkv_w, w.self_qkv_b, normed); @@ -541,30 +488,26 @@ ggml_tensor* K_c = to_heads(C, cross_k, head_dim, n_heads, m_seq); ggml_tensor* V_c = to_heads_v(C, cross_v, head_dim, n_heads, m_seq); - // Cross-attention output length is q_seq, not m_seq. ggml_tensor* att_c = attention(C, Q_c, K_c, V_c, nullptr, scale, hidden, q_seq); ggml_tensor* o_c = linear(C, w.cross_out_w, w.cross_out_b, att_c); h = ggml_add(C, h, o_c); - // FIXED (Point 10): Official PyTorch TransformerDecoderLayer uses ReLU, not GELU + // The upstream TransformerDecoderLayer keeps PyTorch's default ReLU + // activation. See H-EmbodVis/TurboVLA@b29ab142, + // models/action_head.py:39-47. normed = layer_norm(C, h, w.ln3_w, w.ln3_b, 1e-5f); ggml_tensor* ffn_out = ffn_relu(C, w.fc1_w, w.fc1_b, w.fc2_w, w.fc2_b, normed); return ggml_add(C, h, ffn_out); } - // ============================================================================= - // INFERENCE - // ============================================================================= - namespace { static const float IMAGENET_MEAN[3] = {0.485f, 0.456f, 0.406f}; static const float IMAGENET_STD[3] = {0.229f, 0.224f, 0.225f}; - // CHW planar layout matching ggml_conv_2d expectation - // Build the exact flattened Conv2D patch layout expected by a patch-embedding - // weight flattened from [out, C, patch_h, patch_w] to [out, C*patch_h*patch_w]. + // Build the flattened patch matrix expected by a patch-embedding weight + // flattened from [out, C, patch_h, patch_w] to [out, C*patch_h*patch_w]. // Output is GGML matrix layout [patch_dim, n_patches], i.e. each column is one // normalized Cxpatchxpatch image patch. void normalize_imagenet_patches(const ImageView& v, @@ -693,8 +636,10 @@ for (int64_t k = 0; k < seq; ++k) { const bool is_allowed = allowed[(size_t)q * seq + k] != 0; result.bert_self_mask[(size_t)q * seq + k] = is_allowed ? 0.0f : -FLT_MAX; - // The post-BERT enhancer receives both the inverse source mask and - // BERT's key-padding mask. + // The BERT wrapper creates block masks from special tokens; the + // enhancer additionally masks padded keys. See + // H-EmbodVis/TurboVLA@b29ab142, text/bert.py:177-212 and + // models/text_encoder.py:80-94. result.enhancer_self_mask[(size_t)q * seq + k] = is_allowed && valid[(size_t)k] ? 0.0f : -FLT_MAX; } @@ -719,14 +664,13 @@ const int64_t img_size = cfg.image_size; const int64_t n_patches_per_side = img_size / patch_size; const int64_t n_patches = n_patches_per_side * n_patches_per_side; - const int64_t n_reg = cfg.num_register_tokens; // DINOv3 register tokens - const int64_t n_cls = 1; // CLS token - const int64_t internal_vit_seq = n_cls + n_reg + n_patches; // Full internal ViT sequence + const int64_t n_reg = cfg.num_register_tokens; + const int64_t n_cls = 1; + const int64_t internal_vit_seq = n_cls + n_reg + n_patches; const int64_t n_views = cfg.num_views; - // vision_seq is n_patches per view (CLS + register tokens are sliced away) const int64_t vision_seq = n_views * n_patches; - const int64_t hidden = cfg.hidden; // VLA hidden dim (256) - const int64_t vit_dim = cfg.vit_dim; // ViT hidden dim (768) + const int64_t hidden = cfg.hidden; + const int64_t vit_dim = cfg.vit_dim; // Basic shape sanity checks. These fail early with a useful message instead // of failing later inside ggml_reshape_*. @@ -749,10 +693,7 @@ return {}; } - // DINOv3 ViT processes internal sequence: [CLS, register_0..register_N-1, patch_0..patch_255] - // After final DINO norm, we slice away CLS and ALL register tokens - // Only n_patches per view go to vision_projection - // Output per view: [vit_dim, n_patches] (not internal_vit_seq) + // Project only patch tokens; DINOv3's CLS and register tokens stay internal. const int64_t vision_out_size = n_patches * vit_dim; std::vector vision_outputs((size_t)n_views * vision_out_size, 0.0f); @@ -808,22 +749,16 @@ build_dinov3_rope(n_patches_per_side, cfg.vit_head_dim, cfg.rope_theta, rope_cos_host, rope_sin_host); - // ========================================================================= - // DINOv3 Token Flow: Build [CLS, register_tokens..., patches...] - // ========================================================================= - // Concatenate: [vit_dim, 1] + [vit_dim, n_reg] + [vit_dim, n_patches] + // Assemble DINOv3's internal sequence: CLS, register tokens, then patches. ggml_tensor* cls_2d = ggml_reshape_2d(VC, vision_encoder.cls_token, vit_dim, 1); ggml_tensor* vit_out = cls_2d; if (n_reg > 0 && vision_encoder.register_tokens) { - // Add register tokens ggml_tensor* reg_2d = ggml_reshape_2d(VC, vision_encoder.register_tokens, vit_dim, n_reg); vit_out = ggml_concat(VC, vit_out, reg_2d, 1); } - // Add patches vit_out = ggml_concat(VC, vit_out, patches_with_bias, 1); - // Run all DINO layers over FULL internal sequence [CLS, reg, patches] for (int64_t i = 0; i < cfg.vit_layers; ++i) { vit_out = vit_layer(VC, vision_encoder.layers[i], @@ -834,16 +769,10 @@ cfg.vit_heads, cfg.vit_head_dim); } - // ========================================================================= - // DINOv3: Slice away CLS and ALL register tokens - // Only patch tokens (n_patches) go to vision_projection - // ========================================================================= - // The reference TurboVLA encoder uses DINOv3's final hidden state - // (`outputs.hidden_states[-1]`), before the model's final LayerNorm. - // Its shape is [vit_dim, internal_vit_seq]. - // We need to extract only the patch tokens: positions [n_cls + n_reg, internal_vit_seq) - // For ViT-B with 256 patches and 4 register tokens: extract positions [5, 261) - const int64_t patch_token_start = n_cls + n_reg; // First patch position + // TurboVLA selects DINOv3's final hidden-state entry, then removes + // its prefix tokens. See H-EmbodVis/TurboVLA@b29ab142, + // models/vision_encoder.py:109-114. + const int64_t patch_token_start = n_cls + n_reg; ggml_tensor* vit_patches_only = ggml_view_2d(VC, vit_out, @@ -863,7 +792,6 @@ return {}; } - // Only now should patch_input have a backend buffer. if (!patch_input->buffer) { fprintf(stderr, "ERROR: patch_input has no backend buffer after vision graph allocation; " @@ -903,9 +831,7 @@ return {}; } - // Create one graph input per view. Do NOT copy data yet: context is - // no_alloc=true and buffers are created by scratch.alloc(gf) below. - // NOTE: vision_seq = n_views * n_patches (CLS and register tokens were sliced away) + // Buffers are allocated after graph construction because the context uses no_alloc. std::vector view_tokens((size_t)n_views, nullptr); for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { view_tokens[view_idx] = @@ -1037,24 +963,17 @@ ggml_tensor* vl_features = proj; ggml_tensor* text_features = text_tokens; - // TurboVLA fusion: 4 heads × 256 dim (hidden=256, v_proj outputs 4*256=1024) const int64_t vl_heads = cfg.fusion_heads; const int64_t vl_head_dim = cfg.fusion_head_dim; - // TurboVLA text enhancer: 4 heads × 64 dim (hidden=256) const int64_t text_enh_heads = cfg.text_enhancer_heads; const int64_t text_enh_head_dim = cfg.text_enhancer_head_dim; - // FIXED (Point 4): VL interaction is interleaved - // Official forward: - // for i in 0..num_layers-1: - // visual_tokens, text_tokens = fusion_layer[i](visual_tokens, text_tokens) - // text_tokens = text_layer[i](text_tokens) + // Each fusion layer is immediately followed by its paired text enhancer. GGML_ASSERT(cfg.num_fusion_layers == cfg.num_text_layers && "VL fusion and text layers must have same count for interleaved execution"); for (int64_t i = 0; i < cfg.num_fusion_layers; ++i) { - // Fusion layer vl_features = vl_fusion_layer(C, fusion_layers[i], vl_features, @@ -1064,7 +983,6 @@ text_seq, vl_heads, vl_head_dim); - // Text enhancer layer (operates on text_features in-place) text_features = vl_text_layer(C, text_layers[i], text_features, @@ -1089,15 +1007,6 @@ ggml_new_tensor_2d(C, GGML_TYPE_F32, state_dim, 1); ggml_set_input(state_tensor); - // FIXED (Point 9): State projection order - output_norm AFTER position add - // Official order: - // s = LayerNorm(state) - // s = Linear(...) - // s = GELU(s) - // s = Linear(...) - // tokens = reshape(s, num_state_tokens, hidden) - // tokens = tokens + learned_position - // tokens = output_norm(tokens) <- MOVED HERE ggml_tensor* s = layer_norm(C, state_tensor, @@ -1126,16 +1035,15 @@ cfg.num_state_tokens); state_tokens = ggml_add(C, state_tokens, pos_emb_state); - // FIXED: output_norm AFTER position add + // The projection normalizes only after adding learned positions. See + // H-EmbodVis/TurboVLA@b29ab142, models/action_head.py:25-31. state_tokens = layer_norm(C, state_tokens, state_proj.output_norm_w, state_proj.output_norm_b, 1e-5f); - // IMPORTANT: previous code computed state_tokens and then discarded them. - // That made state_tensor a dead graph input and also meant robot state never - // influenced the action decoder. Include state tokens in decoder memory. + // State tokens must be part of the decoder memory. ggml_tensor* condition = ggml_concat(C, vl_condition, state_tokens, 1); const int64_t cond_seq = vl_cond_seq + cfg.num_state_tokens; @@ -1146,7 +1054,6 @@ hidden, cfg.action_horizon); - // TurboVLA ACT decoder: 8 heads × 32 dim const int64_t act_heads = cfg.action_heads; if (hidden % act_heads != 0) { @@ -1275,15 +1182,10 @@ return output; } - // ============================================================================= - // MODEL LOADING - // ============================================================================= - std::unique_ptr turbovla_create(const std::string& /*mmproj_path*/, const std::string& ckpt_path, const std::string& /*config_path*/, const Options& /*opts*/) { - // Open GGUF gguf_reader g("turbovla"); if (!g.open(ckpt_path)) { std::fprintf(stderr, "vla(turbovla): failed to open %s\n", ckpt_path.c_str()); @@ -1292,7 +1194,6 @@ auto* m = new vla::TurboVLA; - // Load config from GGUF m->cfg_.hidden = (int64_t)g.u32("turbovla.hidden"); m->cfg_.vit_dim = (int64_t)g.u32("turbovla.vit_dim"); m->cfg_.vit_layers = (int64_t)g.u32("turbovla.vit_layers"); @@ -1370,7 +1271,6 @@ m->cfg.rope_freq_base = m->cfg_.rope_theta; m->cfg.num_steps = m->cfg_.action_horizon; - // Init backend const Backend b = backend_init("vla(turbovla)", m->n_threads); if (!b.handle) { std::fprintf(stderr, "vla(turbovla): backend_init failed\n"); @@ -1381,10 +1281,8 @@ ggml_init_params wp = { (size_t)16*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); - // Load weights WeightLoader L("turbovla", g, m->ctx_weights, GGML_TYPE_F32); - // Vision encoder m->vision_encoder.cls_token = L.f32("vit.cls_token"); m->vision_encoder.patch_embed_w = L.gemm("vit.patch_embed.weight"); m->vision_encoder.patch_embed_b = L.gemm("vit.patch_embed.bias"); @@ -1421,10 +1319,7 @@ m->vision_encoder.final_ln_w = L.f32("vit.final_norm.weight"); m->vision_encoder.final_ln_b = L.f32("vit.final_norm.bias"); - // Precompute DINOv3 2D RoPE frequencies - // This is done in predict() since we need n_patches_per_side which depends on image_size - // Vision projection m->vision_proj.input_norm_w = L.f32("vit_proj.input_norm.weight"); m->vision_proj.input_norm_b = L.f32("vit_proj.input_norm.bias"); m->vision_proj.mlp_0_w = L.f32("vit_proj.mlp.0.weight"); @@ -1435,7 +1330,6 @@ m->vision_proj.output_norm_w = L.f32("vit_proj.output_norm.weight"); m->vision_proj.output_norm_b = L.f32("vit_proj.output_norm.bias"); - // Text encoder m->text_encoder.word_embed = L.f32("text.embed.word_embeddings"); m->text_encoder.pos_embed = L.f32("text.embed.position_embeddings"); m->text_encoder.tok_type_embed = L.f32("text.embed.token_type_embeddings"); @@ -1467,7 +1361,6 @@ w.ln2_b = L.f32("text.encoder.layer.%lld.output.LayerNorm.bias", (long long) i); } - // VL Fusion layers m->fusion_layers.resize(m->cfg_.num_fusion_layers); for (int64_t i = 0; i < m->cfg_.num_fusion_layers; ++i) { auto& w = m->fusion_layers[i]; @@ -1491,7 +1384,6 @@ w.gamma_l = L.f32("vl_fusion.%lld.gamma_l", (long long) i); } - // VL Text layers m->text_layers.resize(m->cfg_.num_text_layers); for (int64_t i = 0; i < m->cfg_.num_text_layers; ++i) { auto& w = m->text_layers[i]; @@ -1509,7 +1401,6 @@ w.ln2_b = L.f32("vl_text.%lld.ln2.bias", (long long) i); } - // Action decoder m->action_decoder.action_q = L.f32("act.q.weight"); m->action_decoder.layers.resize(m->cfg_.num_action_decoder_layers); for (int64_t i = 0; i < m->cfg_.num_action_decoder_layers; ++i) { @@ -1540,7 +1431,6 @@ m->action_decoder.proj_2_w = L.f32("act.proj.2.weight"); m->action_decoder.proj_2_b = L.f32("act.proj.2.bias"); - // State projection m->state_proj.net_0_w = L.f32("state.proj.0.weight"); m->state_proj.net_0_b = L.f32("state.proj.0.bias"); m->state_proj.net_1_w = L.f32("state.proj.1.weight"); @@ -1551,7 +1441,6 @@ m->state_proj.output_norm_b = L.f32("state.proj.output_norm.bias"); m->state_proj.position = L.f32("state.proj.position"); - // View embeddings m->view_emb = L.f32("view_emb"); if (!L.upload(m->backend, &m->weight_buf)) { @@ -1559,7 +1448,6 @@ return nullptr; } - // Print DINOv3 config std::printf("vla(turbovla): hidden=%lld vit_layers=%lld text_layers=%lld fusion=%lld action_horizon=%lld action_dim=%lld\n", (long long)m->cfg_.hidden, (long long)m->cfg_.vit_layers, diff --git a/tests/py/test_converters.py b/tests/py/test_converters.py index a9467b7..f19bb18 100644 --- a/tests/py/test_converters.py +++ b/tests/py/test_converters.py @@ -61,6 +61,9 @@ def numpy(self): return None def float(self): return _F32Tensor() def to(self, *a): return self def reshape(self, *a): return self + def squeeze(self, *a): return self + def clone(self): return self + def __mul__(self, other): return self class _F32Tensor(_Tensor): @@ -259,6 +262,93 @@ def test_prismatic_tower_and_lm(): assert qsrcs[2] == "language_model.model.layers.0.self_attn.q_proj.weight" +class _SourceTensors(dict): + """Return a tensor stub while recording every converter source lookup.""" + + def __init__(self): + self.keys_read = [] + + def __getitem__(self, key): + self.keys_read.append(key) + return _Tensor() + + def __contains__(self, key): + return True + + +def test_turbovla_converter_remap(): + import importlib + + T = importlib.import_module("convert_turbovla_to_gguf") + tensors = _SourceTensors() + dims = types.SimpleNamespace( + vit_dim=1, + vit_layers=1, + text_layers=1, + num_fusion_layers=1, + num_text_layers=1, + num_action_decoder_layers=1, + ) + writer = _Writer() + + T.write_vision_encoder(writer, tensors, dims) + T.write_text_encoder(writer, tensors, dims) + T.write_vision_projection(writer, tensors) + T.write_vision_language_interaction(writer, tensors, dims) + T.write_action_decoder(writer, tensors, dims) + T.write_state_projection(writer, tensors) + T.write_view_embeddings(writer, tensors) + + assert writer.names == [ + "vit.cls_token", "vit.patch_embed.weight", "vit.patch_embed.bias", "vit.register_tokens", + "vit.blk.0.attn_q.weight", "vit.blk.0.attn_q.bias", "vit.blk.0.attn_k.weight", "vit.blk.0.attn_k.bias", + "vit.blk.0.attn_v.weight", "vit.blk.0.attn_v.bias", "vit.blk.0.attn_o.weight", "vit.blk.0.attn_o.bias", + "vit.blk.0.ln1.weight", "vit.blk.0.ln1.bias", "vit.blk.0.ln2.weight", "vit.blk.0.ln2.bias", + "vit.blk.0.fc1.weight", "vit.blk.0.fc1.bias", "vit.blk.0.fc2.weight", "vit.blk.0.fc2.bias", + "vit.final_norm.weight", "vit.final_norm.bias", + "text.embed.word_embeddings", "text.embed.position_embeddings", "text.embed.token_type_embeddings", + "text.embed.LayerNorm.weight", "text.embed.LayerNorm.bias", + "text.encoder.layer.0.attention.self.query.weight", "text.encoder.layer.0.attention.self.query.bias", + "text.encoder.layer.0.attention.self.key.weight", "text.encoder.layer.0.attention.self.key.bias", + "text.encoder.layer.0.attention.self.value.weight", "text.encoder.layer.0.attention.self.value.bias", + "text.encoder.layer.0.attention.output.dense.weight", "text.encoder.layer.0.attention.output.dense.bias", + "text.encoder.layer.0.attention.output.LayerNorm.weight", "text.encoder.layer.0.attention.output.LayerNorm.bias", + "text.encoder.layer.0.intermediate.dense.weight", "text.encoder.layer.0.intermediate.dense.bias", + "text.encoder.layer.0.output.dense.weight", "text.encoder.layer.0.output.dense.bias", + "text.encoder.layer.0.output.LayerNorm.weight", "text.encoder.layer.0.output.LayerNorm.bias", + "text.pooler.dense.weight", "text.pooler.dense.bias", "text_proj.weight", "text_proj.bias", + "vit_proj.input_norm.weight", "vit_proj.input_norm.bias", "vit_proj.mlp.0.weight", "vit_proj.mlp.0.bias", + "vit_proj.mlp.3.weight", "vit_proj.mlp.3.bias", "vit_proj.skip.weight", "vit_proj.output_norm.weight", + "vit_proj.output_norm.bias", + "vl_fusion.0.v_proj.weight", "vl_fusion.0.v_proj.bias", "vl_fusion.0.l_proj.weight", "vl_fusion.0.l_proj.bias", + "vl_fusion.0.values_v.weight", "vl_fusion.0.values_v.bias", "vl_fusion.0.values_l.weight", "vl_fusion.0.values_l.bias", + "vl_fusion.0.out_v.weight", "vl_fusion.0.out_v.bias", "vl_fusion.0.out_l.weight", "vl_fusion.0.out_l.bias", + "vl_fusion.0.norm_v.weight", "vl_fusion.0.norm_v.bias", "vl_fusion.0.norm_l.weight", "vl_fusion.0.norm_l.bias", + "vl_fusion.0.gamma_v", "vl_fusion.0.gamma_l", + "vl_text.0.attn_qkv.weight", "vl_text.0.attn_qkv.bias", "vl_text.0.attn_o.weight", "vl_text.0.attn_o.bias", + "vl_text.0.ln1.weight", "vl_text.0.ln1.bias", "vl_text.0.fc1.weight", "vl_text.0.fc1.bias", + "vl_text.0.fc2.weight", "vl_text.0.fc2.bias", "vl_text.0.ln2.weight", "vl_text.0.ln2.bias", + "act.q.weight", "act.dec.0.self_qkv.weight", "act.dec.0.self_qkv.bias", "act.dec.0.self_out.weight", + "act.dec.0.self_out.bias", "act.dec.0.cross_qkv.weight", "act.dec.0.cross_qkv.bias", + "act.dec.0.cross_out.weight", "act.dec.0.cross_out.bias", "act.dec.0.ln1.weight", "act.dec.0.ln1.bias", + "act.dec.0.ln2.weight", "act.dec.0.ln2.bias", "act.dec.0.ln3.weight", "act.dec.0.ln3.bias", + "act.dec.0.fc1.weight", "act.dec.0.fc1.bias", "act.dec.0.fc2.weight", "act.dec.0.fc2.bias", + "act.proj.0.weight", "act.proj.0.bias", "act.proj.1.weight", "act.proj.1.bias", "act.proj.2.weight", + "act.proj.2.bias", "state.proj.0.weight", "state.proj.0.bias", "state.proj.1.weight", "state.proj.1.bias", + "state.proj.4.weight", "state.proj.4.bias", "state.proj.output_norm.weight", "state.proj.output_norm.bias", + "state.proj.position", "view_emb", + ] + assert { + "vision_encoder.backbone.embeddings.patch_embeddings.weight", + "text_encoder.bert.encoder.layer.0.attention.self.query.weight", + "vision_language_interaction.fusion_layers.0.attn.v_proj.weight", + "vision_language_interaction.text_layers.0.self_attn.in_proj_weight", + "action_head.decoder.decoder.layers.0.multihead_attn.in_proj_weight", + "action_head.state_projection.net.4.weight", + "view_embedding", + } <= set(tensors.keys_read) + + def test_every_converter_imports(): # every converter must resolve against the two shared modules import importlib From f54dd08fff8d9da0033e37a89beff7a9ad8239f2 Mon Sep 17 00:00:00 2001 From: Khanh Nguyen Date: Thu, 24 Sep 2026 22:06:39 +0700 Subject: [PATCH 3/5] fix(turbovla): match upstream text padding, read the released checkpoint, and run the engine as one cached graph --- eval/client/vla_cpp_client.py | 8 +- scripts/convert_turbovla_to_gguf.py | 317 ++-- src/layers/rope.h | 19 - src/models/turbovla.cpp | 2260 ++++++++++----------------- tests/py/test_converters.py | 3 +- 5 files changed, 927 insertions(+), 1680 deletions(-) diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index d7f2946..b295905 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -331,7 +331,7 @@ def _oft_norm(x, q01=q01, q99=q99, mask=mask): raise ValueError(f"{stats_path} contains invalid TurboVLA normalization ranges") def _state_norm(state_8d, mean=state_mean, std=state_std): - return ((state_8d - mean) / (std + 1e-8)).astype(np.float32) + return ((state_8d - mean) / (std + 1e-6)).astype(np.float32) def _action_unnorm(chunk, mn=action_min, mx=action_max): normalized = np.clip(chunk.astype(np.float32), -1.0, 1.0) @@ -339,9 +339,9 @@ def _action_unnorm(chunk, mn=action_min, mx=action_max): action[..., :6] = ((normalized[..., :6] + 1.0) * 0.5 * (mx[:6] - mn[:6]) + mn[:6]) # TurboVLA thresholds the normalized gripper instead of applying - # continuous min/max denormalization. See H-EmbodVis/TurboVLA@ - # b29ab142, turbovla/evaluation/policy.py:213-221. - action[..., 6] = np.where(normalized[..., 6] > 0.0, 1.0, -1.0) + # continuous min/max denormalization; exactly 0 maps to +1. See + # H-EmbodVis/TurboVLA@b29ab142, turbovla/evaluation/policy.py:205-221. + action[..., 6] = np.where(normalized[..., 6] < 0.0, -1.0, 1.0) return action.astype(np.float32) self._turbovla_state_norm = _state_norm diff --git a/scripts/convert_turbovla_to_gguf.py b/scripts/convert_turbovla_to_gguf.py index 5bee8fd..e910827 100644 --- a/scripts/convert_turbovla_to_gguf.py +++ b/scripts/convert_turbovla_to_gguf.py @@ -13,13 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Convert TurboVLA (LeRobot format) checkpoint to GGUF. +"""Convert a TurboVLA checkpoint (H-EmbodVis/TurboVLA) to GGUF. -Supports any TurboVLA variant with auto-detection of dimensions from checkpoint. -Architecture must be compatible with the LeRobot TurboVLA format. +Reads the released .pth ({model_state_dict, model_config}, e.g. +checkpoints/libero/turbovla_libero.pth) or a directory holding +model.safetensors and config.json. Dimensions come from tensor shapes. Usage: - python convert_turbovla_to_gguf.py /path/to/checkpoint [--out output.gguf] [--verify] + python convert_turbovla_to_gguf.py turbovla_libero.pth -o turbovla-libero.gguf [--verify] """ from __future__ import annotations @@ -83,17 +84,45 @@ def max_layer(keys: set[str], pfx: str) -> int: return m + 1 -def load_safetensors(ckpt: Path) -> dict[str, torch.Tensor]: - """Load all tensors from safetensors file.""" - tensors = {} - safetensors_path = ckpt / "model.safetensors" - if not safetensors_path.exists(): - raise SystemExit(f"model.safetensors not found in {ckpt}") +# facebook/dinov3-vitb16-pretrain-lvd1689m is gated, but only these architecture +# values are needed: the fine-tuned weights ship inside the TurboVLA checkpoint. +DINOV3_VITB16 = {"rope_theta": 100.0, "num_register_tokens": 4} - with safe_open(str(safetensors_path), framework="pt", device="cpu") as f: - for key in f: - tensors[key] = f.get_tensor(key) - return tensors +# Tensors the runtime never reads: DINOv3's final norm (TurboVLA taps +# hidden_states[-1], before it), BERT's pooler, and the MAE mask token. +UNUSED = (f"{PREFIX_VIT}.norm.", f"{PREFIX_TEXT}.pooler.", f"{PREFIX_VIT}.embeddings.mask_token") + + +class TrackedTensors(dict): + """A state dict that records which keys the writers read.""" + + def __init__(self, *a, **k): + super().__init__(*a, **k) + self.read = set() + + def __getitem__(self, key): + self.read.add(key) + return super().__getitem__(key) + + +def load_checkpoint(ckpt: Path) -> tuple[TrackedTensors, dict]: + """Return (state dict, TurboVLA model_config) from a .pth or a directory.""" + if ckpt.is_file(): + blob = torch.load(str(ckpt), map_location="cpu", weights_only=False) + if not isinstance(blob, dict) or "model_state_dict" not in blob: + raise SystemExit(f"{ckpt} is not a TurboVLA checkpoint (no model_state_dict)") + state = blob["model_state_dict"] + cfg = blob.get("model_config") or {} + else: + st_path = ckpt / "model.safetensors" + if not st_path.exists(): + raise SystemExit(f"model.safetensors not found in {ckpt}") + with safe_open(str(st_path), framework="pt", device="cpu") as f: + state = {k: f.get_tensor(k) for k in f.keys()} + cfg_path = ckpt / "config.json" + cfg = json.loads(cfg_path.read_text()) if cfg_path.exists() else {} + state = {k.removeprefix("module."): v for k, v in state.items()} + return TrackedTensors(state), cfg class TurboVLADims: @@ -191,7 +220,10 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d # checkpoint's TurboVLA config names the backbone, while the matching # DINO config supplies rope_theta and coordinate normalization. if "rope_theta" not in self.dinov3_cfg: - raise SystemExit("DINOv3 config with rope_theta is required; pass --dinov3-config") + backbone = self.cfg.get("vision", {}).get("model_name_or_path", "") + if "dinov3-vitb16" not in backbone: + raise SystemExit(f"unknown DINOv3 backbone {backbone!r}; pass --dinov3-config") + self.dinov3_cfg = DINOV3_VITB16 self.rope_theta = float(self.dinov3_cfg["rope_theta"]) self.rope_normalize_coords = str(self.dinov3_cfg.get("rope_normalize_coords", "separate")) if self.rope_normalize_coords != "separate": @@ -204,7 +236,6 @@ def __init__(self, tensors: dict[str, torch.Tensor], keys: set[str], cfg_json: d word_emb = self._get(f"{PREFIX_TEXT}.embeddings.word_embeddings.weight") self.vocab_size = int(word_emb.shape[0]) - self.dropout = float(self.cfg.get("vision", {}).get("dropout", 0.0)) def _get(self, key: str) -> torch.Tensor: if key not in self.tensors: @@ -309,9 +340,6 @@ def write_vision_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLAD add_tensor(writer, f"vit.blk.{i}.fc2.weight", w_fc2) add_tensor(writer, f"vit.blk.{i}.fc2.bias", b_fc2) - add_tensor(writer, "vit.final_norm.weight", tensors[f"{root}.norm.weight"]) - add_tensor(writer, "vit.final_norm.bias", tensors[f"{root}.norm.bias"]) - def write_text_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADims) -> None: """Write BERT text encoder.""" @@ -343,9 +371,6 @@ def write_text_encoder(writer: gguf.GGUFWriter, tensors: dict, dims: TurboVLADim add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.weight", tensors[f"{lr}.output.LayerNorm.weight"]) add_tensor(writer, f"text.encoder.layer.{i}.output.LayerNorm.bias", tensors[f"{lr}.output.LayerNorm.bias"]) - add_tensor(writer, "text.pooler.dense.weight", tensors[f"{root}.pooler.dense.weight"]) - add_tensor(writer, "text.pooler.dense.bias", tensors[f"{root}.pooler.dense.bias"]) - add_tensor(writer, "text_proj.weight", tensors[KEY_TEXT_PROJ + ".weight"]) add_tensor(writer, "text_proj.bias", tensors[KEY_TEXT_PROJ + ".bias"]) @@ -474,181 +499,38 @@ def write_view_embeddings(writer: gguf.GGUFWriter, tensors: dict) -> None: add_tensor(writer, "view_emb", tensors[KEY_VIEW_EMB]) -def verify_consumed_tensors(source_keys: set[str], dims: TurboVLADims) -> None: - """Verify all source tensors are consumed by the converter.""" - consumed = set() - - root_vit = PREFIX_VIT - consumed.add(f"{root_vit}.embeddings.cls_token") - consumed.add(f"{root_vit}.embeddings.patch_embeddings.weight") - consumed.add(f"{root_vit}.embeddings.patch_embeddings.bias") - reg_key = f"{root_vit}.embeddings.register_tokens" - if reg_key in source_keys: - consumed.add(reg_key) - - for i in range(dims.vit_layers): - lr = f"{root_vit}.layer.{i}" - consumed.add(f"{lr}.attention.q_proj.weight") - consumed.add(f"{lr}.attention.q_proj.bias") - consumed.add(f"{lr}.attention.k_proj.weight") - kb = f"{lr}.attention.k_proj.bias" - if kb in source_keys: - consumed.add(kb) - consumed.add(f"{lr}.attention.v_proj.weight") - consumed.add(f"{lr}.attention.v_proj.bias") - consumed.add(f"{lr}.attention.o_proj.weight") - consumed.add(f"{lr}.attention.o_proj.bias") - consumed.add(f"{lr}.norm1.weight") - consumed.add(f"{lr}.norm1.bias") - consumed.add(f"{lr}.norm2.weight") - consumed.add(f"{lr}.norm2.bias") - consumed.add(f"{lr}.mlp.up_proj.weight") - consumed.add(f"{lr}.mlp.up_proj.bias") - consumed.add(f"{lr}.mlp.down_proj.weight") - consumed.add(f"{lr}.mlp.down_proj.bias") - consumed.add(f"{lr}.layer_scale1.lambda1") - consumed.add(f"{lr}.layer_scale2.lambda1") - - consumed.add(f"{root_vit}.norm.weight") - consumed.add(f"{root_vit}.norm.bias") - - root_text = PREFIX_TEXT - consumed.add(f"{root_text}.embeddings.word_embeddings.weight") - consumed.add(f"{root_text}.embeddings.position_embeddings.weight") - consumed.add(f"{root_text}.embeddings.token_type_embeddings.weight") - consumed.add(f"{root_text}.embeddings.LayerNorm.weight") - consumed.add(f"{root_text}.embeddings.LayerNorm.bias") - consumed.add(f"{root_text}.pooler.dense.weight") - consumed.add(f"{root_text}.pooler.dense.bias") - - for i in range(dims.text_layers): - lr = f"{root_text}.encoder.layer.{i}" - consumed.add(f"{lr}.attention.self.query.weight") - consumed.add(f"{lr}.attention.self.query.bias") - consumed.add(f"{lr}.attention.self.key.weight") - consumed.add(f"{lr}.attention.self.key.bias") - consumed.add(f"{lr}.attention.self.value.weight") - consumed.add(f"{lr}.attention.self.value.bias") - consumed.add(f"{lr}.attention.output.dense.weight") - consumed.add(f"{lr}.attention.output.dense.bias") - consumed.add(f"{lr}.attention.output.LayerNorm.weight") - consumed.add(f"{lr}.attention.output.LayerNorm.bias") - consumed.add(f"{lr}.intermediate.dense.weight") - consumed.add(f"{lr}.intermediate.dense.bias") - consumed.add(f"{lr}.output.dense.weight") - consumed.add(f"{lr}.output.dense.bias") - consumed.add(f"{lr}.output.LayerNorm.weight") - consumed.add(f"{lr}.output.LayerNorm.bias") - - consumed.add(KEY_TEXT_PROJ + ".weight") - consumed.add(KEY_TEXT_PROJ + ".bias") - - vp = PREFIX_VIT_PROJ - consumed.add(f"{vp}.input_norm.weight") - consumed.add(f"{vp}.input_norm.bias") - consumed.add(f"{vp}.mlp.0.weight") - consumed.add(f"{vp}.mlp.0.bias") - consumed.add(f"{vp}.mlp.3.weight") - consumed.add(f"{vp}.mlp.3.bias") - consumed.add(f"{vp}.skip.weight") - consumed.add(f"{vp}.output_norm.weight") - consumed.add(f"{vp}.output_norm.bias") +def write_text_groups(writer: gguf.GGUFWriter, text_cfg: dict) -> None: + """Store each training instruction's BERT tokens and padded length. - for i in range(dims.num_fusion_layers): - lr = f"{PREFIX_FUSION}.{i}" - consumed.add(f"{lr}.attn.v_proj.weight") - consumed.add(f"{lr}.attn.v_proj.bias") - consumed.add(f"{lr}.attn.l_proj.weight") - consumed.add(f"{lr}.attn.l_proj.bias") - consumed.add(f"{lr}.attn.values_v_proj.weight") - consumed.add(f"{lr}.attn.values_v_proj.bias") - consumed.add(f"{lr}.attn.values_l_proj.weight") - consumed.add(f"{lr}.attn.values_l_proj.bias") - consumed.add(f"{lr}.attn.out_v_proj.weight") - consumed.add(f"{lr}.attn.out_v_proj.bias") - consumed.add(f"{lr}.attn.out_l_proj.weight") - consumed.add(f"{lr}.attn.out_l_proj.bias") - consumed.add(f"{lr}.layer_norm_v.weight") - consumed.add(f"{lr}.layer_norm_v.bias") - consumed.add(f"{lr}.layer_norm_l.weight") - consumed.add(f"{lr}.layer_norm_l.bias") - consumed.add(f"{lr}.gamma_v") - consumed.add(f"{lr}.gamma_l") - - for i in range(dims.num_text_layers): - lr = f"{PREFIX_VL_TEXT}.{i}" - consumed.add(f"{lr}.self_attn.in_proj_weight") - consumed.add(f"{lr}.self_attn.in_proj_bias") - consumed.add(f"{lr}.self_attn.out_proj.weight") - consumed.add(f"{lr}.self_attn.out_proj.bias") - consumed.add(f"{lr}.norm1.weight") - consumed.add(f"{lr}.norm1.bias") - consumed.add(f"{lr}.linear1.weight") - consumed.add(f"{lr}.linear1.bias") - consumed.add(f"{lr}.linear2.weight") - consumed.add(f"{lr}.linear2.bias") - consumed.add(f"{lr}.norm2.weight") - consumed.add(f"{lr}.norm2.bias") - - consumed.add(f"{PREFIX_ACT_DEC}.action_queries.weight") - for i in range(dims.num_action_decoder_layers): - lr = f"{PREFIX_ACT_DEC}.decoder.layers.{i}" - consumed.add(f"{lr}.self_attn.in_proj_weight") - consumed.add(f"{lr}.self_attn.in_proj_bias") - consumed.add(f"{lr}.self_attn.out_proj.weight") - consumed.add(f"{lr}.self_attn.out_proj.bias") - consumed.add(f"{lr}.multihead_attn.in_proj_weight") - consumed.add(f"{lr}.multihead_attn.in_proj_bias") - consumed.add(f"{lr}.multihead_attn.out_proj.weight") - consumed.add(f"{lr}.multihead_attn.out_proj.bias") - consumed.add(f"{lr}.norm1.weight") - consumed.add(f"{lr}.norm1.bias") - consumed.add(f"{lr}.norm2.weight") - consumed.add(f"{lr}.norm2.bias") - consumed.add(f"{lr}.norm3.weight") - consumed.add(f"{lr}.norm3.bias") - consumed.add(f"{lr}.linear1.weight") - consumed.add(f"{lr}.linear1.bias") - consumed.add(f"{lr}.linear2.weight") - consumed.add(f"{lr}.linear2.bias") - - ap = f"{PREFIX_ACT_DEC}.action_projection" - consumed.add(f"{ap}.layers.0.weight") - consumed.add(f"{ap}.layers.0.bias") - consumed.add(f"{ap}.layers.1.weight") - consumed.add(f"{ap}.layers.1.bias") - consumed.add(f"{ap}.layers.2.weight") - consumed.add(f"{ap}.layers.2.bias") - - sp = PREFIX_STATE_PROJ - consumed.add(f"{sp}.net.0.weight") - consumed.add(f"{sp}.net.0.bias") - consumed.add(f"{sp}.net.1.weight") - consumed.add(f"{sp}.net.1.bias") - consumed.add(f"{sp}.net.4.weight") - consumed.add(f"{sp}.net.4.bias") - consumed.add(f"{sp}.output_norm.weight") - consumed.add(f"{sp}.output_norm.bias") - consumed.add(f"{sp}.position") - - consumed.add(KEY_VIEW_EMB) - - missing = source_keys - consumed - unexpected_missing = [ - k for k in missing - if not any(x in k for x in ["_optim", "_avg", "mask", "running_", ".nbytes"]) - and not k.endswith(".idx") - ] - - if unexpected_missing: - print(f"\n WARNING: {len(unexpected_missing)} keys not consumed:") - for k in sorted(unexpected_missing)[:20]: - print(f" {k}") - if len(unexpected_missing) > 20: - print(f" ... and {len(unexpected_missing) - 20} more") - print(f"\n Consumed: {len(consumed)} / {len(source_keys)} keys (missing {len(unexpected_missing)})") - else: - print(f"\n All {len(consumed)} keys consumed successfully!") + The ACT decoder attends the padded text rows too, so the padded length is + part of the model's output. TurboVLA pads each instruction to the length + its training batch used (padding_length_by_instruction); the runtime only + sees token ids, so the table is keyed by them. + """ + groups = text_cfg.get("padding_length_by_instruction") or {} + writer.add_uint32(kv_prefix("text_groups.count"), len(groups)) + if not groups: + return + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(text_cfg.get("model_name_or_path", "bert-base-uncased")) + ids, lengths, pad_to = [], [], [] + for instruction, length in sorted(groups.items()): + seq = tok(str(instruction), truncation=True, max_length=int(length))["input_ids"] + ids += [int(t) for t in seq] + lengths.append(len(seq)) + pad_to.append(int(length)) + writer.add_array(kv_prefix("text_groups.tokens"), ids) + writer.add_array(kv_prefix("text_groups.lengths"), lengths) + writer.add_array(kv_prefix("text_groups.pad_to"), pad_to) + + +def verify_consumed_tensors(tensors: TrackedTensors) -> None: + """Fail on any checkpoint tensor that no writer read.""" + left = sorted(k for k in tensors if k not in tensors.read and not k.startswith(UNUSED)) + if left: + raise SystemExit(f"{len(left)} checkpoint tensors not converted: {left[:20]}") + print(f" All {len(tensors.read)} used tensors consumed.") def main() -> int: @@ -656,35 +538,26 @@ def main() -> int: parser = argparse.ArgumentParser( description=f"Convert TurboVLA checkpoint to GGUF ({ARCH})" ) - parser.add_argument("ckpt", type=Path, help="Path to TurboVLA checkpoint directory") + parser.add_argument("ckpt", type=Path, + help="TurboVLA .pth, or a directory with model.safetensors + config.json") parser.add_argument("-o", "--out", type=Path, default=None, help="Output GGUF path") - parser.add_argument("--dinov3-config", type=Path, required=True, - help="Validated config.json for the exact DINOv3 backbone") - parser.add_argument("--verify", action="store_true", help="Verify all tensor keys are consumed") + parser.add_argument("--dinov3-config", type=Path, default=None, + help="config.json of the DINOv3 backbone (default: built-in ViT-B/16 values)") + parser.add_argument("--verify", action="store_true", help="Fail if a checkpoint tensor is left unconverted") args = parser.parse_args() ckpt = args.ckpt.resolve() - if not (ckpt / "model.safetensors").exists(): - raise SystemExit(f"model.safetensors not found in {ckpt}") - - out = args.out or (ckpt / f"{ARCH}.gguf") + out = args.out or (ckpt.with_suffix(".gguf") if ckpt.is_file() else ckpt / f"{ARCH}.gguf") print(f"Loading checkpoint from {ckpt}...") - tensors = load_safetensors(ckpt) + tensors, cfg_json = load_checkpoint(ckpt) keys = set(tensors.keys()) print(f" Loaded {len(tensors)} tensors") - cfg_json = {} - cfg_path = ckpt / "config.json" - if cfg_path.exists(): - cfg_json = json.loads(cfg_path.read_text()) - print(f" Config: {cfg_path}") - - dinov3_cfg_path = args.dinov3_config.resolve() - if not dinov3_cfg_path.is_file(): - raise SystemExit(f"DINOv3 config not found: {dinov3_cfg_path}") - dinov3_cfg = json.loads(dinov3_cfg_path.read_text()) - print(f" DINOv3 config: {dinov3_cfg_path}") + dinov3_cfg = {} + if args.dinov3_config: + dinov3_cfg = json.loads(args.dinov3_config.read_text()) + print(f" DINOv3 config: {args.dinov3_config}") print(" Inferring dimensions from checkpoint...") dims = TurboVLADims(tensors, keys, cfg_json, dinov3_cfg) @@ -725,7 +598,13 @@ def main() -> int: writer.add_uint32(kv("num_register_tokens"), dims.num_register_tokens) writer.add_float32(kv("rope_theta"), float(dims.rope_theta)) writer.add_uint32(kv("vocab_size"), dims.vocab_size) - writer.add_uint32(kv("max_text_length"), int(cfg_json.get("text", {}).get("padding_length", 256))) + # The runtime pads to a fixed length; a checkpoint without text.padding_length + # pads each batch to its longest instruction, which it does not implement. + padding_length = cfg_json.get("text", {}).get("padding_length") + if padding_length is None: + raise SystemExit("checkpoint config has no text.padding_length; padding to the longest " + "instruction is not supported") + writer.add_uint32(kv("max_text_length"), int(padding_length)) # TurboVLA's BERT wrapper uses these exact punctuation IDs when it creates # sub-sentence attention masks. Persist them so the GGUF runtime does not # silently depend on a tokenizer installation. @@ -741,7 +620,7 @@ def main() -> int: writer.add_uint32(kv("sep_token_id"), 102) writer.add_uint32(kv("period_token_id"), 1012) writer.add_uint32(kv("question_token_id"), 1029) - writer.add_float32(kv("dropout"), dims.dropout) + write_text_groups(writer, text_cfg) print(" Writing vision encoder...") write_vision_encoder(writer, tensors, dims) @@ -769,7 +648,7 @@ def main() -> int: if args.verify: print(" Verifying tensor consumption...") - verify_consumed_tensors(keys, dims) + verify_consumed_tensors(tensors) writer.write_header_to_file() writer.write_kv_data_to_file() diff --git a/src/layers/rope.h b/src/layers/rope.h index 19d5a2b..ca06156 100644 --- a/src/layers/rope.h +++ b/src/layers/rope.h @@ -59,25 +59,6 @@ inline ggml_tensor * rope_2d(ggml_context * C, ggml_tensor * x, ggml_tensor * co return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, rot, sin_t)); } -// DINOv3 uses axial RoPE. x is [head_dim, patches, heads]; cos/sin are -// [head_dim, patches] generated from the official patch-centre coordinates. -inline ggml_tensor * rope_dinov3_axial(ggml_context * C, ggml_tensor * x, - ggml_tensor * cos_t, ggml_tensor * sin_t) { - GGML_ASSERT(x->ne[0] == cos_t->ne[0]); - GGML_ASSERT(x->ne[1] == cos_t->ne[1]); - ggml_tensor * cos = ggml_reshape_3d(C, cos_t, cos_t->ne[0], cos_t->ne[1], 1); - ggml_tensor * sin = ggml_reshape_3d(C, sin_t, sin_t->ne[0], sin_t->ne[1], 1); - cos = ggml_repeat(C, cos, x); - sin = ggml_repeat(C, sin, x); - const int64_t half = x->ne[0] / 2; - ggml_tensor * x0 = ggml_cont(C, ggml_view_3d(C, x, half, x->ne[1], x->ne[2], - x->nb[1], x->nb[2], 0)); - ggml_tensor * x1 = ggml_cont(C, ggml_view_3d(C, x, half, x->ne[1], x->ne[2], - x->nb[1], x->nb[2], half * x->nb[0])); - ggml_tensor * rotated = ggml_concat(C, ggml_neg(C, x1), x0, 0); - return ggml_add(C, ggml_mul(C, x, cos), ggml_mul(C, rotated, sin)); -} - 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]; diff --git a/src/models/turbovla.cpp b/src/models/turbovla.cpp index 8372e9a..3d3d3ae 100644 --- a/src/models/turbovla.cpp +++ b/src/models/turbovla.cpp @@ -1,1467 +1,855 @@ - // 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 turbovla.cpp - * @brief TurboVLA implementation: DINOv3 ViT + BERT + VL Fusion + ACT decoder. - */ - - #include "arch.h" - #include "loader.h" - #include "model.h" - #include "options.h" - - #include "layers/linear.h" - #include "layers/norm.h" - #include "layers/ffn.h" - #include "layers/attn.h" - #include "layers/rope.h" - #include "modules/preprocess.h" - #include "gguf_reader.h" - #include "scratch_ctx.h" - #include "backend.h" - - #include "ggml.h" - #include "ggml-cpu.h" - #include "ggml-backend.h" - - #include - #include - #include - #include - #include - #include - #include - #include - #include - - namespace vla { - - struct TurboVLAConfig { - int64_t hidden = 256; - int64_t vit_dim = 768; - int64_t vit_layers = 12; - int64_t vit_head_dim = 64; - int64_t vit_heads = 12; - int64_t text_dim = 768; - int64_t text_layers = 12; - int64_t text_head_dim = 64; - int64_t text_heads = 12; - int64_t vocab_size = 30522; - int64_t max_text_length = 256; - int64_t num_fusion_layers = 6; - int64_t fusion_heads = 4; - int64_t fusion_head_dim = 256; - int64_t text_enhancer_heads = 4; - int64_t text_enhancer_head_dim = 64; - int64_t action_heads = 8; - int64_t action_head_dim = 32; - int64_t num_text_layers = 6; - int64_t num_action_decoder_layers = 3; - int64_t action_dim = 7; - int64_t state_dim = 8; - int64_t num_state_tokens = 2; - int64_t action_horizon = 12; - int64_t image_size = 256; - int64_t patch_size = 16; - int64_t num_views = 2; - int64_t num_register_tokens = 0; - float rope_theta = 10000.0f; - float dropout = 0.1f; - int64_t pad_token_id = 0; - int64_t cls_token_id = 101; - int64_t sep_token_id = 102; - int64_t period_token_id = 1012; - int64_t question_token_id = 1029; - }; - - namespace { - - struct VitBlockW { - ggml_tensor* ln1_w, *ln1_b; - ggml_tensor* ln2_w, *ln2_b; - ggml_tensor* attn_q_w, *attn_q_b; - ggml_tensor* attn_k_w, *attn_k_b; - ggml_tensor* attn_v_w, *attn_v_b; - ggml_tensor* attn_o_w, *attn_o_b; - ggml_tensor* fc1_w, *fc1_b; - ggml_tensor* fc2_w, *fc2_b; - }; - - struct VisionEncoderW { - ggml_tensor* cls_token; - ggml_tensor* patch_embed_w; - ggml_tensor* patch_embed_b; - ggml_tensor* register_tokens; - std::vector layers; - ggml_tensor* final_ln_w; - ggml_tensor* final_ln_b; - // One frequency tensor per patch-grid position. - std::vector rope_freqs; - }; - - struct VisionProjectionW { - ggml_tensor* input_norm_w, *input_norm_b; - ggml_tensor* mlp_0_w, *mlp_0_b; - ggml_tensor* mlp_3_w, *mlp_3_b; - ggml_tensor* skip_w; - ggml_tensor* output_norm_w, *output_norm_b; - }; - - struct BertLayerW { - ggml_tensor* attn_q_w, *attn_q_b; - ggml_tensor* attn_k_w, *attn_k_b; - ggml_tensor* attn_v_w, *attn_v_b; - ggml_tensor* attn_o_w, *attn_o_b; - ggml_tensor* ln1_w, *ln1_b; - ggml_tensor* fc1_w, *fc1_b; - ggml_tensor* fc2_w, *fc2_b; - ggml_tensor* ln2_w, *ln2_b; - }; - - struct TextEncoderW { - ggml_tensor* word_embed; - ggml_tensor* pos_embed; - ggml_tensor* tok_type_embed; - ggml_tensor* embed_ln_w, *embed_ln_b; - std::vector layers; - ggml_tensor* pooler_w, *pooler_b; - ggml_tensor* text_proj_w, *text_proj_b; - }; +// 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. + +// TurboVLA (H-EmbodVis/TurboVLA@b29ab142): DINOv3 ViT-B/16 over each camera, +// BERT-base over the instruction, six Grounding-DINO bi-attention fusion layers +// each followed by a text enhancer, and a three-layer ACT decoder whose twelve +// learned queries read the fused tokens plus two state tokens. One forward pass, +// no denoising loop, so the whole model is one cached graph. + +#include "arch.h" +#include "backend.h" +#include "gguf.h" +#include "gguf_reader.h" +#include "loader.h" +#include "model.h" +#include "options.h" +#include "scratch_ctx.h" +#include "layers/attn.h" +#include "layers/ffn.h" +#include "layers/linear.h" +#include "layers/norm.h" +#include "modules/preprocess.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vla { +namespace { + +constexpr float kImagenetMean[3] = {0.485f, 0.456f, 0.406f}; +constexpr float kImagenetStd[3] = {0.229f, 0.224f, 0.225f}; + +// Epsilons are the module defaults upstream: nn.LayerNorm everywhere except +// BERT, whose config sets 1e-12. +constexpr float kLnEps = 1e-5f; +constexpr float kBertEps = 1e-12f; + +struct VitLayerW { + ggml_tensor *ln1_w, *ln1_b, *qkv_w, *qkv_b, *o_w, *o_b; + ggml_tensor *ln2_w, *ln2_b, *fc1_w, *fc1_b, *fc2_w, *fc2_b; +}; + +struct BertLayerW { + ggml_tensor *qkv_w, *qkv_b, *o_w, *o_b, *ln1_w, *ln1_b; + ggml_tensor *fc1_w, *fc1_b, *fc2_w, *fc2_b, *ln2_w, *ln2_b; +}; + +// v_proj/values_v and l_proj/values_l are fused per side: v_proj(v) is the +// vision query and also the key the language side attends, l_proj(l) the +// other way round, so each side needs one GEMM for both of its projections. +struct FusionLayerW { + ggml_tensor *norm_v_w, *norm_v_b, *norm_l_w, *norm_l_b; + ggml_tensor *v_w, *v_b, *l_w, *l_b; + ggml_tensor *out_v_w, *out_v_b, *out_l_w, *out_l_b; + ggml_tensor *gamma_v, *gamma_l; +}; + +struct EnhancerLayerW { + ggml_tensor *qkv_w, *qkv_b, *o_w, *o_b, *ln1_w, *ln1_b; + ggml_tensor *fc1_w, *fc1_b, *fc2_w, *fc2_b, *ln2_w, *ln2_b; +}; + +struct DecoderLayerW { + ggml_tensor *ln1_w, *ln1_b, *self_qkv_w, *self_qkv_b, *self_o_w, *self_o_b; + ggml_tensor *ln2_w, *ln2_b, *cross_qkv_w, *cross_qkv_b, *cross_o_w, *cross_o_b; + ggml_tensor *ln3_w, *ln3_b, *fc1_w, *fc1_b, *fc2_w, *fc2_b; +}; + +// A training instruction's BERT token ids and the length its batch was padded +// to. See text_len() for why the length is part of the model. +struct TextGroup { + std::vector ids; + int64_t pad_to = 0; +}; + +} // namespace + +struct TurboVlaModelArch : public ModelArchBase { + TurboVlaModelArch() : ModelArchBase(Arch::TURBOVLA) {} + ~TurboVlaModelArch() override { + graph.release(); + if (const_buf) ggml_backend_buffer_free(const_buf); + if (ctx_const) ggml_free(ctx_const); + if (weight_buf) ggml_backend_buffer_free(weight_buf); + if (ctx_weights) ggml_free(ctx_weights); + if (backend) ggml_backend_free(backend); + } - struct VlFusionLayerW { - ggml_tensor* v_proj_w, *v_proj_b; - ggml_tensor* l_proj_w, *l_proj_b; - ggml_tensor* values_v_w, *values_v_b; - ggml_tensor* values_l_w, *values_l_b; - ggml_tensor* out_v_w, *out_v_b; - ggml_tensor* out_l_w, *out_l_b; - ggml_tensor* norm_v_w, *norm_v_b; - ggml_tensor* norm_l_w, *norm_l_b; - ggml_tensor* gamma_v; - ggml_tensor* gamma_l; + ggml_backend_t backend = nullptr; + int n_threads = default_cpu_threads(); + ggml_context * ctx_weights = nullptr; + ggml_backend_buffer_t weight_buf = nullptr; + ggml_context * ctx_const = nullptr; + ggml_backend_buffer_t const_buf = nullptr; + ggml_type mt = GGML_TYPE_F32; + + int64_t hidden = 256, n_views = 2, image_size = 256, patch = 16, n_reg = 4; + int64_t vit_dim = 768, vit_layers = 12, vit_heads = 12; + int64_t bert_dim = 768, bert_layers = 12, bert_heads = 12, vocab = 30522, bert_max_pos = 512; + int64_t fusion_layers = 6, fusion_heads = 4, fusion_dim = 1024, enh_heads = 4; + int64_t dec_layers = 3, dec_heads = 8, horizon = 12, action_dim = 7, state_dim = 8, n_state_tok = 2; + int64_t text_len_max = 21; + float rope_theta = 100.0f; + int32_t pad_id = 0, cls_id = 101, sep_id = 102, period_id = 1012, question_id = 1029; + std::vector text_groups; + + ggml_tensor *patch_w = nullptr, *patch_b = nullptr, *cls_tok = nullptr, *reg_tok = nullptr; + std::vector vit; + ggml_tensor *vp_in_w = nullptr, *vp_in_b = nullptr, *vp_fc1_w = nullptr, *vp_fc1_b = nullptr; + ggml_tensor *vp_fc2_w = nullptr, *vp_fc2_b = nullptr, *vp_skip_w = nullptr; + ggml_tensor *vp_out_w = nullptr, *vp_out_b = nullptr, *view_emb = nullptr; + ggml_tensor *word_emb = nullptr, *pos_emb = nullptr, *type_emb = nullptr; + ggml_tensor *emb_ln_w = nullptr, *emb_ln_b = nullptr; + std::vector bert; + ggml_tensor *text_proj_w = nullptr, *text_proj_b = nullptr; + std::vector fusion; + std::vector enhancer; + ggml_tensor *st_ln_w = nullptr, *st_ln_b = nullptr, *st_fc1_w = nullptr, *st_fc1_b = nullptr; + ggml_tensor *st_fc2_w = nullptr, *st_fc2_b = nullptr, *st_pos = nullptr, *st_out_w = nullptr, *st_out_b = nullptr; + ggml_tensor *act_q = nullptr; + std::vector dec; + ggml_tensor *ap0_w = nullptr, *ap0_b = nullptr, *ap1_w = nullptr, *ap1_b = nullptr; + ggml_tensor *ap2_w = nullptr, *ap2_b = nullptr; + + // DINOv3's RoPE depends only on the patch grid, so its tables are computed + // once and live next to the weights, shaped [head_dim, 1, T] to broadcast + // over heads. Rows for the CLS and register tokens are the identity (cos 1, + // sin 0), which lets the rotation cover the whole sequence instead of + // splitting off the prefix in every layer. rope_sin carries rotate_half's + // sign: negative on the first half of the head. + ggml_tensor *rope_cos = nullptr, *rope_sin = nullptr; + + struct Key { + int64_t bert_len = -1; + bool operator==(const Key & o) const { return bert_len == o.bert_len; } }; - - struct VlTextLayerW { - ggml_tensor* attn_qkv_w, *attn_qkv_b; - ggml_tensor* attn_o_w, *attn_o_b; - ggml_tensor* ln1_w, *ln1_b; - ggml_tensor* fc1_w, *fc1_b; - ggml_tensor* fc2_w, *fc2_b; - ggml_tensor* ln2_w, *ln2_b; + struct IO { + ggml_tensor *patches = nullptr, *ids = nullptr, *pos = nullptr, *bert_mask = nullptr; + ggml_tensor *enh_mask = nullptr, *fus_mask = nullptr, *state = nullptr, *actions = nullptr; }; + graph_cache graph; + + int64_t grid() const { return image_size / patch; } + int64_t n_patches() const { return grid() * grid(); } + int64_t vit_seq() const { return 1 + n_reg + n_patches(); } + + int64_t text_len(const int32_t * ids, int64_t n) const; + ggml_cgraph * build(ggml_context * C, IO & io, int64_t bert_len) const; + bool upload_rope_tables(); + + std::vector predict(const Inputs& in) override; +}; + +namespace { + +ggml_tensor * split_heads(ggml_context * C, ggml_tensor * qkv, int64_t hd, int64_t heads, + int64_t T, int64_t nv, int part) { + const size_t es = ggml_element_size(qkv); + return ggml_view_4d(C, qkv, hd, heads, T, nv, (size_t) hd*es, qkv->nb[1], qkv->nb[2], + (size_t) part*hd*heads*es); +} + +// Fused-QKV self-attention over [3*dim, T, nv]; returns [dim, T, nv]. +ggml_tensor * self_attention(ggml_context * C, ggml_tensor * qkv, ggml_tensor * mask, + int64_t hd, int64_t heads, int64_t T, int64_t nv) { + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, split_heads(C, qkv, hd, heads, T, nv, 0), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, split_heads(C, qkv, hd, heads, T, nv, 1), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, split_heads(C, qkv, hd, heads, T, nv, 2), 1, 2, 0, 3)); + return attention(C, Q, K, V, mask, 1.0f / std::sqrt((float) hd), hd*heads, T, nv); +} + +// Half-split RoPE applied in the projection's own [hd, heads, T, nv] layout, so +// no head permute is needed first: rotate_half(x) is [-x1, x0], and the sign +// lives in the sin table, leaving one concat of two views. +ggml_tensor * rope_heads(ggml_context * C, ggml_tensor * x, ggml_tensor * cos_t, ggml_tensor * sin_signed) { + const int64_t half = x->ne[0] / 2; + const size_t es = ggml_element_size(x); + ggml_tensor * x0 = ggml_view_4d(C, x, half, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0); + ggml_tensor * x1 = ggml_view_4d(C, x, half, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], half*es); + return ggml_add(C, ggml_mul(C, x, cos_t), ggml_mul(C, ggml_concat(C, x1, x0, 0), sin_signed)); +} + +// DINOv3 attention. Flash attention takes Q and K as permuted views and casts +// K and V itself, so the rotated projections are never copied again. +ggml_tensor * vit_attention(ggml_context * C, ggml_tensor * qkv, ggml_tensor * cos_t, ggml_tensor * sin_signed, + int64_t hd, int64_t heads, int64_t T, int64_t nv, bool flash) { + const float scale = 1.0f / std::sqrt((float) hd); + ggml_tensor * q = rope_heads(C, split_heads(C, qkv, hd, heads, T, nv, 0), cos_t, sin_signed); + ggml_tensor * k = rope_heads(C, split_heads(C, qkv, hd, heads, T, nv, 1), cos_t, sin_signed); + ggml_tensor * v = split_heads(C, qkv, hd, heads, T, nv, 2); + if (flash) { + ggml_tensor * o = flash_attention(C, ggml_permute(C, q, 0, 2, 1, 3), ggml_permute(C, k, 0, 2, 1, 3), + ggml_permute(C, v, 0, 2, 1, 3), nullptr, scale); + return ggml_reshape_3d(C, o, hd*heads, T, nv); + } + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, q, 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, k, 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, v, 1, 2, 0, 3)); + return attention(C, Q, K, V, nullptr, scale, hd*heads, T, nv); +} + +// The released BERT wrapper (text/bert.py, from Grounding DINO) splits the +// instruction at [CLS], [SEP], '.' and '?': a token attends only inside its +// sub-sentence and its position restarts there. Tokens outside any span (the +// padding, and a special token at either end) see only themselves at position 0. +void special_token_blocks(const std::vector & ids, const TurboVlaModelArch & m, + std::vector & allowed, std::vector & pos) { + const int64_t n = (int64_t) ids.size(); + allowed.assign((size_t) n*n, 0); + pos.assign((size_t) n, 0); + for (int64_t i = 0; i < n; ++i) + allowed[(size_t) i*n + i] = 1; + int64_t prev = 0; + for (int64_t col = 0; col < n; ++col) { + const int32_t id = ids[(size_t) col]; + if (id != m.cls_id && id != m.sep_id && id != m.period_id && id != m.question_id) + continue; + if (col != 0 && col != n - 1) { + for (int64_t q = prev + 1; q <= col; ++q) { + pos[(size_t) q] = (int32_t) (q - prev - 1); + for (int64_t k = prev + 1; k <= col; ++k) + allowed[(size_t) q*n + k] = 1; + } + } + prev = col; + } +} + +} // namespace + +// The ACT decoder attends every text row, padding included (encode_condition +// concatenates them unmasked), so the padded length is part of the output. The +// checkpoint pins it per training instruction (11, 14 or 21 for LIBERO): BERT +// runs over that many tokens and the rows after it, up to text_len_max, are +// zeros that the text projection turns into its bias. Instructions the +// checkpoint does not list are padded to text_len_max. A caller that sends +// trailing [PAD] tokens has chosen the length itself. +int64_t TurboVlaModelArch::text_len(const int32_t * ids, int64_t n) const { + int64_t core = n; + while (core > 0 && ids[core - 1] == pad_id) + --core; + if (core < n) + return n; + for (const TextGroup & g : text_groups) + if ((int64_t) g.ids.size() == n && std::equal(g.ids.begin(), g.ids.end(), ids)) + return g.pad_to; + return text_len_max; +} + +bool TurboVlaModelArch::upload_rope_tables() { + // DINOv3 axial RoPE (HF modeling_dinov3_vit.py): patch centres normalized to + // [-1, 1] per axis, head_dim/4 periods base^(4i/head_dim), angles ordered + // [h, w, h, w] across the head, rotated half-split. See rope_heads. + const int64_t hd = vit_dim / vit_heads, q = hd / 4, T = vit_seq(), g = grid(), pre = 1 + n_reg; + std::vector c((size_t) hd*T, 1.0f), s((size_t) hd*T, 0.0f); + for (int64_t p = 0; p < n_patches(); ++p) { + const float h = 2.0f * ((float) (p / g) + 0.5f) / (float) g - 1.0f; + const float w = 2.0f * ((float) (p % g) + 0.5f) / (float) g - 1.0f; + float * cr = c.data() + (size_t) (pre + p)*hd; + float * sr = s.data() + (size_t) (pre + p)*hd; + for (int64_t i = 0; i < q; ++i) { + const float period = std::pow(rope_theta, 4.0f * (float) i / (float) hd); + const float ah = 6.28318530717958647692f * h / period; + const float aw = 6.28318530717958647692f * w / period; + for (int64_t blk = 0; blk < 4; ++blk) { + const float a = (blk % 2 == 0) ? ah : aw; + cr[blk*q + i] = std::cos(a); + sr[blk*q + i] = blk < 2 ? -std::sin(a) : std::sin(a); + } + } + } - struct ActDecoderLayerW { - ggml_tensor* self_qkv_w, *self_qkv_b; - ggml_tensor* self_out_w, *self_out_b; - ggml_tensor* cross_qkv_w, *cross_qkv_b; - ggml_tensor* cross_out_w, *cross_out_b; - ggml_tensor* ln1_w, *ln1_b; - ggml_tensor* ln2_w, *ln2_b; - ggml_tensor* ln3_w, *ln3_b; - ggml_tensor* fc1_w, *fc1_b; - ggml_tensor* fc2_w, *fc2_b; - }; + ggml_init_params p = { 2*ggml_tensor_overhead(), nullptr, true }; + ctx_const = ggml_init(p); + if (!ctx_const) + return false; + rope_cos = ggml_new_tensor_3d(ctx_const, GGML_TYPE_F32, hd, 1, T); + rope_sin = ggml_new_tensor_3d(ctx_const, GGML_TYPE_F32, hd, 1, T); + const_buf = alloc_weights(ctx_const, backend); + if (!const_buf) + return false; + ggml_backend_tensor_set(rope_cos, c.data(), 0, ggml_nbytes(rope_cos)); + ggml_backend_tensor_set(rope_sin, s.data(), 0, ggml_nbytes(rope_sin)); + return true; +} + +ggml_cgraph * TurboVlaModelArch::build(ggml_context * C, IO & io, int64_t bert_len) const { + const int64_t NP = n_patches(), T = vit_seq(), nv = n_views, VS = nv*NP, LT = text_len_max; + const int64_t pdim = 3*patch*patch; + + // DINOv3. Patches arrive flattened on the host; views are the batch dim. + io.patches = ggml_new_tensor_3d(C, GGML_TYPE_F32, pdim, NP, nv); + ggml_set_input(io.patches); + ggml_tensor * prefix = ggml_reshape_2d(C, cls_tok, vit_dim, 1); + if (reg_tok) + prefix = ggml_concat(C, prefix, reg_tok, 1); + prefix = ggml_repeat_4d(C, prefix, vit_dim, 1 + n_reg, nv, 1); + ggml_tensor * x = ggml_concat(C, prefix, linear(C, patch_w, patch_b, io.patches), 1); + for (const VitLayerW & l : vit) { + ggml_tensor * qkv = linear(C, l.qkv_w, l.qkv_b, layer_norm(C, x, l.ln1_w, l.ln1_b, kLnEps)); + ggml_tensor * att = vit_attention(C, qkv, rope_cos, rope_sin, vit_dim / vit_heads, vit_heads, T, nv, + flash_attn_enabled()); + x = ggml_add(C, x, linear(C, l.o_w, l.o_b, att)); + x = ggml_add(C, x, ffn_gelu_erf(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, + layer_norm(C, x, l.ln2_w, l.ln2_b, kLnEps))); + } + // TurboVLA reads hidden_states[-1], which is before DINOv3's final norm + // (models/vision_encoder.py:109), then drops the CLS and register tokens. + x = ggml_view_3d(C, x, vit_dim, NP, nv, x->nb[1], x->nb[2], (size_t) (1 + n_reg)*x->nb[1]); + x = ggml_reshape_2d(C, ggml_cont(C, x), vit_dim, VS); + + // VisionProjection, then one learned embedding per camera. + ggml_tensor * mlp = linear(C, vp_fc1_w, vp_fc1_b, layer_norm(C, x, vp_in_w, vp_in_b, kLnEps)); + mlp = linear(C, vp_fc2_w, vp_fc2_b, ggml_gelu_erf(C, mlp)); + ggml_tensor * v = layer_norm(C, ggml_add(C, ggml_mul_mat(C, vp_skip_w, x), mlp), vp_out_w, vp_out_b, kLnEps); + v = ggml_reshape_2d(C, ggml_add(C, ggml_reshape_3d(C, v, hidden, NP, nv), + ggml_reshape_3d(C, view_emb, hidden, 1, nv)), hidden, VS); + + // BERT over the checkpoint's padded length, token type 0 throughout. + io.ids = ggml_new_tensor_1d(C, GGML_TYPE_I32, bert_len); + io.pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, bert_len); + io.bert_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, bert_len, bert_len); + ggml_set_input(io.ids); + ggml_set_input(io.pos); + ggml_set_input(io.bert_mask); + ggml_tensor * t = ggml_add(C, ggml_get_rows(C, word_emb, io.ids), ggml_get_rows(C, pos_emb, io.pos)); + t = ggml_add(C, t, ggml_view_1d(C, type_emb, bert_dim, 0)); + t = layer_norm(C, t, emb_ln_w, emb_ln_b, kBertEps); + for (const BertLayerW & l : bert) { + ggml_tensor * att = self_attention(C, linear(C, l.qkv_w, l.qkv_b, t), io.bert_mask, + bert_dim / bert_heads, bert_heads, bert_len, 1); + t = layer_norm(C, ggml_add(C, t, linear(C, l.o_w, l.o_b, att)), l.ln1_w, l.ln1_b, kBertEps); + t = layer_norm(C, ggml_add(C, t, ffn_gelu_erf(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, t)), + l.ln2_w, l.ln2_b, kBertEps); + } + ggml_tensor * lang = linear(C, text_proj_w, text_proj_b, t); + if (bert_len < LT) + lang = ggml_concat(C, lang, ggml_repeat_4d(C, text_proj_b, hidden, LT - bert_len, 1, 1), 1); + + // Fusion. One score matrix serves both directions upstream (the language + // side softmaxes its transpose); two small products are cheaper than a + // transpose here. Only vision queries mask padded language keys. + io.enh_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, LT, LT); + io.fus_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, LT, VS); + ggml_set_input(io.enh_mask); + ggml_set_input(io.fus_mask); + const int64_t fhd = fusion_dim / fusion_heads, ehd = hidden / enh_heads; + const float fscale = 1.0f / std::sqrt((float) fhd); + for (int64_t i = 0; i < fusion_layers; ++i) { + const FusionLayerW & f = fusion[(size_t) i]; + ggml_tensor * vn = layer_norm(C, v, f.norm_v_w, f.norm_v_b, kLnEps); + ggml_tensor * ln = layer_norm(C, lang, f.norm_l_w, f.norm_l_b, kLnEps); + ggml_tensor * vp = linear(C, f.v_w, f.v_b, vn); // [q_v | values_v] + ggml_tensor * lp = linear(C, f.l_w, f.l_b, ln); // [k_l | values_l] + auto heads = [&](ggml_tensor * p, int64_t L, int part) { + return ggml_view_3d(C, p, fhd, fusion_heads, L, (size_t) fhd*ggml_element_size(p), p->nb[1], + (size_t) part*fusion_dim*ggml_element_size(p)); + }; + ggml_tensor * qv = ggml_cont(C, ggml_permute(C, heads(vp, VS, 0), 0, 2, 1, 3)); + ggml_tensor * kl = ggml_cont(C, ggml_permute(C, heads(lp, LT, 0), 0, 2, 1, 3)); + ggml_tensor * vv = ggml_cont(C, ggml_permute(C, heads(vp, VS, 1), 1, 2, 0, 3)); + ggml_tensor * vl = ggml_cont(C, ggml_permute(C, heads(lp, LT, 1), 1, 2, 0, 3)); + ggml_tensor * dv = attention(C, qv, kl, vl, io.fus_mask, fscale, fusion_dim, VS); + ggml_tensor * dl = attention(C, kl, qv, vv, nullptr, fscale, fusion_dim, LT); + // residual_style "normalized": the gated deltas land on the normed tokens. + v = ggml_add(C, vn, ggml_mul(C, linear(C, f.out_v_w, f.out_v_b, dv), f.gamma_v)); + lang = ggml_add(C, ln, ggml_mul(C, linear(C, f.out_l_w, f.out_l_b, dl), f.gamma_l)); + + // Text enhancer: post-norm, ReLU, the sub-sentence mask and no + // key-padding mask (the key_padding_mask argument is commented out + // upstream, models/components/transformer.py). + const EnhancerLayerW & e = enhancer[(size_t) i]; + ggml_tensor * att = self_attention(C, linear(C, e.qkv_w, e.qkv_b, lang), io.enh_mask, ehd, enh_heads, LT, 1); + lang = layer_norm(C, ggml_add(C, lang, linear(C, e.o_w, e.o_b, att)), e.ln1_w, e.ln1_b, kLnEps); + lang = layer_norm(C, ggml_add(C, lang, ffn_relu(C, e.fc1_w, e.fc1_b, e.fc2_w, e.fc2_b, lang)), + e.ln2_w, e.ln2_b, kLnEps); + } - struct ActionDecoderW { - ggml_tensor* action_q; - std::vector layers; - ggml_tensor* proj_0_w, *proj_0_b; - ggml_tensor* proj_1_w, *proj_1_b; - ggml_tensor* proj_2_w, *proj_2_b; + // State tokens: LayerNorm, MLP to num_state_tokens*hidden, + position, norm. + io.state = ggml_new_tensor_1d(C, GGML_TYPE_F32, state_dim); + ggml_set_input(io.state); + ggml_tensor * s = linear(C, st_fc1_w, st_fc1_b, layer_norm(C, io.state, st_ln_w, st_ln_b, kLnEps)); + s = ggml_reshape_2d(C, linear(C, st_fc2_w, st_fc2_b, ggml_gelu_erf(C, s)), hidden, n_state_tok); + s = layer_norm(C, ggml_add(C, s, ggml_reshape_2d(C, st_pos, hidden, n_state_tok)), st_out_w, st_out_b, kLnEps); + + ggml_tensor * mem = ggml_concat(C, ggml_concat(C, v, lang, 1), s, 1); + const int64_t M = VS + LT + n_state_tok; + + // ACT decoder: nn.TransformerDecoderLayer, norm_first, ReLU, no final norm. + const int64_t dhd = hidden / dec_heads; + const float dscale = 1.0f / std::sqrt((float) dhd); + ggml_tensor * a = act_q; + for (const DecoderLayerW & l : dec) { + ggml_tensor * qkv = linear(C, l.self_qkv_w, l.self_qkv_b, layer_norm(C, a, l.ln1_w, l.ln1_b, kLnEps)); + ggml_tensor * att = self_attention(C, qkv, nullptr, dhd, dec_heads, horizon, 1); + a = ggml_add(C, a, linear(C, l.self_o_w, l.self_o_b, att)); + + // in_proj rows are [W_q; W_k; W_v]: the queries take the first block, + // the memory the other two, in one GEMM each. + const size_t wr = l.cross_qkv_w->nb[1], br = ggml_element_size(l.cross_qkv_b); + ggml_tensor * wq = ggml_view_2d(C, l.cross_qkv_w, hidden, hidden, wr, 0); + ggml_tensor * wkv = ggml_view_2d(C, l.cross_qkv_w, hidden, 2*hidden, wr, (size_t) hidden*wr); + ggml_tensor * bq = ggml_view_1d(C, l.cross_qkv_b, hidden, 0); + ggml_tensor * bkv = ggml_view_1d(C, l.cross_qkv_b, 2*hidden, (size_t) hidden*br); + ggml_tensor * q = linear(C, wq, bq, layer_norm(C, a, l.ln2_w, l.ln2_b, kLnEps)); + ggml_tensor * kv = linear(C, wkv, bkv, mem); + ggml_tensor * Q = to_heads(C, q, dhd, dec_heads, horizon); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 0), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 1), 1, 2, 0, 3)); + att = attention(C, Q, K, V, nullptr, dscale, hidden, horizon); + a = ggml_add(C, a, linear(C, l.cross_o_w, l.cross_o_b, att)); + a = ggml_add(C, a, ffn_relu(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, + layer_norm(C, a, l.ln3_w, l.ln3_b, kLnEps))); + } + a = ggml_relu(C, linear(C, ap0_w, ap0_b, a)); + a = ggml_relu(C, linear(C, ap1_w, ap1_b, a)); + io.actions = ggml_tanh(C, linear(C, ap2_w, ap2_b, a)); + ggml_set_output(io.actions); + + ggml_cgraph * gf = ggml_new_graph_custom(C, 8192, false); + ggml_build_forward_expand(gf, io.actions); + return gf; +} + +namespace { + +bool load_config(const gguf_reader & g, TurboVlaModelArch & m) { + auto U = [&](const char * k, int64_t & dst) { + char key[96]; + std::snprintf(key, sizeof(key), "turbovla.%s", k); + if (!g.has(key)) + return; + dst = (int64_t) g.u32(key); }; - - struct StateProjectionW { - ggml_tensor* net_0_w, *net_0_b; - ggml_tensor* net_1_w, *net_1_b; - ggml_tensor* net_4_w, *net_4_b; - ggml_tensor* output_norm_w, *output_norm_b; - ggml_tensor* position; + auto I = [&](const char * k, int32_t & dst) { + int64_t v = dst; + U(k, v); + dst = (int32_t) v; }; - - static void check_heads(const char* tag, ggml_tensor* t, int64_t hd, int64_t heads, int64_t T, int64_t nv = 1) { - (void) tag; - const int64_t got = ggml_nelements(t); - const int64_t expected = hd * heads * T * nv; - GGML_ASSERT(got == expected); + U("hidden", m.hidden); U("num_views", m.n_views); + U("image_size", m.image_size); U("patch_size", m.patch); + U("num_register_tokens", m.n_reg); + U("vit_dim", m.vit_dim); U("vit_layers", m.vit_layers); U("vit_heads", m.vit_heads); + U("text_dim", m.bert_dim); U("text_layers", m.bert_layers); U("text_heads", m.bert_heads); + U("vocab_size", m.vocab); + U("num_fusion_layers", m.fusion_layers); U("fusion_heads", m.fusion_heads); + U("text_enhancer_heads", m.enh_heads); + U("num_action_decoder_layers", m.dec_layers); U("action_heads", m.dec_heads); + U("action_horizon", m.horizon); U("action_dim", m.action_dim); + U("state_dim", m.state_dim); U("num_state_tokens", m.n_state_tok); + U("max_text_length", m.text_len_max); + I("pad_token_id", m.pad_id); I("cls_token_id", m.cls_id); I("sep_token_id", m.sep_id); + I("period_token_id", m.period_id); I("question_token_id", m.question_id); + if (g.has("turbovla.rope_theta")) + m.rope_theta = g.f32("turbovla.rope_theta"); + int64_t fhd = m.fusion_dim / m.fusion_heads; + U("fusion_head_dim", fhd); + m.fusion_dim = fhd * m.fusion_heads; + + if (m.image_size % m.patch || m.vit_dim % m.vit_heads || (m.vit_dim / m.vit_heads) % 4 || + m.bert_dim % m.bert_heads || m.hidden % m.enh_heads || m.hidden % m.dec_heads) { + std::fprintf(stderr, "vla(turbovla): inconsistent dimensions in GGUF metadata\n"); + return false; } - } // anonymous namespace - - struct TurboVLA : public ModelArchBase { - TurboVLA() : ModelArchBase(Arch::TURBOVLA) {} - ~TurboVLA() override; - - ggml_backend_t backend = nullptr; - int n_threads = default_cpu_threads(); - ggml_context* ctx_weights = nullptr; - ggml_backend_buffer_t weight_buf = nullptr; - scratch_ctx scratch; - - TurboVLAConfig cfg_; - - VisionEncoderW vision_encoder; - VisionProjectionW vision_proj; - TextEncoderW text_encoder; - std::vector fusion_layers; - std::vector text_layers; - ActionDecoderW action_decoder; - StateProjectionW state_proj; - ggml_tensor* view_emb = nullptr; - - std::vector predict(const Inputs& in) override; - - private: - ggml_tensor* vit_layer(ggml_context* C, const VitBlockW& w, ggml_tensor* x, - ggml_tensor* rope_cos, ggml_tensor* rope_sin, - int64_t internal_seq, int64_t n_heads, int64_t head_dim) const; - ggml_tensor* bert_layer(ggml_context* C, const BertLayerW& w, ggml_tensor* x, - ggml_tensor* self_mask, - int64_t seq, int64_t n_heads, int64_t head_dim) const; - ggml_tensor* vl_fusion_layer(ggml_context* C, const VlFusionLayerW& w, - ggml_tensor* v, ggml_tensor*& l, - ggml_tensor* language_key_mask, - int64_t v_seq, int64_t l_seq, - int64_t n_heads, int64_t head_dim) const; - ggml_tensor* vl_text_layer(ggml_context* C, const VlTextLayerW& w, ggml_tensor* x, - ggml_tensor* self_mask, - int64_t seq, int64_t n_heads, int64_t head_dim) const; - ggml_tensor* act_decoder_layer(ggml_context* C, const ActDecoderLayerW& w, - ggml_tensor* queries, ggml_tensor* memory, - int64_t q_seq, int64_t m_seq, - int64_t n_heads, int64_t head_dim) const; - }; - - TurboVLA::~TurboVLA() { - if (weight_buf) ggml_backend_buffer_free(weight_buf); - if (ctx_weights) ggml_free(ctx_weights); - if (backend) ggml_backend_free(backend); + // The padding table (see text_len). A GGUF without it predates the table and + // would pad every instruction to max_text_length, which is wrong for the + // shorter LIBERO groups. + if (!g.has("turbovla.text_groups.count")) { + std::fprintf(stderr, "vla(turbovla): GGUF has no text padding table; re-convert it with " + "scripts/convert_turbovla_to_gguf.py\n"); + return false; } - - ggml_tensor* TurboVLA::vit_layer(ggml_context* C, const VitBlockW& w, ggml_tensor* x, - ggml_tensor* rope_cos, ggml_tensor* rope_sin, - int64_t internal_seq, int64_t n_heads, int64_t head_dim) const { - const float scale = 1.0f / std::sqrt((float)head_dim); - const int64_t hidden = n_heads * head_dim; - const int64_t patch_start = 1 + cfg_.num_register_tokens; - const int64_t n_patches = internal_seq - patch_start; - GGML_ASSERT(n_patches > 0 && rope_cos->ne[1] == n_patches && rope_sin->ne[1] == n_patches); - - ggml_tensor* normed = layer_norm(C, x, w.ln1_w, w.ln1_b, 1e-5f); - ggml_tensor* q = linear(C, w.attn_q_w, w.attn_q_b, normed); - ggml_tensor* k = linear(C, w.attn_k_w, w.attn_k_b, normed); - ggml_tensor* v = linear(C, w.attn_v_w, w.attn_v_b, normed); - - ggml_tensor* Q = to_heads(C, q, head_dim, n_heads, internal_seq); - ggml_tensor* K = to_heads(C, k, head_dim, n_heads, internal_seq); - ggml_tensor* V = to_heads_v(C, v, head_dim, n_heads, internal_seq); - const size_t q_offset = (size_t) patch_start * Q->nb[1]; - ggml_tensor* Q_prefix = ggml_cont(C, ggml_view_3d(C, Q, head_dim, patch_start, n_heads, - Q->nb[1], Q->nb[2], 0)); - ggml_tensor* K_prefix = ggml_cont(C, ggml_view_3d(C, K, head_dim, patch_start, n_heads, - K->nb[1], K->nb[2], 0)); - ggml_tensor* Q_patch = ggml_cont(C, ggml_view_3d(C, Q, head_dim, n_patches, n_heads, - Q->nb[1], Q->nb[2], q_offset)); - ggml_tensor* K_patch = ggml_cont(C, ggml_view_3d(C, K, head_dim, n_patches, n_heads, - K->nb[1], K->nb[2], (size_t) patch_start*K->nb[1])); - Q_patch = rope_dinov3_axial(C, Q_patch, rope_cos, rope_sin); - K_patch = rope_dinov3_axial(C, K_patch, rope_cos, rope_sin); - Q = ggml_concat(C, Q_prefix, Q_patch, 1); - K = ggml_concat(C, K_prefix, K_patch, 1); - ggml_tensor* att = attention(C, Q, K, V, nullptr, scale, hidden, internal_seq); - ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); - ggml_tensor* h = ggml_add(C, x, o); - ggml_tensor* mlp = ffn_gelu_erf(C, w.fc1_w, w.fc1_b, w.fc2_w, w.fc2_b, - layer_norm(C, h, w.ln2_w, w.ln2_b, 1e-5f)); - return ggml_add(C, h, mlp); + const size_t n = g.u32("turbovla.text_groups.count"); + if (n == 0) + return true; + const int64_t kt = gguf_find_key(g.gctx, "turbovla.text_groups.tokens"); + const int64_t kl = gguf_find_key(g.gctx, "turbovla.text_groups.lengths"); + const int64_t kp = gguf_find_key(g.gctx, "turbovla.text_groups.pad_to"); + if (kt < 0 || kl < 0 || kp < 0) { + std::fprintf(stderr, "vla(turbovla): malformed turbovla.text_groups metadata\n"); + return false; } - - ggml_tensor* TurboVLA::bert_layer(ggml_context* C, const BertLayerW& w, ggml_tensor* x, - ggml_tensor* self_mask, - int64_t seq, int64_t n_heads, int64_t head_dim) const { - const float scale = 1.0f / std::sqrt((float)head_dim); - const int64_t hidden = n_heads * head_dim; - - // BERT uses post-norm blocks: attention on x, residual + LayerNorm, - // then FFN, residual + LayerNorm. - ggml_tensor* q = linear(C, w.attn_q_w, w.attn_q_b, x); - ggml_tensor* k = linear(C, w.attn_k_w, w.attn_k_b, x); - ggml_tensor* v = linear(C, w.attn_v_w, w.attn_v_b, x); - - ggml_tensor* Q = to_heads(C, q, head_dim, n_heads, seq); - ggml_tensor* K = to_heads(C, k, head_dim, n_heads, seq); - ggml_tensor* V = to_heads_v(C, v, head_dim, n_heads, seq); - - ggml_tensor* att = attention(C, Q, K, V, self_mask, scale, hidden, seq); - ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); - ggml_tensor* h = layer_norm(C, - ggml_add(C, x, o), - w.ln1_w, - w.ln1_b, - 1e-12f); - - ggml_tensor* ffn_out = ffn_gelu_erf(C, - w.fc1_w, - w.fc1_b, - w.fc2_w, - w.fc2_b, - h); - return layer_norm(C, - ggml_add(C, h, ffn_out), - w.ln2_w, - w.ln2_b, - 1e-12f); + if (gguf_get_arr_type(g.gctx, kt) != GGUF_TYPE_INT32 || gguf_get_arr_type(g.gctx, kl) != GGUF_TYPE_INT32 || + gguf_get_arr_type(g.gctx, kp) != GGUF_TYPE_INT32 || gguf_get_arr_n(g.gctx, kl) != n || + gguf_get_arr_n(g.gctx, kp) != n) { + std::fprintf(stderr, "vla(turbovla): malformed turbovla.text_groups metadata\n"); + return false; } - - ggml_tensor* TurboVLA::vl_fusion_layer(ggml_context* C, const VlFusionLayerW& w, - ggml_tensor* v, ggml_tensor*& l, - ggml_tensor* language_key_mask, - int64_t v_seq, int64_t l_seq, - int64_t n_heads, int64_t head_dim) const { - const float scale = 1.0f / std::sqrt((float)head_dim); - - const int64_t att_dim = n_heads * head_dim; - - // TurboVLA's normalized residual style adds the deltas to normalized - // tokens, not the pre-normalization residuals. See - // H-EmbodVis/TurboVLA@b29ab142, models/components/fusion.py:336-347. - ggml_tensor* v_norm = layer_norm(C, v, w.norm_v_w, w.norm_v_b, 1e-5f); - ggml_tensor* l_norm = layer_norm(C, l, w.norm_l_w, w.norm_l_b, 1e-5f); - - // Only vision-to-language attention masks language keys. See - // H-EmbodVis/TurboVLA@b29ab142, models/components/fusion.py:246-267. - ggml_tensor* q_v = linear(C, w.v_proj_w, w.v_proj_b, v_norm); - ggml_tensor* k_l = linear(C, w.l_proj_w, w.l_proj_b, l_norm); - ggml_tensor* v_l = linear(C, w.values_l_w, w.values_l_b, l_norm); - - ggml_tensor* Q = to_heads(C, q_v, head_dim, n_heads, v_seq); - ggml_tensor* K = to_heads(C, k_l, head_dim, n_heads, l_seq); - ggml_tensor* V = to_heads_v(C, v_l, head_dim, n_heads, l_seq); - ggml_tensor* o = attention(C, Q, K, V, language_key_mask, scale, att_dim, v_seq); - ggml_tensor* o_proj = linear(C, w.out_v_w, w.out_v_b, o); - - ggml_tensor* gamma_v_scaled = ggml_repeat(C, w.gamma_v, o_proj); - ggml_tensor* delta_v = ggml_mul(C, gamma_v_scaled, o_proj); - - ggml_tensor* v_out = ggml_add(C, v_norm, delta_v); - - // L attends to V. There is no visual padding, hence no mask on this branch. - ggml_tensor* q_l = linear(C, w.l_proj_w, w.l_proj_b, l_norm); - ggml_tensor* k_v = linear(C, w.v_proj_w, w.v_proj_b, v_norm); - ggml_tensor* v_v = linear(C, w.values_v_w, w.values_v_b, v_norm); - - ggml_tensor* Q_l = to_heads(C, q_l, head_dim, n_heads, l_seq); - ggml_tensor* K_v = to_heads(C, k_v, head_dim, n_heads, v_seq); - ggml_tensor* V_v = to_heads_v(C, v_v, head_dim, n_heads, v_seq); - ggml_tensor* o_l = attention(C, Q_l, K_v, V_v, nullptr, scale, att_dim, l_seq); - ggml_tensor* o_proj_l = linear(C, w.out_l_w, w.out_l_b, o_l); - - ggml_tensor* gamma_l_scaled = ggml_repeat(C, w.gamma_l, o_proj_l); - ggml_tensor* delta_l = ggml_mul(C, gamma_l_scaled, o_proj_l); - - l = ggml_add(C, l_norm, delta_l); - - return v_out; + const int32_t * tok = (const int32_t *) gguf_get_arr_data(g.gctx, kt); + const int32_t * len = (const int32_t *) gguf_get_arr_data(g.gctx, kl); + const int32_t * pad = (const int32_t *) gguf_get_arr_data(g.gctx, kp); + const size_t n_tok = gguf_get_arr_n(g.gctx, kt); + size_t off = 0; + for (size_t i = 0; i < n; ++i) { + if (len[i] <= 0 || off + (size_t) len[i] > n_tok || pad[i] < len[i] || pad[i] > m.text_len_max) { + std::fprintf(stderr, "vla(turbovla): malformed turbovla.text_groups entry %zu\n", i); + return false; + } + m.text_groups.push_back({ std::vector(tok + off, tok + off + len[i]), pad[i] }); + off += (size_t) len[i]; } + return true; +} + +bool load_weights(TurboVlaModelArch & m, gguf_reader & g) { + ggml_init_params wp = { (size_t) 16*1024*1024, nullptr, true }; + m.ctx_weights = ggml_init(wp); + if (!m.ctx_weights) + return false; + WeightLoader L("turbovla", g, m.ctx_weights, m.mt); + char b[128]; + auto N = [&](const char * fmt, int64_t i, const char * s) { + std::snprintf(b, sizeof(b), fmt, (long long) i, s); + return std::string(b); + }; - ggml_tensor* TurboVLA::vl_text_layer(ggml_context* C, const VlTextLayerW& w, ggml_tensor* x, - ggml_tensor* self_mask, - int64_t seq, int64_t n_heads, int64_t head_dim) const { - // TurboVLA's text enhancer uses post-norm residuals and ReLU. - const float scale = 1.0f / std::sqrt((float)head_dim); - const int64_t att_dim = n_heads * head_dim; - const int64_t n = n_heads; - - ggml_tensor* qkv = linear(C, w.attn_qkv_w, w.attn_qkv_b, x); - - // Most checkpoints store fused [Q; K; V]. Keep compatibility with a - // checkpoint that stores a single shared projection as well. - ggml_tensor* q = nullptr; - ggml_tensor* k = nullptr; - ggml_tensor* v = nullptr; - if (qkv->ne[0] == 3 * att_dim) { - const size_t es = ggml_element_size(qkv); - q = ggml_cont(C, ggml_view_2d(C, qkv, att_dim, seq, qkv->nb[1], 0)); - k = ggml_cont(C, ggml_view_2d(C, qkv, att_dim, seq, qkv->nb[1], (size_t)att_dim * es)); - v = ggml_cont(C, ggml_view_2d(C, qkv, att_dim, seq, qkv->nb[1], (size_t)2 * att_dim * es)); - } else { - GGML_ASSERT(qkv->ne[0] == att_dim); - q = qkv; - k = qkv; - v = qkv; - } - - check_heads("VLText.Q", q, head_dim, n, seq); - check_heads("VLText.K", k, head_dim, n, seq); - check_heads("VLText.V", v, head_dim, n, seq); - ggml_tensor* Q = to_heads(C, q, head_dim, n, seq); - ggml_tensor* K = to_heads(C, k, head_dim, n, seq); - ggml_tensor* V = to_heads_v(C, v, head_dim, n, seq); - - ggml_tensor* att = attention(C, Q, K, V, self_mask, scale, att_dim, seq); - ggml_tensor* o = linear(C, w.attn_o_w, w.attn_o_b, att); - ggml_tensor* h = ggml_add(C, x, o); - - h = layer_norm(C, h, w.ln1_w, w.ln1_b, 1e-5f); - - ggml_tensor* ff = linear(C, w.fc1_w, w.fc1_b, h); - ff = ggml_relu(C, ff); - ggml_tensor* ff_out = linear(C, w.fc2_w, w.fc2_b, ff); - - ggml_tensor* out = ggml_add(C, h, ff_out); - return layer_norm(C, out, w.ln2_w, w.ln2_b, 1e-5f); + m.cls_tok = L.f32("vit.cls_token"); + if (m.n_reg > 0) + m.reg_tok = L.f32("vit.register_tokens"); + m.patch_w = L.gemm("vit.patch_embed.weight"); + m.patch_b = L.f32("vit.patch_embed.bias"); + m.vit.resize((size_t) m.vit_layers); + for (int64_t i = 0; i < m.vit_layers; ++i) { + VitLayerW & w = m.vit[(size_t) i]; + const char * f = "vit.blk.%lld.%s"; + w.ln1_w = L.f32("%s", N(f, i, "ln1.weight").c_str()); + w.ln1_b = L.f32("%s", N(f, i, "ln1.bias").c_str()); + w.qkv_w = L.fuse_gemm(N(f, i, "attn_qkv.weight").c_str(), + { N(f, i, "attn_q.weight"), N(f, i, "attn_k.weight"), N(f, i, "attn_v.weight") }); + w.qkv_b = L.fuse_f32(N(f, i, "attn_qkv.bias").c_str(), + { N(f, i, "attn_q.bias"), N(f, i, "attn_k.bias"), N(f, i, "attn_v.bias") }); + w.o_w = L.gemm("%s", N(f, i, "attn_o.weight").c_str()); + w.o_b = L.f32("%s", N(f, i, "attn_o.bias").c_str()); + w.ln2_w = L.f32("%s", N(f, i, "ln2.weight").c_str()); + w.ln2_b = L.f32("%s", N(f, i, "ln2.bias").c_str()); + w.fc1_w = L.gemm("%s", N(f, i, "fc1.weight").c_str()); + w.fc1_b = L.f32("%s", N(f, i, "fc1.bias").c_str()); + w.fc2_w = L.gemm("%s", N(f, i, "fc2.weight").c_str()); + w.fc2_b = L.f32("%s", N(f, i, "fc2.bias").c_str()); } - ggml_tensor* TurboVLA::act_decoder_layer(ggml_context* C, const ActDecoderLayerW& w, - ggml_tensor* queries, ggml_tensor* memory, - int64_t q_seq, int64_t m_seq, - int64_t n_heads, int64_t head_dim) const { - const float scale = 1.0f / std::sqrt((float)head_dim); - const int64_t hidden = n_heads * head_dim; - - ggml_tensor* normed = layer_norm(C, queries, w.ln1_w, w.ln1_b, 1e-5f); - ggml_tensor* self_qkv = linear(C, w.self_qkv_w, w.self_qkv_b, normed); - - ggml_tensor* self_q = nullptr; - ggml_tensor* self_k = nullptr; - ggml_tensor* self_v = nullptr; - if (self_qkv->ne[0] == 3 * hidden) { - const size_t self_es = ggml_element_size(self_qkv); - self_q = ggml_cont(C, ggml_view_2d(C, self_qkv, hidden, q_seq, self_qkv->nb[1], 0)); - self_k = ggml_cont(C, ggml_view_2d(C, self_qkv, hidden, q_seq, self_qkv->nb[1], (size_t)hidden * self_es)); - self_v = ggml_cont(C, ggml_view_2d(C, self_qkv, hidden, q_seq, self_qkv->nb[1], (size_t)2 * hidden * self_es)); - } else { - GGML_ASSERT(self_qkv->ne[0] == hidden); - self_q = self_qkv; - self_k = self_qkv; - self_v = self_qkv; - } - - check_heads("ACT.self.Q", self_q, head_dim, n_heads, q_seq); - check_heads("ACT.self.K", self_k, head_dim, n_heads, q_seq); - check_heads("ACT.self.V", self_v, head_dim, n_heads, q_seq); - ggml_tensor* Q = to_heads(C, self_q, head_dim, n_heads, q_seq); - ggml_tensor* K = to_heads(C, self_k, head_dim, n_heads, q_seq); - ggml_tensor* V = to_heads_v(C, self_v, head_dim, n_heads, q_seq); - - ggml_tensor* att = attention(C, Q, K, V, nullptr, scale, hidden, q_seq); - ggml_tensor* o = linear(C, w.self_out_w, w.self_out_b, att); - ggml_tensor* h = ggml_add(C, queries, o); - - // Cross-attention. If cross_qkv is fused [Q;K;V], project both query and - // memory and select their respective blocks. If the checkpoint stores only - // a single query projection, preserve the original direct-memory K/V path. - normed = layer_norm(C, h, w.ln2_w, w.ln2_b, 1e-5f); - ggml_tensor* cross_qkv_q = linear(C, w.cross_qkv_w, w.cross_qkv_b, normed); - - ggml_tensor* cross_q = nullptr; - ggml_tensor* cross_k = nullptr; - ggml_tensor* cross_v = nullptr; - if (cross_qkv_q->ne[0] == 3 * hidden) { - ggml_tensor* cross_qkv_m = linear(C, w.cross_qkv_w, w.cross_qkv_b, memory); - GGML_ASSERT(cross_qkv_m->ne[0] == 3 * hidden); - const size_t cross_m_es = ggml_element_size(cross_qkv_m); - cross_q = ggml_cont(C, ggml_view_2d(C, cross_qkv_q, hidden, q_seq, cross_qkv_q->nb[1], 0)); - cross_k = ggml_cont(C, ggml_view_2d(C, cross_qkv_m, hidden, m_seq, cross_qkv_m->nb[1], (size_t)hidden * cross_m_es)); - cross_v = ggml_cont(C, ggml_view_2d(C, cross_qkv_m, hidden, m_seq, cross_qkv_m->nb[1], (size_t)2 * hidden * cross_m_es)); - } else { - GGML_ASSERT(cross_qkv_q->ne[0] == hidden); - cross_q = cross_qkv_q; - cross_k = ggml_is_contiguous(memory) ? memory : ggml_cont(C, memory); - cross_v = cross_k; - } - check_heads("ACT.cross.Q", cross_q, head_dim, n_heads, q_seq); - check_heads("ACT.cross.K", cross_k, head_dim, n_heads, m_seq); - check_heads("ACT.cross.V", cross_v, head_dim, n_heads, m_seq); - ggml_tensor* Q_c = to_heads(C, cross_q, head_dim, n_heads, q_seq); - ggml_tensor* K_c = to_heads(C, cross_k, head_dim, n_heads, m_seq); - ggml_tensor* V_c = to_heads_v(C, cross_v, head_dim, n_heads, m_seq); - - ggml_tensor* att_c = attention(C, Q_c, K_c, V_c, nullptr, scale, hidden, q_seq); - ggml_tensor* o_c = linear(C, w.cross_out_w, w.cross_out_b, att_c); - h = ggml_add(C, h, o_c); - - // The upstream TransformerDecoderLayer keeps PyTorch's default ReLU - // activation. See H-EmbodVis/TurboVLA@b29ab142, - // models/action_head.py:39-47. - normed = layer_norm(C, h, w.ln3_w, w.ln3_b, 1e-5f); - ggml_tensor* ffn_out = ffn_relu(C, w.fc1_w, w.fc1_b, w.fc2_w, w.fc2_b, normed); - - return ggml_add(C, h, ffn_out); + m.vp_in_w = L.f32("vit_proj.input_norm.weight"); + m.vp_in_b = L.f32("vit_proj.input_norm.bias"); + m.vp_fc1_w = L.gemm("vit_proj.mlp.0.weight"); + m.vp_fc1_b = L.f32("vit_proj.mlp.0.bias"); + m.vp_fc2_w = L.gemm("vit_proj.mlp.3.weight"); + m.vp_fc2_b = L.f32("vit_proj.mlp.3.bias"); + m.vp_skip_w = L.gemm("vit_proj.skip.weight"); + m.vp_out_w = L.f32("vit_proj.output_norm.weight"); + m.vp_out_b = L.f32("vit_proj.output_norm.bias"); + m.view_emb = L.f32("view_emb"); + + m.word_emb = L.f32("text.embed.word_embeddings"); + m.pos_emb = L.f32("text.embed.position_embeddings"); + m.type_emb = L.f32("text.embed.token_type_embeddings"); + m.emb_ln_w = L.f32("text.embed.LayerNorm.weight"); + m.emb_ln_b = L.f32("text.embed.LayerNorm.bias"); + m.bert.resize((size_t) m.bert_layers); + for (int64_t i = 0; i < m.bert_layers; ++i) { + BertLayerW & w = m.bert[(size_t) i]; + const char * f = "text.encoder.layer.%lld.%s"; + w.qkv_w = L.fuse_gemm(N(f, i, "attention.self.qkv.weight").c_str(), + { N(f, i, "attention.self.query.weight"), N(f, i, "attention.self.key.weight"), + N(f, i, "attention.self.value.weight") }); + w.qkv_b = L.fuse_f32(N(f, i, "attention.self.qkv.bias").c_str(), + { N(f, i, "attention.self.query.bias"), N(f, i, "attention.self.key.bias"), + N(f, i, "attention.self.value.bias") }); + w.o_w = L.gemm("%s", N(f, i, "attention.output.dense.weight").c_str()); + w.o_b = L.f32("%s", N(f, i, "attention.output.dense.bias").c_str()); + w.ln1_w = L.f32("%s", N(f, i, "attention.output.LayerNorm.weight").c_str()); + w.ln1_b = L.f32("%s", N(f, i, "attention.output.LayerNorm.bias").c_str()); + w.fc1_w = L.gemm("%s", N(f, i, "intermediate.dense.weight").c_str()); + w.fc1_b = L.f32("%s", N(f, i, "intermediate.dense.bias").c_str()); + w.fc2_w = L.gemm("%s", N(f, i, "output.dense.weight").c_str()); + w.fc2_b = L.f32("%s", N(f, i, "output.dense.bias").c_str()); + w.ln2_w = L.f32("%s", N(f, i, "output.LayerNorm.weight").c_str()); + w.ln2_b = L.f32("%s", N(f, i, "output.LayerNorm.bias").c_str()); } - - namespace { - - static const float IMAGENET_MEAN[3] = {0.485f, 0.456f, 0.406f}; - static const float IMAGENET_STD[3] = {0.229f, 0.224f, 0.225f}; - - // Build the flattened patch matrix expected by a patch-embedding weight - // flattened from [out, C, patch_h, patch_w] to [out, C*patch_h*patch_w]. - // Output is GGML matrix layout [patch_dim, n_patches], i.e. each column is one - // normalized Cxpatchxpatch image patch. - void normalize_imagenet_patches(const ImageView& v, - int64_t side, - int64_t patch_size, - std::vector& out) { - const int64_t grid = side / patch_size; - const int64_t n_patches = grid * grid; - const int64_t patch_dim = 3 * patch_size * patch_size; - out.assign((size_t)patch_dim * n_patches, 0.0f); - - for (int64_t py = 0; py < grid; ++py) { - for (int64_t px_i = 0; px_i < grid; ++px_i) { - const int64_t patch_idx = py * grid + px_i; - for (int64_t c = 0; c < 3; ++c) { - for (int64_t ky = 0; ky < patch_size; ++ky) { - for (int64_t kx = 0; kx < patch_size; ++kx) { - const int64_t h = py * patch_size + ky; - const int64_t w = px_i * patch_size + kx; - float pixel; - if (v.format == PixelFormat::U8) { - pixel = ((const uint8_t*)v.data)[(h * side + w) * 3 + c] / 255.0f; - } else { - pixel = ((const float*)v.data)[(h * side + w) * 3 + c]; - } - const float norm = (pixel - IMAGENET_MEAN[c]) / IMAGENET_STD[c]; - const int64_t feature_idx = - (c * patch_size + ky) * patch_size + kx; - out[(size_t)feature_idx + (size_t)patch_idx * patch_dim] = norm; - } - } - } - } - } + m.text_proj_w = L.gemm("text_proj.weight"); + m.text_proj_b = L.f32("text_proj.bias"); + + m.fusion.resize((size_t) m.fusion_layers); + m.enhancer.resize((size_t) m.fusion_layers); + for (int64_t i = 0; i < m.fusion_layers; ++i) { + FusionLayerW & w = m.fusion[(size_t) i]; + const char * f = "vl_fusion.%lld.%s"; + w.norm_v_w = L.f32("%s", N(f, i, "norm_v.weight").c_str()); + w.norm_v_b = L.f32("%s", N(f, i, "norm_v.bias").c_str()); + w.norm_l_w = L.f32("%s", N(f, i, "norm_l.weight").c_str()); + w.norm_l_b = L.f32("%s", N(f, i, "norm_l.bias").c_str()); + w.v_w = L.fuse_gemm(N(f, i, "v.weight").c_str(), { N(f, i, "v_proj.weight"), N(f, i, "values_v.weight") }); + w.v_b = L.fuse_f32(N(f, i, "v.bias").c_str(), { N(f, i, "v_proj.bias"), N(f, i, "values_v.bias") }); + w.l_w = L.fuse_gemm(N(f, i, "l.weight").c_str(), { N(f, i, "l_proj.weight"), N(f, i, "values_l.weight") }); + w.l_b = L.fuse_f32(N(f, i, "l.bias").c_str(), { N(f, i, "l_proj.bias"), N(f, i, "values_l.bias") }); + w.out_v_w = L.gemm("%s", N(f, i, "out_v.weight").c_str()); + w.out_v_b = L.f32("%s", N(f, i, "out_v.bias").c_str()); + w.out_l_w = L.gemm("%s", N(f, i, "out_l.weight").c_str()); + w.out_l_b = L.f32("%s", N(f, i, "out_l.bias").c_str()); + w.gamma_v = L.f32("%s", N(f, i, "gamma_v").c_str()); + w.gamma_l = L.f32("%s", N(f, i, "gamma_l").c_str()); + + EnhancerLayerW & e = m.enhancer[(size_t) i]; + const char * t = "vl_text.%lld.%s"; + e.qkv_w = L.gemm("%s", N(t, i, "attn_qkv.weight").c_str()); + e.qkv_b = L.f32("%s", N(t, i, "attn_qkv.bias").c_str()); + e.o_w = L.gemm("%s", N(t, i, "attn_o.weight").c_str()); + e.o_b = L.f32("%s", N(t, i, "attn_o.bias").c_str()); + e.ln1_w = L.f32("%s", N(t, i, "ln1.weight").c_str()); + e.ln1_b = L.f32("%s", N(t, i, "ln1.bias").c_str()); + e.fc1_w = L.gemm("%s", N(t, i, "fc1.weight").c_str()); + e.fc1_b = L.f32("%s", N(t, i, "fc1.bias").c_str()); + e.fc2_w = L.gemm("%s", N(t, i, "fc2.weight").c_str()); + e.fc2_b = L.f32("%s", N(t, i, "fc2.bias").c_str()); + e.ln2_w = L.f32("%s", N(t, i, "ln2.weight").c_str()); + e.ln2_b = L.f32("%s", N(t, i, "ln2.bias").c_str()); } - void build_dinov3_rope(int64_t grid, int64_t head_dim, float base, - std::vector& cos, std::vector& sin) { - const int64_t n_patches = grid * grid; - const int64_t axis_dim = head_dim / 4; - cos.resize((size_t)head_dim * n_patches); - sin.resize((size_t)head_dim * n_patches); - for (int64_t row = 0; row < grid; ++row) { - const float h = 2.0f * ((float(row) + 0.5f) / float(grid)) - 1.0f; - for (int64_t col = 0; col < grid; ++col) { - const float w = 2.0f * ((float(col) + 0.5f) / float(grid)) - 1.0f; - const size_t offset = (size_t)(row * grid + col) * head_dim; - for (int64_t i = 0; i < axis_dim; ++i) { - const float period = std::pow(base, 4.0f * float(i) / float(head_dim)); - constexpr float two_pi = 6.28318530717958647692f; - const float ah = two_pi * h / period; - const float aw = two_pi * w / period; - const int64_t dims[4] = {i, axis_dim + i, 2 * axis_dim + i, 3 * axis_dim + i}; - const float angles[4] = {ah, aw, ah, aw}; - for (int j = 0; j < 4; ++j) { - cos[offset + dims[j]] = std::cos(angles[j]); - sin[offset + dims[j]] = std::sin(angles[j]); - } - } - } - } + m.st_ln_w = L.f32("state.proj.0.weight"); + m.st_ln_b = L.f32("state.proj.0.bias"); + m.st_fc1_w = L.gemm("state.proj.1.weight"); + m.st_fc1_b = L.f32("state.proj.1.bias"); + m.st_fc2_w = L.gemm("state.proj.4.weight"); + m.st_fc2_b = L.f32("state.proj.4.bias"); + m.st_pos = L.f32("state.proj.position"); + m.st_out_w = L.f32("state.proj.output_norm.weight"); + m.st_out_b = L.f32("state.proj.output_norm.bias"); + + m.act_q = L.f32("act.q.weight"); + m.dec.resize((size_t) m.dec_layers); + for (int64_t i = 0; i < m.dec_layers; ++i) { + DecoderLayerW & w = m.dec[(size_t) i]; + const char * f = "act.dec.%lld.%s"; + w.ln1_w = L.f32("%s", N(f, i, "ln1.weight").c_str()); + w.ln1_b = L.f32("%s", N(f, i, "ln1.bias").c_str()); + w.self_qkv_w = L.gemm("%s", N(f, i, "self_qkv.weight").c_str()); + w.self_qkv_b = L.f32("%s", N(f, i, "self_qkv.bias").c_str()); + w.self_o_w = L.gemm("%s", N(f, i, "self_out.weight").c_str()); + w.self_o_b = L.f32("%s", N(f, i, "self_out.bias").c_str()); + w.ln2_w = L.f32("%s", N(f, i, "ln2.weight").c_str()); + w.ln2_b = L.f32("%s", N(f, i, "ln2.bias").c_str()); + w.cross_qkv_w = L.gemm("%s", N(f, i, "cross_qkv.weight").c_str()); + w.cross_qkv_b = L.f32("%s", N(f, i, "cross_qkv.bias").c_str()); + w.cross_o_w = L.gemm("%s", N(f, i, "cross_out.weight").c_str()); + w.cross_o_b = L.f32("%s", N(f, i, "cross_out.bias").c_str()); + w.ln3_w = L.f32("%s", N(f, i, "ln3.weight").c_str()); + w.ln3_b = L.f32("%s", N(f, i, "ln3.bias").c_str()); + w.fc1_w = L.gemm("%s", N(f, i, "fc1.weight").c_str()); + w.fc1_b = L.f32("%s", N(f, i, "fc1.bias").c_str()); + w.fc2_w = L.gemm("%s", N(f, i, "fc2.weight").c_str()); + w.fc2_b = L.f32("%s", N(f, i, "fc2.bias").c_str()); } - - struct TurboVLATextInputs { - std::vector token_ids; - std::vector position_ids; - std::vector bert_self_mask; - std::vector enhancer_self_mask; - std::vector fusion_language_key_mask; - }; - - TurboVLATextInputs build_turbovla_text_inputs(const Inputs& in, - const TurboVLAConfig& cfg, - int64_t vision_seq) { - const int64_t seq = cfg.max_text_length; - TurboVLATextInputs result; - result.token_ids.assign((size_t)seq, (int32_t)cfg.pad_token_id); - result.position_ids.assign((size_t)seq, 0); - - if (in.n_lang > seq) { - throw std::runtime_error("TurboVLA language input exceeds the checkpoint padding length"); - } - if (in.lang_tokens && in.n_lang > 0) { - std::memcpy(result.token_ids.data(), in.lang_tokens, - (size_t)in.n_lang * sizeof(int32_t)); - } - - std::vector valid((size_t)seq, 0); - for (int64_t i = 0; i < seq; ++i) { - const bool from_mask = in.attention_mask && - ((in.attention_mask_n == seq) || (i < in.attention_mask_n)); - valid[(size_t)i] = from_mask ? in.attention_mask[i] != 0 - : result.token_ids[(size_t)i] != cfg.pad_token_id; - } - - // Port of generate_masks_with_special_tokens from TurboVLA's BERT wrapper. - // Its block masks separate sub-sentences at [CLS], [SEP], '.' and '?'. - std::vector allowed((size_t)seq * seq, 0); - for (int64_t i = 0; i < seq; ++i) { - allowed[(size_t)i * seq + i] = 1; - } - int64_t previous_special = 0; - for (int64_t col = 0; col < seq; ++col) { - const int32_t id = result.token_ids[(size_t)col]; - const bool special = id == cfg.cls_token_id || id == cfg.sep_token_id || - id == cfg.period_token_id || id == cfg.question_token_id; - if (!special) { - continue; - } - if (col == 0 || col == seq - 1) { - result.position_ids[(size_t)col] = 0; - } else { - for (int64_t q = previous_special + 1; q <= col; ++q) { - result.position_ids[(size_t)q] = (int32_t)(q - previous_special - 1); - for (int64_t k = previous_special + 1; k <= col; ++k) { - allowed[(size_t)q * seq + k] = 1; - } + m.ap0_w = L.gemm("act.proj.0.weight"); + m.ap0_b = L.f32("act.proj.0.bias"); + m.ap1_w = L.gemm("act.proj.1.weight"); + m.ap1_b = L.f32("act.proj.1.bias"); + m.ap2_w = L.gemm("act.proj.2.weight"); + m.ap2_b = L.f32("act.proj.2.bias"); + + if (!L.upload(m.backend, &m.weight_buf)) + return false; + + m.bert_max_pos = m.pos_emb->ne[1]; + if (m.vocab != m.word_emb->ne[1] || m.text_len_max > m.bert_max_pos || + m.patch_w->ne[0] != 3*m.patch*m.patch || (m.reg_tok && m.reg_tok->ne[1] != m.n_reg) || + ggml_nelements(m.view_emb) != m.hidden*m.n_views || m.act_q->ne[1] != m.horizon || + m.dec[0].cross_qkv_w->ne[1] != 3*m.hidden) { + std::fprintf(stderr, "vla(turbovla): tensor shapes disagree with GGUF metadata\n"); + return false; + } + return true; +} + +// Flattens one view into DINOv3's patch-embedding input: column p is patch p +// (row-major over the grid), rows follow the conv weight's (c, ky, kx) order. +void patchify(const ImageView & v, int64_t side, int64_t P, float * out) { + const int64_t g = side / P, pdim = 3*P*P; + for (int64_t p = 0; p < g*g; ++p) { + const int64_t y0 = (p / g)*P, x0 = (p % g)*P; + float * col = out + p*pdim; + for (int64_t c = 0; c < 3; ++c) + for (int64_t ky = 0; ky < P; ++ky) + for (int64_t kx = 0; kx < P; ++kx) { + const size_t idx = (size_t) ((y0 + ky)*side + x0 + kx)*3 + c; + const float px = v.format == PixelFormat::U8 ? ((const uint8_t *) v.data)[idx] / 255.0f + : ((const float *) v.data)[idx]; + col[(c*P + ky)*P + kx] = (px - kImagenetMean[c]) / kImagenetStd[c]; } - } - previous_special = col; - } - - result.bert_self_mask.resize((size_t)seq * seq); - result.enhancer_self_mask.resize((size_t)seq * seq); - for (int64_t q = 0; q < seq; ++q) { - for (int64_t k = 0; k < seq; ++k) { - const bool is_allowed = allowed[(size_t)q * seq + k] != 0; - result.bert_self_mask[(size_t)q * seq + k] = is_allowed ? 0.0f : -FLT_MAX; - // The BERT wrapper creates block masks from special tokens; the - // enhancer additionally masks padded keys. See - // H-EmbodVis/TurboVLA@b29ab142, text/bert.py:177-212 and - // models/text_encoder.py:80-94. - result.enhancer_self_mask[(size_t)q * seq + k] = - is_allowed && valid[(size_t)k] ? 0.0f : -FLT_MAX; - } - } - - result.fusion_language_key_mask.resize((size_t)vision_seq * seq); - for (int64_t q = 0; q < vision_seq; ++q) { - for (int64_t k = 0; k < seq; ++k) { - result.fusion_language_key_mask[(size_t)q * seq + k] = - valid[(size_t)k] ? 0.0f : -FLT_MAX; - } - } - return result; } - - } // anonymous namespace - - std::vector TurboVLA::predict(const Inputs& in) { - const auto& cfg = cfg_; - - const int64_t patch_size = cfg.patch_size; - const int64_t img_size = cfg.image_size; - const int64_t n_patches_per_side = img_size / patch_size; - const int64_t n_patches = n_patches_per_side * n_patches_per_side; - const int64_t n_reg = cfg.num_register_tokens; - const int64_t n_cls = 1; - const int64_t internal_vit_seq = n_cls + n_reg + n_patches; - const int64_t n_views = cfg.num_views; - const int64_t vision_seq = n_views * n_patches; - const int64_t hidden = cfg.hidden; - const int64_t vit_dim = cfg.vit_dim; - - // Basic shape sanity checks. These fail early with a useful message instead - // of failing later inside ggml_reshape_*. - if (patch_size <= 0 || img_size <= 0 || img_size % patch_size != 0) { - fprintf(stderr, "ERROR: invalid image/patch configuration\n"); - return {}; - } - if (cfg.vit_heads * cfg.vit_head_dim != vit_dim) { - fprintf(stderr, - "ERROR: vit_heads * vit_head_dim = %lld, expected vit_dim = %lld\n", - (long long)(cfg.vit_heads * cfg.vit_head_dim), - (long long)vit_dim); - return {}; - } - if (n_views <= 0 || !in.images || in.n_images < n_views) { - fprintf(stderr, - "ERROR: model expects %lld image view(s), but input has %lld\n", - (long long)n_views, - (long long)in.n_images); - return {}; - } - - // Project only patch tokens; DINOv3's CLS and register tokens stay internal. - const int64_t vision_out_size = n_patches * vit_dim; - std::vector vision_outputs((size_t)n_views * vision_out_size, 0.0f); - - for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { - ggml_context* VC = scratch.reset(256 * 1024 * 1024); // metadata arena - if (!VC) { - fprintf(stderr, "ERROR: vision scratch reset failed\n"); - return {}; - } - - // Patch-embedding input [patch_dim, n_patches]. Patch extraction is - // performed on the host so that each column is exactly one spatial - // Cxpatchxpatch block; a plain reshape of CHW image memory does not - // create spatial patches correctly. - const int64_t patch_dim = 3 * patch_size * patch_size; - ggml_tensor* patch_input = - ggml_new_tensor_2d(VC, GGML_TYPE_F32, patch_dim, n_patches); - ggml_set_input(patch_input); - - // patch_embed_w is a flattened Conv2D patch projection. GGML linear - // layout is [in=patch_dim, out=vit_dim], so mul_mat produces - // [vit_dim, n_patches]. - if (vision_encoder.patch_embed_w->ne[0] != patch_dim) { - fprintf(stderr, - "ERROR: patch_embed_w input dim=%lld, expected patch_dim=%lld\n", - (long long)vision_encoder.patch_embed_w->ne[0], - (long long)patch_dim); - return {}; - } - - ggml_tensor* patches_emb = - ggml_mul_mat(VC, vision_encoder.patch_embed_w, patch_input); - ggml_tensor* patches_2d = ggml_cont(VC, patches_emb); - - if (patches_2d->ne[0] != vit_dim || patches_2d->ne[1] != n_patches) { - fprintf(stderr, - "ERROR: patch embedding shape is [%lld,%lld], expected [%lld,%lld]\n", - (long long)patches_2d->ne[0], - (long long)patches_2d->ne[1], - (long long)vit_dim, - (long long)n_patches); - return {}; - } - - ggml_tensor* patches_with_bias = - ggml_add(VC, patches_2d, vision_encoder.patch_embed_b); - - ggml_tensor* rope_cos = ggml_new_tensor_2d(VC, GGML_TYPE_F32, cfg.vit_head_dim, n_patches); - ggml_tensor* rope_sin = ggml_new_tensor_2d(VC, GGML_TYPE_F32, cfg.vit_head_dim, n_patches); - ggml_set_input(rope_cos); - ggml_set_input(rope_sin); - std::vector rope_cos_host, rope_sin_host; - build_dinov3_rope(n_patches_per_side, cfg.vit_head_dim, cfg.rope_theta, - rope_cos_host, rope_sin_host); - - // Assemble DINOv3's internal sequence: CLS, register tokens, then patches. - ggml_tensor* cls_2d = ggml_reshape_2d(VC, vision_encoder.cls_token, vit_dim, 1); - - ggml_tensor* vit_out = cls_2d; - if (n_reg > 0 && vision_encoder.register_tokens) { - ggml_tensor* reg_2d = ggml_reshape_2d(VC, vision_encoder.register_tokens, vit_dim, n_reg); - vit_out = ggml_concat(VC, vit_out, reg_2d, 1); - } - vit_out = ggml_concat(VC, vit_out, patches_with_bias, 1); - - for (int64_t i = 0; i < cfg.vit_layers; ++i) { - vit_out = vit_layer(VC, - vision_encoder.layers[i], - vit_out, - rope_cos, - rope_sin, - internal_vit_seq, - cfg.vit_heads, - cfg.vit_head_dim); - } - // TurboVLA selects DINOv3's final hidden-state entry, then removes - // its prefix tokens. See H-EmbodVis/TurboVLA@b29ab142, - // models/vision_encoder.py:109-114. - const int64_t patch_token_start = n_cls + n_reg; - - ggml_tensor* vit_patches_only = ggml_view_2d(VC, - vit_out, - vit_dim, - n_patches, - vit_out->nb[1], - patch_token_start * vit_out->nb[1]); - vit_patches_only = ggml_cont(VC, vit_patches_only); - - ggml_set_output(vit_patches_only); - - ggml_cgraph* vg = ggml_new_graph_custom(VC, 8192, false); - ggml_build_forward_expand(vg, vit_patches_only); - - if (!scratch.alloc(backend, vg)) { - fprintf(stderr, "ERROR: vision gallocr alloc failed\n"); - return {}; - } - - if (!patch_input->buffer) { - fprintf(stderr, - "ERROR: patch_input has no backend buffer after vision graph allocation; " - "vision graph is disconnected\n"); - return {}; - } - - std::vector patch_buf; - normalize_imagenet_patches(in.images[view_idx], - img_size, - patch_size, - patch_buf); - ggml_backend_tensor_set(patch_input, - patch_buf.data(), - 0, - patch_buf.size() * sizeof(float)); - ggml_backend_tensor_set(rope_cos, rope_cos_host.data(), 0, - rope_cos_host.size() * sizeof(float)); - ggml_backend_tensor_set(rope_sin, rope_sin_host.data(), 0, - rope_sin_host.size() * sizeof(float)); - - if (ggml_backend_graph_compute(backend, vg) != GGML_STATUS_SUCCESS) { - fprintf(stderr, "ERROR: vision compute failed\n"); - return {}; - } - - float* out_ptr = vision_outputs.data() + view_idx * vision_out_size; - ggml_backend_tensor_get(vit_patches_only, - out_ptr, - 0, - vision_out_size * sizeof(float)); - } - - ggml_context* C = scratch.reset(512 * 1024 * 1024); // metadata arena - if (!C) { - fprintf(stderr, "ERROR: main scratch reset failed\n"); - return {}; - } - - // Buffers are allocated after graph construction because the context uses no_alloc. - std::vector view_tokens((size_t)n_views, nullptr); - for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { - view_tokens[view_idx] = - ggml_new_tensor_2d(C, GGML_TYPE_F32, vit_dim, n_patches); - ggml_set_input(view_tokens[view_idx]); - } - - ggml_tensor* vit_tokens = view_tokens[0]; - for (int64_t view_idx = 1; view_idx < n_views; ++view_idx) { - vit_tokens = ggml_concat(C, vit_tokens, view_tokens[view_idx], 1); - } - - // Vision projection: [vit_dim, vision_seq] -> [hidden, vision_seq]. - ggml_tensor* proj = - layer_norm(C, - vit_tokens, - vision_proj.input_norm_w, - vision_proj.input_norm_b, - 1e-5f); - ggml_tensor* mlp_out = - linear(C, vision_proj.mlp_0_w, vision_proj.mlp_0_b, proj); - mlp_out = ggml_gelu_erf(C, mlp_out); - mlp_out = - linear(C, vision_proj.mlp_3_w, vision_proj.mlp_3_b, mlp_out); - - ggml_tensor* skip_out = - ggml_mul_mat(C, vision_proj.skip_w, vit_tokens); - proj = ggml_add(C, mlp_out, skip_out); - proj = layer_norm(C, - proj, - vision_proj.output_norm_w, - vision_proj.output_norm_b, - 1e-5f); - - // The checkpoint stores one embedding per camera. Each embedding must form - // a contiguous patch block: [view0 x 256 patches, view1 x 256 patches]. - const int64_t view_emb_elems = ggml_nelements(view_emb); - ggml_tensor* view_emb_used = nullptr; - if (view_emb_elems == (int64_t)hidden * n_views) { - ggml_tensor* by_view = ggml_reshape_2d(C, view_emb, hidden, n_views); - for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { - ggml_tensor* one_view = ggml_cont(C, ggml_view_2d( - C, by_view, hidden, 1, by_view->nb[1], view_idx * by_view->nb[1])); - ggml_tensor* view_block = ggml_repeat( - C, one_view, ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden, n_patches)); - view_emb_used = view_emb_used ? ggml_concat(C, view_emb_used, view_block, 1) : view_block; - } - } else if (view_emb_elems == (int64_t)hidden) { - // Single shared [hidden] embedding: repeat to [hidden, vision_seq] - view_emb_used = ggml_repeat(C, view_emb, - ggml_new_tensor_2d(C, GGML_TYPE_F32, hidden, vision_seq)); - } else if (view_emb_elems == (int64_t)hidden * vision_seq) { - view_emb_used = ggml_reshape_2d(C, view_emb, (int64_t)hidden, vision_seq); - } else { - fprintf(stderr, - "ERROR: view_emb has %lld values; expected %lld (n_views,hidden), " - "%lld (shared), or %lld (per-pos)\n", - (long long)view_emb_elems, - (long long)((int64_t)hidden * n_views), - (long long)(int64_t)hidden, - (long long)((int64_t)hidden * vision_seq)); - return {}; - } - proj = ggml_add(C, proj, view_emb_used); - - const int64_t text_seq = cfg.max_text_length; - if (text_seq <= 0) { - fprintf(stderr, "ERROR: text_seq must be positive\n"); - return {}; - } - TurboVLATextInputs text_inputs; - try { - text_inputs = build_turbovla_text_inputs(in, cfg, vision_seq); - } catch (const std::exception& e) { - fprintf(stderr, "ERROR: %s\n", e.what()); - return {}; - } - std::vector tok_type_ids_host((size_t)text_seq, 0); - - ggml_tensor* token_ids = - ggml_new_tensor_1d(C, GGML_TYPE_I32, text_seq); - ggml_set_input(token_ids); - - ggml_tensor* pos_ids = - ggml_new_tensor_1d(C, GGML_TYPE_I32, text_seq); - ggml_set_input(pos_ids); - - ggml_tensor* tok_type_ids = - ggml_new_tensor_1d(C, GGML_TYPE_I32, text_seq); - ggml_set_input(tok_type_ids); - - ggml_tensor* bert_self_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, text_seq, text_seq); - ggml_tensor* enhancer_self_mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, text_seq, text_seq); - ggml_tensor* fusion_language_key_mask = - ggml_new_tensor_2d(C, GGML_TYPE_F32, text_seq, vision_seq); - ggml_set_input(bert_self_mask); - ggml_set_input(enhancer_self_mask); - ggml_set_input(fusion_language_key_mask); - - ggml_tensor* pos_emb = - ggml_get_rows(C, text_encoder.pos_embed, pos_ids); - ggml_tensor* text_emb = - ggml_get_rows(C, text_encoder.word_embed, token_ids); - ggml_tensor* tok_type_emb = - ggml_get_rows(C, text_encoder.tok_type_embed, tok_type_ids); - - ggml_tensor* text_tokens = - ggml_add(C, ggml_add(C, text_emb, pos_emb), tok_type_emb); - text_tokens = layer_norm(C, - text_tokens, - text_encoder.embed_ln_w, - text_encoder.embed_ln_b, - 1e-12f); - - for (int64_t i = 0; i < cfg.text_layers; ++i) { - text_tokens = bert_layer(C, - text_encoder.layers[i], - text_tokens, - bert_self_mask, - text_seq, - cfg.text_heads, - cfg.text_head_dim); - } - text_tokens = - linear(C, - text_encoder.text_proj_w, - text_encoder.text_proj_b, - text_tokens); - ggml_tensor* vl_features = proj; - ggml_tensor* text_features = text_tokens; - - const int64_t vl_heads = cfg.fusion_heads; - const int64_t vl_head_dim = cfg.fusion_head_dim; - - const int64_t text_enh_heads = cfg.text_enhancer_heads; - const int64_t text_enh_head_dim = cfg.text_enhancer_head_dim; - - // Each fusion layer is immediately followed by its paired text enhancer. - GGML_ASSERT(cfg.num_fusion_layers == cfg.num_text_layers && - "VL fusion and text layers must have same count for interleaved execution"); - - for (int64_t i = 0; i < cfg.num_fusion_layers; ++i) { - vl_features = vl_fusion_layer(C, - fusion_layers[i], - vl_features, - text_features, - fusion_language_key_mask, - vision_seq, - text_seq, - vl_heads, - vl_head_dim); - text_features = vl_text_layer(C, - text_layers[i], - text_features, - enhancer_self_mask, - text_seq, - text_enh_heads, - text_enh_head_dim); - } - - ggml_tensor* vl_condition = - ggml_concat(C, vl_features, text_features, 1); - const int64_t vl_cond_seq = vision_seq + text_seq; - - const int64_t state_dim = cfg.state_dim; - - std::vector state_host((size_t)state_dim, 0.0f); - if (in.state) { - std::copy(in.state, in.state + state_dim, state_host.begin()); - } - - ggml_tensor* state_tensor = - ggml_new_tensor_2d(C, GGML_TYPE_F32, state_dim, 1); - ggml_set_input(state_tensor); - - ggml_tensor* s = - layer_norm(C, - state_tensor, - state_proj.net_0_w, - state_proj.net_0_b, - 1e-5f); - s = linear(C, state_proj.net_1_w, state_proj.net_1_b, s); - s = ggml_gelu_erf(C, s); - s = linear(C, state_proj.net_4_w, state_proj.net_4_b, s); - - // state projection is expected to produce hidden * num_state_tokens values. - if (ggml_nelements(s) != hidden * cfg.num_state_tokens) { - fprintf(stderr, - "ERROR: state projection produced %lld values; expected hidden*num_state_tokens=%lld\n", - (long long)ggml_nelements(s), - (long long)(hidden * cfg.num_state_tokens)); - return {}; - } - - ggml_tensor* state_tokens = - ggml_reshape_2d(C, s, hidden, cfg.num_state_tokens); - ggml_tensor* pos_emb_state = - ggml_reshape_2d(C, - state_proj.position, - hidden, - cfg.num_state_tokens); - state_tokens = ggml_add(C, state_tokens, pos_emb_state); - - // The projection normalizes only after adding learned positions. See - // H-EmbodVis/TurboVLA@b29ab142, models/action_head.py:25-31. - state_tokens = layer_norm(C, - state_tokens, - state_proj.output_norm_w, - state_proj.output_norm_b, - 1e-5f); - - // State tokens must be part of the decoder memory. - ggml_tensor* condition = - ggml_concat(C, vl_condition, state_tokens, 1); - const int64_t cond_seq = vl_cond_seq + cfg.num_state_tokens; - - ggml_tensor* queries = - ggml_reshape_2d(C, - action_decoder.action_q, - hidden, - cfg.action_horizon); - - const int64_t act_heads = cfg.action_heads; - - if (hidden % act_heads != 0) { - fprintf(stderr, - "ERROR: hidden=%lld is not divisible by act_heads=%lld\n", - (long long)hidden, - (long long)act_heads); - return {}; - } - const int64_t act_head_dim = cfg.action_head_dim; - - for (int64_t i = 0; i < cfg.num_action_decoder_layers; ++i) { - queries = act_decoder_layer(C, - action_decoder.layers[i], - queries, - condition, - cfg.action_horizon, - cond_seq, - act_heads, - act_head_dim); - } - - ggml_tensor* actions = - linear(C, action_decoder.proj_0_w, action_decoder.proj_0_b, queries); - actions = ggml_relu(C, actions); - actions = - linear(C, action_decoder.proj_1_w, action_decoder.proj_1_b, actions); - actions = ggml_relu(C, actions); - actions = - linear(C, action_decoder.proj_2_w, action_decoder.proj_2_b, actions); - actions = ggml_tanh(C, actions); - ggml_set_output(actions); - - ggml_cgraph* gf = ggml_new_graph_custom(C, 16384, false); - ggml_build_forward_expand(gf, actions); - - if (!scratch.alloc(backend, gf)) { - fprintf(stderr, "ERROR: scratch alloc failed\n"); - return {}; - } - - // All graph inputs now have backend buffers. Copy them only at this point. - for (int64_t view_idx = 0; view_idx < n_views; ++view_idx) { - if (!view_tokens[view_idx]->buffer) { - fprintf(stderr, - "ERROR: view_tokens[%lld] has no backend buffer after main graph allocation\n", - (long long)view_idx); - return {}; - } - const float* out_ptr = - vision_outputs.data() + view_idx * vision_out_size; - ggml_backend_tensor_set(view_tokens[view_idx], - out_ptr, - 0, - vision_out_size * sizeof(float)); - } - - if (!token_ids->buffer || !pos_ids->buffer || !tok_type_ids->buffer || - !bert_self_mask->buffer || !enhancer_self_mask->buffer || - !fusion_language_key_mask->buffer || !state_tensor->buffer) { - fprintf(stderr, - "ERROR: one or more main graph inputs have no backend buffer " - "(token=%p pos=%p type=%p bert=%p enhancer=%p fusion=%p state=%p)\n", - (void*)token_ids->buffer, - (void*)pos_ids->buffer, - (void*)tok_type_ids->buffer, - (void*)bert_self_mask->buffer, - (void*)enhancer_self_mask->buffer, - (void*)fusion_language_key_mask->buffer, - (void*)state_tensor->buffer); - return {}; - } - - ggml_backend_tensor_set(token_ids, - text_inputs.token_ids.data(), - 0, - text_inputs.token_ids.size() * sizeof(int32_t)); - ggml_backend_tensor_set(pos_ids, - text_inputs.position_ids.data(), - 0, - text_inputs.position_ids.size() * sizeof(int32_t)); - ggml_backend_tensor_set(tok_type_ids, - tok_type_ids_host.data(), - 0, - tok_type_ids_host.size() * sizeof(int32_t)); - ggml_backend_tensor_set(bert_self_mask, - text_inputs.bert_self_mask.data(), - 0, - text_inputs.bert_self_mask.size() * sizeof(float)); - ggml_backend_tensor_set(enhancer_self_mask, - text_inputs.enhancer_self_mask.data(), - 0, - text_inputs.enhancer_self_mask.size() * sizeof(float)); - ggml_backend_tensor_set(fusion_language_key_mask, - text_inputs.fusion_language_key_mask.data(), - 0, - text_inputs.fusion_language_key_mask.size() * sizeof(float)); - ggml_backend_tensor_set(state_tensor, - state_host.data(), - 0, - state_host.size() * sizeof(float)); - - if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { - fprintf(stderr, "ERROR: compute failed\n"); +} + +} // namespace + +std::unique_ptr turbovla_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string&, + const Options& opts) { + if (!mmproj_path.empty()) + std::printf("vla(turbovla): note - mmproj '%s' is ignored (vision is baked into the GGUF)\n", + mmproj_path.c_str()); + + auto m = std::make_unique(); + m->mt = opts.weight_dtype.value_or(GGML_TYPE_F32); + + gguf_reader g("turbovla"); + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("turbovla.architecture")) { + std::fprintf(stderr, "vla(turbovla): %s is not a TurboVLA GGUF\n", ckpt_path.c_str()); + return nullptr; + } + if (!load_config(g, *m)) + return nullptr; + + const Backend b = backend_init("vla(turbovla)", m->n_threads); + if (!b.handle) + return nullptr; + m->backend = b.handle; + + if (!load_weights(*m, g) || !m->upload_rope_tables()) + return nullptr; + + m->cfg.n_img = m->n_views * m->n_patches(); + m->cfg.n_lang = m->text_len_max; + m->cfg.n_state = m->n_state_tok; + m->cfg.n_suffix = m->horizon; + m->cfg.hidden = m->hidden; + m->cfg.max_state_dim = m->state_dim; + m->cfg.real_state_dim = m->state_dim; + m->cfg.max_action_dim = m->action_dim; + m->cfg.real_action_dim = m->action_dim; + + std::printf("vla(turbovla): weights resident %.2f GiB (%s) - DINOv3 ViT-B/16 x%lld views + BERT + " + "%lld fusion layers + ACT decoder, horizon %lld, %zu text groups\n", + ggml_backend_buffer_get_size(m->weight_buf)/(1024.0*1024.0*1024.0), dtype_name(m->mt), + (long long) m->n_views, (long long) m->fusion_layers, (long long) m->horizon, + m->text_groups.size()); + return m; +} + +std::vector TurboVlaModelArch::predict(const Inputs& in) { + using clock = std::chrono::steady_clock; + const auto t0 = clock::now(); + stats = Stats{}; + + if (in.precomputed_img_emb) { + std::fprintf(stderr, "vla(turbovla): precomputed_img_emb is not supported; pass raw images\n"); + return {}; + } + if (!in.images || in.n_images != n_views) { + std::fprintf(stderr, "vla(turbovla): expected %lld views (agentview, wrist), got %d\n", + (long long) n_views, in.n_images); + return {}; + } + for (int64_t i = 0; i < n_views; ++i) { + const ImageView & iv = in.images[i]; + if (!view_is_side(iv.data, iv.w, iv.h, image_size)) { + std::fprintf(stderr, "vla(turbovla): view %lld is %dx%d, expected %lldx%lld\n", + (long long) i, iv.w, iv.h, (long long) image_size, (long long) image_size); return {}; } - - const int64_t expected_action_values = - cfg.action_horizon * cfg.action_dim; - if (ggml_nelements(actions) != expected_action_values) { - fprintf(stderr, - "ERROR: actions has %lld values, expected %lld (= horizon %lld * dim %lld)\n", - (long long)ggml_nelements(actions), - (long long)expected_action_values, - (long long)cfg.action_horizon, - (long long)cfg.action_dim); + } + if (!in.lang_tokens || in.n_lang < 1 || in.n_lang > text_len_max) { + std::fprintf(stderr, "vla(turbovla): %d language tokens; need 1..%lld (truncate when tokenizing)\n", + in.n_lang, (long long) text_len_max); + return {}; + } + // ggml_get_rows does not bound-check. + for (int i = 0; i < in.n_lang; ++i) + if (in.lang_tokens[i] < 0 || in.lang_tokens[i] >= vocab) { + std::fprintf(stderr, "vla(turbovla): token %d out of vocab\n", in.lang_tokens[i]); return {}; } - std::vector output((size_t)expected_action_values, 0.0f); - ggml_backend_tensor_get(actions, - output.data(), - 0, - output.size() * sizeof(float)); - - return output; + const int64_t LT = text_len_max, VS = n_views*n_patches(); + const int64_t bert_len = text_len(in.lang_tokens, in.n_lang); + std::vector ids((size_t) bert_len, pad_id); + std::copy(in.lang_tokens, in.lang_tokens + in.n_lang, ids.begin()); + + std::vector allowed; + std::vector pos; + special_token_blocks(ids, *this, allowed, pos); + const float NEG = -INFINITY; + std::vector bert_mask((size_t) bert_len*bert_len), enh_mask((size_t) LT*LT, NEG), fus_mask((size_t) LT*VS); + for (int64_t q = 0; q < bert_len; ++q) + for (int64_t k = 0; k < bert_len; ++k) { + bert_mask[(size_t) q*bert_len + k] = allowed[(size_t) q*bert_len + k] ? 0.0f : NEG; + enh_mask[(size_t) q*LT + k] = bert_mask[(size_t) q*bert_len + k]; + } + for (int64_t q = bert_len; q < LT; ++q) + enh_mask[(size_t) q*LT + q] = 0.0f; + for (int64_t k = 0; k < LT; ++k) { + const float val = (k < bert_len && ids[(size_t) k] != pad_id) ? 0.0f : NEG; + for (int64_t q = 0; q < VS; ++q) + fus_mask[(size_t) q*LT + k] = val; } - std::unique_ptr turbovla_create(const std::string& /*mmproj_path*/, - const std::string& ckpt_path, - const std::string& /*config_path*/, - const Options& /*opts*/) { - gguf_reader g("turbovla"); - if (!g.open(ckpt_path)) { - std::fprintf(stderr, "vla(turbovla): failed to open %s\n", ckpt_path.c_str()); - return nullptr; - } - - auto* m = new vla::TurboVLA; - - m->cfg_.hidden = (int64_t)g.u32("turbovla.hidden"); - m->cfg_.vit_dim = (int64_t)g.u32("turbovla.vit_dim"); - m->cfg_.vit_layers = (int64_t)g.u32("turbovla.vit_layers"); - m->cfg_.vit_head_dim = (int64_t)g.u32("turbovla.vit_head_dim"); - m->cfg_.vit_heads = (int64_t)g.u32("turbovla.vit_heads"); - m->cfg_.text_dim = (int64_t)g.u32("turbovla.text_dim"); - m->cfg_.text_layers = (int64_t)g.u32("turbovla.text_layers"); - m->cfg_.text_head_dim = (int64_t)g.u32("turbovla.text_head_dim"); - m->cfg_.text_heads = (int64_t)g.u32("turbovla.text_heads"); - m->cfg_.num_fusion_layers = (int64_t)g.u32("turbovla.num_fusion_layers"); - m->cfg_.fusion_heads = (int64_t)g.u32("turbovla.fusion_heads"); - m->cfg_.fusion_head_dim = (int64_t)g.u32("turbovla.fusion_head_dim"); - m->cfg_.text_enhancer_heads = (int64_t)g.u32("turbovla.text_enhancer_heads"); - m->cfg_.text_enhancer_head_dim = (int64_t)g.u32("turbovla.text_enhancer_head_dim"); - m->cfg_.action_heads = (int64_t)g.u32("turbovla.action_heads"); - m->cfg_.action_head_dim = (int64_t)g.u32("turbovla.action_head_dim"); - m->cfg_.num_text_layers = (int64_t)g.u32("turbovla.num_text_layers"); - m->cfg_.num_action_decoder_layers = (int64_t)g.u32("turbovla.num_action_decoder_layers"); - m->cfg_.action_dim = (int64_t)g.u32("turbovla.action_dim"); - m->cfg_.state_dim = (int64_t)g.u32("turbovla.state_dim"); - m->cfg_.num_state_tokens = (int64_t)g.u32("turbovla.num_state_tokens"); - m->cfg_.action_horizon = (int64_t)g.u32("turbovla.action_horizon"); - m->cfg_.image_size = (int64_t)g.u32("turbovla.image_size"); - m->cfg_.patch_size = (int64_t)g.u32("turbovla.patch_size"); - m->cfg_.num_views = (int64_t)g.u32("turbovla.num_views"); - m->cfg_.num_register_tokens = (int64_t)g.u32("turbovla.num_register_tokens"); - if (!g.has("turbovla.rope_theta")) { - std::fprintf(stderr, "vla(turbovla): GGUF is missing required DINOv3 rope_theta metadata\n"); - return nullptr; - } - m->cfg_.rope_theta = g.f32("turbovla.rope_theta"); - m->cfg_.max_text_length = (int64_t)g.u32("turbovla.max_text_length"); - if (!g.has("turbovla.pad_token_id") || !g.has("turbovla.cls_token_id") || - !g.has("turbovla.sep_token_id") || !g.has("turbovla.period_token_id") || - !g.has("turbovla.question_token_id")) { - std::fprintf(stderr, "vla(turbovla): GGUF is missing required BERT special-token metadata\n"); - return nullptr; - } - m->cfg_.pad_token_id = (int64_t)g.u32("turbovla.pad_token_id"); - m->cfg_.cls_token_id = (int64_t)g.u32("turbovla.cls_token_id"); - m->cfg_.sep_token_id = (int64_t)g.u32("turbovla.sep_token_id"); - m->cfg_.period_token_id = (int64_t)g.u32("turbovla.period_token_id"); - m->cfg_.question_token_id = (int64_t)g.u32("turbovla.question_token_id"); - if (m->cfg_.fusion_heads * m->cfg_.fusion_head_dim != 1024 || - m->cfg_.text_enhancer_heads * m->cfg_.text_enhancer_head_dim != m->cfg_.hidden || - m->cfg_.action_heads * m->cfg_.action_head_dim != m->cfg_.hidden) { - std::fprintf(stderr, "vla(turbovla): incompatible attention metadata\n"); - return nullptr; - } - - // Publish the dimensions used by serving tools. Without this, callers - // allocate an empty state vector from ModelArchBase::cfg while predict() - // reads the checkpoint's state_dim values. - const int64_t patches_per_side = m->cfg_.image_size / m->cfg_.patch_size; - m->cfg.n_img = m->cfg_.num_views * patches_per_side * patches_per_side; - m->cfg.n_lang = m->cfg_.max_text_length; - m->cfg.n_state = m->cfg_.num_state_tokens; - m->cfg.n_prefix = m->cfg.n_img + m->cfg.n_lang + m->cfg.n_state; - m->cfg.n_suffix = m->cfg_.action_horizon; - m->cfg.n_full = m->cfg.n_prefix + m->cfg.n_suffix; - m->cfg.hidden = m->cfg_.hidden; - m->cfg.n_q_heads = m->cfg_.action_heads; - m->cfg.n_kv_heads = m->cfg_.action_heads; - m->cfg.head_dim = m->cfg_.action_head_dim; - m->cfg.q_full_dim = m->cfg_.hidden; - m->cfg.kv_full_dim = m->cfg_.hidden; - m->cfg.n_layers = m->cfg_.num_action_decoder_layers; - m->cfg.max_state_dim = m->cfg_.state_dim; - m->cfg.real_state_dim = m->cfg_.state_dim; - m->cfg.max_action_dim = m->cfg_.action_dim; - m->cfg.real_action_dim = m->cfg_.action_dim; - m->cfg.norm_eps = 1e-5f; - m->cfg.rms_eps = 1e-5f; - m->cfg.rope_n_dims = m->cfg_.vit_head_dim; - m->cfg.rope_freq_base = m->cfg_.rope_theta; - m->cfg.num_steps = m->cfg_.action_horizon; - - const Backend b = backend_init("vla(turbovla)", m->n_threads); - if (!b.handle) { - std::fprintf(stderr, "vla(turbovla): backend_init failed\n"); - return nullptr; - } - m->backend = b.handle; - - ggml_init_params wp = { (size_t)16*1024*1024, nullptr, true }; - m->ctx_weights = ggml_init(wp); - - WeightLoader L("turbovla", g, m->ctx_weights, GGML_TYPE_F32); - - m->vision_encoder.cls_token = L.f32("vit.cls_token"); - m->vision_encoder.patch_embed_w = L.gemm("vit.patch_embed.weight"); - m->vision_encoder.patch_embed_b = L.gemm("vit.patch_embed.bias"); - - // DINOv3 register tokens: load only when num_register_tokens > 0 - if (m->cfg_.num_register_tokens > 0) { - m->vision_encoder.register_tokens = L.f32("vit.register_tokens"); - if (!m->vision_encoder.register_tokens) { - std::fprintf(stderr, "vla(turbovla): WARNING: num_register_tokens=%lld but vit.register_tokens not found\n", - (long long)m->cfg_.num_register_tokens); - } - } - - m->vision_encoder.layers.resize(m->cfg_.vit_layers); - for (int64_t i = 0; i < m->cfg_.vit_layers; ++i) { - auto& w = m->vision_encoder.layers[i]; - w.ln1_w = L.f32("vit.blk.%lld.ln1.weight", (long long) i); - w.ln1_b = L.f32("vit.blk.%lld.ln1.bias", (long long) i); - w.attn_q_w = L.f32("vit.blk.%lld.attn_q.weight", (long long) i); - w.attn_q_b = L.f32("vit.blk.%lld.attn_q.bias", (long long) i); - w.attn_k_w = L.f32("vit.blk.%lld.attn_k.weight", (long long) i); - w.attn_k_b = L.f32("vit.blk.%lld.attn_k.bias", (long long) i); - w.attn_v_w = L.f32("vit.blk.%lld.attn_v.weight", (long long) i); - w.attn_v_b = L.f32("vit.blk.%lld.attn_v.bias", (long long) i); - w.attn_o_w = L.f32("vit.blk.%lld.attn_o.weight", (long long) i); - w.attn_o_b = L.f32("vit.blk.%lld.attn_o.bias", (long long) i); - w.ln2_w = L.f32("vit.blk.%lld.ln2.weight", (long long) i); - w.ln2_b = L.f32("vit.blk.%lld.ln2.bias", (long long) i); - w.fc1_w = L.f32("vit.blk.%lld.fc1.weight", (long long) i); - w.fc1_b = L.f32("vit.blk.%lld.fc1.bias", (long long) i); - w.fc2_w = L.f32("vit.blk.%lld.fc2.weight", (long long) i); - w.fc2_b = L.f32("vit.blk.%lld.fc2.bias", (long long) i); - } - m->vision_encoder.final_ln_w = L.f32("vit.final_norm.weight"); - m->vision_encoder.final_ln_b = L.f32("vit.final_norm.bias"); - - - m->vision_proj.input_norm_w = L.f32("vit_proj.input_norm.weight"); - m->vision_proj.input_norm_b = L.f32("vit_proj.input_norm.bias"); - m->vision_proj.mlp_0_w = L.f32("vit_proj.mlp.0.weight"); - m->vision_proj.mlp_0_b = L.f32("vit_proj.mlp.0.bias"); - m->vision_proj.mlp_3_w = L.f32("vit_proj.mlp.3.weight"); - m->vision_proj.mlp_3_b = L.f32("vit_proj.mlp.3.bias"); - m->vision_proj.skip_w = L.f32("vit_proj.skip.weight"); - m->vision_proj.output_norm_w = L.f32("vit_proj.output_norm.weight"); - m->vision_proj.output_norm_b = L.f32("vit_proj.output_norm.bias"); - - m->text_encoder.word_embed = L.f32("text.embed.word_embeddings"); - m->text_encoder.pos_embed = L.f32("text.embed.position_embeddings"); - m->text_encoder.tok_type_embed = L.f32("text.embed.token_type_embeddings"); - m->text_encoder.embed_ln_w = L.f32("text.embed.LayerNorm.weight"); - m->text_encoder.embed_ln_b = L.f32("text.embed.LayerNorm.bias"); - m->text_encoder.pooler_w = L.f32("text.pooler.dense.weight"); - m->text_encoder.pooler_b = L.f32("text.pooler.dense.bias"); - m->text_encoder.text_proj_w = L.f32("text_proj.weight"); - m->text_encoder.text_proj_b = L.f32("text_proj.bias"); - - m->text_encoder.layers.resize(m->cfg_.text_layers); - for (int64_t i = 0; i < m->cfg_.text_layers; ++i) { - auto& w = m->text_encoder.layers[i]; - w.attn_q_w = L.f32("text.encoder.layer.%lld.attention.self.query.weight", (long long) i); - w.attn_q_b = L.f32("text.encoder.layer.%lld.attention.self.query.bias", (long long) i); - w.attn_k_w = L.f32("text.encoder.layer.%lld.attention.self.key.weight", (long long) i); - w.attn_k_b = L.f32("text.encoder.layer.%lld.attention.self.key.bias", (long long) i); - w.attn_v_w = L.f32("text.encoder.layer.%lld.attention.self.value.weight", (long long) i); - w.attn_v_b = L.f32("text.encoder.layer.%lld.attention.self.value.bias", (long long) i); - w.attn_o_w = L.f32("text.encoder.layer.%lld.attention.output.dense.weight", (long long) i); - w.attn_o_b = L.f32("text.encoder.layer.%lld.attention.output.dense.bias", (long long) i); - w.ln1_w = L.f32("text.encoder.layer.%lld.attention.output.LayerNorm.weight", (long long) i); - w.ln1_b = L.f32("text.encoder.layer.%lld.attention.output.LayerNorm.bias", (long long) i); - w.fc1_w = L.f32("text.encoder.layer.%lld.intermediate.dense.weight", (long long) i); - w.fc1_b = L.f32("text.encoder.layer.%lld.intermediate.dense.bias", (long long) i); - w.fc2_w = L.f32("text.encoder.layer.%lld.output.dense.weight", (long long) i); - w.fc2_b = L.f32("text.encoder.layer.%lld.output.dense.bias", (long long) i); - w.ln2_w = L.f32("text.encoder.layer.%lld.output.LayerNorm.weight", (long long) i); - w.ln2_b = L.f32("text.encoder.layer.%lld.output.LayerNorm.bias", (long long) i); - } - - m->fusion_layers.resize(m->cfg_.num_fusion_layers); - for (int64_t i = 0; i < m->cfg_.num_fusion_layers; ++i) { - auto& w = m->fusion_layers[i]; - w.v_proj_w = L.f32("vl_fusion.%lld.v_proj.weight", (long long) i); - w.v_proj_b = L.f32("vl_fusion.%lld.v_proj.bias", (long long) i); - w.l_proj_w = L.f32("vl_fusion.%lld.l_proj.weight", (long long) i); - w.l_proj_b = L.f32("vl_fusion.%lld.l_proj.bias", (long long) i); - w.values_v_w = L.f32("vl_fusion.%lld.values_v.weight", (long long) i); - w.values_v_b = L.f32("vl_fusion.%lld.values_v.bias", (long long) i); - w.values_l_w = L.f32("vl_fusion.%lld.values_l.weight", (long long) i); - w.values_l_b = L.f32("vl_fusion.%lld.values_l.bias", (long long) i); - w.out_v_w = L.f32("vl_fusion.%lld.out_v.weight", (long long) i); - w.out_v_b = L.f32("vl_fusion.%lld.out_v.bias", (long long) i); - w.out_l_w = L.f32("vl_fusion.%lld.out_l.weight", (long long) i); - w.out_l_b = L.f32("vl_fusion.%lld.out_l.bias", (long long) i); - w.norm_v_w = L.f32("vl_fusion.%lld.norm_v.weight", (long long) i); - w.norm_v_b = L.f32("vl_fusion.%lld.norm_v.bias", (long long) i); - w.norm_l_w = L.f32("vl_fusion.%lld.norm_l.weight", (long long) i); - w.norm_l_b = L.f32("vl_fusion.%lld.norm_l.bias", (long long) i); - w.gamma_v = L.f32("vl_fusion.%lld.gamma_v", (long long) i); - w.gamma_l = L.f32("vl_fusion.%lld.gamma_l", (long long) i); - } - - m->text_layers.resize(m->cfg_.num_text_layers); - for (int64_t i = 0; i < m->cfg_.num_text_layers; ++i) { - auto& w = m->text_layers[i]; - w.attn_qkv_w = L.f32("vl_text.%lld.attn_qkv.weight", (long long) i); - w.attn_qkv_b = L.f32("vl_text.%lld.attn_qkv.bias", (long long) i); - w.attn_o_w = L.f32("vl_text.%lld.attn_o.weight", (long long) i); - w.attn_o_b = L.f32("vl_text.%lld.attn_o.bias", (long long) i); - w.ln1_w = L.f32("vl_text.%lld.ln1.weight", (long long) i); - w.ln1_b = L.f32("vl_text.%lld.ln1.bias", (long long) i); - w.fc1_w = L.f32("vl_text.%lld.fc1.weight", (long long) i); - w.fc1_b = L.f32("vl_text.%lld.fc1.bias", (long long) i); - w.fc2_w = L.f32("vl_text.%lld.fc2.weight", (long long) i); - w.fc2_b = L.f32("vl_text.%lld.fc2.bias", (long long) i); - w.ln2_w = L.f32("vl_text.%lld.ln2.weight", (long long) i); - w.ln2_b = L.f32("vl_text.%lld.ln2.bias", (long long) i); - } - - m->action_decoder.action_q = L.f32("act.q.weight"); - m->action_decoder.layers.resize(m->cfg_.num_action_decoder_layers); - for (int64_t i = 0; i < m->cfg_.num_action_decoder_layers; ++i) { - auto& w = m->action_decoder.layers[i]; - w.self_qkv_w = L.f32("act.dec.%lld.self_qkv.weight", (long long) i); - w.self_qkv_b = L.f32("act.dec.%lld.self_qkv.bias", (long long) i); - w.self_out_w = L.f32("act.dec.%lld.self_out.weight", (long long) i); - w.self_out_b = L.f32("act.dec.%lld.self_out.bias", (long long) i); - w.cross_qkv_w = L.f32("act.dec.%lld.cross_qkv.weight", (long long) i); - w.cross_qkv_b = L.f32("act.dec.%lld.cross_qkv.bias", (long long) i); - w.cross_out_w = L.f32("act.dec.%lld.cross_out.weight", (long long) i); - w.cross_out_b = L.f32("act.dec.%lld.cross_out.bias", (long long) i); - w.ln1_w = L.f32("act.dec.%lld.ln1.weight", (long long) i); - w.ln1_b = L.f32("act.dec.%lld.ln1.bias", (long long) i); - w.ln2_w = L.f32("act.dec.%lld.ln2.weight", (long long) i); - w.ln2_b = L.f32("act.dec.%lld.ln2.bias", (long long) i); - w.ln3_w = L.f32("act.dec.%lld.ln3.weight", (long long) i); - w.ln3_b = L.f32("act.dec.%lld.ln3.bias", (long long) i); - w.fc1_w = L.f32("act.dec.%lld.fc1.weight", (long long) i); - w.fc1_b = L.f32("act.dec.%lld.fc1.bias", (long long) i); - w.fc2_w = L.f32("act.dec.%lld.fc2.weight", (long long) i); - w.fc2_b = L.f32("act.dec.%lld.fc2.bias", (long long) i); - } - m->action_decoder.proj_0_w = L.f32("act.proj.0.weight"); - m->action_decoder.proj_0_b = L.f32("act.proj.0.bias"); - m->action_decoder.proj_1_w = L.f32("act.proj.1.weight"); - m->action_decoder.proj_1_b = L.f32("act.proj.1.bias"); - m->action_decoder.proj_2_w = L.f32("act.proj.2.weight"); - m->action_decoder.proj_2_b = L.f32("act.proj.2.bias"); - - m->state_proj.net_0_w = L.f32("state.proj.0.weight"); - m->state_proj.net_0_b = L.f32("state.proj.0.bias"); - m->state_proj.net_1_w = L.f32("state.proj.1.weight"); - m->state_proj.net_1_b = L.f32("state.proj.1.bias"); - m->state_proj.net_4_w = L.f32("state.proj.4.weight"); - m->state_proj.net_4_b = L.f32("state.proj.4.bias"); - m->state_proj.output_norm_w = L.f32("state.proj.output_norm.weight"); - m->state_proj.output_norm_b = L.f32("state.proj.output_norm.bias"); - m->state_proj.position = L.f32("state.proj.position"); - - m->view_emb = L.f32("view_emb"); - - if (!L.upload(m->backend, &m->weight_buf)) { - std::fprintf(stderr, "vla(turbovla): weight upload failed\n"); - return nullptr; - } - - std::printf("vla(turbovla): hidden=%lld vit_layers=%lld text_layers=%lld fusion=%lld action_horizon=%lld action_dim=%lld\n", - (long long)m->cfg_.hidden, - (long long)m->cfg_.vit_layers, - (long long)m->cfg_.text_layers, - (long long)m->cfg_.num_fusion_layers, - (long long)m->cfg_.action_horizon, - (long long)m->cfg_.action_dim); - std::printf("vla(turbovla): DINOv3 config - image_size=%lld patch_size=%lld num_register_tokens=%lld rope_theta=%.1f\n", - (long long)m->cfg_.image_size, - (long long)m->cfg_.patch_size, - (long long)m->cfg_.num_register_tokens, - (double)m->cfg_.rope_theta); - - return std::unique_ptr(m); + std::vector patches((size_t) 3*patch*patch*VS); + for (int64_t i = 0; i < n_views; ++i) + patchify(in.images[i], image_size, patch, patches.data() + (size_t) i*3*patch*patch*n_patches()); + std::vector state((size_t) state_dim, 0.0f); + if (in.state) + std::copy(in.state, in.state + state_dim, state.begin()); + + const size_t arena = ggml_tensor_overhead()*8192 + ggml_graph_overhead_custom(8192, false); + if (!graph.ensure(backend, Key{bert_len}, arena, + [&](ggml_context * C, IO & io) { return build(C, io, bert_len); })) { + std::fprintf(stderr, "vla(turbovla): graph build/alloc failed\n"); + return {}; } + IO & io = graph.io(); + ggml_backend_tensor_set(io.patches, patches.data(), 0, ggml_nbytes(io.patches)); + ggml_backend_tensor_set(io.ids, ids.data(), 0, ggml_nbytes(io.ids)); + ggml_backend_tensor_set(io.pos, pos.data(), 0, ggml_nbytes(io.pos)); + ggml_backend_tensor_set(io.bert_mask, bert_mask.data(), 0, ggml_nbytes(io.bert_mask)); + ggml_backend_tensor_set(io.enh_mask, enh_mask.data(), 0, ggml_nbytes(io.enh_mask)); + ggml_backend_tensor_set(io.fus_mask, fus_mask.data(), 0, ggml_nbytes(io.fus_mask)); + ggml_backend_tensor_set(io.state, state.data(), 0, ggml_nbytes(io.state)); + + const auto tc = clock::now(); + graph_unique_names(graph.graph()); + if (ggml_backend_graph_compute(backend, graph.graph()) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "vla(turbovla): compute failed\n"); + return {}; + } + std::vector out((size_t) (horizon*action_dim)); + ggml_backend_tensor_get(io.actions, out.data(), 0, out.size()*sizeof(float)); + + stats.ms_inference = std::chrono::duration(clock::now() - tc).count(); + stats.ms_total = std::chrono::duration(clock::now() - t0).count(); + return out; +} - } // namespace vla +} // namespace vla diff --git a/tests/py/test_converters.py b/tests/py/test_converters.py index f19bb18..ba403e2 100644 --- a/tests/py/test_converters.py +++ b/tests/py/test_converters.py @@ -305,7 +305,6 @@ def test_turbovla_converter_remap(): "vit.blk.0.attn_v.weight", "vit.blk.0.attn_v.bias", "vit.blk.0.attn_o.weight", "vit.blk.0.attn_o.bias", "vit.blk.0.ln1.weight", "vit.blk.0.ln1.bias", "vit.blk.0.ln2.weight", "vit.blk.0.ln2.bias", "vit.blk.0.fc1.weight", "vit.blk.0.fc1.bias", "vit.blk.0.fc2.weight", "vit.blk.0.fc2.bias", - "vit.final_norm.weight", "vit.final_norm.bias", "text.embed.word_embeddings", "text.embed.position_embeddings", "text.embed.token_type_embeddings", "text.embed.LayerNorm.weight", "text.embed.LayerNorm.bias", "text.encoder.layer.0.attention.self.query.weight", "text.encoder.layer.0.attention.self.query.bias", @@ -316,7 +315,7 @@ def test_turbovla_converter_remap(): "text.encoder.layer.0.intermediate.dense.weight", "text.encoder.layer.0.intermediate.dense.bias", "text.encoder.layer.0.output.dense.weight", "text.encoder.layer.0.output.dense.bias", "text.encoder.layer.0.output.LayerNorm.weight", "text.encoder.layer.0.output.LayerNorm.bias", - "text.pooler.dense.weight", "text.pooler.dense.bias", "text_proj.weight", "text_proj.bias", + "text_proj.weight", "text_proj.bias", "vit_proj.input_norm.weight", "vit_proj.input_norm.bias", "vit_proj.mlp.0.weight", "vit_proj.mlp.0.bias", "vit_proj.mlp.3.weight", "vit_proj.mlp.3.bias", "vit_proj.skip.weight", "vit_proj.output_norm.weight", "vit_proj.output_norm.bias", From 8ae670dbbb518e5f05097584f973a24c0e6868ff Mon Sep 17 00:00:00 2001 From: Khanh Nguyen Date: Fri, 25 Sep 2026 10:45:36 +0700 Subject: [PATCH 4/5] fix(turbovla): slice the cross-attention projection output, not its weight, so bf16 runs on ggml-openvino --- src/models/turbovla.cpp | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/models/turbovla.cpp b/src/models/turbovla.cpp index 3d3d3ae..236aa75 100644 --- a/src/models/turbovla.cpp +++ b/src/models/turbovla.cpp @@ -416,18 +416,14 @@ ggml_cgraph * TurboVlaModelArch::build(ggml_context * C, IO & io, int64_t bert_l ggml_tensor * att = self_attention(C, qkv, nullptr, dhd, dec_heads, horizon, 1); a = ggml_add(C, a, linear(C, l.self_o_w, l.self_o_b, att)); - // in_proj rows are [W_q; W_k; W_v]: the queries take the first block, - // the memory the other two, in one GEMM each. - const size_t wr = l.cross_qkv_w->nb[1], br = ggml_element_size(l.cross_qkv_b); - ggml_tensor * wq = ggml_view_2d(C, l.cross_qkv_w, hidden, hidden, wr, 0); - ggml_tensor * wkv = ggml_view_2d(C, l.cross_qkv_w, hidden, 2*hidden, wr, (size_t) hidden*wr); - ggml_tensor * bq = ggml_view_1d(C, l.cross_qkv_b, hidden, 0); - ggml_tensor * bkv = ggml_view_1d(C, l.cross_qkv_b, 2*hidden, (size_t) hidden*br); - ggml_tensor * q = linear(C, wq, bq, layer_norm(C, a, l.ln2_w, l.ln2_b, kLnEps)); - ggml_tensor * kv = linear(C, wkv, bkv, mem); - ggml_tensor * Q = to_heads(C, q, dhd, dec_heads, horizon); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 0), 0, 2, 1, 3)); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 1), 1, 2, 0, 3)); + // in_proj is [W_q; W_k; W_v]: the queries keep the Q block of their + // projection, the memory its K and V blocks. Slicing the weight instead + // saves little and breaks ggml-openvino once the weights are BF16. + ggml_tensor * q = linear(C, l.cross_qkv_w, l.cross_qkv_b, layer_norm(C, a, l.ln2_w, l.ln2_b, kLnEps)); + ggml_tensor * kv = linear(C, l.cross_qkv_w, l.cross_qkv_b, mem); + ggml_tensor * Q = ggml_cont(C, ggml_permute(C, split_heads(C, q, dhd, dec_heads, horizon, 1, 0), 0, 2, 1, 3)); + ggml_tensor * K = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 1), 0, 2, 1, 3)); + ggml_tensor * V = ggml_cont(C, ggml_permute(C, split_heads(C, kv, dhd, dec_heads, M, 1, 2), 1, 2, 0, 3)); att = attention(C, Q, K, V, nullptr, dscale, hidden, horizon); a = ggml_add(C, a, linear(C, l.cross_o_w, l.cross_o_b, att)); a = ggml_add(C, a, ffn_relu(C, l.fc1_w, l.fc1_b, l.fc2_w, l.fc2_b, From 0d70b424e81a09f9b2ee31392a69fba62676cf28 Mon Sep 17 00:00:00 2001 From: Khanh Nguyen Date: Fri, 25 Sep 2026 11:26:21 +0700 Subject: [PATCH 5/5] update support matrix with TurboVLA --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d00a12c..69a9434 100644 --- a/README.md +++ b/README.md @@ -326,7 +326,7 @@ supported (released and benchmarked), `~` = in progress, `-` = planned. | [OpenVLA-OFT](https://hf.co/vrfai/openvla-oft-libero-gguf) | Y | Y | - | Y | Y | - | | [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | - | Y | Y | - | | [Octo-Small](https://hf.co/vrfai/octo-small-libero-gguf) | Y | Y | Y | Y | - | - | -| [TurboVLA](https://github.com/H-EmbodVis/TurboVLA) | ~ | ~ | - | - | - | - | +| [TurboVLA](https://hf.co/vrfai/turbovla-libero-gguf) | Y | Y | Y | Y | Y | - | ---