# Where Is the KV Cache Implemented in Robbyant/lingbot-map?

> Find the KV cache implementation in Robbyant/lingbot-map within lingbot_map/layers/flashinfer_cache.py. Learn about the FlashInferKVCacheManager class for streaming inference.

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

---

**The KV cache implementation in Robbyant/lingbot-map is located in [`lingbot_map/layers/flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/flashinfer_cache.py), specifically within the `FlashInferKVCacheManager` class that provides a two-stream paged cache for streaming inference.**

The LingBot-Map project implements a specialized key-value (KV) cache to enable efficient streaming inference across video frames. This implementation utilizes a two-stream paged architecture that separates patch pages from special token pages, integrating directly with the FlashInfer library for optimized attention computation. Understanding this KV cache implementation is essential for developers working with the streaming GCT models or customizing the inference pipeline.

## Core KV Cache Implementation

The primary KV cache implementation resides in [`lingbot_map/layers/flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/flashinfer_cache.py). This file contains the `FlashInferKVCacheManager` class, which serves as the central management system for allocating, writing, and evicting cached key-value tensors.

### FlashInferKVCacheManager Architecture

The `FlashInferKVCacheManager` implements a two-stream paged KV cache design. It maintains separate memory pools for patch tokens and special tokens, allowing the model to handle varying sequence lengths efficiently during streaming inference. The manager is instantiated per transformer block, with each block maintaining its own independent cache state.

### Public API Methods

The cache exposes several critical methods used throughout the model:

- `append_frame(block_idx, k, v)` - Appends KV tensors for a new frame to the cache
- `evict_frames(block_idx, scale_frames, sliding_window, ...)` - Recycles old patch pages while preserving scale and special pages
- `compute_attention(block_idx, q)` - Executes cross-frame attention using FlashInfer or an fp32 fallback
- `reset()` - Clears all per-block state for fresh sequences
- `get_cache_stats()` - Returns cache occupancy metrics for debugging

## Integration with Streaming Models

The KV cache manager integrates with the streaming inference pipeline through several model files:

- [`lingbot_map/models/gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream.py) - Constructs the manager at line 118 and delegates KV handling for continuous streaming
- [`lingbot_map/models/gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream_window.py) - Similar usage at line 117 for windowed streaming scenarios
- [`lingbot_map/models/gct_stream_window_v2.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream_window_v2.py) - Enhanced version with additional KV-debug utilities

These files demonstrate how the cache manager is wired into the aggregation pipeline, handling frame-by-frame processing while maintaining context across the sliding window.

## Understanding the Two-Stream Paged Design

The implementation distinguishes between two types of cache entries:

**Patch Pages**: Store visual token embeddings that are subject to eviction based on the sliding window configuration.

**Special Pages**: Store special tokens (6 tokens per frame) that are retained longer and managed separately from the patch pool.

This architecture allows the model to maintain a fixed memory footprint while processing arbitrarily long video streams, evicting older visual patches while preserving critical metadata tokens.

## Practical Usage Example

Here's how to instantiate and use the KV cache manager:

```python
import torch
from lingbot_map.layers.flashinfer_cache import FlashInferKVCacheManager

# Initialize the manager (one cache per transformer block)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
kv_mgr = FlashInferKVCacheManager(
    num_blocks=12,
    max_num_frames=88,
    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,
)

# Append frame KV tensors

k = torch.randn(262, 16, 64, dtype=torch.bfloat16, device=device)
v = torch.randn(262, 16, 64, dtype=torch.bfloat16, device=device)
kv_mgr.append_frame(block_idx=0, k=k, v=v)

# Evict old frames based on sliding window

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

# Compute attention

q = torch.randn(262, 16, 64, dtype=torch.bfloat16, device=device)
out = kv_mgr.compute_attention(block_idx=0, q=q)

# Debug and reset

stats = kv_mgr.get_cache_stats()
kv_mgr.reset()

```

## Summary

- The KV cache implementation is located in [`lingbot_map/layers/flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/flashinfer_cache.py) within the `FlashInferKVCacheManager` class
- It uses a two-stream paged architecture separating patch pages from special token pages
- Key methods include `append_frame()`, `evict_frames()`, and `compute_attention()` for managing streaming inference
- The cache integrates with [`gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_stream.py) and windowed variants to process video frames efficiently
- Configuration parameters like `scale_frames` and `sliding_window` control memory usage and context retention

## Frequently Asked Questions

### What is the purpose of the two-stream design in the KV cache?

The two-stream design separates patch tokens (visual features) from special tokens (metadata), allowing the cache to evict older visual patches while preserving critical structural tokens. This maintains a fixed memory footprint during long video sequences while retaining essential context for attention computation.

### How does the KV cache handle memory management during streaming?

The `evict_frames()` method implements a sliding window mechanism that recycles patch pages based on the `sliding_window` parameter, while `scale_frames` determines how many frames remain resident in memory. Special tokens are managed separately to ensure proper attention computation across the entire sequence.

### Which files should I modify to adjust KV cache behavior for custom models?

Start with [`lingbot_map/layers/flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/flashinfer_cache.py) to modify the core cache logic. For integration changes, examine [`lingbot_map/models/gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream.py) (continuous streaming) or [`lingbot_map/models/gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream_window.py) (windowed processing), where the `FlashInferKVCacheManager` is instantiated and configured.

### Does the implementation support CPU inference or only CUDA?

The `FlashInferKVCacheManager` accepts a `device` parameter and can fall back to fp32 computation when FlashInfer is unavailable, though optimal performance requires CUDA. The `compute_attention()` method automatically handles the fallback based on hardware availability.