How the Paged KV Cache Attention Mechanism Enables Long Sequence Inference in LingBot-Map

The paged KV cache attention mechanism uses a two-stream memory architecture that bounds GPU memory to a fixed sliding window plus persistent special tokens, enabling processing of arbitrarily long video sequences without quadratic memory growth.

LingBot-Map is an open-source video processing framework designed for streaming inference over long visual sequences. Traditional transformer architectures require storing all past key-value (KV) pairs, causing memory consumption to explode as O(L²) with sequence length. This repository solves the scalability challenge through a paged KV cache attention mechanism that implements logical memory paging directly inside the FlashInfer attention kernel.

Two-Stream Page Design

The system partitions KV cache memory into two distinct streams with different eviction policies. This separation allows the model to discard old visual information while preserving critical metadata required for coherent video understanding.

Patch Stream (Recyclable)

The patch stream stores dense visual tokens extracted from each video frame. According to the implementation in lingbot_map/layers/flashinfer_cache.py, one patch page is allocated per frame with a size equal to the number of patches per frame. The system maintains only the most recent scale frames plus a configurable sliding window of frames in live memory. When frames exceed this window, their patch pages are returned to a free list and physically reused for incoming frames.

Special Stream (Append-Only)

The special stream handles a small, fixed set of non-visual tokens including camera pose embeddings, register tokens, and scale indicators. These tokens are packed continuously into special pages that persist for the entire video duration. Unlike patch pages, special pages are never evicted and remain visible to attention computations at every timestep.

Physical Memory Layout

In lingbot_map/layers/flashinfer_cache.py, the physical storage is organized as a 5-D tensor with shape [max_num_pages, 2, page_size, H, D]. The second dimension indexes the two streams (0 for patch, 1 for special), while the remaining dimensions handle page sizing, attention heads, and head dimension. This layout allows the FlashInferAttention class to address both streams through a unified indexing scheme.

Runtime Flow for Frame Processing

The interplay between page management and kernel execution follows a strict four-phase pipeline that executes for every new frame during streaming inference.

Appending New Frames

When a new frame arrives, the kv_cache.append_frame(block_idx, k, v) method splits incoming K/V tensors into special and patch components. Special tokens are appended to the continuous special pages, while patch tokens write into newly allocated or recycled patch pages. This method is called once per transformer block, as shown in lingbot_map/layers/block.py.

Evicting Stale Frames

Memory pressure is controlled through kv_cache.evict_frames, which identifies live_window_patch_pages exceeding the configured sliding window size. The method recycles only patch pages; special pages remain untouched regardless of frame age. The eviction logic supports cross-frame special token retention and configurable scale frame buffers.

Building the Visible Page Table

Before attention computation, kv_cache.build_visible_page_table constructs a compact list of currently resident page IDs. This includes all special pages, active scale frames, and the most recent sliding window of patch pages. The method calculates paged_kv_last_page_len to handle partially filled final pages, communicating exact token counts to the kernel.

FlashInfer Kernel Execution

The FlashInferAttention.compute_attention method in lingbot_map/layers/attention.py invokes flashinfer.BatchPrefillWithPagedKVCacheWrapper with the visible page table. The kernel operates exclusively on the listed pages, ensuring computational work scales with O(window · patches) rather than O(total frames). The first block builds an execution plan that subsequent layers reuse across all transformer blocks.

Memory Efficiency Analysis

The paged approach fundamentally changes the resource requirements for long-sequence modeling.

Memory growth is bounded to O(window · patch + special) rather than O(L²), remaining constant regardless of total video length. Data movement is minimized because only newly written pages require GPU transfer; evicted pages are overwritten in place. Kernel flexibility supports arbitrary token counts per frame, including non-power-of-2 page sizes for FlashAttention-2 compatibility. Precision modes include an optional force_fp32 path that gathers paged K/V into dense tensors and runs PyTorch's scaled_dot_product_attention for numerical validation.

Implementation Example

The following pattern demonstrates stream processing using the paged cache manager:

from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager
from lingbot_map.layers.attention import FlashInferAttention
import torch

# Initialize cache manager for a 12-layer transformer

kv_cache = FlashInferKVCacheManager(
    num_blocks=12,               # One cache per transformer block

    max_num_frames=2000,         # Pre-allocation upper bound

    tokens_per_frame=262,        # 256 visual patches + 6 special tokens

    num_heads=8,
    head_dim=64,
    dtype=torch.bfloat16,
    device=torch.device("cuda"),
    num_special_tokens=6,
    scale_frames=8,              # Recent frames kept at full resolution

    sliding_window=64,           # Maximum history window

)

# Streaming inference loop

def process_video_stream(model, video_frames):
    for frame in video_frames:
        q, k, v = model.attn.prepare_qkv(frame)
        
        # Process each transformer block

        for block_idx in range(12):
            kv_cache.append_frame(block_idx, k[block_idx], v[block_idx])
            kv_cache.evict_frames(
                block_idx=block_idx,
                scale_frames=8,
                sliding_window=64,
                cross_frame_special=True,
                include_scale_frames=True
            )
            
            # Compute attention using only visible pages

            out = model.attn.compute_attention(block_idx, q[block_idx])

For a complete model implementation, see lingbot_map/models/gct_stream.py, which orchestrates frame-wise inference through the paged cache.

Summary

• Two-stream architecture separates recyclable visual patches from persistent special tokens, enabling aggressive memory reuse without losing critical metadata.
• Physical paging stores KV pairs in a 5-D tensor [max_num_pages, 2, page_size, H, D] managed by FlashInferKVCacheManager in lingbot_map/layers/flashinfer_cache.py.
• Constant memory bound replaces quadratic growth by limiting visible context to a configurable sliding window plus special tokens.
• FlashInfer integration ensures efficient kernel execution through BatchPrefillWithPagedKVCacheWrapper with reusable execution plans across transformer blocks.

Frequently Asked Questions

How does the two-stream design differ from a standard KV cache?

A standard KV cache treats all tokens uniformly, requiring storage of the entire sequence history. LingBot-Map's two-stream design in lingbot_map/layers/flashinfer_cache.py separates frame-specific visual patches (which can be discarded when old) from global special tokens like camera pose (which must persist). This allows the system to recycle memory for old frames while maintaining critical cross-frame information.

What determines which frames are evicted from the cache?

Eviction is controlled by the sliding_window and scale_frames parameters passed to kv_cache.evict_frames. The method calculates a live window based on these values and returns patch pages for frames outside this window to the free list. Special pages are excluded from eviction regardless of their age, ensuring persistent tokens remain available.

Why use FlashInfer instead of standard PyTorch attention?

FlashInfer provides BatchPrefillWithPagedKVCacheWrapper, which natively supports non-contiguous memory access through page tables. As implemented in lingbot_map/layers/attention.py, this allows the kernel to operate on scattered physical pages without gathering them into a dense tensor first, reducing memory bandwidth and enabling the O(1) memory scaling that makes long-sequence inference feasible.

How does the system handle variable numbers of tokens per frame?

The build_visible_page_table method calculates paged_kv_last_page_len to track the exact number of valid tokens in the final page of each sequence. This value is passed to the FlashInfer kernel, which adjusts its internal loops to ignore padding. The design supports non-power-of-2 page sizes, accommodating frames with different patch counts without wasting memory on alignment padding.

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