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

> Discover how LingBot-Map leverages FlashInfer KV cache for efficient streaming 3D reconstruction. Explore its innovative two-stream architecture for constant-time per-frame attention.

- Repository: [Robbyant/lingbot-map](https://github.com/Robbyant/lingbot-map)
- Tags: deep-dive
- Published: 2026-06-29

---

**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`](https://github.com/Robbyant/lingbot-map/blob/main/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:

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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:

```python
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:

```python

# 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:

```python

# 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:

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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.