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

> Discover how LingBot-Map optimizes streaming inference with a paged KV cache and FlashInfer integration for efficient attention computation in a single kernel launch.

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

---

**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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/flashinfer_cache.py) (lines 40-44):

```python

# 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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/flashinfer_cache.py)), the system builds a plan once per frame step that describes the visible page indices and last-page lengths:

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/aggregator/stream.py) (lines 79-103) creates the cache manager lazily to defer memory allocation until the first forward pass:

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/block.py) demonstrates how a single frame flows through the system:

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

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