Paged KV Cache Design in LingBot-Map: Two-Stream Sliding-Window Architecture

The paged KV cache in LingBot-Map implements a two-stream memory architecture managed by FlashInferKVCacheManager that separates recyclable patch tokens from permanent special tokens, using sliding-window eviction to maintain constant memory bounds during streaming video inference.

The LingBot-Map repository optimizes transformer inference for long video sequences through a sophisticated key-value caching mechanism. This design centers on the FlashInferKVCacheManager class in lingbot_map/layers/flashinfer_cache.py, which organizes KV tensors into distinct logical streams backed by a unified physical page pool to enable efficient cross-frame attention with bounded memory consumption.

Two-Stream Logical Architecture

The cache divides tokens into two logical streams that share a single physical page pool per transformer block. This separation allows the system to recycle high-volume patch data while preserving low-volume metadata essential for every frame.

The Patch Stream (Recyclable)

The patch stream stores per-frame image tokens—the bulk of the visual data. According to the module docstring at lines 4‑15, this stream operates with the following characteristics:

• Page allocation: One page per frame, with page_size equal to patches_per_frame (e.g., 256 tokens for 224×224 inputs)
• Scale frames: A small fixed set (default 8) stored in scale_patch_pages that are never evicted
• Live window: Recent frames stored in live_window_patch_pages that are evicted when exceeding the sliding_window limit (default 64)

The Special Stream (Append-Only)

The special stream holds a fixed set of metadata tokens per frame (camera parameters, registers, scale indicators—default 6 tokens). Unlike patch pages, special pages are never recycled:

• Storage: Packed continuously where one special page holds floor(page_size/6) frames
• Lifetime: Persists for the entire sequence, ensuring metadata remains accessible even after patch eviction
• Growth: Monotonically increasing, avoiding the rewrite overhead that would occur if mixed with recyclable patch data

Physical Page Pool Layout

For each transformer block, the KV cache allocates a 5‑D tensor in __init__ (lines 31‑41):

kv_caches[block_idx]: [max_num_pages, 2, page_size, H, D]

• Dimension 0: Page index (0 to max_patch_pages‑1 for patch pages; max_patch_pages to max_num_pages‑1 for special pages)
• Dimension 1: KV split (0 = Key, 1 = Value)
• Dimensions 2‑4: page_size, number of heads (H), and head dimension (D)

This physical layout enables FlashInfer to access both streams through a unified page table while maintaining logical separation.

Core Cache Operations

The manager exposes several key methods that streaming models call every frame:

Appending Frames

The append_frame method (lines 202‑214) splits incoming K/V tensors into special and patch components. It writes patch pages to either the scale set or live window, then appends special tokens to the special stream:

for blk in range(manager.num_blocks):
    k, v = make_kv()  # Shape: [262, 16, 64] for 256 patches + 6 specials

    manager.append_frame(blk, k, v)

Sliding-Window Eviction

The evict_frames method (lines 229‑245) implements the cache's memory bounding strategy. It recycles patch pages older than the sliding window while preserving scale frames and all special pages:

manager.evict_frames(
    block_idx=0,
    scale_frames=8,      # Permanent reference frames

    sliding_window=64    # Maximum recent frames to retain

)

Attention Computation with FlashInfer

The compute_attention method interfaces with FlashInfer's BatchPrefillWithPagedKVCacheWrapper. For block 0, it creates a plan using build_visible_page_table (lines 386‑393), which returns the ordered list scale → live window → specials. Subsequent blocks reuse this plan. The method also handles partial final pages through compute_last_page_len (lines 352‑362) to determine valid token counts in non-full special pages.

Integration in Streaming Inference

Streaming models in lingbot_map/models/gct_stream.py and its windowed variants instantiate the manager and invoke it within the per-frame loop:

from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager

# Initialization

self.kv_cache = FlashInferKVCacheManager(
    num_blocks=num_transformer_blocks,
    max_num_frames=88,              # scale + window + headroom

    tokens_per_frame=262,           # patches + specials

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

# Per-frame processing

self.kv_cache.append_frame(block_idx, k, v)
self.kv_cache.evict_frames(
    block_idx,
    scale_frames=self.cfg.kv_cache_scale_frames,
    sliding_window=self.cfg.kv_cache_sliding_window
)
output = self.kv_cache.compute_attention(block_idx, q)

Practical Code Examples

Instantiating the Cache Manager

import torch
from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

mgr = FlashInferKVCacheManager(
    num_blocks=2,
    max_num_frames=88,
    tokens_per_frame=262,  # 256 patch + 6 special

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

Processing Frames with Automatic Eviction

def process_sequence(frames):
    for frame_idx, (k, v) in enumerate(frames):
        for blk in range(mgr.num_blocks):
            mgr.append_frame(blk, k, v)
            mgr.evict_frames(blk, scale_frames=8, sliding_window=64)

Computing Cross-Frame Attention

# Query tensor for current frame

q = torch.randn(262, 16, 64, dtype=torch.bfloat16, device=device)

# Compute attention using cached KV from all visible frames

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

Resetting Between Video Sequences

# Clear all page deques and counters for new video

mgr.reset()

Monitoring Cache Statistics

stats = mgr.get_cache_stats(block_idx=0)
print(stats)

# Output: {'frame_count': 10, 'scale_pages': 8, 'live_pages': 2,

#          'free_pages': 78, 'special_tokens': 60}

Summary

• Two-stream design: Separates high-volume recyclable patch tokens from low-volume permanent special tokens, minimizing data movement during eviction.
• Physical layout: 5‑D tensor [max_num_pages, 2, page_size, H, D] with partitioned indices for patch and special pages.
• Sliding-window eviction: Bounds memory usage to scale_frames + sliding_window patch pages while preserving all special tokens.
• FlashInfer integration: Uses build_visible_page_table to present contiguous logical views of non-contiguous physical pages.
• Streaming optimization: Designed for per-frame processing in gct_stream.py variants with constant-time append and eviction operations.

Frequently Asked Questions

How does the two-stream design optimize memory usage?

The two-stream design prevents fragmentation and unnecessary copying by separating tokens with different lifetimes. Patch tokens are large (hundreds per frame) and can be safely discarded when outside the sliding window, while special tokens are tiny (six per frame) but needed for every historical frame. By keeping specials in an append-only stream, the system avoids rewriting metadata after every eviction, reducing memory bandwidth and simplifying the page table.

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

Scale frames constitute a small fixed set (default 8) of reference frames that remain permanently in the patch cache to provide stable spatial anchors for the model. Sliding window frames (default max 64) are the most recent frames kept for temporal context. When new frames arrive, the oldest sliding window pages are recycled, while scale frames persist indefinitely alongside the growing special token stream.

How does FlashInfer integrate with the page table?

FlashInfer receives a logical page table built by build_visible_page_table (lines 386‑393), which concatenates scale pages, live window pages, and special pages into a contiguous addressing scheme. The compute_attention method creates a FlashInfer plan on the first transformer block that maps these logical indices to the physical 5‑D tensor. Subsequent blocks reuse this plan, avoiding redundant metadata computation during multi-layer inference.

What are the default configuration parameters for video inference?

The default configuration assumes 224×224 image inputs with 256 patch tokens per frame plus 6 special tokens. The cache reserves 8 scale frames permanently and maintains a sliding window of 64 recent frames, with the physical pool sized to scale_frames + sliding_window + headroom (typically 88 total pages) to prevent allocation failures during peak usage.

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