# How FlashInfer's Paged KV Cache Attention Outperforms SDPA Fallback in Streaming Inference

> FlashInfer's paged KV cache attention speeds up streaming inference by reducing memory bandwidth and kernel launches outperforming SDPA fallback.

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

---

**FlashInfer's paged KV cache attention reduces memory bandwidth and kernel launch overhead by reading only visible cache pages rather than dense tensors, delivering substantial speed-ups over PyTorch's native SDPA fallback for long streaming sequences.**

LingBot-Map implements a two-stream paged KV-cache system using FlashInfer to accelerate transformer attention in streaming applications. According to the Robbyant/lingbot-map source code, the `FlashInferAttention` layer leverages the `BatchPrefillWithPagedKVCacheWrapper` kernel to process frames efficiently, falling back to native PyTorch Scaled Dot-Product Attention (SDPA) only when FlashInfer is unavailable or when `use_sdpa=True` is explicitly set. This architecture specifically optimizes for long sequences where frames share KV-cache entries, minimizing both memory traffic and computational overhead compared to dense attention approaches.

## Architecture of the Paged KV Cache

### Two-Stream Design

The cache splits storage across **two distinct streams** managed by `FlashInferKVCacheManager` in [`lingbot_map/layers/flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/flashinfer_cache.py):

- **Patch stream (recyclable):** Holds KV entries for scale frames and sliding-window live frames. Each frame occupies exactly one patch page (`self.patches_per_frame` tokens) with zero padding. Pages are allocated from `self.free_patch_pages` and recycled after exiting the sliding window via `evict_frames`.
- **Special stream (append-only):** Stores KV for special tokens (camera, register, scale token) that appear in every frame. Tokens pack continuously across pages, which are never evicted. Allocation occurs in `_write_special_tokens` using `self.all_special_pages` and `self.free_special_pages`.

### Physical Memory Layout

The cache stores data as a 5-D tensor per block with shape `[max_num_pages, 2, page_size, H, D]`, where dimension 1 separates keys and values. This layout is instantiated in `FlashInferKVCacheManager.__init__` (lines 42-48 of [`flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/flashinfer_cache.py)).

### Plan-Run Execution Workflow

The attention mechanism uses an efficient plan-run pattern implemented in `compute_attention` (lines 85-115 of [`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/attention.py)):

1. **Plan (once per frame step):** Only the first block (`block_idx == 0`) builds the visible page table via `build_visible_page_table` (lines 38-44), computes the last-page length with `compute_last_page_len` (lines 52-63), and calls `self.prefill_wrapper.plan(...)`.
2. **Run (once per block):** Each block reuses the cached plan and executes the kernel via `self.prefill_wrapper.run(...)`.

This amortizes expensive planning work across all blocks, reducing per-block overhead to a lightweight kernel dispatch.

## Performance Comparison: FlashInfer vs SDPA Fallback

### Memory Bandwidth Efficiency

**FlashInfer** reads only the *visible* pages (scale + window + special) as determined by `build_visible_page_table`. This sparse access pattern avoids materializing dense KV tensors for the entire sliding window. In contrast, **SDPA fallback** operates on a dense tensor of shape `[B, N, H, D]`, forcing the kernel to load the complete window (often thousands of tokens) at every step regardless of actual dependencies.

### Kernel Launch Overhead

The FlashInfer path executes one plan kernel plus one run kernel per block, with planning cost shared across blocks. The SDPA fallback requires a separate `torch.nn.functional.scaled_dot_product_attention` call (line 59 of [`attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/attention.py)) per forward pass, processing the entire dense KV tensor each time without amortization.

### Cache Reuse and Page Recycling

FlashInfer stores KV pages once and reuses them across frames; only new pages are written via `append_frame`, while evicted pages are recycled through `free_patch_pages` without memory copies. SDPA recomputes the dense KV representation each step, incurring extra allocation overhead and potential fragmentation.

### Precision and Hardware Utilization

FlashInfer kernels operate natively in fp16/bf16 (with fp32 support via `force_fp32`), keeping data on-device without conversion and enabling tensor core usage. The SDPA fallback may up-cast to fp32 internally, and when forced to fp32 mode loses tensor core acceleration entirely.

### Scalability for Streaming

The page-based design supports arbitrarily long streaming sequences bounded only by the sliding-window size. SDPA memory consumption grows linearly with total window length, quickly exhausting GPU memory for extended streams.

## Implementation in the LingBot-Map Codebase

### Cache Manager Integration

The `AggregatorStream` class in [`lingbot_map/aggregator/stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/aggregator/stream.py) instantiates either `FlashInferBlock` or `SDPABlock` based on the `use_flashinfer` flag (lines 91-99). It lazily creates a `FlashInferKVCacheManager` on the first forward pass (lines 199-215), ensuring the cache is properly sized for the model configuration.

### Attention Layer Logic

`FlashInferBlock` (defined in [`lingbot_map/layers/block.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/block.py), lines 151-180) preprocesses QKV projections, appends KV to the cache via the manager's methods, and invokes `compute_attention`. If the cache is not a `FlashInferKVCacheManager` instance (line 46 of [`attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/attention.py)), or if `use_sdpa=True` is passed to model constructors, the code bypasses the paged cache entirely and routes to the SDPA path.

### Model Configuration

Model constructors in [`lingbot_map/models/gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream.py) (line 125) and [`lingbot_map/models/gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream_window.py) (line 166) expose the `use_sdpa` argument. When `False` (default), FlashInfer serves as the primary backend, as documented in [`gct_stream_window_v2.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_stream_window_v2.py) (line 183).

## Practical Usage Examples

### Instantiating a Model with FlashInfer (Default)

```python
from lingbot_map.models.gct_stream import GCTStream

model = GCTStream(
    num_blocks=12,
    dim=1024,
    num_heads=16,
    use_sdpa=False,  # FlashInfer is default; set True to force SDPA

    kv_cache_sliding_window=64,
    kv_cache_scale_frames=8,
)

```

### Running a Single Streaming Step

```python
import torch

# Dummy input for one frame (B=1, N=262 tokens, C=dim)

x = torch.randn(1, 262, 1024, device="cuda")

# Forward pass – lazily creates KV-cache manager

out = model(x)  # Uses FlashInferAttention internally

```

### Forcing the SDPA Fallback

```python
model_sdpa = GCTStream(
    num_blocks=12,
    dim=1024,
    num_heads=16,
    use_sdpa=True,  # Disables FlashInfer completely

)

out = model_sdpa(x)  # Routes through SDPA path in FlashInferAttention.forward

```

### Inspecting Cache Statistics

```python

# After processing frames

kv_manager = model.blocks[0].attn.kv_cache  # FlashInferKVCacheManager instance

stats = kv_manager.get_cache_stats()  # Defined in flashinfer_cache.py lines 303-322

print(stats)

# Output: {'frame_count': 42, 'scale_pages': 8, 'live_pages': 34,

#          'free_pages': 46, 'special_tokens': 252}

```

## Summary

- **FlashInfer's paged KV cache** minimizes memory traffic by reading only visible pages rather than dense tensors, dramatically reducing bandwidth requirements for long sequences.
- The **plan-run workflow** amortizes planning overhead across transformer blocks, while SDPA requires full dense tensor processing on every forward pass.
- **Two-stream architecture** separates recyclable patch pages from append-only special tokens, enabling efficient page reuse without memory copies.
- Implementation spans [`lingbot_map/layers/flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/flashinfer_cache.py) (cache management), [`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/attention.py) (kernel execution), and [`lingbot_map/layers/block.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/block.py) (integration), with fallback logic triggered when `use_sdpa=True` or when FlashInfer is unavailable.
- The system automatically handles variable-length sequences via `compute_last_page_len` and `build_visible_page_table`, ensuring correctness while maintaining performance.

## Frequently Asked Questions

### When does LingBot-Map fall back to SDPA instead of using FlashInfer?

The code falls back to SDPA when FlashInfer is not installed, when the model runs in fp32 mode without tensor core support, or when the user explicitly sets `use_sdpa=True` in the model constructor. This check occurs in `FlashInferAttention.forward` (around line 46 of [`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/attention.py)), which verifies whether the provided `kv_cache` is an instance of `FlashInferKVCacheManager`; if not, it routes to `torch.nn.functional.scaled_dot_product_attention` (line 59).

### How does the paged cache handle variable sequence lengths across frames?

The cache handles variable lengths through the `compute_last_page_len` method (lines 52-63 of [`flashinfer_cache.py`](https://github.com/Robbyant/lingbot-map/blob/main/flashinfer_cache.py)), which calculates how many tokens are valid on the final special page since special tokens append continuously and may not fill a page completely. The `build_visible_page_table` function (lines 38-44) constructs the ordered list of pages (scale → live → special) for each timestep, ensuring the attention kernel only processes valid entries regardless of how many frames are currently in the sliding window.

### What is the physical memory format of the KV cache?

The KV cache stores data as a 5-D tensor with shape `[max_num_pages, 2, page_size, H, D]`, where the second dimension (size 2) separates keys and values, `page_size` is the number of tokens per physical page, `H` is the number of heads, and `D` is the head dimension. This layout is allocated in `FlashInferKVCacheManager.__init__` (lines 42-48) and enables FlashInfer to index pages directly without additional memory copies.

### Can FlashInfer be used with fp32 precision?

Yes, FlashInfer supports fp32 through the `force_fp32` path, though it is optimized for fp16 and bf16 to leverage tensor cores. When the model is forced to fp32, the kernels can still execute, but users may experience reduced performance compared to lower precision modes. The SDPA fallback is typically used only when FlashInfer is unavailable, not as a precision-specific optimization, since FlashInfer handles the necessary type conversions internally.