Quartz 5

Realtime-VLA

Properties1
tagspaper, vla, m-pi0, i-perf, i-cuda-opt, i-latency, k-dynamic, c-impl

8 min read

Realtime-VLA

Abstract

  • π0レベルのマルチビューVLAを単一コンシューマGPU (RTX 4090) で30Hz フレームレート、最大480Hzの軌道周波数で実行する方法を示す
  • モデル推論のオーバーヘッドを排除する一連の戦略を導入
  • 実世界実験で落下するペンの把持タスクにおいてπ0ポリシーが100%の成功率を達成
  • Full Streaming Inferenceフレームワークを提案

1 Introduction

  • VLAモデル (数十億パラメータ) は汎化性能に優れるがレイテンシが課題
  • 1 forward passに数百msかかるため、動的タスクへの対応が困難
  • 33ms (≈1/30s) 以下がリアルタイム動作の転換点：30FPS映像の全フレームを処理可能
  • 最適化後、RTX 4090で2ビュー入力に対し27.3msのレイテンシを達成

推論速度比較 (RTX 4090, 空プロンプト, chunk_length=63)

実装1 view2 views3 views
naive torch105.0ms106.5ms113.9ms
openpi/jax43.8ms53.7ms67.6ms
ours20.0ms27.3ms36.8ms
  • 落下ペン把持タスクで100%成功率を実証（30cm落下、反応時間200ms以内、人間と同等）
  • Full Streaming Inferenceモード：最大480Hzの制御信号生成が可能

2 Preliminary on π0-level Model

  • VLM backbone: PaliGemma (3Bパラメータ) で初期化
    • Vision Encoder: SigLIP (400Mパラメータ)
    • LLM: Gemma (2.6Bパラメータ)
  • Action Expert (AE): VLMとMoEアーキテクチャで結合
    • Gemmaから縮小 (幅・MLP次元を削減)、300Mパラメータ
    • Flow Matching によるaction chunkingの予測
    • VLMの入力: マルチビュー画像 + タスクプロンプト
    • AEの入力: 状態 + アクションノイズ

3 Eliminating the Overheads

3.1 Removing CPU Overhead

  • Python → CUDAカーネル起動のオーバーヘッドが深刻（π0モデルで1000以上のカーネル起動）
  • CUDA Graph を使用してカーネルストリームを記録・リプレイ
    • リプレイ時はGPUとドライバのみでカーネル起動（Python実行オーバーヘッド排除）
    • VLAではtransformerブロックに動的分岐がないため適用可能
    • 約2倍の高速化を達成

3.2 Simplifying the Graph

  • RMSNorm affine吸収: RMS normのアフィンパラメータを後続の線形層に融合（線形性の結合則）
  • Action Time Encoder折り畳み: 2つの連続カーネルを1つに統合。時間ステップは10種のみなので結果をテーブル化
  • QKV融合: Q, K, Vの射影行列を1つの大行列に統合。RoPEもmatmulに融合し、重みを事前計算
  • これらの変換で7-8msの推論時間削減

3.3 Avoiding Other Overheads

  • 画像リサイズ: カメラISPの出力解像度を224×224に近い値（240×320等）に設定。手書きリサイズで60us以下
  • ピンドメモリ使用、静的CPUバッファ、ゼロコピーカメラフレーム処理

4 In-Depth Optimization of the Kernels

  • 簡略化後の計算グラフは24のGEMM的演算とスカラー演算で構成

4.1 Tuning Tile Parameters of GEMM

  • cuBLAS (torch.matmul) のデフォルト設定が最適でないケースあり
  • Triton実装でタイリング戦略を手動チューニング → 約1.5ms改善
  • LLMのtransformerは17回（18ではなく）のattention/FFN実行（最終層のKVキャッシュのみAEに渡すため）→ 約0.7ms節約

4.2 Fusing Gated Linear Layers

  • FFNのGated up-projection: FC1(x,w1) · GELU(FC2(x,w2))
  • 2つのmatmulを並列実行し、ロード/ストア操作を結合
  • 入力タイル1つのロード後に2つの重みタイルを処理、結合結果のみ書き戻し
  • 1.7msの改善

4.3 Partial Split-k

  • 512×1152×1152のGEMM: 64×64タイルで144ブロック → RTX 4090の128SMに不均等分配
  • 解決策: 512×1152×1024 (64×64タイル、均等分配) + 512×1152×128 (32×32ブロック、split-2) に分割
  • 1カーネルで記述可能（独立な計算）、改善は0.1ms未満だが興味深い事例

4.4 Fusing the Scalar Operations

  • bias、残差接続、活性化関数をGEMMに統合
  • RMS Normは先にトークンレベル統計を計算し、次のGEMMで累積後に対応ファクタで除算
  • 推定約4msの改善

カーネル詳細 (2 views, RTX 4090)

コンポーネントRooflinecuBLASTriton実測
Vision Encoder2.485ms4.334ms3.957ms4.059ms
LLM10.727ms12.875ms12.732ms12.503ms
Action Expert6.486ms10.808ms9.831ms11.001ms
合計19.698ms28.017ms26.520ms27.299ms

5 Establishing Lower Bound

5.1 Roofline of the GEMMs

  • BF16 GEMM (N×K×M) の下限: t_roofline = max(2KM / T_bandwidth, NKM / T_compute)
  • RTX 4090: メモリ帯域幅 1.01 TB/s、BF16 MAC/s 91.4T (ブースト周波数2.79GHz考慮)
  • Roofline下限:
1 view2 views3 views
12.8ms19.7ms26.7ms

5.2 Synchronization Overhead

  • 1378回のmatmulに対する同期オーバーヘッドを測定
方法時間オーバーヘッド
Pytorch (逐次起動)13.81ms+12.92ms
CUDA Graph2.61ms+1.72ms
Software Barrier1.75ms+0.86ms
Fused (同期なし)0.89ms+0ms
  • Software Barrier: ブロック数=SM数のとき、グローバルメモリでバリア実装可能
  • 同期込み下限:
1 view2 views3 views
13.7ms20.6ms27.6ms
  • 現在の実装は最適実装まで残り約30%の余地

6 Full Streaming Inference

6.1 Overlapped Streaming Inference

  • AEカーネル (IO律速) とVLMカーネル (計算律速) を並行実行でスループット向上
構成推論時間
Sequential VLM + 10 AE27.3ms
Concurrent VLM + 10 AE26.3ms
Concurrent VLM + 16 AE32.7ms
  • 並行実行なら1秒あたり30 VLM + 480 AE実行が可能

6.2 Re-thinking Action Expert: Up to 480 Hz

  • 480Hzの制御周波数実現のために、高周波サンプリング可能なセンサ（力センサ 2kHz+等）をAEに入力
  • Flow Matchingを「漸次的」生成に書き換え：各ステップで部分アクションリストを生成（RTC類似）
  • 2つのフィードバックループ:
    • Quick loop: 高周波信号入力 → AE処理 → 反応軌道生成（最速2ms）
    • Slow loop: 画像駆動ループ、VLM処理後AEに反映（最速1/30s）

6.3 Fusing VLMs: Going below 1Hz

  • VLMの30Hzフレームエンコーディングにテキスト推論をピギーバック
  • 1つの重み行列ロード後にVLM計算とテキストデータ推論を順次実行
  • 追加のオートリグレッシブテキストストリーム（30 token/s）で対話・推論が可能

6.4 Summary

  • 3つのフィードバックループ:
    • Force loop (480Hz): AEによる高周波入出力
    • Visual loop (30Hz): VLMによる画像反応
    • Textual loop (<1Hz): テキストベースの対話・推論

7 Real World Validation

7.1 Setup and Data Collection

  • 垂直配置の2つのグラバーによるペン把持タスク
  • 30FPS 720P USBカメラ（遅延約2フレーム）
  • ペン落下距離約30cm、200msで把持する600エピソード収集
  • 2ビュー入力（現在フレーム + 前フレーム）で速度情報も取得

7.2 Training and Inference

  • openpiリポジトリで訓練、空プロンプト、数エポック
  • マルチスレッド構成: カメラスレッド → リングバッファ → 推論スレッド → 出力バッファ → 送信スレッド
  • GPU暖機後に安定速度で実験開始

7.3 Result

  • 10回連続の把持実験で100%成功率
  • タスク自体は単純だが、システム視点ではVLA実装の低レイテンシを実証
  • 人間の反応（30cm落下距離）と同等の反応速度
  • 手動設計カーネル: XLA, TensorRT, CUTLASS, Triton, FlashAttention, FlashDecoding
  • スケジューリングベース: TVM（アルゴリズムと実行スケジュールの分離）
  • 超最適化: Mirage, Neptune（階層的最適化だが検索時間長い）

9 Future Directions

9.1 Visual-based Latency: Going to 60-120 FPS

  • 低精度（8bit）量子化や適応的ビュー選択で高FPS化
  • 人間は30FPSと60FPSの差を認識可能、60FPSが次の目標

9.2 Larger Models: Going to 7B and Beyond

  • RTX 5090は帯域幅 (1.79TB/s) が大幅向上、AEスケーリングに有利
  • 7Bパラメータが次のマイルストーン

9.3 More Fine-grained Feed-back Loop

  • AE内の各層レベルでの信号注入による更なる高周波制御の可能性（現状未解決）

Graph View

Backlinks