vllm-project/vllm-omni

[RFC]: Support 3D World Model (VGGT, etc.)

Open

#2,727 opened on Apr 13, 2026

View on GitHub
 (6 comments) (9 reactions) (2 assignees)Python (1,067 forks)github user discovery
help wantednew model

Repository metrics

Stars
 (4,990 stars)
PR merge metrics
 (PR metrics pending)

Description

1. Summary

This RFC proposes introducing a third engine type in vLLM-Omni — the Reconstruction Engine (stage_type: "reconstruct"), with VGGT (CVPR 2025 Best Paper, 1.26B) as the first supported model.

Engine stage_type Computation Mode
AR Engine llm Autoregressive token-by-token generation
Diffusion Engine diffusion Iterative T-step denoising generation
Reconstruction Engine reconstruct Chunked feed-forward reconstruction

The Reconstruction Engine adopts a dLLM-style chunked computation + sliding window KV Cache architecture: input frames are grouped by chunk_size, and global attention is computed only over the KV of the most recent sliding_window_num chunks. The effective sequence length is sequence_length = sliding_window_num × chunk_size, which strictly bounds peak memory independent of total frame count N. This design naturally supports streaming video input, KV compression, Ring Attention parallelism, and quantization, while paving the way for future dLLM integration.


2. Background

2.1 Why a Reconstruction Engine?

vLLM-Omni currently covers generation and understanding of text, image, video, and audio, but lacks 3D spatial reconstruction capabilities. VGGT infers camera poses, depth maps, point clouds, and 3D point tracks from multiple 2D images in a single forward pass — this "recovering structure from input" computation pattern is neither autoregressive generation nor iterative denoising, requiring a dedicated engine.

Moreover, VGGT's outputs naturally feed into downstream 3D generation tasks (3DGS, NeRF, conditional diffusion), serving as a Reconstruct → Diffusion preceding stage in vLLM-Omni's multi-stage pipeline.

2.2 VGGT Architecture

Image1  Image2  ...  ImageN  (hundreds to thousands)
  │       │           │
  ▼       ▼           ▼
┌──────────────────────────────┐
│ Patch Embed (DINOv2-L frozen) │  ~86M params
│ Each image → ~1374 tokens     │
└──────────┬───────────────────┘
           │
  ╔════════▼═══════════════╗
  ║  Aggregator × 24 layers ║  ~1.1B params
  ║  Frame self-attn ↕ alt.  ║
  ║  Global self-attn (all)  ║  ← Bottleneck O((N×L)²)
  ╚════════╤═══════════════╝
           │ Multi-scale features
    ┌──────┼──────┬──────────┐
    ▼      ▼      ▼          ▼
 Camera  Point  Depth     Track
  Head    Head   Head      Head

Key properties: Purely feed-forward (no iterative loop), multi-task shared backbone, global attention as the compute/memory bottleneck.

2.3 Core Challenges and Existing Solutions

Global attention performs self-attention over all frame tokens — at 100 frames the attention matrix reaches ~35 GB, causing single-GPU OOM. VGGT can only handle 60–80 frames on a 24 GB RTX 4090; KITTI Seq 00 (4600 frames) would require ~1400 GB of memory.

The research community has tackled this bottleneck from two directions:

VGGT-Long (ICRA 2026) — Chunking + Alignment + Loop Closure:

  • Splits long sequences into overlapping chunks, each independently processed by vanilla VGGT
  • Adjacent chunks are aligned via overlapping frames' 3D point clouds using confidence-weighted Sim(3) alignment (IRLS + Huber loss), where high-confidence static structures dominate alignment, filtering out sky/dynamic objects
  • Global loop closure detection + Sim(3) LM optimization corrects accumulated drift
  • No retraining required; successfully extends VGGT to kilometer-scale autonomous driving scenarios (KITTI/Waymo)

FastVGGT (ICLR 2026) — Token Sparsification:

  • Identifies a token collapse phenomenon in global attention: many tokens have nearly identical attention distributions, causing redundant computation
  • Proposes a three-class token partitioning strategy: initial-frame tokens (global reference anchors, non-mergeable), salient tokens (top-k by norm, preserving details), uniform tokens (region-based random sampling for spatial coverage)
  • Training-free, achieves 4× speedup on 1000-frame scenarios while suppressing long-sequence error accumulation

Sparse-VGGT — Block-Sparse Attention (concurrent work):

  • Empirically finds that global attention probability mass concentrates on cross-view geometric matching patch-patch interactions
  • Replaces dense global attention with a block-sparse attention kernel, achieving ~4× speedup
  • Applicable to both VGGT and π³

These three works validate the feasibility of the chunking + sparsification design in this RFC, and provide direct implementation references for the Reconstruction Engine's acceleration system:

                    VGGT-Long          FastVGGT / Sparse-VGGT
                    (chunk alignment)    (token sparsification)
                        │                       │
                        ▼                       ▼
    ReconEngine maps:  ChunkManager           KV compression strategies
                      + Sim(3) post-alignment  (AVGGT / BlockSparse / TokenMerge)
                      + loop closure OutputProcessor

Relationship to dLLM / VGGT-Long:

ReconEngine's chunking strategy fuses dLLM chunked prefill and VGGT-Long concepts, but introduces a sliding window KV Cache to enforce a strict memory upper bound:

Key parameters:
  chunk_size         = number of frames per chunk (e.g., 8)
  sliding_window_num = number of chunks retained in the window (e.g., 5)
  sequence_length    = sliding_window_num × chunk_size (e.g., 40 frames)
  ↑ Global attention is computed only over tokens within sequence_length

VGGT-Long:   Independent chunk inference (sliding_window_num = 1)
             → Sim(3) post-alignment → loop closure optimization
             Memory = chunk_size × L attention (fixed)

ReconEngine: Sliding window retains KV of the most recent sliding_window_num chunks
             → Current chunk attends to all KV within the window (approx. global attn)
             → KV of chunks outside the window is evicted; their outputs are
               stitched via Sim(3) alignment
             Memory bound = sliding_window_num × chunk_size × L (fixed)

dLLM analogy: sliding window attention (Mistral-style)
              sliding_window_num chunks ≈ window_size token groups
sliding_window_num Behavior Memory
Full global attention (small-scale, all KV retained) Grows linearly with N
> 1 Sliding window + Sim(3) alignment outside window sliding_window_num × chunk_size × L (fixed)
1 Degenerates to VGGT-Long (fully independent chunks + post-alignment) chunk_size × L (minimum)

3. Design Overview

3.1 Architecture Overview

3.2 Component Responsibilities

Component Responsibility
Orchestrator Request entry point. Contains InputPreprocessor for image decoding/resizing/normalization; pushes raw pixel data to VerticalConnector and passes only metadata (shape, request ID, frame count) to ReconEngine
VerticalConnector Cross-process/cross-node large tensor transport channel (reuses vLLM-Omni's SharedMemory/Mooncake/Yuanrong), avoiding raw pixel transfer through Engine internal queues
ReconEngine Engine main body, orchestrates the full inference workflow
Scheduler Chunk-level scheduling: manages request queues, determines batching combinations and execution order
ChunkManager Internal Scheduler component: groups each request's N frames by chunk_size, maintains sliding window state (current window range, KV Cache eviction), tracks chunk progress
InputProcessor Pulls pixel data from VerticalConnector, performs patch embedding (DINOv2), builds chunk tensors and initializes KV Cache
Executor Core execution loop: drives Worker chunk-by-chunk through frame attention + global attention (sliding window KV Cache); KV of chunks outside the window is evicted (their outputs can be cached for subsequent Sim(3) alignment); triggers prediction heads after all chunks complete
Worker GPU compute unit: executes Aggregator forward_chunk (FlashAttention + KV Cache append) and prediction head inference; supports Ring Attention for multi-GPU distribution
OutputProcessor Takes Worker output (prediction head results), splits by request; when sliding_window_num < ∞, performs inter-window Sim(3) alignment (VGGT-Long style); filters/serializes results based on user-selected output_heads; pushes to VerticalConnector

3.3 Data Flow

Request arrives
  │
  ▼
Orchestrator.InputPreprocessor
  ├─ Real Data (pixel tensor) ──▶ VerticalConnector ──▶ InputProcessor
  └─ Meta Data (req_id, N frames, heads) ──▶ Scheduler
                                            │
                                    ChunkManager grouping:
                                    req_1: [chunk_0, chunk_1, ..., chunk_k]
                                    req_2: [chunk_0, chunk_1, ..., chunk_j]
                                            │
                                            ▼
                                    Executor chunk loop (sliding window):
                                    ┌─ for chunk_idx in range(max_chunks):
                                    │    InputProcessor.get_chunk(chunk_idx)
                                    │         │
                                    │         ▼
                                    │    Worker.forward_chunk()
                                    │    ├─ DINOv2 patch embed (CPU offloadable)
                                    │    ├─ Frame attention (FlashAttn, per-frame)
                                    │    ├─ Global attention (attend in-window KV + current)
                                    │    ├─ KV Cache append
                                    │    └─ Evict oldest chunk KV (if len > W)
                                    └─ end loop
                                            │
                                            ▼
                                    Worker.execute_heads()
                                    ├─ CameraHead → pose
                                    ├─ DepthHead  → depth map
                                    ├─ PointHead  → point cloud
                                    └─ TrackHead  → 3D tracks (optional)
                                            │
                                            ▼
                                    OutputProcessor
                                    ├─ Filter (return only requested heads)
                                    ├─ Serialize (tensor → base64/npz/ply)
                                    └─ ▶ VerticalConnector ──▶ Response

3.4 Acceleration System

Five orthogonal dimensions, independently enableable and combinable:

# Dimension Mechanism Reference Effect
dLLM-style chunked computation Global attention computed incrementally per chunk + sliding window KV Cache dLLM chunked prefill + sliding window attn; VGGT-Long chunk concept Peak memory = sliding_window_num × chunk_size × L (strict bound, independent of total N)
Streaming video input Process frames as they arrive, incrementally append KV Cache, periodically refresh prediction heads VGGT-Long sliding window; SLAM real-time requirements Online SLAM / real-time reconstruction
KV compression / token sparsification Reduce effective sequence length of global attention FastVGGT (token partitioning); Sparse-VGGT (block-sparse kernel); AVGGT (random KV subset) Reduce O(n²), up to 4× speedup
Ring Attention Global attention KV sharded across GPU ring Ring Attention paper Linear scaling to 500+ frames
Quantization Weight INT8/INT4/FP8 + KV Cache FP8/INT8 Mature LLM quantization methods ~50% memory reduction

③ KV Compression Strategy Details (based on FastVGGT / Sparse-VGGT papers):

Strategy Source Mechanism Speedup
FastVGGT Token Partition FastVGGT (ICLR 2026) Three-class partition: initial-frame tokens as global anchors (non-mergeable) + top-k salient tokens (preserve details) + remaining tokens via region-based random sampling ~4× (1000 frames)
Block-Sparse Attention Sparse-VGGT Global attention probability mass concentrates on cross-view geometric matches; block-sparse kernel skips low-weight blocks ~4×
AVGGT AVGGT paper Per-layer random KV subset sampling; different layers sample different subsets for coverage 8–10×
Token Merge FastVGGT variant Merge high cosine-similarity KV token pairs, addressing token collapse redundancy 2–4×

Typical combined configurations:

Scenario Chunking KV Compression Ring Attn Quantization
Entry (< 20 frames) Full None None None
Medium (20–100 frames) chunk=8 AVGGT 0.25 None W-INT8
Large-scale (100–500 frames) chunk=8 FastVGGT partition 4 GPU W-INT8 + KV-FP8
Kilometer-scale (VGGT-Long mode) chunk=60 + overlap AVGGT + Sim(3) alignment None W-INT8
Real-time streaming chunk=4 Block-Sparse 2 GPU W-FP8

3.5 OmniStage Integration

As a new stage type in vLLM-Omni's multi-stage pipeline, alongside existing llm / diffusion:

# omni_stage.py
if stage_config.stage_type == "llm":
    return OmniLLM(...)
elif stage_config.stage_type == "diffusion":
    return OmniDiffusion(...)
elif stage_config.stage_type == "reconstruct":
    return OmniReconstruct(...)  # New

Multi-stage pipeline example:

# Reconstruct → Diffusion (VGGT output as diffusion model conditioning)
stage_args:
  - stage_id: 0
    stage_type: reconstruct
    engine_args:
      model: facebook/VGGT-1B-Commercial
      chunk_size: 8
      kv_compression: avggt
      enable_depth: true
      enable_point: true
    final_output: false

  - stage_id: 1
    stage_type: diffusion
    engine_args:
      model: your-org/geo-conditioned-nvs
    engine_input_source: [0]
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.vggt.reconstruct2diffusion
    final_output: true

3.6 dLLM Reuse Path

The Reconstruction Engine's chunked computation is architecturally isomorphic to dLLM chunked prefill:

Dimension dLLM Chunked Prefill ReconEngine Chunked Attn
Input unit text token chunk frame token chunk (1 frame ≈ 1374 tokens)
KV Cache Incremental K/V append Sliding window KV Cache (retain most recent sliding_window_num chunks)
Memory bound context_window × head_dim sliding_window_num × chunk_size × L × head_dim
Beyond-window handling Discard / sink token Sim(3) alignment (VGGT-Long style)
Scheduling Chunk-batched prefill Frame-group chunk-batched global attn
Ring Attention Long-context across GPUs In-window KV across GPUs

Reuse layer: ChunkManager (sliding_window_num scheduling) + Sliding Window KV Cache Pool + Ring Attention + KV Compression Divergence layer: InputProcessor (patch embed vs tokenizer) + OutputProcessor (prediction heads + Sim(3) alignment vs decode)

Future dLLM models can reuse ReconEngine's core scheduling and execution path by swapping Input/OutputProcessor.


References

  • VGGT — CVPR 2025 Best Paper, 1.26B feed-forward 3D reconstruction
  • VGGT-Long — ICRA 2026, chunk + overlap + Sim(3) alignment, kilometer-scale long-sequence extension
  • FastVGGT — ICLR 2026, training-free three-class token partitioning sparsification, 4× speedup
  • Sparse-VGGT — Block-sparse global attention kernel, 4× speedup
  • AVGGT — Per-layer random KV subset sampling, 8–10× speedup
  • Ring Attention — Sequence parallelism
  • Sarathi-Serve — dLLM chunked prefill
  • vLLM-Omni — Omni-modality inference framework

Feedback Period.

No response

CC List.

@ywang96 @hsliuustc0106 @Gaohan123 @princepride @lishunyang12 @fake0fan @natureofnature @wtomin @SamitHuang @gcanlin

Any Other Things.

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Contributor guide