How the KV Cache Impacts Performance in Robbyant/lingbot-map

The KV cache in lingbot-map reduces attention complexity from O(T²) to O(T·W) by storing key and value tensors in a two-stream paged memory system, enabling real-time streaming inference while maintaining bounded memory usage.

The lingbot-map repository implements a specialized KV cache mechanism that fundamentally changes how transformer models process sequential video frames. By leveraging a two-stream paged architecture implemented in lingbot_map/layers/flashinfer_cache.py, the system eliminates redundant computation of key and value tensors for previously processed frames. This design directly impacts inference speed and memory consumption, making it critical for streaming video applications.

Architecture of the Two-Stream Paged KV Cache

The cache splits memory management between two distinct streams to optimize for different token types. This separation is crucial for performance because it allows aggressive recycling of video patch memory while preserving special metadata tokens indefinitely.

Patch Stream (Recyclable Pages)

The patch stream holds regular patch tokens for each video frame. It allocates a fixed number of pages calculated as scale_frames + sliding_window + 16 headroom, where pages are reused when older frames leave the sliding window. This circular buffer approach prevents unbounded memory growth during streaming inference.

The implementation handles page allocation in _write_patch_page (lines ~74-87) through direct slice assignments:

self.kv_caches[block_idx][page_id, 0, :P] = patch_k

Special Stream (Append-Only Pages)

The special stream stores six persistent tokens per frame—camera parameters, register tokens, and scale information. Unlike the patch stream, these pages are never evicted, causing the special-token stream to grow monotonically with the sequence length. This ensures temporal consistency for critical metadata across all frames.

Page allocation occurs in _write_special_tokens (lines ~140-156), maintaining separate indices from the recyclable patch pages.

Physical Memory Layout

The cache stores tensors in a contiguous CUDA memory block to minimize GPU overhead. In __init__ (lines ~41-44), the system initializes:

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

Where dimension 1 distinguishes between K (index 0) and V (index 1) tensors. This layout enables coalesced memory access patterns during attention computation.

Visibility Table and FlashInfer Integration

The build_visible_page_table method (lines ~384-390) concatenates scale pages, window pages, and special pages into the exact order expected by FlashInfer's kernels. The compute_last_page_len function (lines ~452-470) tracks valid token counts in the final page, which is essential for the append-only special stream where pages may be partially filled.

Performance Optimizations in the KV Cache

The KV cache delivers performance gains through several architectural decisions that minimize computation and memory overhead.

Reduced Recomputation via Sliding Window

Without caching, self-attention recomputes K and V tensors for all previous frames at every step, resulting in O(T²) complexity. The KV cache reduces this to O(T·W), where W is the sliding window size (default 64). Only the current frame's K/V tensors are computed; all prior frames within the window are read from cache.

GPU-Friendly Paging

The cache uses a contiguous CUDA tensor allocated with torch.zeros and performs writes through direct slice assignments. This eliminates Python-to-CUDA overhead that would arise from per-token insertion operations, keeping the GPU execution pipeline saturated.

Single Planning Step for Multi-Layer Execution

In compute_attention (lines ~485-526), the system calls prefill_wrapper.plan once per frame step when processing block 0:

if block_idx == 0:
    prefill_wrapper.plan(...)

All subsequent transformer layers reuse this plan through prefill_wrapper.run, avoiding repeated index-building work for each of the 12+ layers in the model.

Deferred Eviction for Key-Frame Selection

The cache supports execute_deferred_eviction and rollback_last_frame (lines ~522-572), allowing frames to be temporarily retained beyond the sliding window. This enables key-frame selection algorithms to roll back frames without recomputation, trading temporary memory for algorithmic flexibility.

Memory-Efficient Separation

By isolating recyclable patch pages from never-evicted special pages, the cache bounds patch memory to the sliding window size while allowing special tokens to accumulate linearly. With 256 tokens per page, the special stream's growth remains modest even for long sequences.

Implementing the KV Cache in Practice

Initializing the KV Cache Manager

Create a FlashInferKVCacheManager for each transformer block:

import torch
from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
kv_manager = FlashInferKVCacheManager(
    num_blocks=12,                 # one manager per transformer block

    max_num_frames=88,             # scale + sliding window + headroom

    tokens_per_frame=262,          # 256 patches + 6 special tokens

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

Appending a Frame's K/V Tensors

Write new frame data to the cache:

# k, v are [tokens_per_frame, H, D] tensors produced by the model

k, v = make_kv()                         # user-defined helper

block_idx = 0                             # choose a transformer block

kv_manager.append_frame(block_idx, k, v) # writes patch & special pages

The append_frame method (lines ~200-227) automatically routes patch tokens to recyclable pages and special tokens to the append-only stream.

Evicting Old Frames

Maintain the sliding window by evicting stale frames:

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

The evict_frames method (lines ~260-267) recycles only the oldest window patch pages, preserving special tokens and scale frames.

Computing Attention with Cached KV

Query the cache during forward passes:

# q: [tokens_per_frame, H, D] for the current frame

out = kv_manager.compute_attention(block_idx=0, q=current_q)

The compute_attention method (lines ~485-531) builds the visibility table once per frame and executes FlashInfer's optimized kernels.

Integration in the Attention Layer

Use the cache within the model's attention mechanism:

from lingbot_map.layers.attention import FlashInferAttention

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

# In the model forward loop (streaming mode):

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

    kv_cache=kv_manager,   # pass the manager

    global_idx=0,
    num_frames=1,
)

The FlashInferAttention.forward method (lines ~460-680) detects the manager, appends K/V tensors, and executes attention using the cached history.

Summary

• Complexity reduction: The KV cache transforms attention from O(T²) to O(T·W) by caching historical K/V tensors, where W is the configurable sliding window size.
• Dual-stream architecture: Separating recyclable patch pages from append-only special pages optimizes memory usage while preserving critical metadata.
• GPU optimization: Contiguous CUDA tensors and single-plan execution minimize Python overhead and kernel launch costs.
• Flexible eviction: Deferred eviction supports advanced features like key-frame selection without sacrificing the caching benefits.
• Real-time streaming: Bounded memory growth and reduced computation enable processing of arbitrarily long video sequences in near real-time.

Frequently Asked Questions

What is the memory complexity of the KV cache in lingbot-map?

The patch stream maintains O(W) memory complexity, where W is the sliding window size (default 64), because pages are recycled when frames leave the window. The special stream grows linearly O(T) with total frames but uses large pages (256 tokens) to minimize overhead. This contrasts with standard transformers that require O(T) memory for all layers.

How does the two-stream design improve performance compared to single-stream caching?

The two-stream design allows the system to aggressively recycle memory for video patches (which constitute ~98% of tokens) while permanently retaining special tokens (camera, registers, scale) that must remain visible to all future frames. Without this separation, either special tokens would be incorrectly evicted breaking temporal consistency, or patch memory would grow unbounded, causing OOM errors during long videos.

Can the KV cache handle variable frame sizes or dynamic insertion?

The cache handles fixed-size frames efficiently through pre-allocated pages, but supports variable valid lengths within the final page via compute_last_page_len. For dynamic insertion or out-of-order processing, the deferred eviction mechanism (rollback_last_frame) allows temporary retention of frames beyond the sliding window, though standard inference assumes sequential frame arrival.

What happens when the sliding window exceeds the allocated headroom?

When the sliding window moves forward, evict_frames identifies the oldest patch pages for recycling. The system reserves 16 pages of headroom beyond scale_frames + sliding_window to accommodate this. If the sequence exceeds max_total_frames (default 2048), the special stream would theoretically grow too large, but in practice, the recyclable patch stream maintains constant memory regardless of video length.

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