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.