# How torch.compile with reduce-overhead Mode Works in the LingBot-Map Streaming Pipeline

> Learn how torch.compile reduce-overhead mode optimizes the LingBot-Map streaming pipeline by compiling transformer blocks into fast CUDA kernels for faster inference without Python overhead.

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

---

**`torch.compile(..., mode="reduce-overhead")` compiles fixed-shape transformer blocks into static CUDA kernels that are replayed for every frame, eliminating Python dispatch overhead during the streaming inference loop.**

The LingBot-Map repository implements a high-throughput video understanding system that processes frames incrementally using a streaming transformer architecture. By applying **torch.compile** with **reduce-overhead** mode to hot sub-modules, the system converts dynamic PyTorch operations into pre-compiled kernels that minimize latency during per-frame inference while maintaining compatibility with FlashInfer's paged KV-cache. This article examines how the implementation in Robbyant/lingbot-map combines module-level compilation with CUDA graph capture to achieve low-latency streaming.

## Understanding the Streaming Pipeline Architecture

The streaming pipeline centers on the `AggregatorStream` class defined in [`lingbot_map/aggregator/stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/aggregator/stream.py). This architecture separates processing into **frame-wise** transformers (`frame_blocks`) and **global** transformers (`global_blocks`), supporting both FlashInfer's paged KV-cache and pure-PyTorch SDPA implementations.

The model handles two distinct inference paths:
- **Keyframe processing**: Full attention updates with KV-cache appends
- **Non-keyframe processing**: Skip-append mode that changes the KV-cache update path without altering tensor shapes

Because the compiled modules remain shape-stable across these paths, the streaming loop can toggle between keyframe and non-keyframe logic without recompiling the graph.

## Compiling Hot Sub-Modules with reduce-overhead

Before streaming begins, the `compile_model` helper function (located in [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) at lines 66-80) walks through the aggregator's modules and replaces each with a compiled version. The `mode="reduce-overhead"` flag instructs PyTorch to generate fast, static graphs that minimize Python-side overhead while preserving the original module semantics.

```python
def compile_model(model):
    """Compile fixed-shape modules with reduce-overhead."""
    agg = model.aggregator

    # Frame-wise transformer blocks

    for i, block in enumerate(agg.frame_blocks):
        agg.frame_blocks[i] = torch.compile(block, mode="reduce-overhead")

    # Patch-embedding blocks (used before frame blocks)

    for i, block in enumerate(agg.patch_embed.blocks):
        agg.patch_embed.blocks[i] = torch.compile(block, mode="reduce-overhead")

    # Global attention blocks (FlashInfer or SDPA)

    for block in agg.global_blocks:
        if hasattr(block, "attn_pre"):
            block.attn_pre = torch.compile(block.attn_pre, mode="reduce-overhead")
        if hasattr(block, "ffn_residual"):
            block.ffn_residual = torch.compile(block.ffn_residual,
                                               mode="reduce-overhead")
        block.attn.proj = torch.compile(block.attn.proj, mode="reduce-overhead")

```

This compilation targets the fixed-shape components: frame blocks, patch-embedding blocks, and global attention projections. The compiled kernels remain valid across the streaming sequence because the tensor shapes for these blocks do not change between frames.

## CUDA Graph Warm-Up and Step Markers

Before the actual stream runs, the pipeline executes a warm-up phase that captures a CUDA graph containing the compiled blocks. The warm-up uses `torch.compiler.cudagraph_mark_step_begin()` to mark the start of each step that will later be replayed by the CUDA graph engine.

The `_warm_streaming` function in [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) (lines 88-105) implements this preparation:

```python
def _warm_streaming(model, images, scale_frames, warm_stream_n,
                    dtype, passes=1, keyframe_interval=1):
    # … (prepare scale and streaming tensors) …

    for _ in range(passes):
        model.clean_kv_cache()
        torch.compiler.cudagraph_mark_step_begin()
        with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
            model.forward(
                warm_scale,
                num_frame_for_scale=scale_frames,
                num_frame_per_block=scale_frames,
                causal_inference=True,
            )
        # Run a few non-keyframe steps to capture the skip-append path

        for i in range(warm_stream_n):
            is_keyframe = (keyframe_interval <= 1) or (i % keyframe_interval == 0)
            if not is_keyframe:
                model._set_skip_append(True)
            torch.compiler.cudagraph_mark_step_begin()
            with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
                model.forward(
                    warm_stream[:, i:i + 1],
                    num_frame_for_scale=scale_frames,
                    num_frame_per_block=1,
                    causal_inference=True,
                )
            if not is_keyframe:
                model._set_skip_append(False)
    torch.cuda.synchronize()
    model.clean_kv_cache()

```

This warm-up captures both the initial scale batch processing and the single-frame streaming paths, including the non-keyframe skip-append logic, ensuring the CUDA graph contains all necessary kernel configurations.

## The Streaming Inference Loop

During actual inference, the `run_streaming_once` function in [`scripts/benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py) (lines 98-125) executes the streaming loop. Each iteration marks a CUDA-graph step, processes either a scale batch or single frame, and toggles the skip-append flag for non-keyframes without disturbing the compiled graph.

```python
@torch.no_grad()
def run_streaming_once(*, model, images, frame_count, dtype,
                      device, num_scale_frames, keyframe_interval):
    model.clean_kv_cache()
    torch.cuda.synchronize(device)
    start = time.perf_counter()

    # --- initial scale batch (warm-up) ---

    scale_batch = torch.cat(
        [get_frame(images, i, dtype, device) for i in range(num_scale_frames)],
        dim=1,
    )
    torch.compiler.cudagraph_mark_step_begin()
    with make_autocast(dtype):
        model.forward(
            scale_batch,
            num_frame_for_scale=num_scale_frames,
            num_frame_per_block=num_scale_frames,
            causal_inference=True,
        )
    del scale_batch

    # --- streaming frames ---

    for i in range(num_scale_frames, frame_count):
        is_keyframe = (keyframe_interval <= 1) or ((i - num_scale_frames) % keyframe_interval == 0)
        if not is_keyframe:
            model._set_skip_append(True)          # non-keyframe path

        frame = get_frame(images, i, dtype, device)
        torch.compiler.cudagraph_mark_step_begin()
        with make_autocast(dtype):
            model.forward(
                frame,
                num_frame_for_scale=num_scale_frames,
                num_frame_per_block=1,
                causal_inference=True,
            )
        del frame
        if not is_keyframe:
            model._set_skip_append(False)
    torch.cuda.synchronize(device)
    elapsed = time.perf_counter() - start
    # … (report memory, timing) …

```

The `torch.compiler.cudagraph_mark_step_begin()` calls ensure that the CUDA graph engine replays the pre-compiled kernels for each frame, while `model.clean_kv_cache()` resets the FlashInfer cache between runs.

## Why reduce-overhead Mode is Essential for Streaming

The streaming pipeline processes many **tiny** forward passes (one per frame). Without compilation, each call would incur Python dispatch and kernel launch overhead, which dominates latency when processing individual frames. 

**`torch.compile(..., mode="reduce-overhead")`** creates a single, ahead-of-time compiled kernel for each block, so every frame only costs the data-movement and KV-cache lookup time. The CUDA-graph warm-up guarantees that the runtime replays the same kernel configuration each iteration, eliminating repeated graph-building overhead. This combination allows LingBot-Map to maintain complex FlashInfer KV-cache logic while achieving the performance characteristics of a static inference engine.

## Summary

- **`torch.compile` with reduce-overhead** compiles fixed-shape transformer blocks in `frame_blocks`, `patch_embed.blocks`, and `global_blocks` to eliminate Python overhead.
- **CUDA graph warm-up** in `_warm_streaming` captures the compiled kernels and skip-append paths before the actual streaming begins.
- **Step markers** (`torch.compiler.cudagraph_mark_step_begin()`) in the streaming loop trigger replay of the pre-compiled graphs for each frame.
- **Shape-stable compilation** allows the pipeline to toggle between keyframe and non-keyframe processing without recompiling.
- **FlashInfer compatibility** is maintained because the compiled modules operate independently of the KV-cache management logic.

## Frequently Asked Questions

### What is the difference between reduce-overhead and other torch.compile modes in streaming contexts?

**`reduce-overhead`** mode specifically targets the elimination of Python-side dispatch costs and framework overhead, making it ideal for streaming pipelines that execute many small forward passes. Other modes like `max-autotune` focus on kernel selection for throughput, while `default` provides a balance that may retain more Python overhead than necessary for latency-sensitive per-frame processing.

### Why does the LingBot-Map streaming pipeline require CUDA graph warm-up?

The warm-up phase captures the **compiled kernel configurations** and memory allocations into a CUDA graph that can be replayed without re-executing Python code. According to the implementation in [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py), the `_warm_streaming` function exercises both the initial scale batch and single-frame paths (including non-keyframe skip-append logic) to ensure the graph contains all necessary operations before the timing-critical streaming loop begins.

### How does torch.compile interact with the FlashInfer KV-cache?

The compiled modules (frame blocks and attention projections) operate on **fixed-shape tensors** independently of the KV-cache implementation. The FlashInfer paged KV-cache management in [`lingbot_map/aggregator/stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/aggregator/stream.py) handles memory addressing and cache updates outside the compiled regions, while the compiled kernels perform the actual attention computations. This separation allows the pipeline to use FlashInfer's efficient memory management without invalidating the compiled graphs.

### What happens to the compiled graph when skip_append is toggled for non-keyframes?

Toggling `skip_append` changes the KV-cache update path but **does not alter the tensor shapes** processed by the compiled modules. Because `torch.compile` with reduce-overhead mode generates static graphs based on shape stability, the compiled attention and FFN kernels remain valid for both keyframe and non-keyframe iterations. The warm-up phase explicitly captures both paths to ensure the CUDA graph contains the necessary kernel launches for each case.