How LingBot-Map Implements the KV Cache Mechanism for Streaming Inference with FlashInfer Integration

LingBot-Map uses a two-stream paged KV cache that separates recyclable patch pages from append-only special pages, integrated with FlashInfer's BatchPrefillWithPagedKVCacheWrapper to execute attention in a single kernel launch during streaming inference.

LingBot-Map is a causal transformer architecture designed for streaming video processing where each new frame must attend only to past frames while maintaining strict memory bounds. The repository implements a sophisticated KV cache mechanism for streaming inference that combines a custom FlashInferKVCacheManager with FlashInfer integration to eliminate key-value re-materialization and minimize GPU memory footprint.

The Two-Stream Paged KV Cache Architecture

The core of the system is the FlashInferKVCacheManager class defined in lingbot_map/layers/flashinfer_cache.py (lines 55-71). This manager implements a novel two-stream design that treats visual patch tokens and metadata special tokens differently to optimize for both memory efficiency and attention correctness.

Patch Pages vs. Special Pages

The cache maintains two distinct page pools per transformer block:

• Patch pages: Recyclable pages that store the visual patch tokens of each video frame. These follow a sliding window eviction policy to bound memory usage.
• Special pages: An append-only stream that stores the six "special" tokens (camera parameters, registers, and scale embeddings). These persist across frames to maintain architectural state.

This separation is crucial because it allows the build_visible_page_table() method (lines 386-400) to construct a deterministic visibility order: scale frames → live sliding window → special tokens. The special tokens always appear at the end of the sequence, ensuring they are attended to correctly without shifting indices during frame eviction.

Append and Evict API

The manager provides explicit control over cache lifecycle through methods implemented in flashinfer_cache.py (lines 199-260):

• append_frame(block_idx, K, V): Writes the key/value tensors of a new frame into the appropriate page pool. For non-keyframes, deferred eviction is supported via the _skip_append flag and rollback_last_frame() mechanism.
• evict_frames(...): Recycles the oldest patch pages from the live sliding window while optionally preserving special tokens. This runs in O(1) amortized time per frame.

The eviction logic respects three memory regions: the fixed scale window (recent history), the live sliding window (active context), and the free pool. By tracking which pages belong to which region, the system maintains bounded memory of scale_frames + sliding_window + headroom patch pages plus a constant number of special pages.

Physical Page Layout and Indexing

The physical storage uses a compact tensor layout defined in flashinfer_cache.py (lines 40-44):

# Shape: [max_num_pages, 2, page_size, num_heads, head_dim]

# Where index 0 = Key, index 1 = Value

self.kv_caches = torch.empty(
    max_num_pages, 2, page_size, num_heads, head_dim,
    dtype=dtype, device=device
)

Indices 0 through max_patch_pages-1 are reserved for patch pages, while the remainder store special pages. The compute_last_page_len() method (lines 520-562) calculates how many valid tokens reside in the final (potentially partially-filled) special page, allowing FlashInfer to avoid custom masking logic.

FlashInfer Integration for Attention Computation

FlashInfer is integrated as a drop-in replacement for standard attention through three key integration points, enabling single-kernel-launch attention on the paged KV cache.

FlashInferAttention Layer Implementation

The FlashInferAttention class in lingbot_map/layers/attention.py (lines 352-498) implements a dual-mode operation:

Phase 1 (Batch Mode): During initialization or bulk processing, the layer executes standard dense attention on all frames, then sequentially calls manager.append_frame() to populate the cache.

Phase 2 (Streaming Mode): For live inference, the layer prepares Q/K/V in NHD format via prepare_qkv(), appends them to the cache, triggers eviction, and delegates computation to manager.compute_attention().

Kernel-Level Execution Flow

Inside FlashInferKVCacheManager.compute_attention() (lines 186-194 in flashinfer_cache.py), the system builds a plan once per frame step that describes the visible page indices and last-page lengths:

self.prefill_wrapper.plan(
    self._qo_indptr,
    paged_kv_indptr,
    paged_kv_indices,
    paged_kv_last_page_len,
    num_qo_heads=self.num_heads,
    num_kv_heads=self.num_heads,
    head_dim_qk=self.head_dim,
    page_size=self.page_size,
    causal=False,
    pos_encoding_mode="NONE",
    q_data_type=self.dtype,
)

This plan is reused across all transformer blocks. The actual computation invokes self.prefill_wrapper.run(), which maps directly to FlashInfer's BatchPrefillWithPagedKVCacheWrapper. This executes scaled dot-product attention on the GPU by fetching only the required K/V pages, avoiding unnecessary data movement.

Streaming Block Integration

The FlashInferBlock class in lingbot_map/layers/block.py (lines 51-78) orchestrates the per-block operations. In streaming mode, it calls self.attn_pre() to apply normalization, linear projections, and RoPE (Rotary Position Embedding), yielding NHD-formatted tensors. It then follows the append-evict-compute pattern, ensuring that each transformer block operates on the same consistent view of the cached history.

Implementation Examples

Lazy Initialization of the KV Cache

The AggregatorStream class in lingbot_map/aggregator/stream.py (lines 79-103) creates the cache manager lazily to defer memory allocation until the first forward pass:

def _get_flashinfer_manager(self, device, dtype, tokens_per_frame=None):
    if self.kv_cache_manager is None:
        tokens_per_frame = tokens_per_frame or self.num_patches + self.num_special_tokens
        self.kv_cache_manager = FlashInferKVCacheManager(
            num_blocks=self.depth,
            max_num_frames=self.max_frame_num,
            tokens_per_frame=tokens_per_frame,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            dtype=dtype,
            device=device,
            num_special_tokens=self.num_special_tokens,
            scale_frames=self.kv_cache_scale_frames,
            sliding_window=self.kv_cache_sliding_window,
        )
    return self.kv_cache_manager

Streaming Forward Pass

The following pattern from lingbot_map/layers/block.py demonstrates how a single frame flows through the system:

def forward(self, x, pos=None, kv_cache=None, global_idx=0, ...):
    is_streaming = kv_cache is not None and (num_frames is None or num_frames <= 1)
    if is_streaming:
        manager = kv_cache
        q_nhd, k_nhd, v_nhd = self.attn_pre(x, pos=pos, enable_3d_rope=True)
        
        manager.append_frame(global_idx, k_nhd, v_nhd)
        manager.evict_frames(
            block_idx=global_idx,
            scale_frames=self.attn.kv_cache_scale_frames,
            sliding_window=self.attn.kv_cache_sliding_window,
            cross_frame_special=self.attn.kv_cache_cross_frame_special,
            camera_only=self.attn.kv_cache_camera_only,
        )
        attn_x = manager.compute_attention(global_idx, q_nhd)
        
        # Reshape and residual connection

        attn_x = attn_x.reshape(x.shape[0], q_nhd.shape[0], -1)
        x = x + self.ls1(self.attn.proj(attn_x))

End-to-End Usage

To use the streaming aggregator with FlashInfer enabled:

from lingbot_map.aggregator.stream import AggregatorStream
import torch

model = AggregatorStream(
    depth=12,
    embed_dim=768,
    num_heads=12,
    use_sdpa=False,  # Force FlashInfer path

    kv_cache_sliding_window=64,
    kv_cache_scale_frames=8,
).cuda()

# Process frames sequentially

frame = torch.randn(1, 262, 768).cuda()  # 256 patches + 6 specials

output = model(frame, num_frames=1)  # KV cache initialized automatically

next_frame = torch.randn(1, 262, 768).cuda()
output = model(next_frame, num_frames=1)  # Reuses existing cache

Summary

• Two-stream architecture: Separates recyclable patch pages from persistent special tokens, enabling efficient sliding window eviction while preserving camera and register states.
• O(1) operations: Appending and evicting frames execute in constant amortized time through direct page table manipulation in FlashInferKVCacheManager.
• FlashInfer integration: Uses BatchPrefillWithPagedKVCacheWrapper to execute attention in a single CUDA kernel launch, avoiding Python overhead and minimizing memory bandwidth.
• Lazy initialization: The AggregatorStream defers cache allocation until the first forward pass, allowing dynamic configuration of frame dimensions.
• Fallback support: When FlashInfer is unavailable, the system gracefully degrades to SDPAAttention using standard PyTorch scaled dot-product attention.

Frequently Asked Questions

What is the difference between patch pages and special pages in the KV cache?

Patch pages store the visual tokens extracted from video frames and are recycled according to the sliding window policy to maintain bounded memory. Special pages store the six architectural tokens (camera, registers, scale) and use an append-only strategy, ensuring these metadata tokens persist across all frames and appear at the end of the attention sequence.

How does FlashInfer reduce latency compared to standard attention?

FlashInfer eliminates the need to reconstruct the full key-value tensor from Python data structures by operating directly on the paged physical layout. The BatchPrefillWithPagedKVCacheWrapper fetches only the required pages and executes attention in a single kernel launch, reducing kernel dispatch overhead and memory bandwidth usage compared to standard PyTorch attention that requires contiguous tensor reassembly.

Can LingBot-Map run without FlashInfer installed?

Yes. The repository provides a fallback SDPAAttention implementation in lingbot_map/layers/attention.py that uses PyTorch's native scaled dot-product attention. Set use_sdpa=True when instantiating AggregatorStream to disable the FlashInfer code path, though this will result in higher memory usage and latency during streaming inference.

How is memory bounded in streaming mode?

Memory is strictly bounded by the configuration parameters scale_frames, sliding_window, and num_special_tokens. The FlashInferKVCacheManager pre-allocates max_patch_pages (equal to scale_frames + sliding_window + headroom) and a fixed number of special pages. When the live window exceeds its limit, evict_frames() recycles the oldest patch pages back to the free pool, ensuring GPU memory usage remains constant regardless of stream duration.

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