[Model] Add InternVLA-A1 offline inference support (#2737)

Signed-off-by: zhuyangkun <zyk_120403228@163.com>
Signed-off-by: bl-wang1 <918395762@qq.com>
Co-authored-by: wangbolun <918395762@qq.com>

docs/.nav.yml

@@ -22,6 +22,7 @@ nav:
       - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
       - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md

docs/models/supported_models.md

@@ -20,6 +20,7 @@ th {
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
+| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/user_guide/examples/offline_inference/internvla_a1.md

@@ -0,0 +1,205 @@
+# InternVLA-A1
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/internvla_a1>.
+
+
+## Setup
+This example is adapted from: https://github.com/InternRobotics/InternVLA-A1/blob/master/tests/policies/internvla_a1_3b/open_loop_genie1_real.ipynb
+
+This example runs the single-path vLLM offline inference workflow for InternVLA-A1 open-loop action prediction.
+
+Before running the script, export the required local paths:
+
+```bash
+# hf Jia-Zeng/InternVLA-A1-3B-FineTuned-Place_Markpen
+export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
+# hf download InternRobotics/InternData-A1 real_lerobotv30/genie1/Genie1-Place_Markpen.tar.gz --repo-type dataset --local-dir /path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
+# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
+export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
+```
+
+The shell entrypoint also accepts these variables as defaults, so you can keep the command line itself short.
+
+`INTERNVLA_A1_COSMOS_DIR` is expected to contain:
+
+- `encoder.safetensors`
+- `decoder.safetensors`
+
+Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
+
+## Run examples
+
+Get into the example folder:
+
+```bash
+cd examples/offline_inference/internvla_a1
+```
+
+### Run one sample
+
+This runs a single sample through the maintained vLLM registry path.
+
+```bash
+bash run.sh --num-samples 1 --num-episodes 0
+```
+
+### Run open-loop evaluation against GT
+
+This runs one episode, writes evaluation metrics, and saves prediction-vs-GT plots when `matplotlib` is available.
+
+```bash
+bash run.sh --num-episodes 1
+```
+
+Outputs are written under:
+
+```bash
+outputs/internvla_a1/vllm_infer/
+```
+
+Typical files:
+
+- `summary.json`: top-level run summary
+- `registry/log.json`: per-episode GT comparison metrics
+- `registry/plots/*.jpg`: prediction-vs-GT figures
+
+### Eager float32 validation
+
+Use eager attention and float32 when you want the most stable numerical comparison baseline.
+
+```bash
+bash run.sh \
+  --num-episodes 1 \
+  --attn-implementation eager \
+  --dtype float32
+```
+
+### Optional runtime switches
+
+- `--dtype {bfloat16,float32}`: choose inference dtype
+- `--attn-implementation {eager,sdpa}`: switch attention backend
+- `--enable-regional-compile`: enable regional `torch.compile`
+- `--enable-warmup`: run pipeline warmup in initialization
+- `--skip-plots`: skip plot generation even if `matplotlib` is installed
+
+### Collect results and performance logs
+
+Use the helper below when you want to collect one-sample output, forward latency, GT comparison, plots, timing logs, and optional pytest output into a single result directory:
+
+```bash
+bash collect_results.sh
+```
+
+The script writes a timestamped directory under:
+
+```bash
+outputs/internvla_a1/collected_results/
+```
+
+Typical result files:
+
+- `env_summary.txt`: environment and path summary
+- `sample_run.log`: one-sample inference log
+- `forward_benchmark/forward_latency.json`: isolated `pipeline.forward` latency summary
+- `eval_run.log`: open-loop GT evaluation log
+- `eval_outputs/summary.json`: top-level output summary
+- `eval_outputs/registry/log.json`: GT comparison metrics
+- `eval_outputs/registry/plots/*.jpg`: prediction-vs-GT figures
+- `*_time.txt`: timing output
+- `*_gpu.csv`: sampled GPU usage
+
+### Benchmark forward latency
+
+Use the dedicated mode below when you want the isolated `pipeline.forward` latency instead of end-to-end script time:
+
+```bash
+python end2end.py \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --benchmark-forward \
+  --dtype bfloat16 \
+  --attn-implementation eager \
+  --warmup-iters 3 \
+  --benchmark-iters 10 \
+  --output-dir outputs/internvla_a1/forward_benchmark
+```
+
+The output JSON is written to `outputs/internvla_a1/forward_benchmark/forward_latency.json` and contains:
+
+- `cold_start_ms`
+- `warmup_summary`
+- `benchmark_summary`
+- `benchmark_samples_ms`
+
+### Reference results
+
+Reference run collected on `1x NVIDIA H200`, `bfloat16`, `eager`:
+
+- one-sample end-to-end run: `38s`
+- one-episode GT evaluation run: `45s`
+- `average_mse = 1.7173260857816786e-05`
+- `average_mae = 0.0011860118247568607`
+- `average_mse_joint = 7.42028441891307e-06`
+- `average_mae_joint = 0.0010777723509818316`
+- `average_mse_gripper = 8.544408774469048e-05`
+- `average_mae_gripper = 0.0019436875591054559`
+
+The reference plot filename is:
+
+- `eval_outputs/registry/plots/vllm_registry_open_loop_ep0.jpg`
+
+## Measure latency and VRAM
+
+You can collect end-to-end runtime and peak host memory with:
+
+```bash
+/usr/bin/time -v bash run.sh --num-episodes 1
+```
+
+For GPU memory snapshots, use:
+
+```bash
+nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv
+```
+
+Run `nvidia-smi` before and during inference to record the peak device memory for your PR report.
+
+## FAQ
+
+If `matplotlib` is missing, evaluation still runs and only plot generation is skipped. To enable plots:
+
+```bash
+pip install matplotlib
+```
+
+## Example materials
+
+- `end2end.py`: main offline inference and GT evaluation entrypoint
+- `run.sh`: shell wrapper with local path env vars
+- `collect_results.sh`: helper to gather result summaries and performance logs into one directory
+- `internvla_a1_common.py`: dataset, evaluation, and plotting helpers
+
+## Embedded source listings
+
+??? abstract "collect_results.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/internvla_a1/collect_results.sh"
+    ``````
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/internvla_a1/end2end.py"
+    ``````
+??? abstract "internvla_a1_common.py"
+    ``````py
+    --8<-- "examples/offline_inference/internvla_a1/internvla_a1_common.py"
+    ``````
+??? abstract "run.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/internvla_a1/run.sh"
+    ``````
+??? abstract "standalone/outputs/qwena1/standalone_infer/log.json"
+    ``````json
+    --8<-- "examples/offline_inference/internvla_a1/standalone/outputs/qwena1/standalone_infer/log.json"
+    ``````

examples/offline_inference/internvla_a1/README.md

@@ -0,0 +1,30 @@
+# InternVLA-A1
+
+Full usage and result-reporting guidance lives in [docs/user_guide/examples/offline_inference/internvla_a1.md](../../../docs/user_guide/examples/offline_inference/internvla_a1.md).
+
+Quick start:
+
+```bash
+export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
+export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
+# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
+export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
+
+bash run.sh --num-samples 1 --num-episodes 0
+bash run.sh --num-episodes 1
+bash collect_results.sh
+```
+
+Expected files under `INTERNVLA_A1_COSMOS_DIR`:
+
+- `encoder.safetensors`
+- `decoder.safetensors`
+
+Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
+
+Key entrypoints:
+
+- `run.sh`: wrapper for offline inference and GT evaluation
+- `collect_results.sh`: collect sample output, latency, metrics, plots, and logs
+- `end2end.py`: underlying Python entrypoint

examples/offline_inference/internvla_a1/collect_results.sh

@@ -0,0 +1,220 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+RESULT_ROOT="${INTERNVLA_A1_RESULT_DIR:-$REPO_ROOT/outputs/internvla_a1/collected_results/$TIMESTAMP}"
+
+: "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
+: "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
+: "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
+: "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
+
+EVAL_OUTPUT_DIR="$RESULT_ROOT/eval_outputs"
+GPU_MONITOR_INTERVAL="${GPU_MONITOR_INTERVAL:-1}"
+
+mkdir -p "$RESULT_ROOT"
+
+export INTERNVLA_A1_MODEL_DIR
+export INTERNVLA_A1_DATASET_DIR
+export INTERNVLA_A1_PROCESSOR_DIR
+export INTERNVLA_A1_COSMOS_DIR
+
+write_env_summary() {
+  cat >"$RESULT_ROOT/env_summary.txt" <<EOF
+timestamp=$TIMESTAMP
+repo_root=$REPO_ROOT
+result_root=$RESULT_ROOT
+python=$(command -v python || true)
+python_version=$(python --version 2>&1 || true)
+model_dir=$INTERNVLA_A1_MODEL_DIR
+dataset_dir=$INTERNVLA_A1_DATASET_DIR
+processor_dir=$INTERNVLA_A1_PROCESSOR_DIR
+cosmos_dir=$INTERNVLA_A1_COSMOS_DIR
+pwd=$(pwd)
+EOF
+}
+
+capture_gpu_snapshot() {
+  local output_file="$1"
+  if command -v nvidia-smi >/dev/null 2>&1; then
+    nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv >"$output_file" 2>&1 || true
+  else
+    echo "nvidia-smi not found" >"$output_file"
+  fi
+}
+
+start_gpu_monitor() {
+  local output_file="$1"
+  if ! command -v nvidia-smi >/dev/null 2>&1; then
+    echo "nvidia-smi not found" >"$output_file"
+    return 1
+  fi
+
+  (
+    echo "timestamp,index,name,memory.total [MiB],memory.used [MiB],utilization.gpu [%]"
+    while true; do
+      local now
+      now="$(date '+%Y-%m-%d %H:%M:%S')"
+      nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv,noheader,nounits \
+        | awk -F',' -v ts="$now" '{gsub(/^[ \t]+|[ \t]+$/, "", $0); print ts "," $0}'
+      sleep "$GPU_MONITOR_INTERVAL"
+    done
+  ) >"$output_file" 2>/dev/null &
+  echo $!
+}
+
+stop_gpu_monitor() {
+  local monitor_pid="${1:-}"
+  if [[ -n "$monitor_pid" ]] && kill -0 "$monitor_pid" 2>/dev/null; then
+    kill "$monitor_pid" 2>/dev/null || true
+    wait "$monitor_pid" 2>/dev/null || true
+  fi
+}
+
+copy_eval_outputs() {
+  local source_dir="$1"
+  local target_dir="$2"
+  local source_real=""
+  local target_real=""
+  source_real="$(realpath "$source_dir" 2>/dev/null || echo "$source_dir")"
+  target_real="$(realpath "$target_dir" 2>/dev/null || echo "$target_dir")"
+  if [[ "$source_real" == "$target_real" ]]; then
+    return 0
+  fi
+  mkdir -p "$target_dir"
+  if [[ -d "$source_dir" ]]; then
+    cp -r "$source_dir"/. "$target_dir"/
+  fi
+}
+
+run_with_artifacts() {
+  local name="$1"
+  shift
+
+  local log_file="$RESULT_ROOT/${name}.log"
+  local time_file="$RESULT_ROOT/${name}_time.txt"
+  local gpu_file="$RESULT_ROOT/${name}_gpu.csv"
+  local status_file="$RESULT_ROOT/${name}_status.txt"
+  local monitor_pid=""
+
+  monitor_pid="$(start_gpu_monitor "$gpu_file" || true)"
+  set +e
+  local exit_code=0
+  local start_ts=0
+  local end_ts=0
+  start_ts="$(date +%s)"
+  if [[ -x /usr/bin/time ]]; then
+    /usr/bin/time -v -o "$time_file" "$@" >"$log_file" 2>&1
+    exit_code=$?
+    if [[ ! -s "$time_file" ]]; then
+      end_ts="$(date +%s)"
+      {
+        echo "timing_mode=usr_bin_time_empty_fallback"
+        echo "start_ts=$start_ts"
+        echo "end_ts=$end_ts"
+        echo "elapsed_seconds=$((end_ts - start_ts))"
+      } >"$time_file"
+    fi
+  else
+    "$@" >"$log_file" 2>&1
+    exit_code=$?
+    end_ts="$(date +%s)"
+    {
+      echo "timing_mode=shell_date"
+      echo "start_ts=$start_ts"
+      echo "end_ts=$end_ts"
+      echo "elapsed_seconds=$((end_ts - start_ts))"
+    } >"$time_file"
+  fi
+  set -e
+  stop_gpu_monitor "$monitor_pid"
+
+  echo "exit_code=$exit_code" >"$status_file"
+  if [[ $exit_code -ne 0 ]]; then
+    echo "[error] command failed for $name, see $log_file"
+    return $exit_code
+  fi
+}
+
+write_manifest() {
+  cat >"$RESULT_ROOT/README.txt" <<EOF
+InternVLA-A1 collected results
+
+Key files:
+- env_summary.txt: environment and path summary
+- sample_run.log / sample_run_time.txt: one-sample functional run
+- forward_benchmark.log / forward_benchmark_time.txt: pure pipeline.forward latency benchmark
+- eval_run.log / eval_run_time.txt: GT evaluation run
+- pytest_e2e.log / pytest_e2e_time.txt: offline e2e pytest result
+- gpu_info_before.csv / gpu_info_after.csv: point-in-time GPU snapshots
+- *_gpu.csv: sampled GPU usage during each run
+- eval_outputs/: copied output directory from the GT evaluation run
+
+Important outputs:
+- forward_benchmark/forward_latency.json
+- eval_outputs/summary.json
+- eval_outputs/registry/log.json
+- eval_outputs/registry/plots/
+EOF
+}
+
+write_skip_artifact() {
+  local name="$1"
+  local reason="$2"
+  echo "$reason" >"$RESULT_ROOT/${name}.log"
+  echo "timing_mode=skipped" >"$RESULT_ROOT/${name}_time.txt"
+  echo "exit_code=0" >"$RESULT_ROOT/${name}_status.txt"
+  echo "skipped_reason=$reason" >>"$RESULT_ROOT/${name}_status.txt"
+  if [[ ! -f "$RESULT_ROOT/${name}_gpu.csv" ]]; then
+    echo "skipped,$reason" >"$RESULT_ROOT/${name}_gpu.csv"
+  fi
+}
+
+write_env_summary
+write_manifest
+capture_gpu_snapshot "$RESULT_ROOT/gpu_info_before.csv"
+
+run_with_artifacts \
+  "sample_run" \
+  bash "$ROOT_DIR/run.sh" \
+  --output-dir "$RESULT_ROOT/sample_outputs" \
+  --num-samples 1 \
+  --num-episodes 0
+
+run_with_artifacts \
+  "forward_benchmark" \
+  python "$ROOT_DIR/end2end.py" \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --benchmark-forward \
+  --dtype bfloat16 \
+  --attn-implementation eager \
+  --warmup-iters 3 \
+  --benchmark-iters 10 \
+  --output-dir "$RESULT_ROOT/forward_benchmark"
+
+run_with_artifacts \
+  "eval_run" \
+  bash "$ROOT_DIR/run.sh" \
+  --output-dir "$EVAL_OUTPUT_DIR" \
+  --num-episodes 1
+
+copy_eval_outputs "$EVAL_OUTPUT_DIR" "$RESULT_ROOT/eval_outputs"
+
+if python - <<'PY' >/dev/null 2>&1
+import importlib.util
+raise SystemExit(0 if importlib.util.find_spec("pytest") else 1)
+PY
+then
+  run_with_artifacts \
+    "pytest_e2e" \
+    python -m pytest -sv tests/e2e/offline_inference/test_internvla_a1.py -m advanced_model
+else
+  write_skip_artifact "pytest_e2e" "pytest is not installed in the current python environment"
+fi
+
+capture_gpu_snapshot "$RESULT_ROOT/gpu_info_after.csv"
+
+echo "Results written to: $RESULT_ROOT"

examples/offline_inference/internvla_a1/end2end.py

@@ -0,0 +1,327 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+_WORKSPACE_ROOT = Path(__file__).resolve().parent
+_REPO_ROOT = _WORKSPACE_ROOT.parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    # Keep the shared helper import clean by doing script-only path bootstrap here.
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from internvla_a1_common import (  # noqa: E402
+    A2DOpenLoopDataset,
+    collate_open_loop_samples,
+    make_shared_noise,
+    run_open_loop_evaluation,
+    select_indices,
+    tensor_dtype,
+    tensor_sha256,
+)
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig  # noqa: E402
+from vllm_omni.diffusion.models.internvla_a1 import (  # noqa: E402
+    InternVLAA1Config,
+    InternVLAA1TrainMetadata,
+)
+from vllm_omni.diffusion.models.internvla_a1.config import OBS_STATE  # noqa: E402
+from vllm_omni.diffusion.registry import initialize_model  # noqa: E402
+from vllm_omni.diffusion.request import OmniDiffusionRequest  # noqa: E402
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams  # noqa: E402
+
+
+def _required_path_arg(env_name: str, cli_value: str | None) -> str:
+    value = cli_value or os.getenv(env_name)
+    if not value:
+        raise ValueError(f"Missing required path: set --{env_name.lower().replace('_', '-')} or {env_name}.")
+    return value
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run vLLM InternVLA-A1 inference on a few samples.")
+    parser.add_argument("--model-dir")
+    parser.add_argument("--dataset-dir")
+    parser.add_argument("--num-samples", type=int, default=1)
+    parser.add_argument("--num-episodes", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--compile-model", action="store_true")
+    parser.add_argument("--attn-implementation", default="eager")
+    parser.add_argument("--enable-regional-compile", action="store_true")
+    parser.add_argument("--enable-warmup", action="store_true")
+    parser.add_argument("--strict-load", action="store_true")
+    parser.add_argument("--output-dir", default="outputs/internvla_a1/vllm_infer")
+    parser.add_argument("--skip-plots", action="store_true")
+    parser.add_argument("--benchmark-forward", action="store_true")
+    parser.add_argument("--warmup-iters", type=int, default=3)
+    parser.add_argument("--benchmark-iters", type=int, default=10)
+    return parser.parse_args()
+
+
+def build_od_config(args: argparse.Namespace, processor_model_name: str) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model=str(Path(args.model_dir).resolve()),
+        model_class_name="InternVLAA1Pipeline",
+        dtype=tensor_dtype(args.dtype),
+        custom_pipeline_args={
+            "device": args.device,
+            "dtype": args.dtype,
+            "compile_model": args.compile_model,
+            "attn_implementation": args.attn_implementation,
+            "enable_regional_compile": args.enable_regional_compile,
+            "enable_warmup": args.enable_warmup,
+            "strict_load": args.strict_load,
+            "processor_model_name": processor_model_name,
+        },
+    )
+
+
+def build_dataset(args: argparse.Namespace) -> tuple[A2DOpenLoopDataset, InternVLAA1Config, InternVLAA1TrainMetadata]:
+    model_dir = Path(args.model_dir)
+    config = InternVLAA1Config.from_pretrained(model_dir)
+    config.device = args.device
+    config.dtype = args.dtype
+    config.compile_model = args.compile_model
+    config.attn_implementation = args.attn_implementation
+    config.enable_regional_compile = args.enable_regional_compile
+
+    train_meta = InternVLAA1TrainMetadata.from_pretrained(model_dir)
+    with open(model_dir / "stats.json", encoding="utf-8") as f:
+        train_stats = json.load(f)["a2d"]
+
+    dataset = A2DOpenLoopDataset(
+        args.dataset_dir,
+        config=config,
+        train_stats=train_stats,
+    )
+    return dataset, config, train_meta
+
+
+def run_one_path(
+    pipeline,
+    dataset: A2DOpenLoopDataset,
+    config: InternVLAA1Config,
+    args: argparse.Namespace,
+    indices: list[int],
+) -> list[dict[str, object]]:
+    results: list[dict[str, object]] = []
+    for index in indices:
+        sample = dataset.get_sample(index)
+        batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
+        noise = make_shared_noise(
+            args.seed,
+            index,
+            (
+                batch_inputs[OBS_STATE].shape[0],
+                config.chunk_size,
+                config.max_action_dim,
+            ),
+            args.device,
+        )
+        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+        pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+        results.append(
+            {
+                "path": "registry",
+                "index": index,
+                "episode_index": sample.episode_index,
+                "task": sample.task,
+                "seed": args.seed,
+                "shape": list(pred.shape),
+                "mean": float(pred.mean().item()),
+                "std": float(pred.std().item()),
+                "action_sha256": tensor_sha256(pred),
+                "first_action_prefix": pred[0, 0, :8].tolist(),
+            }
+        )
+    return results
+
+
+def run_pipeline_forward(
+    pipeline,
+    batch_inputs: dict[str, torch.Tensor],
+    noise: torch.Tensor,
+) -> torch.Tensor:
+    output = pipeline.forward(
+        OmniDiffusionRequest(
+            prompts=[""],
+            sampling_params=OmniDiffusionSamplingParams(
+                extra_args={
+                    "batch_inputs": batch_inputs,
+                    "noise": noise,
+                    "decode_image": False,
+                }
+            ),
+        )
+    )
+    if output.error:
+        raise RuntimeError(output.error)
+    if output.output is None:
+        raise RuntimeError("InternVLAA1Pipeline.forward returned no output tensor.")
+    return output.output
+
+
+def _synchronize(device: str) -> None:
+    if device.startswith("cuda") and torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def _latency_summary(values_ms: list[float]) -> dict[str, float]:
+    sorted_values = sorted(values_ms)
+    p50_index = min(len(sorted_values) - 1, round(0.50 * (len(sorted_values) - 1)))
+    p90_index = min(len(sorted_values) - 1, round(0.90 * (len(sorted_values) - 1)))
+    return {
+        "mean_ms": float(statistics.mean(sorted_values)),
+        "stdev_ms": float(statistics.pstdev(sorted_values)) if len(sorted_values) > 1 else 0.0,
+        "min_ms": float(sorted_values[0]),
+        "max_ms": float(sorted_values[-1]),
+        "p50_ms": float(sorted_values[p50_index]),
+        "p90_ms": float(sorted_values[p90_index]),
+    }
+
+
+def benchmark_forward(
+    pipeline,
+    dataset: A2DOpenLoopDataset,
+    config: InternVLAA1Config,
+    args: argparse.Namespace,
+    index: int,
+    output_dir: Path,
+) -> dict[str, object]:
+    sample = dataset.get_sample(index)
+    batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
+    noise = make_shared_noise(
+        args.seed,
+        index,
+        (
+            batch_inputs[OBS_STATE].shape[0],
+            config.chunk_size,
+            config.max_action_dim,
+        ),
+        args.device,
+    )
+
+    _synchronize(args.device)
+    cold_start_begin = time.perf_counter()
+    pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+    _synchronize(args.device)
+    cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
+
+    warmup_ms: list[float] = []
+    for _ in range(args.warmup_iters):
+        _synchronize(args.device)
+        begin = time.perf_counter()
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _synchronize(args.device)
+        warmup_ms.append((time.perf_counter() - begin) * 1000.0)
+
+    benchmark_ms: list[float] = []
+    for _ in range(args.benchmark_iters):
+        _synchronize(args.device)
+        begin = time.perf_counter()
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _synchronize(args.device)
+        benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
+
+    pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+    summary = {
+        "mode": "forward_latency",
+        "model_dir": str(Path(args.model_dir).resolve()),
+        "dataset_dir": str(Path(args.dataset_dir).resolve()),
+        "sample_index": index,
+        "episode_index": sample.episode_index,
+        "task": sample.task,
+        "device": args.device,
+        "dtype": args.dtype,
+        "attn_implementation": args.attn_implementation,
+        "enable_regional_compile": args.enable_regional_compile,
+        "warmup_iters": args.warmup_iters,
+        "benchmark_iters": args.benchmark_iters,
+        "output_shape": list(pred.shape),
+        "cold_start_ms": cold_start_ms,
+        "warmup_summary": _latency_summary(warmup_ms) if warmup_ms else {},
+        "benchmark_summary": _latency_summary(benchmark_ms) if benchmark_ms else {},
+        "benchmark_samples_ms": benchmark_ms,
+    }
+    with open(output_dir / "forward_latency.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+    return summary
+
+
+def main() -> None:
+    args = parse_args()
+    args.model_dir = _required_path_arg("INTERNVLA_A1_MODEL_DIR", args.model_dir)
+    args.dataset_dir = _required_path_arg("INTERNVLA_A1_DATASET_DIR", args.dataset_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    dataset, config, train_meta = build_dataset(args)
+    indices = select_indices(dataset, args.num_samples)
+    od_config = build_od_config(args, train_meta.processor_model_name)
+
+    eval_summaries: dict[str, object] = {}
+    pipeline = initialize_model(od_config)
+    if args.benchmark_forward:
+        benchmark_forward(
+            pipeline=pipeline,
+            dataset=dataset,
+            config=config,
+            args=args,
+            index=indices[0],
+            output_dir=output_dir,
+        )
+        return
+    results = run_one_path(
+        pipeline=pipeline,
+        dataset=dataset,
+        config=config,
+        args=args,
+        indices=indices,
+    )
+    if args.num_episodes > 0:
+        eval_summaries["registry"] = run_open_loop_evaluation(
+            mode="vllm_registry",
+            policy=pipeline,
+            config=config,
+            dataset=dataset,
+            train_meta=train_meta,
+            collate_samples=collate_open_loop_samples,
+            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
+            num_episodes=args.num_episodes,
+            seed=args.seed,
+            device=args.device,
+            dtype=tensor_dtype(args.dtype),
+            output_dir=output_dir / "registry",
+            skip_plots=args.skip_plots,
+        )
+
+    summary = {
+        "mode": "registry",
+        "model_dir": str(Path(args.model_dir).resolve()),
+        "dataset_dir": str(Path(args.dataset_dir).resolve()),
+        "device": args.device,
+        "dtype": args.dtype,
+        "attn_implementation": args.attn_implementation,
+        "enable_regional_compile": args.enable_regional_compile,
+        "seed": args.seed,
+        "indices": indices,
+        "results": results,
+        "output_dir": str(output_dir.resolve()),
+        "eval_summaries": eval_summaries,
+    }
+    with open(output_dir / "summary.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/internvla_a1/internvla_a1_common.py

@@ -0,0 +1,471 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import torch.nn.functional as F
+import torchvision
+
+from vllm_omni.diffusion.models.internvla_a1.config import (  # noqa: E402
+    OBS_IMAGES,
+    OBS_STATE,
+    OBS_TASK,
+    InternVLAA1Config,
+)
+
+
+def tensor_dtype(name: str) -> torch.dtype:
+    return torch.bfloat16 if name == "bfloat16" else torch.float32
+
+
+def select_indices(dataset, num_samples: int) -> list[int]:
+    indices: list[int] = []
+    for _, episode_indices in dataset.episode_start_indices(max_episodes=num_samples):
+        if episode_indices:
+            indices.append(episode_indices[0])
+    return indices[:num_samples]
+
+
+def make_shared_noise(seed: int, sample_index: int, shape: tuple[int, ...], device: str) -> torch.Tensor:
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed + sample_index)
+    noise = torch.randn(shape, generator=generator, dtype=torch.float32)
+    return noise.to(device=device)
+
+
+def tensor_sha256(tensor: torch.Tensor) -> str:
+    array = tensor.detach().contiguous().cpu().numpy()
+    return hashlib.sha256(array.tobytes()).hexdigest()
+
+
+def _clamp_index(index: int, start: int, end: int) -> int:
+    return max(start, min(end - 1, index))
+
+
+def _load_parquet_rows(path: Path) -> list[dict[str, Any]]:
+    return pq.read_table(path).to_pylist()
+
+
+def _load_json(path: Path) -> dict[str, Any]:
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _stack_stats(stats: dict[str, Any], keys: list[str]) -> dict[str, torch.Tensor]:
+    result = {}
+    for stat_name in ("mean", "std"):
+        values = []
+        for key in keys:
+            values.extend(stats[key][stat_name])
+        result[stat_name] = torch.tensor(values, dtype=torch.float32)
+    return result
+
+
+def normalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    denom = torch.where(stats["std"] == 0, torch.ones_like(stats["std"]), stats["std"])
+    return (values - stats["mean"]) / denom
+
+
+def unnormalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    return values * stats["std"] + stats["mean"]
+
+
+class TorchvisionVideoReaderCache:
+    def __init__(self, backend: str = "pyav") -> None:
+        self.backend = backend
+        self._readers: dict[str, Any] = {}
+        torchvision.set_video_backend(backend)
+
+    def get(self, path: str) -> Any:
+        reader = self._readers.get(path)
+        if reader is None:
+            reader = torchvision.io.VideoReader(path, "video")
+            self._readers[path] = reader
+        return reader
+
+    def decode_frames(self, path: str, timestamps: list[float], tolerance_s: float = 1e-4) -> torch.Tensor:
+        reader = self.get(path)
+        first_ts = min(timestamps)
+        last_ts = max(timestamps)
+        reader.seek(first_ts, keyframes_only=self.backend == "pyav")
+
+        loaded_frames: list[torch.Tensor] = []
+        loaded_ts: list[float] = []
+        for frame in reader:
+            current_ts = float(frame["pts"])
+            loaded_frames.append(frame["data"])
+            loaded_ts.append(current_ts)
+            if current_ts >= last_ts:
+                break
+
+        query_ts = torch.tensor(timestamps, dtype=torch.float32)
+        loaded_ts_tensor = torch.tensor(loaded_ts, dtype=torch.float32)
+        distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1)
+        min_dist, argmin = distances.min(dim=1)
+        if not torch.all(min_dist < tolerance_s):
+            raise RuntimeError(
+                f"Video timestamps are outside tolerance: query={query_ts.tolist()} "
+                f"loaded={loaded_ts_tensor.tolist()} path={path}"
+            )
+        return torch.stack([loaded_frames[i] for i in argmin]).float() / 255.0
+
+
+@dataclass
+class A2DOpenLoopSample:
+    index: int
+    episode_index: int
+    task: str
+    state_raw: torch.Tensor
+    action_raw: torch.Tensor
+    inputs: dict[str, torch.Tensor]
+
+
+class A2DOpenLoopDataset:
+    image_keys = [
+        "observation.images.head",
+        "observation.images.hand_left",
+        "observation.images.hand_right",
+    ]
+    state_keys = [
+        "observation.states.joint.position",
+        "observation.states.effector.position",
+    ]
+    action_keys = [
+        "actions.joint.position",
+        "actions.effector.position",
+    ]
+
+    def __init__(
+        self,
+        dataset_root: str | Path,
+        *,
+        config: InternVLAA1Config,
+        train_stats: dict[str, Any],
+        image_offsets: tuple[int, int] = (-15, 0),
+        tolerance_s: float = 1e-4,
+    ) -> None:
+        self.root = Path(dataset_root)
+        self.config = config
+        self.info = _load_json(self.root / "meta" / "info.json")
+        self.dataset_stats = _load_json(self.root / "meta" / "stats.json")
+        self.data_rows = _load_parquet_rows(self.root / "data" / "chunk-000" / "file-000.parquet")
+        self.episode_rows = _load_parquet_rows(self.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
+        self.task_rows = _load_parquet_rows(self.root / "meta" / "tasks.parquet")
+
+        self.state_stats = _stack_stats(train_stats, self.state_keys)
+        self.action_stats = _stack_stats(train_stats, self.action_keys)
+        self.image_offsets = image_offsets
+        self.tolerance_s = tolerance_s
+        self.video_reader = TorchvisionVideoReaderCache(backend="pyav")
+
+    @property
+    def num_episodes(self) -> int:
+        return len(self.episode_rows)
+
+    @property
+    def physical_action_dim(self) -> int:
+        return 16
+
+    def episode_start_indices(self, max_episodes: int | None = None) -> list[tuple[int, list[int]]]:
+        rows = self.episode_rows if max_episodes is None else self.episode_rows[:max_episodes]
+        result = []
+        for ep in rows:
+            start = int(ep["dataset_from_index"])
+            end = int(ep["dataset_to_index"])
+            result.append((int(ep["episode_index"]), list(range(start, end, self.config.chunk_size))))
+        return result
+
+    def _task_text(self, task_index: int) -> str:
+        return self.task_rows[task_index]["__index_level_0__"]
+
+    def _episode_for_index(self, row: dict[str, Any]) -> dict[str, Any]:
+        return self.episode_rows[int(row["episode_index"])]
+
+    def _state_vector(self, row: dict[str, Any]) -> torch.Tensor:
+        return torch.tensor(row[self.state_keys[0]] + row[self.state_keys[1]], dtype=torch.float32)
+
+    def _action_vector(self, row: dict[str, Any]) -> torch.Tensor:
+        return torch.tensor(row[self.action_keys[0]] + row[self.action_keys[1]], dtype=torch.float32)
+
+    def _query_rows(self, idx: int, deltas: list[int]) -> list[dict[str, Any]]:
+        row = self.data_rows[idx]
+        episode = self._episode_for_index(row)
+        start = int(episode["dataset_from_index"])
+        end = int(episode["dataset_to_index"])
+        return [self.data_rows[_clamp_index(idx + delta, start, end)] for delta in deltas]
+
+    def _decode_camera_history(
+        self, episode: dict[str, Any], camera_key: str, rows: list[dict[str, Any]]
+    ) -> torch.Tensor:
+        timestamps = [float(r["timestamp"]) for r in rows]
+        shifted = [float(episode[f"videos/{camera_key}/from_timestamp"]) + ts for ts in timestamps]
+        chunk_idx = int(episode[f"videos/{camera_key}/chunk_index"])
+        file_idx = int(episode[f"videos/{camera_key}/file_index"])
+        path = self.root / self.info["video_path"].format(
+            video_key=camera_key,
+            chunk_index=chunk_idx,
+            file_index=file_idx,
+        )
+        frames = self.video_reader.decode_frames(str(path), shifted, tolerance_s=self.tolerance_s)
+        return frames
+
+    def get_sample(self, idx: int) -> A2DOpenLoopSample:
+        row = self.data_rows[idx]
+        episode = self._episode_for_index(row)
+        image_rows = self._query_rows(idx, list((-15, 0)))
+        action_rows = self._query_rows(idx, list(range(self.config.chunk_size)))
+        camera_images = [self._decode_camera_history(episode, camera_key, image_rows) for camera_key in self.image_keys]
+        state_raw = self._state_vector(row)
+        state_norm = normalize_vector(state_raw, self.state_stats)
+        action_raw = torch.stack([self._action_vector(action_row) for action_row in action_rows], dim=0)
+        task = self._task_text(int(row["task_index"]))
+        inputs = {
+            OBS_STATE: state_norm,
+            OBS_TASK: task,
+            f"{OBS_IMAGES}.image0": camera_images[0],
+            f"{OBS_IMAGES}.image1": camera_images[1],
+            f"{OBS_IMAGES}.image2": camera_images[2],
+            f"{OBS_IMAGES}.image0_mask": torch.tensor(True),
+            f"{OBS_IMAGES}.image1_mask": torch.tensor(True),
+            f"{OBS_IMAGES}.image2_mask": torch.tensor(True),
+        }
+        return A2DOpenLoopSample(
+            index=idx,
+            episode_index=int(row["episode_index"]),
+            task=task,
+            state_raw=state_raw,
+            action_raw=action_raw,
+            inputs=inputs,
+        )
+
+
+def collate_open_loop_samples(
+    samples: list[A2DOpenLoopSample],
+    *,
+    device: str,
+    dtype: torch.dtype,
+) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+    first = samples[0]
+    batch_inputs: dict[str, torch.Tensor] = {}
+    for key in first.inputs:
+        values = [sample.inputs[key] for sample in samples]
+        if isinstance(values[0], torch.Tensor):
+            tensor = torch.stack(values, dim=0)
+            if tensor.dtype in (torch.int64, torch.bool):
+                batch_inputs[key] = tensor.to(device=device)
+            else:
+                batch_inputs[key] = tensor.to(device=device, dtype=dtype)
+        else:
+            batch_inputs[key] = values
+
+    metadata = {
+        "indices": [sample.index for sample in samples],
+        "episode_indices": [sample.episode_index for sample in samples],
+        "tasks": [sample.task for sample in samples],
+        "state_raw": torch.stack([sample.state_raw for sample in samples], dim=0),
+        "action_raw": torch.stack([sample.action_raw for sample in samples], dim=0),
+    }
+    return batch_inputs, metadata
+
+
+def plot_prediction_series(
+    *,
+    series: dict[str, np.ndarray],
+    output_path: Path,
+    title: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    first_series = next(iter(series.values()))
+    num_dims = int(first_series.shape[1])
+    num_cols = 2
+    num_rows = math.ceil(num_dims / num_cols)
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(16, max(6, 3 * num_rows)))
+    axs = np.array(axs).reshape(-1)
+    x_values = np.arange(first_series.shape[0])
+    styles = [
+        ("Ground Truth", "blue", "-", 1.5),
+        ("VLLM", "green", "-.", 1.4),
+        ("Predicted", "red", "--", 1.4),
+    ]
+    style_map = {name: (color, linestyle, linewidth) for name, color, linestyle, linewidth in styles}
+
+    for dim in range(num_dims):
+        ax = axs[dim]
+        for name, values in series.items():
+            color, linestyle, linewidth = style_map.get(name, (None, "-", 1.2))
+            ax.plot(
+                x_values,
+                values[:, dim],
+                label=name,
+                color=color,
+                linestyle=linestyle,
+                linewidth=linewidth,
+            )
+        ax.set_title(f"Dimension {dim + 1}")
+        ax.set_xlabel("Time Step / Sample Index")
+        ax.set_ylabel(f"Value Dim {dim + 1}")
+        ax.legend(loc="upper right")
+        ax.grid(True, linestyle="--", alpha=0.7)
+
+    for dim in range(num_dims, len(axs)):
+        axs[dim].axis("off")
+
+    fig.suptitle(title, fontsize=16)
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path)
+    plt.close(fig)
+
+
+def summarize_prediction_metrics(
+    *,
+    gt_tensor: torch.Tensor,
+    pred_tensor: torch.Tensor,
+    joint_dims: int = 14,
+) -> dict[str, float]:
+    return {
+        "mse": float(F.mse_loss(gt_tensor, pred_tensor, reduction="mean").item()),
+        "mae": float(F.l1_loss(gt_tensor, pred_tensor, reduction="mean").item()),
+        "mse_joint": float(F.mse_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
+        "mae_joint": float(F.l1_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
+        "mse_gripper": float(
+            F.mse_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
+        ),
+        "mae_gripper": float(
+            F.l1_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
+        ),
+    }
+
+
+def summarize_metric_list(values: list[float]) -> float | None:
+    return float(np.mean(values)) if values else None
+
+
+def run_open_loop_evaluation(
+    *,
+    mode: str,
+    policy,
+    config,
+    dataset,
+    train_meta,
+    collate_samples: Callable[..., tuple[dict[str, torch.Tensor], dict[str, Any]]],
+    run_sample_actions: Callable[[Any, dict[str, torch.Tensor], torch.Tensor], torch.Tensor],
+    num_episodes: int,
+    seed: int,
+    device: str,
+    dtype: torch.dtype,
+    output_dir: str | Path,
+    skip_plots: bool,
+) -> dict[str, Any]:
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    plots_dir = output_dir / "plots"
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    mse_values: list[float] = []
+    mae_values: list[float] = []
+    mse_joint_values: list[float] = []
+    mae_joint_values: list[float] = []
+    mse_gripper_values: list[float] = []
+    mae_gripper_values: list[float] = []
+    per_episode: list[dict[str, Any]] = []
+
+    plotting_available = True
+    if not skip_plots:
+        try:
+            import matplotlib.pyplot as _plt  # noqa: F401
+        except ModuleNotFoundError:
+            plotting_available = False
+            print("[warning] matplotlib is not installed; plots will be skipped.")
+
+    episode_specs = dataset.episode_start_indices(max_episodes=num_episodes)
+    for episode_index, indices in episode_specs:
+        print(f"[{mode}] episode: {episode_index}")
+        pred_chunks = []
+        gt_chunks = []
+        visited_indices = []
+        task = None
+
+        for index in indices:
+            sample = dataset.get_sample(index)
+            task = sample.task
+            batch_inputs, meta = collate_samples([sample], device=device, dtype=dtype)
+            noise = make_shared_noise(
+                seed,
+                index,
+                (batch_inputs[OBS_STATE].shape[0], config.chunk_size, config.max_action_dim),
+                device,
+            )
+            with torch.no_grad():
+                pred = run_sample_actions(policy, batch_inputs, noise)
+
+            pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+            pred_phys = unnormalize_vector(pred, dataset.action_stats)
+            if train_meta.action_mode == "delta":
+                pred_phys[:, :, :14] += meta["state_raw"][:, None, :14]
+
+            gt_phys = meta["action_raw"][:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+            pred_chunks.append(pred_phys[0])
+            gt_chunks.append(gt_phys[0])
+            visited_indices.append(index)
+
+        pred_tensor = torch.cat(pred_chunks, dim=0)
+        gt_tensor = torch.cat(gt_chunks, dim=0)
+        metrics = summarize_prediction_metrics(gt_tensor=gt_tensor, pred_tensor=pred_tensor)
+
+        mse_values.append(metrics["mse"])
+        mae_values.append(metrics["mae"])
+        mse_joint_values.append(metrics["mse_joint"])
+        mae_joint_values.append(metrics["mae_joint"])
+        mse_gripper_values.append(metrics["mse_gripper"])
+        mae_gripper_values.append(metrics["mae_gripper"])
+
+        if not skip_plots and plotting_available:
+            plot_prediction_series(
+                series={
+                    "Ground Truth": gt_tensor.numpy(),
+                    "VLLM": pred_tensor.numpy(),
+                },
+                output_path=plots_dir / f"{mode}_open_loop_ep{episode_index}.jpg",
+                title=f"{mode} Ground Truth vs VLLM",
+            )
+
+        episode_log = {
+            "episode_id": int(episode_index),
+            "task": task,
+            "visited_indices": visited_indices,
+            "num_pred_steps": int(pred_tensor.shape[0]),
+            **metrics,
+        }
+        per_episode.append(episode_log)
+        print(json.dumps({"mode": mode, **episode_log}, ensure_ascii=False))
+
+    summary = {
+        "mode": mode,
+        "num_episodes": len(per_episode),
+        "mse": mse_values,
+        "mae": mae_values,
+        "average_mse": summarize_metric_list(mse_values),
+        "average_mae": summarize_metric_list(mae_values),
+        "mse_joint": mse_joint_values,
+        "mae_joint": mae_joint_values,
+        "average_mse_joint": summarize_metric_list(mse_joint_values),
+        "average_mae_joint": summarize_metric_list(mae_joint_values),
+        "mse_gripper": mse_gripper_values,
+        "mae_gripper": mae_gripper_values,
+        "average_mse_gripper": summarize_metric_list(mse_gripper_values),
+        "average_mae_gripper": summarize_metric_list(mae_gripper_values),
+        "episodes": per_episode,
+    }
+    with open(output_dir / "log.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, indent=2, ensure_ascii=False)
+    return summary

examples/offline_inference/internvla_a1/run.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
+ : "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
+ : "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
+ : "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
+ : "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
+INTERNVLA_A1_OUTPUT_DIR="${INTERNVLA_A1_OUTPUT_DIR:-$REPO_ROOT/outputs/internvla_a1/vllm_infer}"
+
+export INTERNVLA_A1_MODEL_DIR
+export INTERNVLA_A1_DATASET_DIR
+export INTERNVLA_A1_PROCESSOR_DIR
+export INTERNVLA_A1_COSMOS_DIR
+
+python "$ROOT_DIR/end2end.py" \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --output-dir "$INTERNVLA_A1_OUTPUT_DIR" \
+  --num-episodes "${INTERNVLA_A1_NUM_EPISODES:-1}" \
+  --attn-implementation eager \
+  "$@"

tests/e2e/offline_inference/test_internvla_a1.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Offline E2E test for InternVLA-A1 open-loop inference."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+EXAMPLE_SCRIPT = REPO_ROOT / "examples" / "offline_inference" / "internvla_a1" / "end2end.py"
+
+
+def _required_env(name: str) -> str:
+    value = os.getenv(name)
+    if not value:
+        pytest.skip(f"{name} is not set")
+    return value
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_internvla_a1_offline_open_loop(run_level: str) -> None:
+    if run_level != "advanced_model":
+        pytest.skip("InternVLA-A1 offline evaluation requires real local checkpoints and dataset.")
+
+    model_dir = _required_env("INTERNVLA_A1_MODEL_DIR")
+    dataset_dir = _required_env("INTERNVLA_A1_DATASET_DIR")
+    processor_dir = _required_env("INTERNVLA_A1_PROCESSOR_DIR")
+    cosmos_dir = _required_env("INTERNVLA_A1_COSMOS_DIR")
+
+    env = os.environ.copy()
+    env["INTERNVLA_A1_PROCESSOR_DIR"] = processor_dir
+    env["INTERNVLA_A1_COSMOS_DIR"] = cosmos_dir
+
+    with tempfile.TemporaryDirectory(prefix="internvla_a1_e2e_") as tmpdir:
+        output_dir = Path(tmpdir) / "outputs"
+        cmd = [
+            sys.executable,
+            str(EXAMPLE_SCRIPT),
+            "--model-dir",
+            model_dir,
+            "--dataset-dir",
+            dataset_dir,
+            "--output-dir",
+            str(output_dir),
+            "--num-samples",
+            "1",
+            "--num-episodes",
+            "1",
+            "--dtype",
+            "float32",
+            "--attn-implementation",
+            "eager",
+            "--skip-plots",
+        ]
+        subprocess.run(cmd, check=True, env=env)
+
+        summary_path = output_dir / "summary.json"
+        assert summary_path.exists(), f"Missing summary file: {summary_path}"
+
+        with open(summary_path, encoding="utf-8") as f:
+            summary = json.load(f)
+
+        assert summary["mode"] == "registry"
+        assert summary["attn_implementation"] == "eager"
+        assert summary["dtype"] == "float32"
+        assert len(summary["results"]) == 1
+        assert "registry" in summary["eval_summaries"]
+
+        eval_summary = summary["eval_summaries"]["registry"]
+        assert eval_summary["num_episodes"] == 1
+        assert eval_summary["average_mse"] is not None
+        assert eval_summary["average_mae"] is not None
+
+        log_path = output_dir / "registry" / "log.json"
+        assert log_path.exists(), f"Missing evaluation log: {log_path}"

vllm_omni/diffusion/models/internvla_a1/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""InternVLA-A1 diffusion model components."""
+
+from vllm_omni.diffusion.models.internvla_a1.config import (
+    InternVLAA1Config,
+    InternVLAA1TrainMetadata,
+)
+from vllm_omni.diffusion.models.internvla_a1.model_internvla_a1 import (
+    InternVLAA1,
+    InternVLAA1Policy,
+)
+from vllm_omni.diffusion.models.internvla_a1.pipeline_internvla_a1 import (
+    InternVLAA1Pipeline,
+    get_internvla_a1_post_process_func,
+)
+
+__all__ = [
+    "InternVLAA1",
+    "InternVLAA1Config",
+    "InternVLAA1Pipeline",
+    "InternVLAA1TrainMetadata",
+    "InternVLAA1Policy",
+    "get_internvla_a1_post_process_func",
+]

vllm_omni/diffusion/models/internvla_a1/adapter_qwen3_vl.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    ALL_ATTENTION_FUNCTIONS,
+    BaseModelOutputWithPast,
+    Cache,
+    DynamicCache,
+    FlashAttentionKwargs,
+    Qwen3VLConfig,
+    Qwen3VLPreTrainedModel,
+    Qwen3VLTextConfig,
+    Qwen3VLTextRotaryEmbedding,
+    Qwen3VLVisionModel,
+    Unpack,
+    apply_rotary_pos_emb,
+    check_model_inputs,
+    create_causal_mask,
+    deprecate_kwarg,
+    eager_attention_forward,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLForConditionalGeneration as HFQwen3VLForConditionalGeneration,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLModel as HFQwen3VLModel,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextAttention as HFQwen3VLTextAttention,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextDecoderLayer as HFQwen3VLTextDecoderLayer,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextModel as HFQwen3VLTextModel,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextRMSNorm as HFQwen3VLTextRMSNorm,
+)
+
+
+class Qwen3VLTextRMSNorm(HFQwen3VLTextRMSNorm):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+class Qwen3VLTextAttention(HFQwen3VLTextAttention):
+    def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.q_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            if use_cache:
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_values.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs
+                )
+            else:
+                key_states = torch.cat([past_key_values[self.layer_idx][0], key_states], dim=2)
+                value_states = torch.cat([past_key_values[self.layer_idx][1], value_states], dim=2)
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Qwen3VLTextDecoderLayer(HFQwen3VLTextDecoderLayer):
+    def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx)
+        self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+class Qwen3VLTextModel(HFQwen3VLTextModel):
+    config: Qwen3VLTextConfig
+    _no_split_modules = ["Qwen3VLTextDecoderLayer"]
+    _repeated_blocks = ["Qwen3VLTextDecoderLayer"]
+
+    def __init__(self, config: Qwen3VLTextConfig):
+        Qwen3VLPreTrainedModel.__init__(self, config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    @check_model_inputs
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        visual_pos_masks: torch.Tensor | None = None,
+        deepstack_visual_embeds: list[torch.Tensor] | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple | BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if use_cache and past_key_values is None and not torch.jit.is_tracing():
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+        elif position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+            text_position_ids = position_ids[0]
+            position_ids = position_ids[1:]
+        else:
+            text_position_ids = position_ids[0]
+
+        attention_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=text_position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+            if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+                hidden_states = self._deepstack_process(
+                    hidden_states,
+                    visual_pos_masks,
+                    deepstack_visual_embeds[layer_idx],
+                )
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class Qwen3VLModel(HFQwen3VLModel):
+    base_model_prefix = ""
+    _checkpoint_conversion_mapping = {}
+    accepts_loss_kwargs = False
+    config: Qwen3VLConfig
+    _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
+
+    def __init__(self, config):
+        Qwen3VLPreTrainedModel.__init__(self, config)
+        self.visual = Qwen3VLVisionModel._from_config(config.vision_config)
+        self.language_model = Qwen3VLTextModel._from_config(config.text_config)
+        self.rope_deltas = None
+        self.post_init()
+
+
+class Qwen3VLForConditionalGeneration(HFQwen3VLForConditionalGeneration):
+    _checkpoint_conversion_mapping = {}
+    _tied_weights_keys = ["lm_head.weight"]
+    accepts_loss_kwargs = False
+    config: Qwen3VLConfig
+
+    def __init__(self, config):
+        Qwen3VLPreTrainedModel.__init__(self, config)
+        self.model = Qwen3VLModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+
+__all__ = [
+    "Qwen3VLVisionModel",
+    "Qwen3VLForConditionalGeneration",
+    "Qwen3VLModel",
+    "Qwen3VLPreTrainedModel",
+    "Qwen3VLTextModel",
+    "apply_rotary_pos_emb",
+    "eager_attention_forward",
+]

vllm_omni/diffusion/models/internvla_a1/config.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from dataclasses import fields as dataclass_fields
+from pathlib import Path
+from typing import Any
+
+from huggingface_hub.constants import HF_HOME
+
+ACTION = "action"
+OBS_STR = "observation"
+OBS_PREFIX = OBS_STR + "."
+OBS_STATE = OBS_STR + ".state"
+OBS_TASK = OBS_STR + ".task"
+OBS_IMAGE = OBS_STR + ".image"
+OBS_IMAGES = OBS_IMAGE + "s"
+OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38
+
+DEFAULT_QWEN3_VL_MODEL = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", "Qwen/Qwen3-VL-2B-Instruct")
+DEFAULT_COSMOS_REPO = "nvidia/Cosmos-Tokenizer-CI8x8"
+DEFAULT_COSMOS_DIR = Path(
+    os.getenv(
+        "INTERNVLA_A1_COSMOS_DIR",
+        str(Path(HF_HOME) / "hub" / "Cosmos-Tokenizer-CI8x8"),
+    )
+)
+
+
+@dataclass
+class InternVLAA1Config:
+    """Standalone-compatible InternVLA-A1 config with a few fake-smoke defaults."""
+
+    type: str = "internvla_a1"
+    qwen3_vl_variant: str = "qwen3_vl_28l"
+    action_expert_variant: str = "qwen3_28l"
+    dtype: str = "bfloat16"
+    device: str = "cuda"
+    chunk_size: int = 50
+    n_action_steps: int = 50
+    max_state_dim: int = 32
+    max_action_dim: int = 32
+    num_inference_steps: int = 10
+    time_sampling_beta_alpha: float = 1.5
+    time_sampling_beta_beta: float = 1.0
+    time_sampling_scale: float = 0.999
+    time_sampling_offset: float = 0.001
+    min_period: float = 4e-3
+    max_period: float = 4.0
+    image_resolution: tuple[int, int] = (224, 224)
+    empty_cameras: int = 0
+    gradient_checkpointing: bool = False
+    compile_model: bool = False
+    compile_mode: str = "max-autotune"
+    enable_regional_compile: bool = False
+    regional_compile_dynamic: bool = True
+    enable_suffix_static_context_optimization: bool = False
+    attn_implementation: str = "eager"
+    tokenizer_max_length: int = 48
+    freeze_vision_encoder: bool = False
+    train_expert_only: bool = False
+    train_vlm_only: bool = False
+    scale_factor: int = 8
+    lambda_gen: float = 0.01
+    input_features: dict[str, Any] = field(default_factory=dict)
+    output_features: dict[str, Any] = field(default_factory=dict)
+
+    # Fake/smoke-only knobs used by the lightweight pipeline path.
+    vocab_size: int = 256
+    hidden_size: int = 128
+    intermediate_size: int = 256
+    num_attention_heads: int = 4
+    num_hidden_layers: int = 2
+    num_cameras: int = 3
+    image_history: int = 2
+    pixel_feature_dim: int = 48
+
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str | Path) -> InternVLAA1Config:
+        checkpoint_dir = Path(checkpoint_dir)
+        with open(checkpoint_dir / "config.json", encoding="utf-8") as f:
+            raw = json.load(f)
+        return cls.from_model_config(raw)
+
+    @classmethod
+    def from_model_config(cls, model_config: dict[str, Any] | None) -> InternVLAA1Config:
+        if not model_config:
+            return cls()
+
+        raw = dict(model_config)
+        raw["image_resolution"] = tuple(raw.get("image_resolution", [224, 224]))
+        allowed = {item.name for item in dataclass_fields(cls)}
+        filtered = {key: value for key, value in raw.items() if key in allowed}
+        return cls(**filtered)
+
+
+@dataclass
+class InternVLAA1TrainMetadata:
+    action_mode: str = "delta"
+    processor_model_name: str = DEFAULT_QWEN3_VL_MODEL
+
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str | Path) -> InternVLAA1TrainMetadata:
+        checkpoint_dir = Path(checkpoint_dir)
+        path = checkpoint_dir / "train_config.json"
+        if not path.exists():
+            return cls()
+
+        with open(path, encoding="utf-8") as f:
+            raw = json.load(f)
+
+        dataset_cfg = raw.get("dataset", {})
+        processor_model_name = DEFAULT_QWEN3_VL_MODEL
+        for transform in dataset_cfg.get("data_transforms", {}).get("inputs", []):
+            if transform.get("type") == "internvla_a1_processor":
+                processor_model_name = transform.get(
+                    "pretrained_model_name_or_path",
+                    processor_model_name,
+                )
+                break
+
+        return cls(
+            action_mode=dataset_cfg.get("action_mode", "delta"),
+            processor_model_name=processor_model_name,
+        )

vllm_omni/diffusion/models/internvla_a1/cosmos_ci_torch.py

@@ -0,0 +1,416 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file vendors the minimal Continuous Image (CI) tokenizer implementation
+# needed by InternVLA-A1 from NVIDIA/Cosmos-Tokenizer. The code is intentionally
+# reduced to the CI image path only so the runtime does not depend on the
+# external `cosmos_tokenizer` package.
+
+from __future__ import annotations
+
+import math
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+_WAVELETS = {
+    "haar": torch.tensor([0.7071067811865476, 0.7071067811865476]),
+    "rearrange": torch.tensor([1.0, 1.0]),
+}
+_PERSISTENT = False
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    return x * torch.sigmoid(x)
+
+
+def normalize(in_channels: int, num_groups: int = 32) -> nn.GroupNorm:
+    return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class IdentityDistribution(nn.Module):
+    def forward(self, parameters: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        zeros = torch.tensor([0.0], device=parameters.device, dtype=parameters.dtype)
+        return parameters, (zeros, zeros)
+
+
+class Patcher(nn.Module):
+    def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_method = patch_method
+        self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT)
+        self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item()))
+        self.register_buffer(
+            "_arange",
+            torch.arange(_WAVELETS[patch_method].shape[0]),
+            persistent=_PERSISTENT,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_method == "haar":
+            return self._haar(x)
+        if self.patch_method == "rearrange":
+            return self._arrange(x)
+        raise ValueError(f"Unknown patch method: {self.patch_method}")
+
+    def _dwt(self, x: torch.Tensor, mode: str = "reflect", *, rescale: bool = False) -> torch.Tensor:
+        dtype = x.dtype
+        h = self.wavelets
+        n = h.shape[0]
+        g = x.shape[1]
+        hl = h.flip(0).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype)
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype)
+
+        x = F.pad(x, pad=(n - 2, n - 1, n - 2, n - 1), mode=mode).to(dtype)
+        xl = F.conv2d(x, hl.unsqueeze(2), groups=g, stride=(1, 2))
+        xh = F.conv2d(x, hh.unsqueeze(2), groups=g, stride=(1, 2))
+        xll = F.conv2d(xl, hl.unsqueeze(3), groups=g, stride=(2, 1))
+        xlh = F.conv2d(xl, hh.unsqueeze(3), groups=g, stride=(2, 1))
+        xhl = F.conv2d(xh, hl.unsqueeze(3), groups=g, stride=(2, 1))
+        xhh = F.conv2d(xh, hh.unsqueeze(3), groups=g, stride=(2, 1))
+
+        out = torch.cat([xll, xlh, xhl, xhh], dim=1)
+        if rescale:
+            out = out / 2
+        return out
+
+    def _haar(self, x: torch.Tensor) -> torch.Tensor:
+        for _ in self.range:
+            x = self._dwt(x, rescale=True)
+        return x
+
+    def _arrange(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b c (h p1) (w p2) -> b (c p1 p2) h w",
+            p1=self.patch_size,
+            p2=self.patch_size,
+        ).contiguous()
+
+
+class UnPatcher(nn.Module):
+    def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_method = patch_method
+        self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT)
+        self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item()))
+        self.register_buffer(
+            "_arange",
+            torch.arange(_WAVELETS[patch_method].shape[0]),
+            persistent=_PERSISTENT,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_method == "haar":
+            return self._ihaar(x)
+        if self.patch_method == "rearrange":
+            return self._iarrange(x)
+        raise ValueError(f"Unknown patch method: {self.patch_method}")
+
+    def _idwt(self, x: torch.Tensor, *, rescale: bool = False) -> torch.Tensor:
+        dtype = x.dtype
+        h = self.wavelets
+        n = h.shape[0]
+        g = x.shape[1] // 4
+        hl = h.flip([0]).reshape(1, 1, -1).repeat([g, 1, 1]).to(dtype=dtype)
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype)
+
+        xll, xlh, xhl, xhh = torch.chunk(x.to(dtype), 4, dim=1)
+        yl = F.conv_transpose2d(xll, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yl += F.conv_transpose2d(xlh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yh = F.conv_transpose2d(xhl, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yh += F.conv_transpose2d(xhh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        y = F.conv_transpose2d(yl, hl.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2))
+        y += F.conv_transpose2d(yh, hh.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2))
+
+        if rescale:
+            y = y * 2
+        return y
+
+    def _ihaar(self, x: torch.Tensor) -> torch.Tensor:
+        for _ in self.range:
+            x = self._idwt(x, rescale=True)
+        return x
+
+    def _iarrange(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (c p1 p2) h w -> b c (h p1) (w p2)",
+            p1=self.patch_size,
+            p2=self.patch_size,
+        )
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
+        return self.conv(x)
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels: int, out_channels: int | None = None, dropout: float, **_: object) -> None:
+        super().__init__()
+        out_channels = in_channels if out_channels is None else out_channels
+        self.norm1 = normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.nin_shortcut = (
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.conv1(nonlinearity(self.norm1(x)))
+        h = self.conv2(self.dropout(nonlinearity(self.norm2(h))))
+        return self.nin_shortcut(x) + h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.norm = normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        q = self.q(h)
+        k = self.k(h)
+        v = self.v(h)
+        b, c, height, width = q.shape
+        q = q.reshape(b, c, height * width).permute(0, 2, 1)
+        k = k.reshape(b, c, height * width)
+        attn = torch.bmm(q, k) * (int(c) ** -0.5)
+        attn = F.softmax(attn, dim=2)
+        v = v.reshape(b, c, height * width)
+        attn = attn.permute(0, 2, 1)
+        h = torch.bmm(v, attn).reshape(b, c, height, width)
+        return x + self.proj_out(h)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        channels: int,
+        channels_mult: list[int],
+        num_res_blocks: int,
+        attn_resolutions: list[int],
+        dropout: float,
+        resolution: int,
+        z_channels: int,
+        spatial_compression: int,
+        **kwargs: object,
+    ) -> None:
+        super().__init__()
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+
+        patch_size = int(kwargs.get("patch_size", 1))
+        self.patcher = Patcher(patch_size, str(kwargs.get("patch_method", "rearrange")))
+        in_channels = in_channels * patch_size * patch_size
+        self.num_downsamples = int(math.log2(spatial_compression)) - int(math.log2(patch_size))
+        assert self.num_downsamples <= self.num_resolutions
+
+        self.conv_in = nn.Conv2d(in_channels, channels, kernel_size=3, stride=1, padding=1)
+        curr_res = resolution // patch_size
+        in_ch_mult = (1,) + tuple(channels_mult)
+        self.down = nn.ModuleList()
+        block_in = channels
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = channels * in_ch_mult[i_level]
+            block_out = channels * channels_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level < self.num_downsamples:
+                down.downsample = Downsample(block_in)
+                curr_res //= 2
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+        self.norm_out = normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patcher(x)
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level < self.num_downsamples:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = nonlinearity(self.norm_out(h))
+        return self.conv_out(h)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        channels: int,
+        channels_mult: list[int],
+        num_res_blocks: int,
+        attn_resolutions: list[int],
+        dropout: float,
+        resolution: int,
+        z_channels: int,
+        spatial_compression: int,
+        **kwargs: object,
+    ) -> None:
+        super().__init__()
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+
+        patch_size = int(kwargs.get("patch_size", 1))
+        self.unpatcher = UnPatcher(patch_size, str(kwargs.get("patch_method", "rearrange")))
+        out_ch = out_channels * patch_size * patch_size
+        self.num_upsamples = int(math.log2(spatial_compression)) - int(math.log2(patch_size))
+        assert self.num_upsamples <= self.num_resolutions
+
+        block_in = channels * channels_mult[self.num_resolutions - 1]
+        curr_res = (resolution // patch_size) // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = channels * channels_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level >= (self.num_resolutions - self.num_upsamples):
+                up.upsample = Upsample(block_in)
+                curr_res *= 2
+            self.up.insert(0, up)
+
+        self.norm_out = normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        h = self.conv_in(z)
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level >= (self.num_resolutions - self.num_upsamples):
+                h = self.up[i_level].upsample(h)
+        h = self.conv_out(nonlinearity(self.norm_out(h)))
+        return self.unpatcher(h)
+
+
+class ContinuousImageTokenizer(nn.Module):
+    def __init__(self, z_channels: int, z_factor: int, latent_channels: int, **kwargs: object) -> None:
+        super().__init__()
+        self.latent_channels = latent_channels
+        self.encoder = Encoder(z_channels=z_factor * z_channels, **kwargs)
+        self.decoder = Decoder(z_channels=z_channels, **kwargs)
+        self.quant_conv = nn.Conv2d(z_factor * z_channels, z_factor * latent_channels, 1)
+        self.post_quant_conv = nn.Conv2d(latent_channels, z_channels, 1)
+        self.distribution = IdentityDistribution()
+
+    def encoder_module(self) -> nn.Sequential:
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    ("encoder", self.encoder),
+                    ("quant_conv", self.quant_conv),
+                    ("distribution", self.distribution),
+                ]
+            )
+        )
+
+    def decoder_module(self) -> nn.Sequential:
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    ("post_quant_conv", self.post_quant_conv),
+                    ("decoder", self.decoder),
+                ]
+            )
+        )
+
+    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        return self.distribution(moments)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.post_quant_conv(z))
+
+
+def build_cosmos_ci_torch_model(spatial_compression: int) -> ContinuousImageTokenizer:
+    return ContinuousImageTokenizer(
+        attn_resolutions=[32],
+        channels=128,
+        channels_mult=[2, 4, 4],
+        dropout=0.0,
+        in_channels=3,
+        spatial_compression=spatial_compression,
+        num_res_blocks=2,
+        out_channels=3,
+        resolution=1024,
+        patch_size=4,
+        patch_method="haar",
+        latent_channels=16,
+        z_channels=16,
+        z_factor=1,
+        name="CI",
+    )

vllm_omni/diffusion/models/internvla_a1/model_cosmos.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from .cosmos_ci_torch import build_cosmos_ci_torch_model
+
+_SAFETENSORS_EXTENSION = ".safetensors"
+
+
+def is_safetensors_checkpoint(path: str | Path) -> bool:
+    return Path(path).suffix.lower() == _SAFETENSORS_EXTENSION
+
+
+def infer_cosmos_ci_spatial_compression(path: str | Path) -> int:
+    normalized = str(path).upper()
+    if "CI8X8" in normalized or "8X8" in normalized:
+        return 8
+    if "CI16X16" in normalized or "16X16" in normalized:
+        return 16
+    raise ValueError(f"Unable to infer Cosmos CI spatial compression from checkpoint path: {path}")
+
+
+def _load_safetensors_checkpoint(path: str | Path, device: str) -> dict[str, torch.Tensor]:
+    return load_file(str(path), device=device)
+
+
+def _load_native_cosmos_component(
+    checkpoint_filepath: str | Path,
+    *,
+    component: str,
+    device: str,
+) -> torch.nn.Module:
+    full_model = build_cosmos_ci_torch_model(
+        spatial_compression=infer_cosmos_ci_spatial_compression(checkpoint_filepath)
+    )
+    if component == "encoder":
+        model = full_model.encoder_module()
+    elif component == "decoder":
+        model = full_model.decoder_module()
+    else:
+        raise ValueError(f"Unsupported Cosmos component: {component}")
+
+    if not is_safetensors_checkpoint(checkpoint_filepath):
+        raise ValueError(f"Native Cosmos torch backend expects `.safetensors` checkpoints. Got: {checkpoint_filepath}")
+    state_dict = _load_safetensors_checkpoint(checkpoint_filepath, device)
+    model.load_state_dict(state_dict, strict=False)
+    return model.eval().to(device)
+
+
+def load_cosmos_component(
+    checkpoint_filepath: str | Path,
+    *,
+    component: str,
+    device: str = "cuda",
+) -> torch.nn.Module:
+    return _load_native_cosmos_component(
+        checkpoint_filepath,
+        component=component,
+        device=device,
+    )
+
+
+class ImageTokenizer(torch.nn.Module):
+    def __init__(
+        self,
+        checkpoint_enc: str,
+        checkpoint_dec: str,
+        device: str = "cuda",
+    ) -> None:
+        super().__init__()
+        self._enc_model = load_cosmos_component(
+            checkpoint_enc,
+            component="encoder",
+            device=device,
+        )
+        self._dec_model = load_cosmos_component(
+            checkpoint_dec,
+            component="decoder",
+            device=device,
+        )
+
+    def decode(self, input_latent: torch.Tensor) -> torch.Tensor:
+        original_dtype = input_latent.dtype
+        model_dtype = next(self.parameters()).dtype
+        output_tensor = self._dec_model(input_latent.to(model_dtype))
+        return output_tensor.to(original_dtype)
+
+    def encode(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        original_dtype = input_tensor.dtype
+        model_dtype = next(self.parameters()).dtype
+        output_latent = self._enc_model(input_tensor.to(model_dtype))
+        if isinstance(output_latent, torch.Tensor):
+            return output_latent.to(original_dtype)
+        return output_latent[0].to(original_dtype)

vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py

@@ -0,0 +1,962 @@
+from __future__ import annotations
+
+import math
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+import torch
+import torch._dynamo as dynamo
+import torch.nn.functional as F
+from einops import rearrange
+from safetensors import safe_open
+from torch import Tensor, nn
+from transformers import Qwen3VLProcessor
+from transformers.models.auto import CONFIG_MAPPING
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from .adapter_qwen3_vl import (
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLTextModel,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+from .config import (
+    DEFAULT_COSMOS_DIR,
+    DEFAULT_COSMOS_REPO,
+    DEFAULT_QWEN3_VL_MODEL,
+    OBS_IMAGES,
+    OBS_PREFIX,
+    OBS_STATE,
+    OBS_TASK,
+    OPENPI_ATTENTION_MASK_VALUE,
+    InternVLAA1Config,
+)
+from .model_cosmos import ImageTokenizer
+
+
+def get_safe_dtype(target_dtype: torch.dtype, device_type: str) -> torch.dtype:
+    if device_type == "mps" and target_dtype == torch.float64:
+        return torch.float32
+    if device_type == "cpu" and target_dtype == torch.bfloat16:
+        return torch.float32
+    return target_dtype
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.Tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device: torch.device,
+) -> Tensor:
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+    dtype = get_safe_dtype(torch.float64, device.type)
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None]
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+
+
+def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    return att_2d_masks & pad_2d_masks
+
+
+def pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
+    if vector.shape[-1] >= new_dim:
+        return vector
+    return F.pad(vector, (0, new_dim - vector.shape[-1]))
+
+
+def resize_with_pad(images: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
+    target_h, target_w = size
+    if images.ndim != 4:
+        raise ValueError(f"Expected [T, C, H, W], got {tuple(images.shape)}")
+    _, _, src_h, src_w = images.shape
+    scale = min(target_h / src_h, target_w / src_w)
+    resized_h = max(1, int(round(src_h * scale)))
+    resized_w = max(1, int(round(src_w * scale)))
+    resized = F.interpolate(images, size=(resized_h, resized_w), mode="bilinear", align_corners=False)
+    pad_h = target_h - resized_h
+    pad_w = target_w - resized_w
+    pad_top = pad_h // 2
+    pad_bottom = pad_h - pad_top
+    pad_left = pad_w // 2
+    pad_right = pad_w - pad_left
+    return F.pad(resized, (pad_left, pad_right, pad_top, pad_bottom))
+
+
+class Qwen3VLInputBuilder:
+    def __init__(
+        self,
+        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
+        max_length: int = 48,
+        spatial_merge_size: int = 2,
+    ) -> None:
+        processor_model_name = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", processor_model_name)
+        local_path = Path(processor_model_name).expanduser()
+        if local_path.exists():
+            self.processor = Qwen3VLProcessor.from_pretrained(str(local_path), local_files_only=True)
+        else:
+            if os.getenv("INTERNVLA_A1_PROCESSOR_DIR"):
+                raise FileNotFoundError(f"INTERNVLA_A1_PROCESSOR_DIR points to a missing path: {local_path}")
+            self.processor = Qwen3VLProcessor.from_pretrained(processor_model_name)
+        self.max_length = max_length
+        self.spatial_merge_size = spatial_merge_size
+
+    def build(self, camera_images: list[torch.Tensor], task: str) -> dict[str, torch.Tensor]:
+        input_ids: list[int] = []
+        attention_mask: list[int] = []
+        pixel_values: list[torch.Tensor] = []
+        image_grid_thw: list[torch.Tensor] = []
+
+        for image_history in camera_images:
+            current_image = image_history[1]
+            image_inputs = self.processor.image_processor(current_image, do_rescale=False)
+            pixel_values.append(image_inputs.pixel_values)
+            image_grid_thw.append(image_inputs.image_grid_thw)
+            num_img_token = int(torch.prod(image_inputs.image_grid_thw[0]).item()) // (self.spatial_merge_size**2)
+            input_ids.extend(
+                [self.processor.vision_start_token_id]
+                + [self.processor.image_token_id] * num_img_token
+                + [self.processor.vision_end_token_id]
+            )
+            attention_mask.extend([1] * (num_img_token + 2))
+
+        text_inputs = self.processor.tokenizer(
+            task,
+            max_length=self.max_length,
+            padding="max_length",
+            padding_side="right",
+            truncation=True,
+        )
+        input_ids.extend(text_inputs.input_ids)
+        attention_mask.extend(text_inputs.attention_mask)
+
+        return {
+            f"{OBS_PREFIX}pixel_values": torch.cat(pixel_values, dim=0),
+            f"{OBS_PREFIX}image_grid_thw": torch.cat(image_grid_thw, dim=0),
+            f"{OBS_PREFIX}input_ids": torch.tensor(input_ids, dtype=torch.long),
+            f"{OBS_PREFIX}attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+        }
+
+
+@dataclass
+class SuffixStaticContext:
+    state_emb: torch.Tensor
+    full_att_2d_masks_4d: torch.Tensor
+    position_ids: torch.Tensor
+
+
+def compute_layer_complete(
+    layer_idx: int,
+    inputs_embeds: list[torch.Tensor],
+    attention_mask: torch.Tensor,
+    position_ids: torch.Tensor,
+    und_expert: nn.Module,
+    gen_expert: nn.Module,
+    act_expert: nn.Module,
+) -> list[torch.Tensor]:
+    models = [und_expert.language_model, gen_expert, act_expert]
+    query_states = []
+    key_states = []
+    value_states = []
+    for i, hidden_states in enumerate(inputs_embeds):
+        layer = models[i].layers[layer_idx]
+        hidden_states = layer.input_layernorm(hidden_states)
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+        if layer.self_attn.q_proj.weight.dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype=torch.bfloat16)
+        query_state = layer.self_attn.q_norm(layer.self_attn.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_state = layer.self_attn.k_norm(layer.self_attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        query_states.append(query_state)
+        key_states.append(key_state)
+        value_states.append(value_state)
+
+    query_states = torch.cat(query_states, dim=2)
+    key_states = torch.cat(key_states, dim=2)
+    value_states = torch.cat(value_states, dim=2)
+    dummy_tensor = torch.zeros(
+        query_states.shape[0],
+        query_states.shape[2],
+        query_states.shape[-1],
+        device=query_states.device,
+        dtype=query_states.dtype,
+    )
+    cos, sin = und_expert.model.language_model.rotary_emb(dummy_tensor, position_ids)
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states,
+        key_states,
+        cos,
+        sin,
+        unsqueeze_dim=1,
+    )
+    scaling = und_expert.language_model.layers[layer_idx].self_attn.scaling
+    att_output, _ = eager_attention_forward(
+        und_expert.language_model.layers[layer_idx].self_attn,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        scaling,
+    )
+
+    head_dim = und_expert.language_model.layers[layer_idx].self_attn.head_dim
+    num_attention_heads = und_expert.language_model.layers[layer_idx].self_attn.config.num_attention_heads
+    batch_size = query_states.shape[0]
+    att_output = att_output.reshape(batch_size, -1, num_attention_heads * head_dim)
+
+    outputs_embeds = []
+    start_pos = 0
+    for i, hidden_states in enumerate(inputs_embeds):
+        layer = models[i].layers[layer_idx]
+        end_pos = start_pos + hidden_states.shape[1]
+        att_chunk = att_output[:, start_pos:end_pos]
+        if att_chunk.dtype != layer.self_attn.o_proj.weight.dtype:
+            att_chunk = att_chunk.to(layer.self_attn.o_proj.weight.dtype)
+        out_emb = layer.self_attn.o_proj(att_chunk)
+        out_emb = out_emb + hidden_states
+        residual = out_emb
+        out_emb = layer.post_attention_layernorm(out_emb)
+        if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
+            out_emb = out_emb.to(dtype=torch.bfloat16)
+        out_emb = layer.mlp(out_emb)
+        out_emb = out_emb + residual
+        outputs_embeds.append(out_emb)
+        start_pos = end_pos
+    return outputs_embeds
+
+
+class QwenConfig:
+    def __init__(
+        self,
+        head_dim: int,
+        hidden_size: int,
+        intermediate_size: int,
+        num_attention_heads: int,
+        num_hidden_layers: int,
+        num_key_value_heads: int,
+    ) -> None:
+        self.head_dim = head_dim
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.num_key_value_heads = num_key_value_heads
+
+
+def get_qwen_config(variant: str) -> QwenConfig:
+    num_hidden_layers = int(variant.split("_")[-1][:-1])
+    if variant.startswith("qwen3_vl"):
+        return QwenConfig(128, 2048, 6144, 16, num_hidden_layers, 8)
+    if variant.startswith("qwen3"):
+        return QwenConfig(128, 1024, 3072, 16, num_hidden_layers, 8)
+    raise ValueError(f"Unknown variant: {variant}")
+
+
+def resolve_cosmos_checkpoint_paths(
+    *,
+    encoder_path: str | Path | None = None,
+    decoder_path: str | Path | None = None,
+) -> tuple[Path, Path]:
+    encoder_override = encoder_path or os.getenv("INTERNVLA_A1_COSMOS_ENCODER_PATH")
+    decoder_override = decoder_path or os.getenv("INTERNVLA_A1_COSMOS_DECODER_PATH")
+    default_encoder_name = "encoder.safetensors"
+    default_decoder_name = "decoder.safetensors"
+    checkpoint_enc = (
+        Path(encoder_override).expanduser()
+        if encoder_override is not None
+        else DEFAULT_COSMOS_DIR / default_encoder_name
+    )
+    checkpoint_dec = (
+        Path(decoder_override).expanduser()
+        if decoder_override is not None
+        else DEFAULT_COSMOS_DIR / default_decoder_name
+    )
+
+    missing = [str(path) for path in (checkpoint_enc, checkpoint_dec) if not path.exists()]
+    if missing:
+        missing_display = ", ".join(missing)
+        raise FileNotFoundError(
+            "InternVLA-A1 requires local Cosmos tokenizer checkpoints. "
+            f"Missing: {missing_display}. "
+            f"Download {DEFAULT_COSMOS_REPO} and set INTERNVLA_A1_COSMOS_DIR to that directory, "
+            "or point INTERNVLA_A1_COSMOS_ENCODER_PATH / INTERNVLA_A1_COSMOS_DECODER_PATH "
+            "to explicit checkpoint files."
+        )
+
+    return checkpoint_enc, checkpoint_dec
+
+
+class Qwen3VLWithExpertModel(nn.Module):
+    def __init__(
+        self,
+        vlm_config: QwenConfig,
+        action_expert_config: QwenConfig,
+        precision: Literal["bfloat16", "float32"] = "bfloat16",
+    ) -> None:
+        super().__init__()
+
+        vlm_config_hf = CONFIG_MAPPING["qwen3_vl"]()
+        vlm_config_hf.text_config.hidden_size = vlm_config.hidden_size
+        vlm_config_hf.text_config.intermediate_size = vlm_config.intermediate_size
+        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_attention_heads
+        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
+        vlm_config_hf.text_config.num_hidden_layers = vlm_config.num_hidden_layers
+        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_key_value_heads
+        vlm_config_hf.text_config.max_position_embeddings = 262144
+        vlm_config_hf.text_config.rope_scaling = {
+            "mrope_interleaved": True,
+            "mrope_section": [24, 20, 20],
+            "rope_type": "default",
+        }
+        vlm_config_hf.text_config.tie_word_embeddings = True
+        vlm_config_hf.tie_word_embeddings = True
+        vlm_config_hf.vision_config.deepstack_visual_indexes = [5, 11, 17]
+        vlm_config_hf.vision_config.depth = 24
+        vlm_config_hf.vision_config.hidden_size = 1024
+        vlm_config_hf.vision_config.intermediate_size = 4096
+        vlm_config_hf.vision_config.out_hidden_size = 2048
+
+        self.und_expert = Qwen3VLForConditionalGeneration(config=vlm_config_hf)
+
+        gen_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]()
+        gen_expert_config_hf.head_dim = action_expert_config.head_dim
+        gen_expert_config_hf.hidden_size = action_expert_config.hidden_size
+        gen_expert_config_hf.intermediate_size = action_expert_config.intermediate_size
+        gen_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads
+        gen_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers
+        gen_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads
+        gen_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings
+        gen_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling
+        self.gen_expert = Qwen3VLTextModel(config=gen_expert_config_hf)
+        self.gen_expert.embed_tokens = None
+        self.gen_expert.lm_head = None
+
+        act_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]()
+        act_expert_config_hf.head_dim = action_expert_config.head_dim
+        act_expert_config_hf.hidden_size = action_expert_config.hidden_size
+        act_expert_config_hf.intermediate_size = action_expert_config.intermediate_size
+        act_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads
+        act_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers
+        act_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads
+        act_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings
+        act_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling
+        self.act_expert = Qwen3VLTextModel(config=act_expert_config_hf)
+        self.act_expert.embed_tokens = None
+        self.act_expert.lm_head = None
+
+        self.to_selected_precision(precision)
+
+    def to_selected_precision(self, precision: Literal["bfloat16", "float32"]) -> None:
+        if precision == "bfloat16":
+            self.to(dtype=torch.bfloat16)
+        elif precision == "float32":
+            # The fp32 path keeps the whole module in full precision, so no
+            # selective post-conversion fixups are required.
+            self.to(dtype=torch.float32)
+            return
+        else:
+            raise ValueError(f"Unsupported precision: {precision}")
+
+        # Match the original runtime numerics by keeping normalization layers
+        # in fp32 even when the rest of the module runs in bf16.
+        keep_fp32 = ["input_layernorm", "post_attention_layernorm", "model.norm"]
+        for name, param in self.named_parameters():
+            if any(key in name for key in keep_fp32):
+                param.data = param.data.to(dtype=torch.float32)
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.LongTensor | None,
+        past_key_values: Any,
+        inputs_embeds: list[torch.Tensor | None],
+        use_cache: bool,
+    ) -> tuple[list[torch.Tensor | None], Any]:
+        if inputs_embeds[1] is None and inputs_embeds[2] is None:
+            prefix_output = self.und_expert.language_model.forward(
+                inputs_embeds=inputs_embeds[0],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            return [prefix_output.last_hidden_state, None, None], prefix_output.past_key_values
+
+        if inputs_embeds[0] is None and inputs_embeds[2] is None:
+            middle_output = self.gen_expert.forward(
+                inputs_embeds=inputs_embeds[1],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            return [None, middle_output.last_hidden_state, None], middle_output.past_key_values
+
+        if inputs_embeds[0] is None and inputs_embeds[1] is None:
+            suffix_output = self.act_expert.forward(
+                inputs_embeds=inputs_embeds[2],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            return [None, None, suffix_output.last_hidden_state], None
+
+        models = [self.und_expert.language_model, self.gen_expert, self.act_expert]
+        stacked_inputs = [inputs_embeds[0], inputs_embeds[1], inputs_embeds[2]]
+        for layer_idx in range(self.und_expert.config.text_config.num_hidden_layers):
+            stacked_inputs = compute_layer_complete(
+                layer_idx,
+                stacked_inputs,
+                attention_mask,
+                position_ids,
+                und_expert=self.und_expert,
+                gen_expert=self.gen_expert,
+                act_expert=self.act_expert,
+            )
+        outputs = [models[i].norm(stacked_inputs[i]) for i in range(3)]
+        return outputs, None
+
+
+class InternVLAA1(nn.Module):
+    def __init__(
+        self,
+        config: InternVLAA1Config,
+        *,
+        cosmos_encoder_path: str | Path | None = None,
+        cosmos_decoder_path: str | Path | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+
+        vlm_config = get_qwen_config(config.qwen3_vl_variant)
+        action_expert_config = get_qwen_config(config.action_expert_variant)
+        self.qwen3_vl_with_expert = Qwen3VLWithExpertModel(vlm_config, action_expert_config, precision=config.dtype)
+
+        cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths(
+            encoder_path=cosmos_encoder_path,
+            decoder_path=cosmos_decoder_path,
+        )
+        self.cosmos = ImageTokenizer(
+            checkpoint_enc=str(cosmos_encoder_path),
+            checkpoint_dec=str(cosmos_decoder_path),
+            device=config.device,
+        )
+
+        hidden_size = action_expert_config.hidden_size
+        vae_dim = 16
+        ds = config.scale_factor
+        self.cosmos_in_proj = nn.Conv2d(vae_dim, hidden_size, kernel_size=1, stride=1, padding=0)
+        self.downsample_conv = nn.Conv2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0)
+        self.upsample_conv = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0)
+        self.cosmos_out_proj = nn.Linear(hidden_size, vae_dim)
+        self.cosmos_out_layer_norm = nn.LayerNorm(hidden_size)
+
+        self.action_in_proj = nn.Linear(config.max_action_dim, hidden_size)
+        self.action_out_proj = nn.Linear(hidden_size, config.max_action_dim)
+        self.state_proj = nn.Linear(config.max_state_dim, hidden_size)
+        self.action_time_mlp_in = nn.Linear(2 * hidden_size, hidden_size)
+        self.action_time_mlp_out = nn.Linear(hidden_size, hidden_size)
+
+        if config.compile_model:
+            torch.set_float32_matmul_precision("high")
+            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
+
+        self.cosmos.eval()
+        for param in self.cosmos.parameters():
+            param.requires_grad = False
+
+        self.set_attention_implementation(config.attn_implementation)
+
+    def set_attention_implementation(self, attn_implementation: str) -> None:
+        self.config.attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.und_expert.config.text_config._attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.und_expert.language_model.config._attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.gen_expert.config._attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.act_expert.config._attn_implementation = attn_implementation
+
+    def _prepare_attention_masks_4d(
+        self,
+        att_2d_masks: torch.Tensor,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        att_2d_masks_4d = att_2d_masks[:, None, :, :]
+        bias_dtype = torch.float32 if self.config.attn_implementation == "eager" else dtype
+        valid_bias = torch.zeros((), dtype=bias_dtype, device=att_2d_masks.device)
+        invalid_bias = torch.tensor(
+            OPENPI_ATTENTION_MASK_VALUE,
+            dtype=bias_dtype,
+            device=att_2d_masks.device,
+        )
+        return torch.where(att_2d_masks_4d, valid_bias, invalid_bias)
+
+    def sample_noise(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor:
+        return torch.normal(0.0, 1.0, shape, dtype=torch.float32, device=device)
+
+    @dynamo.disable
+    def embed_prefix(
+        self,
+        pixel_values: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        lang_tokens: torch.Tensor,
+        lang_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        image_token_id = self.qwen3_vl_with_expert.und_expert.config.image_token_id
+        pixel_values = pixel_values.view(-1, pixel_values.shape[-1])
+        image_grid_thw = image_grid_thw.view(-1, 3)
+        image_embs, _ = self.qwen3_vl_with_expert.und_expert.visual(pixel_values, image_grid_thw)
+
+        embs = self.qwen3_vl_with_expert.und_expert.get_input_embeddings()(lang_tokens)
+        batch_size, seq_len, hidden_dim = embs.shape
+        embs = embs.view(-1, hidden_dim)
+        lang_tokens_flat = lang_tokens.view(-1)
+        embs[lang_tokens_flat == image_token_id] = image_embs
+        embs = embs.view(batch_size, seq_len, hidden_dim)
+
+        pad_masks = lang_masks.to(torch.bool)
+        att_masks = torch.zeros_like(pad_masks, dtype=torch.bool, device=pad_masks.device)
+        return embs, pad_masks, att_masks
+
+    def get_cosmos_features(self, images: torch.Tensor) -> torch.Tensor:
+        shape = images.shape[:-3]
+        channels, height, width = images.shape[-3:]
+        images = images.reshape(-1, channels, height, width)
+        if (height, width) != (256, 256):
+            images = F.interpolate(images, size=(256, 256), mode="bilinear", align_corners=False)
+        images = images * 2 - 1
+        features = self.cosmos.encode(images)
+        channels, height, width = features.shape[-3:]
+        return features.view(*shape, channels, height, width)
+
+    def embed_middle(
+        self,
+        images: torch.Tensor,
+        img_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        device = images.device
+        batch_size, n_view, timesteps = images.shape[:3]
+        features = self.get_cosmos_features(images)
+        features = rearrange(features, "b n t c h w -> (b n t) c h w")
+        features = self.cosmos_in_proj(features)
+        features = self.downsample_conv(features)
+        features = rearrange(features, "(b n t) c h w -> b n t c h w", b=batch_size, n=n_view, t=timesteps)
+
+        _, _, _, _, height, width = features.shape
+        embs = rearrange(features, "b n t c h w -> b (n t h w) c")
+        pad_masks = torch.zeros((batch_size, n_view, timesteps, height, width), dtype=torch.bool, device=device)
+        pad_masks[img_masks] = True
+        pad_masks = rearrange(pad_masks, "b n t h w -> b (n t h w)")
+        att_masks = torch.tensor([1] + [0] * (embs.shape[1] - 1), dtype=torch.bool, device=device)
+        att_masks = att_masks[None, :].expand(batch_size, att_masks.shape[0])
+        return embs, pad_masks, att_masks
+
+    def prepare_suffix_static_context(
+        self,
+        state: torch.Tensor,
+        prefix_pad_masks: torch.Tensor,
+        max_prefix_position_ids: torch.Tensor,
+    ) -> SuffixStaticContext:
+        if self.state_proj.weight.dtype == torch.float32:
+            state = state.to(torch.float32)
+
+        state_emb = self.state_proj(state)
+        device = state_emb.device
+        batch_size = state_emb.shape[0]
+        suffix_len = self.config.chunk_size + 1
+        prefix_len = prefix_pad_masks.shape[1]
+
+        suffix_pad_masks = torch.ones(batch_size, suffix_len, dtype=torch.bool, device=device)
+        suffix_att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=torch.bool, device=device)
+        suffix_att_masks = suffix_att_masks[None, :].expand(batch_size, suffix_len)
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2),
+            dtype=state_emb.dtype,
+        )
+        position_ids = (
+            torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device)
+            .repeat(3, 1, 1)
+            .to(max_prefix_position_ids)
+            + max_prefix_position_ids
+        )
+        return SuffixStaticContext(
+            state_emb=state_emb,
+            full_att_2d_masks_4d=full_att_2d_masks_4d,
+            position_ids=position_ids,
+        )
+
+    def embed_suffix(
+        self,
+        state: torch.Tensor,
+        noisy_actions: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.state_proj.weight.dtype == torch.float32:
+            state = state.to(torch.float32)
+
+        state_emb = self.state_proj(state)
+        batch_size = state_emb.shape[0]
+        device = state_emb.device
+        state_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=device)
+
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.action_in_proj.out_features,
+            min_period=self.config.min_period,
+            max_period=self.config.max_period,
+            device=timestep.device,
+        ).to(dtype=timestep.dtype)
+        action_emb = self.action_in_proj(noisy_actions)
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+        action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb)))
+        action_mask = torch.ones(batch_size, action_time_emb.shape[1], dtype=torch.bool, device=device)
+
+        embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1)
+        pad_masks = torch.cat([state_mask, action_mask], dim=1)
+        att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=embs.dtype, device=device)
+        att_masks = att_masks[None, :].expand(batch_size, self.config.chunk_size + 1)
+        return embs, pad_masks, att_masks
+
+    def get_position_ids(
+        self,
+        lang_tokens: torch.Tensor,
+        image_grid_thw: torch.Tensor | None,
+        pad_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, Any]:
+        seq_len = lang_tokens.shape[1]
+        padded_lang_tokens = torch.ones_like(pad_masks).to(lang_tokens) * 777
+        padded_lang_tokens[:, :seq_len] = lang_tokens
+        attention_mask = pad_masks.to(lang_tokens)
+        if image_grid_thw is not None:
+            image_grid_thw = image_grid_thw.view(-1, 3)
+        return self.qwen3_vl_with_expert.und_expert.model.get_rope_index(
+            padded_lang_tokens,
+            image_grid_thw,
+            attention_mask=attention_mask,
+        )
+
+    @torch.no_grad()
+    def sample_actions(
+        self,
+        images: torch.Tensor,
+        img_masks: torch.Tensor,
+        pixel_values: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        lang_tokens: torch.Tensor,
+        lang_masks: torch.Tensor,
+        state: torch.Tensor,
+        *,
+        noise: torch.Tensor | None = None,
+        num_steps: int | None = None,
+        decode_image: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if num_steps is None:
+            num_steps = self.config.num_inference_steps
+        batch_size = state.shape[0]
+        device = state.device
+        dtype = state.dtype
+        if noise is None:
+            noise = self.sample_noise((batch_size, self.config.chunk_size, self.config.max_action_dim), device)
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            pixel_values,
+            image_grid_thw,
+            lang_tokens,
+            lang_masks,
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids, _ = self.get_position_ids(lang_tokens, image_grid_thw, prefix_pad_masks)
+        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            prefix_att_2d_masks,
+            dtype=prefix_embs.dtype,
+        )
+        _, past_key_values = self.qwen3_vl_with_expert.forward(
+            attention_mask=prefix_att_2d_masks_4d,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None, None],
+            use_cache=True,
+        )
+        max_prefix_position_ids = prefix_position_ids.max(dim=-1, keepdim=True).values
+
+        middle_embs, middle_pad_masks, middle_att_masks = self.embed_middle(images[:, :, :2], img_masks)
+        middle_len = middle_pad_masks.shape[1]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, middle_len, prefix_len)
+        middle_att_2d_masks = make_att_2d_masks(middle_pad_masks, middle_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, middle_att_2d_masks], dim=2)
+        middle_position_ids = (
+            torch.arange(1, middle_len + 1, device=max_prefix_position_ids.device)
+            .repeat(3, 1, 1)
+            .to(max_prefix_position_ids)
+            + max_prefix_position_ids
+        )
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            full_att_2d_masks,
+            dtype=middle_embs.dtype,
+        )
+        (_, middle_out, _), past_key_values = self.qwen3_vl_with_expert.forward(
+            attention_mask=full_att_2d_masks_4d,
+            position_ids=middle_position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, middle_embs, None],
+            use_cache=True,
+        )
+
+        max_position_ids = middle_position_ids.max(dim=-1, keepdim=True).values
+        curr_pad_masks = torch.cat([prefix_pad_masks, middle_pad_masks], dim=1)
+        dt = torch.tensor(-1.0 / num_steps, dtype=torch.float32, device=device)
+        x_t = noise
+        time = torch.tensor(1.0, dtype=torch.float32, device=device)
+        suffix_static = None
+        if self.config.enable_suffix_static_context_optimization:
+            suffix_static = self.prepare_suffix_static_context(state, curr_pad_masks, max_position_ids)
+        while time >= -dt / 2:
+            expanded_time = time.expand(batch_size)
+            if suffix_static is None:
+                v_t = self.denoise_step(
+                    state,
+                    curr_pad_masks,
+                    past_key_values,
+                    max_position_ids,
+                    x_t.to(dtype),
+                    expanded_time.to(dtype),
+                )
+            else:
+                v_t = self.denoise_step_optimized(
+                    suffix_static, past_key_values, x_t.to(dtype), expanded_time.to(dtype)
+                )
+            x_t = x_t + dt * v_t
+            time += dt
+
+        return x_t, None if not decode_image else middle_out
+
+    def denoise_step(
+        self,
+        state: torch.Tensor,
+        prefix_pad_masks: torch.Tensor,
+        past_key_values: Any,
+        max_prefix_position_ids: torch.Tensor,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep)
+        suffix_len = suffix_pad_masks.shape[1]
+        batch_size = prefix_pad_masks.shape[0]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+        position_ids = (
+            torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device)
+            .repeat(3, 1, 1)
+            .to(max_prefix_position_ids)
+            + max_prefix_position_ids
+        )
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            full_att_2d_masks,
+            dtype=suffix_embs.dtype,
+        )
+        outputs_embeds, _ = self.qwen3_vl_with_expert.forward(
+            attention_mask=full_att_2d_masks_4d,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, None, suffix_embs],
+            use_cache=False,
+        )
+        suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)
+
+    def denoise_step_optimized(
+        self,
+        suffix_static: SuffixStaticContext,
+        past_key_values: Any,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        state_emb = suffix_static.state_emb
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.action_in_proj.out_features,
+            min_period=self.config.min_period,
+            max_period=self.config.max_period,
+            device=timestep.device,
+        ).to(dtype=timestep.dtype)
+        action_emb = self.action_in_proj(x_t)
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+        action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb)))
+        suffix_embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1)
+        outputs_embeds, _ = self.qwen3_vl_with_expert.forward(
+            attention_mask=suffix_static.full_att_2d_masks_4d,
+            position_ids=suffix_static.position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, None, suffix_embs],
+            use_cache=False,
+        )
+        suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)
+
+
+class InternVLAA1Policy(nn.Module):
+    _AUTO_WEIGHTS_IGNORE_PREFIXES = ["model.cosmos."]
+
+    def __init__(
+        self,
+        config: InternVLAA1Config,
+        *,
+        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.input_builder = Qwen3VLInputBuilder(
+            processor_model_name=processor_model_name,
+            max_length=config.tokenizer_max_length,
+        )
+        cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths()
+        self.model = InternVLAA1(
+            config,
+            cosmos_encoder_path=cosmos_encoder_path,
+            cosmos_decoder_path=cosmos_decoder_path,
+        )
+        self.model.to(config.device)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        checkpoint_dir: str | Path,
+        *,
+        config: InternVLAA1Config | None = None,
+        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
+        strict: bool = False,
+    ) -> InternVLAA1Policy:
+        checkpoint_dir = Path(checkpoint_dir)
+        if config is None:
+            config = InternVLAA1Config.from_pretrained(checkpoint_dir)
+        instance = cls(config, processor_model_name=processor_model_name)
+        checkpoint_path = checkpoint_dir / "model.safetensors"
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"InternVLAA1 checkpoint not found: {checkpoint_path}")
+
+        loader = AutoWeightsLoader(
+            instance,
+            ignore_unexpected_prefixes=cls._AUTO_WEIGHTS_IGNORE_PREFIXES,
+        )
+        with safe_open(str(checkpoint_path), framework="pt", device=str(config.device)) as f:
+            file_keys = list(f.keys())
+            loaded = loader.load_weights((name, f.get_tensor(name)) for name in file_keys)
+
+        if strict:
+            expected = {
+                name
+                for name, _ in instance.state_dict().items()
+                if not any(name.startswith(prefix) for prefix in cls._AUTO_WEIGHTS_IGNORE_PREFIXES)
+            }
+            missing = sorted(expected - loaded)
+            if missing:
+                preview = ", ".join(missing[:10])
+                suffix = " ..." if len(missing) > 10 else ""
+                raise RuntimeError(f"Missing weights after AutoWeightsLoader load: {preview}{suffix}")
+
+        instance.eval()
+        return instance
+
+    def to(self, *args, **kwargs):  # type: ignore[override]
+        super().to(*args, **kwargs)
+        self.model.cosmos.to(torch.bfloat16)
+        self.model.action_out_proj.to(torch.float32)
+        return self
+
+    def _prepare_resized_histories(self, batch: dict[str, Any]) -> list[torch.Tensor]:
+        resized_images: list[torch.Tensor] = []
+        for i in range(3):
+            image_history = batch[f"{OBS_IMAGES}.image{i}"]
+            if image_history.shape[0] != 1:
+                raise ValueError(f"InternVLAA1Policy only supports bs=1, got image batch {tuple(image_history.shape)}")
+            resized_images.append(resize_with_pad(image_history[0], self.config.image_resolution))
+        return resized_images
+
+    def _preprocess_images(
+        self,
+        batch: dict[str, torch.Tensor],
+        *,
+        resized_histories: list[torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if resized_histories is None:
+            resized_histories = self._prepare_resized_histories(batch)
+        images = torch.stack(resized_histories, dim=0).unsqueeze(0)
+        img_masks = torch.stack([batch[f"{OBS_IMAGES}.image{i}_mask"] for i in range(3)], dim=1)
+        return images, img_masks
+
+    def prepare_state(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+        return pad_vector(batch[OBS_STATE], self.config.max_state_dim)
+
+    def _get_task(self, batch: dict[str, Any]) -> str:
+        task = batch[OBS_TASK]
+        if isinstance(task, str):
+            return task
+        if isinstance(task, (list, tuple)):
+            if len(task) != 1:
+                raise ValueError(f"InternVLAA1Policy only supports bs=1, got {len(task)} tasks")
+            item = task[0]
+            if not isinstance(item, str):
+                raise TypeError(f"Expected task string, got {type(item)!r}")
+            return item
+        raise TypeError(f"Unsupported task payload type: {type(task)!r}")
+
+    def _prepare_qwen_prefix_inputs(
+        self,
+        resized_histories: list[torch.Tensor],
+        task: str,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        qwen_inputs = self.input_builder.build(resized_histories, task)
+        device = self.model.action_in_proj.weight.device
+        return (
+            qwen_inputs[f"{OBS_PREFIX}pixel_values"]
+            .unsqueeze(0)
+            .to(device=device, dtype=self.model.action_in_proj.weight.dtype),
+            qwen_inputs[f"{OBS_PREFIX}image_grid_thw"].unsqueeze(0).to(device=device),
+            qwen_inputs[f"{OBS_PREFIX}input_ids"].unsqueeze(0).to(device=device),
+            qwen_inputs[f"{OBS_PREFIX}attention_mask"].unsqueeze(0).to(device=device),
+        )
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        batch: dict[str, Any],
+        *,
+        noise: torch.Tensor | None = None,
+        decode_image: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        task = self._get_task(batch)
+        resized_histories = self._prepare_resized_histories(batch)
+        pixel_values, image_grid_thw, lang_tokens, lang_masks = self._prepare_qwen_prefix_inputs(
+            resized_histories,
+            task,
+        )
+        state = self.prepare_state(batch)
+        images, img_masks = self._preprocess_images(batch, resized_histories=resized_histories)
+        return self.model.sample_actions(
+            images,
+            img_masks,
+            pixel_values,
+            image_grid_thw,
+            lang_tokens,
+            lang_masks,
+            state,
+            noise=noise,
+            decode_image=decode_image,
+        )

vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+    wrap_methods_by_paths,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+from .config import (
+    DEFAULT_QWEN3_VL_MODEL,
+    OBS_IMAGES,
+    OBS_STATE,
+    OBS_TASK,
+    InternVLAA1Config,
+)
+from .model_internvla_a1 import InternVLAA1Policy, resolve_cosmos_checkpoint_paths
+
+logger = init_logger(__name__)
+
+
+def get_internvla_a1_post_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def post_process_func(x):
+        return x
+
+    return post_process_func
+
+
+class InternVLAA1Pipeline(nn.Module, DiffusionPipelineProfilerMixin):
+    """InternVLA-A1 pipeline wrapper for the policy implementation."""
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.prefix = prefix
+        self.model_dir = od_config.model
+        self.config = self._build_config(od_config)
+        custom_args = od_config.custom_pipeline_args or {}
+        self.strict_load = bool(custom_args.get("strict_load", False))
+        self.processor_model_name = str(custom_args.get("processor_model_name", DEFAULT_QWEN3_VL_MODEL))
+        enable_warmup = custom_args.get("enable_warmup")
+        self.enable_warmup = bool(enable_warmup) if isinstance(enable_warmup, bool) else False
+
+        self.setup_diffusion_pipeline_profiler(
+            profiler_targets=["forward"],
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
+        )
+        self.policy = self._initialize_policy()
+        self._setup_policy_profiler_targets()
+        if self.enable_warmup:
+            self._warmup()
+
+    def _build_config(self, od_config: OmniDiffusionConfig) -> InternVLAA1Config:
+        config_dict = self._load_config_dict(od_config)
+        config = InternVLAA1Config.from_model_config(config_dict)
+
+        custom_args = od_config.custom_pipeline_args or {}
+        device = custom_args.get("device")
+        if isinstance(device, str):
+            config.device = device
+
+        dtype = custom_args.get("dtype")
+        if isinstance(dtype, str):
+            config.dtype = dtype
+        elif od_config.dtype is not None:
+            config.dtype = str(od_config.dtype).split(".")[-1]
+
+        compile_model = custom_args.get("compile_model")
+        if isinstance(compile_model, bool):
+            config.compile_model = compile_model
+
+        attn_implementation = custom_args.get("attn_implementation")
+        if isinstance(attn_implementation, str):
+            config.attn_implementation = attn_implementation
+
+        enable_regional_compile = custom_args.get("enable_regional_compile")
+        if isinstance(enable_regional_compile, bool):
+            config.enable_regional_compile = enable_regional_compile
+
+        regional_compile_dynamic = custom_args.get("regional_compile_dynamic")
+        if isinstance(regional_compile_dynamic, bool):
+            config.regional_compile_dynamic = regional_compile_dynamic
+
+        return config
+
+    def _load_config_dict(self, od_config: OmniDiffusionConfig) -> dict[str, Any]:
+        if od_config.model_config:
+            return dict(od_config.model_config)
+
+        model_path = od_config.model
+        if not model_path:
+            return {}
+
+        config_path = os.path.join(model_path, "config.json")
+        if not os.path.exists(config_path):
+            logger.info("InternVLAA1Pipeline config.json not found under %s; using defaults.", model_path)
+            return {}
+
+        with open(config_path, encoding="utf-8") as f:
+            return json.load(f)
+
+    def has_real_checkpoint(self) -> bool:
+        return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors"))
+
+    def runtime_mode(self) -> str:
+        return "real_checkpoint_loaded" if self.has_real_checkpoint() else "no_checkpoint_policy"
+
+    def _setup_policy_profiler_targets(self) -> None:
+        if not self.od_config.enable_diffusion_pipeline_profiler:
+            return
+
+        wrap_methods_by_paths(
+            self,
+            [
+                "policy.model.sample_actions",
+                "policy.model.embed_prefix",
+                "policy.model.embed_middle",
+                "policy.model.denoise_step",
+            ],
+        )
+
+    def _apply_policy_optimizations(self, policy: InternVLAA1Policy) -> None:
+        policy.model.set_attention_implementation(self.config.attn_implementation)
+        if not self.config.enable_regional_compile:
+            return
+
+        compile_targets = [
+            "qwen3_vl_with_expert.und_expert.visual",
+            "qwen3_vl_with_expert.und_expert.language_model",
+            "qwen3_vl_with_expert.gen_expert",
+            "qwen3_vl_with_expert.act_expert",
+        ]
+
+        for path in compile_targets:
+            current = policy.model
+            for part in path.split("."):
+                current = getattr(current, part, None)
+                if current is None:
+                    break
+            if current is None:
+                continue
+            try:
+                regionally_compile(current, dynamic=self.config.regional_compile_dynamic)
+                logger.info("InternVLAA1Pipeline regional compile applied to %s", path)
+            except Exception as exc:
+                logger.warning("InternVLAA1Pipeline regional compile failed for %s: %s", path, exc)
+
+    def _initialize_policy(self) -> InternVLAA1Policy:
+        resolve_cosmos_checkpoint_paths()
+        if self.has_real_checkpoint():
+            logger.info("Loading InternVLA-A1 weights from %s", self.model_dir)
+            policy = InternVLAA1Policy.from_pretrained(
+                self.model_dir,
+                config=self.config,
+                processor_model_name=self.processor_model_name,
+                strict=self.strict_load,
+            )
+        else:
+            logger.info("Initializing InternVLA-A1 policy without checkpoint weights.")
+            policy = InternVLAA1Policy(self.config, processor_model_name=self.processor_model_name)
+
+        policy.to(self.config.device)
+        policy.to(getattr(torch, self.config.dtype))
+        policy.eval()
+        self._apply_policy_optimizations(policy)
+        return policy
+
+    def _build_fake_batch_inputs(self) -> dict[str, torch.Tensor]:
+        device = torch.device(self.config.device)
+        dtype = getattr(torch, self.config.dtype)
+        history = 2
+        channels = 3
+        image_size = self.config.image_resolution
+        fake_image = torch.zeros((1, history, channels, image_size[0], image_size[1]), device=device, dtype=dtype)
+        fake_mask = torch.ones((1,), device=device, dtype=torch.bool)
+        return {
+            OBS_STATE: torch.zeros((1, self.config.max_state_dim), device=device, dtype=dtype),
+            OBS_TASK: [""],
+            f"{OBS_IMAGES}.image0": fake_image.clone(),
+            f"{OBS_IMAGES}.image1": fake_image.clone(),
+            f"{OBS_IMAGES}.image2": fake_image.clone(),
+            f"{OBS_IMAGES}.image0_mask": fake_mask.clone(),
+            f"{OBS_IMAGES}.image1_mask": fake_mask.clone(),
+            f"{OBS_IMAGES}.image2_mask": fake_mask.clone(),
+        }
+
+    def _warmup(self) -> None:
+        logger.info("InternVLAA1Pipeline warmup started")
+        try:
+            batch_inputs = self._build_fake_batch_inputs()
+            noise = torch.zeros(
+                (1, self.config.chunk_size, self.config.max_action_dim),
+                device=self.config.device,
+                dtype=torch.float32,
+            )
+            with torch.inference_mode():
+                self.policy.forward(batch_inputs, noise=noise, decode_image=False)
+        except Exception as exc:
+            logger.warning("InternVLAA1Pipeline warmup failed: %s", exc)
+            return
+        logger.info("InternVLAA1Pipeline warmup finished")
+
+    def _predict_actions(
+        self,
+        batch_inputs: dict[str, Any],
+        *,
+        noise: torch.Tensor | None = None,
+        decode_image: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        logger.debug("InternVLAA1Pipeline forward mode=%s", self.runtime_mode())
+        return self.policy.forward(
+            batch_inputs,
+            noise=noise,
+            decode_image=decode_image,
+        )
+
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        if len(req.prompts) > 1:
+            logger.warning("InternVLAA1Pipeline only supports a single prompt/request; taking the first sample.")
+        extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
+        batch_inputs = extra_args.get("batch_inputs")
+        if batch_inputs is None:
+            return DiffusionOutput(
+                error=(
+                    "InternVLAA1Pipeline.forward expects sampling_params.extra_args['batch_inputs'] "
+                    "with pre-built repo-side inputs."
+                ),
+                post_process_func=get_internvla_a1_post_process_func(self.od_config),
+            )
+
+        output, decoded = self._predict_actions(
+            batch_inputs,
+            noise=extra_args.get("noise"),
+            decode_image=bool(extra_args.get("decode_image", False)),
+        )
+        custom_output: dict[str, Any] = {}
+        if decoded is not None:
+            custom_output["decoded"] = decoded
+        return DiffusionOutput(
+            output=output,
+            custom_output=custom_output,
+            post_process_func=get_internvla_a1_post_process_func(self.od_config),
+        )

vllm_omni/diffusion/registry.py

@@ -135,6 +135,11 @@ _DIFFUSION_MODELS = {
         "pipeline_bagel",
         "BagelPipeline",
     ),
+    "InternVLAA1Pipeline": (
+        "internvla_a1",
+        "pipeline_internvla_a1",
+        "InternVLAA1Pipeline",
+    ),
     "LongCatImageEditPipeline": (
         "longcat_image",
         "pipeline_longcat_image_edit",
@@ -421,6 +426,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
+    "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
     "FluxKontextPipeline": "get_flux_kontext_post_process_func",