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

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. 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 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.

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 (lines 88-105) implements this preparation:

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 (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.

@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, 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 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.

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