[[Realtime-VLA]] [[Realtime-VLA-v2]]

はじめに

VLAモデルは推論に数百msかかり、リアルタイム制御の障壁になっています。Realtime-VLA [1] はpi0系モデルの推論をTritonカーネルで再実装することでGPU処理を最適化し、RTX 4090で27.3msの推論速度を達成しました。

この記事ではRealtime-VLAの手法を解説し、手元のRTX 3060でpi0.5を計測した結果を報告します。また、モデルだけでなく推論パイプライン全体の最適化を行ったRealtime-VLA V2 [2] についても概説します。

LeRobotとRealtime-VLAのpi0.5実装の比較

pi0.5の推論フロー

pi0.5 は Vision-Language Model (VLM) と Action Expert (AE) の2つのTransformerで構成されています。推論時はVLMを１回だけ実行してKVキャッシュを作り、AEがそのキャッシュを参照しながらFlow Matchingのdenoiseループを回します。

モデル構成:

Vision EncoderにはSigLIP (400M) を使用し、224×224画像を256個のトークンに変換します。

Encoder (prefix) フェーズ — 1回実行:

1. SigLIPで各ビュー画像をエンコードし256トークンのembeddingを得る（pooler_output * sqrt(2048) でスケーリング）
2. タスクプロンプトの言語トークンをembedding層で変換し * sqrt(2048) でスケーリング
3. 画像トークン列と言語トークン列を結合してprefixシーケンスを構成
4. VLM (Gemma 2B) の18層を通し、全層のKey/Valueを past_key_values (KVキャッシュ) として保存

Decoder (suffix) フェーズ — 10ステップ繰り返し:

各denoiseステップで以下を実行:

1. 入力準備: noisy_actions (chunk_size × 32) を action_in_proj (Linear 32→1024) でAEの次元に射影。timestep $t$ をsinusoidal embeddingで1024次元に変換し、2層MLP (SiLU → Linear → SiLU → Linear) で処理して AdaRMS Norm の条件ベクトル adarms_cond とする
2. AE 18層の実行: 各層で:
   • AdaRMS Norm: dense (Linear 1024→3072) が adarms_cond から scale, shift, gate の3つを生成。x_norm * (1 + scale) + shift で変調
   • Q, K, V 射影 + RoPE → prefixのKVキャッシュと結合してattention
   • O射影 → gated residual: residual + output * gate
   • AdaRMS Norm → GatedMLP (gate_proj * up_proj → down_proj) → gated residual
3. velocity予測: action_out_proj (Linear 1024→32) で velocity $v_t$ を出力
4. Euler積分: $x_{t+dt} = x_t + dt \cdot v_t$（$dt = -1/\text{num_steps}$）

LeRobot実装

LeRobotのpi0.5実装 (modeling_pi05.py) はOpenPIのPyTorch移植で、HuggingFace Transformersの GemmaModel / PaliGemmaForConditionalGeneration 上に構築されています。

クラス階層:

PI05Policy
 └─ PI05Pytorch (sample_actions / denoise_step)
     ├─ PaliGemmaWithExpertModel
     │   ├─ paligemma: PaliGemmaForConditionalGenerationWithPiGemma
     │   │   └─ model: PaliGemmaModelWithPiGemma
     │   │       ├─ vision_tower: SiglipVisionTransformer (SigLIP, 27層)
     │   │       ├─ multi_modal_projector: Linear (1152→2048)
     │   │       └─ language_model: PiGemmaModel (Gemma 2B, 18層)
     │   └─ gemma_expert: PiGemmaForCausalLM
     │       └─ model: PiGemmaModel (Gemma 300M, 18層)
     ├─ action_in_proj: Linear(32, 1024)
     ├─ action_out_proj: Linear(1024, 32)
     ├─ time_mlp_in: Linear(1024, 1024)
     └─ time_mlp_out: Linear(1024, 1024)

推論の分岐構造 (PaliGemmaWithExpertModel.forward):

inputs_embeds はVLMとExpertの2要素リストで、どちらかが None かどうかで3つの処理パスに分岐します:

• inputs_embeds=[prefix, None] — Encoder-only: VLMのみforwardし use_cache=True でKVキャッシュを返す。sample_actions の冒頭で1回実行
• inputs_embeds=[None, suffix] — Decoder-only: Expert のみ forward。KVキャッシュ付きで呼ばれる。推論時の denoise_step がこのパスを使用
• inputs_embeds=[prefix, suffix] — Joint: 訓練時に使用。compute_layer_complete() が各層で両モデルのQ,K,Vを結合して単一のattentionを計算

Decoder-onlyパスの詳細 (denoise_step):

1. embed_suffix() でnoisy_actionsとtimestepをembedding
2. prefix の pad_mask と suffix の attention mask を結合して4Dマスクを構築
3. copy.deepcopy(past_key_values) でKVキャッシュを複製 — Expertが各層でKVキャッシュに自身のK,Vを追記するため、毎ステップコピーが必要
4. gemma_expert.model.forward() → PiGemmaModel.forward() が呼ばれ、18層のTransformerを逐次実行

推論パイプライン全体の流れ:

flowchart TB
    subgraph Encoder["Encoder — 1回実行"]
        direction TB
        IMG["画像 (N views)"] --> SigLIP["SigLIP\n27層 Vision Encoder"]
        SigLIP --> PROJ["multi_modal_projector\nLinear 1152→2048"]
        LANG["言語トークン"] --> EMB["embed_tokens\n× √2048"]
        PROJ --> CAT["concat"]
        EMB --> CAT
        CAT --> VLM["Gemma 2B\n18層 Transformer"]
        VLM --> KV["past_key_values\n18層分 KVキャッシュ"]
    end

    subgraph Denoise["Denoise Loop — 10回繰り返し"]
        direction TB
        NOISE["noisy_actions\nchunk_size × 32"] --> AIN["action_in_proj\nLinear 32→1024"]
        TIME["timestep t"] --> SINCOS["sinusoidal emb"] --> TMLP["time_mlp\nSiLU→Linear→SiLU→Linear"]
        TMLP --> COND["adarms_cond"]

        AIN --> DEEPCOPY
        KV2["past_key_values"] --> DEEPCOPY["copy.deepcopy\n毎ステップ複製"]
        DEEPCOPY --> EXPERT["Gemma 300M Expert\n18層 Transformer\n各層 ~21カーネル起動"]
        COND --> EXPERT
        EXPERT --> AOUT["action_out_proj\nLinear 1024→32"]
        AOUT --> EULER["Euler積分\nx_t += dt × v_t"]
        EULER -.->|"次ステップ"| NOISE
    end

    KV --> KV2
    EULER --> ACTIONS["predicted actions"]

    style Encoder fill:#e8f4fd,stroke:#4a90d9
    style Denoise fill:#fdf2e8,stroke:#d99a4a
    style DEEPCOPY fill:#fdd,stroke:#d44
    style EXPERT fill:#fdd,stroke:#d44

各Transformer層の処理 (PiGemmaDecoderLayer.forward):

各ボックスが独立したCUDAカーネル起動に対応する:

flowchart LR
    subgraph attn_norm["AdaRMS Norm"]
        direction LR
        D1["dense\n1024→3072"] --> N1["RMS\nNorm"] --> SS1["scale/shift\n+ gate"]
    end
    subgraph qkv["QKV射影"]
        direction LR
        Q["q_proj\n1024→2048"] ~~~ K["k_proj\n1024→256"] ~~~ V["v_proj\n1024→256"]
    end
    subgraph rope_attn["Attention"]
        direction LR
        ROPE["RoPE"] --> REP["repeat\n_kv"] --> QKT["Q@Kᵀ"] --> SM["softmax"] --> AV["@V"]
    end
    subgraph o_res["O + Res"]
        direction LR
        OP["o_proj\n2048→1024"] --> GR1["gated\nres"]
    end
    subgraph ffn_norm["AdaRMS Norm"]
        direction LR
        D2["dense\n1024→3072"] --> N2["RMS\nNorm"] --> SS2["scale/shift\n+ gate"]
    end
    subgraph mlp["GatedMLP"]
        direction LR
        GP["gate_proj\n+ GELU"] --> MUL["×"] --> DP["down_proj\n4096→1024"]
        UP["up_proj"] --> MUL
    end
    subgraph res2["Res"]
        GR2["gated\nres"]
    end

    attn_norm --> qkv --> rope_attn --> o_res --> ffn_norm --> mlp --> res2

    style attn_norm fill:#fee,stroke:#c66
    style qkv fill:#fee,stroke:#c66
    style rope_attn fill:#fee,stroke:#c66
    style o_res fill:#fee,stroke:#c66
    style ffn_norm fill:#fee,stroke:#c66
    style mlp fill:#fee,stroke:#c66
    style res2 fill:#fee,stroke:#c66

1層あたり約21カーネル × 18層 × 10ステップ = 約3,780回のカーネル起動

hidden_states
  → PiGemmaRMSNorm (input_layernorm, cond=adarms_cond)
     → dense(cond): Linear(1024, 3072) → scale, shift, gate に分割
     → x_norm * (1 + scale) + shift  /  gate を返却
  → q_proj: Linear(1024, 2048)  # 8 heads × 256 dim
  → k_proj: Linear(1024, 256)   # 1 kv_head × 256 dim
  → v_proj: Linear(1024, 256)   # 1 kv_head × 256 dim
  → apply_rotary_pos_emb (RoPE)
  → repeat_kv (1→8 heads に展開)
  → Q @ K^T * scaling → softmax → @ V  (eager attention)
  → o_proj: Linear(2048, 1024)
  → _gated_residual: residual + output * gate

  → PiGemmaRMSNorm (post_attention_layernorm, cond=adarms_cond) → gate
  → GemmaMLP:
     → gate_proj: Linear(1024, 4096, bias=False) → GELU
     → up_proj:   Linear(1024, 4096, bias=False)
     → GELU(gate_proj(x)) * up_proj(x)
     → down_proj: Linear(4096, 1024, bias=False)
  → _gated_residual: residual + output * gate

オーバーヘッドの要因:

各Transformer層で発生する独立CUDAカーネル起動:

これが18層 × 10 denoiseステップで繰り返され、Decoderだけで約3,780回のカーネル起動が発生します。毎回Python → CUDAドライバのディスパッチを経由し、さらに copy.deepcopy(past_key_values) による18層分のKVテンソルの複製が10回行われます。GPUの実計算時間に対してCPUオーバーヘッドとメモリコピーが支配的です。

Realtime-VLA実装

GitHub: https://github.com/Dexmal/realtime-vla

pi0_infer.py に共通Tritonカーネル群を定義し、pi05_infer.py / dm0_infer.py がそれを import して各モデル固有の推論を構築しています。依存は PyTorch + Triton のみです。

計算グラフの簡略化 (§3.2)

カーネル融合に先立ち、計算グラフ自体を等価変換で簡略化します。

• RMSNorm affine吸収: RMSNorm の学習可能スケールパラメータ $\gamma$ を後続の線形層の重みに事前乗算して融合。$\text{Linear}(\gamma \cdot \hat{x}) = \text{Linear}'(\hat{x})$ と等価変換でき、RMSNorm は正規化のみに簡略化される
• Action Time Encoder折り畳み: denoiseの時間ステップは $t \in {1.0, 0.9, \ldots, 0.1}$ の10通りのみ。time embedding と style projection の結果を __init__ 時に全計算してテーブル化
• QKV融合: Q, K, V の3つの射影行列を1つの大行列 (hidden, (num_heads+2)*head_dim) に結合。RoPEの cos/sin も重みに事前融合し、射影と位置エンコーディングを1つのmatmulで実行

これらの変換で7–8msの推論時間を削減しています。

I/O最適化 (§3.3)

GPU以外の部分でも以下の最適化を行っています。

• 画像リサイズ: カメラISPの出力解像度を224×224に近い値 (240×320等) に設定し、手書きのリサイズカーネルで60μs以下に抑える
• ピンドメモリ / ゼロコピー: 静的CPUバッファをピンドメモリで確保し、カメラフレームのCPU→GPU転送をゼロコピーで処理

カーネル融合と計算グラフの分解 (§4)

計算グラフを簡略化した上で、24の行列積演算に分解し、各演算を専用Tritonカーネルで実装しています。

最適化段階ごとの推論時間 (論文 Figure 2):

naive PyTorch (105ms) → CUDA Graph適用 → 計算グラフ簡略化 → カーネル最適化と段階的に高速化し、Roofline下限に接近しています。

計算グラフの簡略化 (論文 Figure 3):

π0モデルの計算フロー (論文 Figure 4): Vision Encoder (左)、LLM (中央)、Action Expert (右) が24のGEMM演算に分解されています。

Encoder (VLM/LLM) 18層:

• rms_matmul_n_2048_2560_qkv_rope: RMSNorm + QKV射影 + RoPE を1カーネルに融合
• matmul_n_2048_2048_res: O射影 + 残差接続を融合
• rms_matmul_n_2048_16384_gate: RMSNorm + Gated up-projection を融合
• matmul_n_16384_2048_res: MLP down-projection + 残差接続を融合

Decoder (Action Expert) 18層 × 10ステップ:

• adarms_norm_style_proj: AdaRMS正規化 + スタイル射影を融合。RMS正規化後に (1 + scale) * x_norm + shift を適用し、gate値を別バッファに出力
• matmul_rope_qkv: QKV射影 + RoPE融合。Q, K, V を1つの重み行列 (1024, 2560) として射影し、ヘッド次元の偶奇ペアに cos/sin を適用。出力先を Q/K/V の3ポインタに分岐書き込み
• softmax_kernel_prefix_suffix: Encoder KV (prefix) と Decoder self KV (suffix) を結合した attention mask 付き softmax。prefix の有効長を動的に参照し、無効キーに -2.38e38 をマスク
• matmul_small_res_gate: matmul + 残差 + gate を1カーネルで実行。out = res + (x @ W) * gate を1回のロード/ストアで計算

Decoder 1層のカーネル融合 (LeRobotとの比較):

LeRobot: ~21 kernels × 18層 × 10ステップ = ~3,780回
Realtime-VLA: ~5 kernels × 18層 × 10ステップ = ~900回 + CUDA Graphで同期オーバーヘッド最小化

カーネル内部の最適化 (§4)

融合カーネルの内部でも以下の最適化を適用しています。

• GEMMタイルチューニング (§4.1): cuBLAS のデフォルトタイル設定は VLA の行列サイズに最適でない場合がある。Triton実装でタイルサイズ (BLOCK_SIZE_M/N/K) を手動チューニングし約1.5ms改善
• LLM 17層実行 (§4.1): Encoderの最終 (18番目) 層はKVキャッシュのみAEに渡すため、attention/FFNの実行を省略して約0.7ms節約
• Gated Linear Layer融合 (§4.2): FFNの gate_proj と up_proj は同じ入力に対する2つの独立matmul。入力タイルを1回ロードして2つの重みタイルを処理し、GELU(gate) * up の結果のみ書き戻す。1.7msの改善
• Partial Split-k (§4.3): SigLIPの 512×1152×1152 GEMM は64×64タイルで144ブロックとなりRTX 4090の128SMに不均等分配される。512×1152×1024 (均等) + 512×1152×128 (split-2) に分割して1カーネルで記述
• スカラー演算のGEMM統合 (§4.4): bias加算、残差接続、活性化関数をGEMMのエピローグに統合。RMSNormはトークンレベルの統計量を先に計算し、次のGEMMで累積後に正規化ファクタで除算。約4msの改善

Roofline下限 (§5)

RTX 4090 (メモリ帯域幅 1.01 TB/s, BF16 MAC 91.4 TFLOPS) での理論下限は以下のとおりです。

現在の実装は理論下限まで約30%の余地があります。

事前計算 (__init__)

推論時のオーバーヘッドを削減するため、初期化時に以下の事前計算を行っています。

• time embedding: 10ステップ分の SiLU(Linear(SiLU(Linear(t)))) を初期化時に全計算し decoder_time_emb[10, seq, 1024] に格納
• style projection: 10ステップ × 18層分の time_emb @ mod_w + mod_b を事前計算し decoder_style_attn[10, 18, seq, 3072] と decoder_style_ffn[10, 18, seq, 3072] に格納
• action_out_proj の重みとバイアスに -1/num_steps を乗算し、Euler積分（x += -1/T * v）をmatmulに融合。ループ内では単純な加算で済みます
• RoPE テーブル: inv_freq = 1/10000^(2i/256) から cos/sin を事前計算し、encoder/decoder の位置に応じてスライス

CUDA Graph

record_infer_graph() で vision_encoder → transformer_encoder → transformer_decoder の全 forward を1グラフにキャプチャします。forward() は入力バッファへの copy_() 後に infer_graph.replay() の1コールで完結し、Python の制御フローは一切介入しません。

静的バッファ

全中間テンソルを __init__ で事前確保します（torch.empty(..., device="cuda")）。forward 中に動的メモリ割り当ては発生せず、vision_x, encoder_K/V, decoder_q_buf 等すべてのバッファが固定サイズで、CUDA Graph との互換性を保証しています。

RTX3060による計測実験

実験設定

推論実験に用いたHWのスペックは以下のとおりです。

実験に関連するパラメータは以下のとおりです。

結果

推論パイプライン全体を最適化するRealtime-VLA V2 [2]

V1がGPU計算の高速化に焦点を当てていたのに対し、V2は実ロボットで Fast / Smooth / Accurate を同時に達成するための手法です。推論パイプライン全体を高速化するため、以下の手法を用いています。

• 遅延キャリブレーション: t_camera, t_readout, t_proprio, t_motion を計測し、入力の時間整合とコマンドのプリアンプリファイで補償します
• 軌道後処理: Speed Adaptation Model → Temporal Optimization (OSQP) → Spatial Optimization (acados MPC) の3段パイプラインで軌道を最適化します
• 速度適応の学習: Human-in-the-loop でオペレータのスロットル入力を回帰モデルに蒸留し、タスク中の精度要求区間で自動減速を学習します

shirt-folding 等3タスクで人間操作と同等の実行速度を達成しています。

まとめ

Realtime-VLA はカーネル融合・CUDA Graph・静的バッファ等の最適化により、LeRobot実装と比べ推論速度が最大5倍に向上しています。手元のRTX 3060でも107.8msで推論を実行できました。

計測結果からは、pi0.5ではデノイズステップよりもVLMの処理時間が支配的であることがわかります。また、Realtime-VLA V2ではGPU計算だけでなく推論パイプライン全体を最適化するアプローチが提案されており、実ロボットでの高速かつ滑らかな動作を実現しています。

参考文献

• [[experiment]]
• [1] Y. Ma, Y. Zhou, Y. Yang, T. Wang, and H. Fan, "Running VLAs at Real-time Speed." 2025. https://arxiv.org/abs/2510.26742
• [2] Realtime-VLA V2. 2025. https://arxiv.org/abs/2603.26360
• [3] K. Black, M. Y. Galliker, and S. Levine, "Real-time Execution of Action Chunking Flow Policies." 2025. https://arxiv.org/abs/2512.05964