How LingBot-Map Uses FlashInfer KV Cache for Streaming 3D Reconstruction

LingBot-Map implements a two-stream paged KV cache manager on top of FlashInfer that separates recyclable patch tokens from append-only special tokens, enabling constant-time per-frame attention for continuous 3D reconstruction from video streams.

LingBot-Map is an open-source framework for streaming 3D reconstruction that processes video frames incrementally using transformer-based architectures. At its core lies a sophisticated KV cache architecture integrated with FlashInfer, a high-performance kernel library designed for paged attention. This architecture allows the system to maintain temporal context across arbitrarily long sequences without linear memory growth, storing both fixed "scale" frames and a sliding window of recent patch tokens while keeping special tokens permanently resident.

Two-Stream Paginated KV Cache Architecture

The FlashInferKVCacheManager class in lingbot_map/layers/flashinfer_cache.py implements a novel two-stream design that partitions the transformer's key-value tensors into logically separate but physically co-located storage. This design optimizes memory usage for streaming inference where different token types have distinct lifecycles.

Patch Stream and Special Stream

The manager maintains two logical streams that share a single physical page pool per transformer block:

Stream	Purpose	Allocation Strategy
Patch stream (recyclable)	Stores regular image patch tokens for each frame.	Pages 0 to max_patch_pages-1. The first scale_frames pages are kept permanently (scale frames), while the next sliding_window pages form the live window that is recycled as older frames exit.
Special stream (append-only)	Holds fixed-size special tokens (camera token, register tokens, scale token).	Pages max_patch_pages to max_num_pages-1. These are never evicted; the manager appends tokens until the pre-allocated pool is exhausted.

Physical Memory Layout

Per transformer block, the physical tensor layout follows:

kv_caches[block]: [max_num_pages, 2 (K/V), page_size, num_heads, head_dim]

Where max_num_pages = max_patch_pages + max_special_pages and page_size equals patches_per_frame (or the next power-of-two for SM90 kernels). All caches are per-block (one per transformer layer), but page IDs remain identical across blocks, enabling a single execution plan to be reused for all layers.

FlashInfer Integration Mechanics

The architecture leverages BatchPrefillWithPagedKVCacheWrapper from the FlashInfer library to execute attention kernels efficiently over the paged memory structure.

Page Table Planning and Reuse

For each frame step, the manager calls plan() once to build the visible page table using IDs from block 0 (scale_patch_pages + live_window_patch_pages + all_special_pages). This plan is then reused across all layers via run(). The visible page table is ordered scale → window → specials, with paged_kv_last_page_len specifying the exact token count for the final partially-filled page.

Attention Kernel Execution

When compute_attention(block_idx, q) is invoked, the manager executes the FlashInfer prefill kernel using the pre-built plan. For numerical accuracy checks, setting force_fp32=True triggers a fallback that gathers visible K/V into dense tensors and calls PyTorch's scaled_dot_product_attention instead of the FlashInfer kernel.

Streaming 3D Reconstruction Workflow

FlashInferAttention in lingbot_map/layers/attention.py orchestrates the two-phase streaming process that drives the reconstruction pipeline.

Scale Frames vs Sliding Window

The system distinguishes between two types of temporal context:

• Scale frames (kv_cache_scale_frames, default 8): These are never evicted and provide stable global context for camera motion estimation.
• Sliding-window frames (kv_cache_sliding_window, default 64): These constitute the live working memory; older pages beyond this window are recycled via evict_frames().
• Special tokens: Accumulate for the entire sequence and are never removed.

The manager also supports deferred eviction (_defer_eviction) for algorithms requiring frame rollback capabilities, such as key-frame selection heuristics.

Frame Processing Loop

The streaming forward method automatically selects between batch processing and single-frame streaming:

1. Scale Phase (num_frames > 1): Performs a full SDPA pass over multiple initialization frames, then appends each frame's K/V to the cache via append_frame().
2. Streaming Phase (num_frames == 1): Feeds the current query tensor to compute_attention(), which performs cross-frame attention against previously cached K/V using the pre-planned page table.

3D RoPE Position Encoding

Streaming reconstruction requires positional encodings that evolve across time. The StreamAggregator class in lingbot_map/aggregator/stream.py lazily initializes a 3-D rotary embedding (WanRotaryPosEmbed) via _init_3d_rope().

During each forward pass, _get_3d_positions_streaming() constructs a position tensor incorporating the global frame indices (f_start, f_end). These positions are applied to query and key tensors before they enter the KV cache, ensuring that static special tokens (camera, register, scale) carry correct temporal information while maintaining cache compatibility.

Implementation Guide

Initializing the Cache Manager

Create a FlashInferKVCacheManager instance for each video stream:

from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokens_per_frame = 262  # 256 patches + 6 specials

mgr = FlashInferKVCacheManager(
    num_blocks=12,                   # One per transformer block

    max_num_frames=88,               # scale + window + headroom

    tokens_per_frame=tokens_per_frame,
    num_heads=16,
    head_dim=64,
    dtype=torch.bfloat16,
    device=device,
    num_special_tokens=6,
    scale_frames=8,
    sliding_window=64,
    max_total_frames=200,
)

Appending Frames and Eviction

Write new frame data and recycle old pages:

# k_frame and v_frame have shape [tokens_per_frame, num_heads, head_dim]

mgr.append_frame(block_idx=0, k=k_frame, v=v_frame)

# Evict oldest live-window pages beyond the 64-frame sliding window

mgr.evict_frames(block_idx=0, scale_frames=8, sliding_window=64)

Running Attention in Streaming Mode

Compute attention for the current frame using cached history:

# q has shape [tokens_per_frame, num_heads, head_dim]

output = mgr.compute_attention(block_idx=0, q=q)

High-Level Attention Layer

Use the FlashInferAttention wrapper for integration with the reconstruction model:

from lingbot_map.layers.attention import FlashInferAttention

attn = FlashInferAttention(
    dim=1024,
    num_heads=16,
    kv_cache_sliding_window=64,
    kv_cache_scale_frames=8,
)

# Streaming inference for a single frame

output = attn(
    x,                    # [B, N, C] current frame tokens

    kv_cache=mgr,         # The FlashInferKVCacheManager instance

    num_frames=1,         # Streaming mode

)

Summary

• Two-stream architecture: Separates recyclable patch tokens from permanent special tokens, optimizing memory usage in lingbot_map/layers/flashinfer_cache.py.
• FlashInfer integration: Uses BatchPrefillWithPagedKVCacheWrapper with reusable page tables for constant-time attention regardless of sequence length.
• Temporal memory management: Combines fixed scale frames (default 8) with a sliding window (default 64) to balance global context against memory constraints.
• 3D RoPE encoding: Applies temporal positions before caching via StreamAggregator._get_3d_positions_streaming() to maintain spatiotemporal coherence.
• Deferred eviction: Supports rollback mechanisms for key-frame selection algorithms without invalidating the cache.

Frequently Asked Questions

What is the difference between scale frames and sliding window frames?

Scale frames are permanently resident pages (defaulting to 8 frames) that provide stable global context for camera motion estimation throughout the entire sequence. Sliding window frames constitute a recyclable buffer (default 64 frames) representing the recent temporal context; as new frames arrive, the oldest pages in this window are evicted to maintain constant memory usage. This hybrid approach allows LingBot-Map to maintain both long-term stability and recent detail without unbounded memory growth.

How does FlashInfer's BatchPrefillWithPagedKVCacheWrapper integrate with LingBot-Map?

The FlashInferKVCacheManager uses BatchPrefillWithPagedKVCacheWrapper to execute attention kernels over non-contiguous memory. It calls plan() once per frame step to build a view of the current page table (combining scale frames, live window, and special tokens), then reuses this plan across all transformer layers via run(). The paged_kv_last_page_len parameter ensures correct handling of the final partially-filled page, while the identical page IDs across blocks enable single-plan reuse for multi-layer architectures.

Why are special tokens kept in a separate append-only stream?

Special tokens (camera, register, and scale tokens) are kept in an append-only stream because they contain critical metadata that must persist for the entire video sequence and must never be evicted. Separating them from patch tokens prevents the eviction logic from accidentally discarding structural information required for 3D reconstruction, while allowing the patch stream to be aggressively recycled as the sliding window advances.

How is temporal information preserved in the KV cache?

Temporal information is encoded via 3-D Rotary Position Embeddings (RoPE) before tokens enter the cache. The StreamAggregator in lingbot_map/aggregator/stream.py computes positions incorporating global frame indices (f_start, f_end) and applies them to queries and keys via _get_3d_positions_streaming(). Because positions are baked into the tensors before the append_frame() call, the cached values carry correct temporal coordinates without requiring position lookup tables during attention computation.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}