How to Use `torch.compile` to Accelerate LingBot-Map Inference: A Complete Guide

Add the --compile flag when running demo.py or benchmark scripts to enable torch.compile with "reduce-overhead" mode on the GCT model's hot sub-modules, followed by a CUDA graph warm-up for optimal streaming performance.

LingBot-Map is an open-source streaming vision model that processes sequences of images through a Global Context Transformer (GCT). The repository ships with a compact, opt-in workflow that uses torch.compile to fuse heavy linear and attention kernels, significantly reducing latency during streaming inference without increasing memory usage.

How LingBot-Map Implements torch.compile

The compilation strategy targets fixed-shape sub-modules within the GCT aggregator while preserving the dynamic control flow required for KV-cache management. This approach minimizes Just-In-Time (JIT) compilation overhead while maximizing kernel fusion for the compute-intensive parts of the model.

The compile_model Helper Function

All entry-point scripts—including demo.py, scripts/benchmark_gct_memory.py, and gct_profile.py—accept a --compile argument. When present, the model is passed through the compile_model helper before streaming starts. This function iterates over specific transformer blocks and replaces them with compiled versions using mode="reduce-overhead".

Targeted Sub-Module Compilation

Rather than compiling the entire model, LingBot-Map selectively compiles the hot execution paths:

• Frame blocks: The per-frame processing modules (agg.frame_blocks)
• Patch embedding blocks: The initial vision encoding layers (agg.patch_embed.blocks)
• Attention projections: The output projections within attention mechanisms (block.attn.proj)
• Pre-attention and residual branches: Optional attn_pre and ffn_residual sub-modules when present in global blocks

This selective approach is implemented in demo.py and scripts/benchmark_gct_memory.py within their respective compile_model functions.

Step-by-Step Compilation Workflow

1. Enable the Compile Flag

Activate compilation by appending --compile to any supported script. This flag triggers the compile_model helper before the inference loop begins.

python demo.py \
    --image_folder path/to/images \
    --backend flashinfer \
    --dtype auto \
    --compile

The scripts automatically configure PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to prevent out-of-memory errors during the compilation phase.

2. Compile Hot Blocks

The compile_model function applies torch.compile to specific sub-modules. In demo.py, the compilation logic explicitly targets:

agg = model.aggregator

# Compile frame-level blocks

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

# Compile patch embedding blocks

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

# Compile global block components

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")

3. Warm-Up with CUDA Graph Capture

After compilation, the demo executes _warm_streaming to eliminate first-frame overhead. This routine runs a short warm-up pass that calls torch.compiler.cudagraph_mark_step_begin() before each forward operation, recording a CUDA graph that matches the exact tensor shapes used during real inference.

def warmup(model, images, dtype):
    torch.compiler.cudagraph_mark_step_begin()
    with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
        model.forward(images[:1])  # Dummy forward to capture graph

This step ensures the compiled kernels remain "hot" and avoids graph capture latency during actual streaming.

4. Run Streaming Inference

Once warmed up, the normal streaming loop (inference_streaming in lingbot_map/models/gct_stream.py) executes entirely under the compiled graph. The model processes incoming frames with reduced kernel launch overhead while maintaining dynamic control flow for KV-cache management outside the compiled regions.

Why "reduce-overhead" Mode?

LingBot-Map fixes mode="reduce-overhead" because the GCT architecture contains many small, dynamic control-flow steps for KV-cache management and per-frame skips. This compilation mode reduces JIT compilation costs while still fusing heavy linear and attention kernels, making it ideal for streaming workloads where latency consistency matters more than peak throughput.

Code Examples

Running the Demo with Compilation

Execute the streaming demo with full compilation support:

python demo.py \
    --image_folder /path/to/frames \
    --backend flashinfer \
    --dtype auto \
    --compile

Minimal Programmatic Usage

Integrate compilation into custom inference scripts:

import torch
from lingbot_map.models.gct_stream import GCTStream

# Initialize model

model = GCTStream(...)
model.eval().cuda()

# Compile hot modules

agg = model.aggregator
for i, blk in enumerate(agg.frame_blocks):
    agg.frame_blocks[i] = torch.compile(blk, mode="reduce-overhead")
for i, blk in enumerate(agg.patch_embed.blocks):
    agg.patch_embed.blocks[i] = torch.compile(blk, mode="reduce-overhead")
for blk in agg.global_blocks:
    if hasattr(blk, "attn_pre"):
        blk.attn_pre = torch.compile(blk.attn_pre, mode="reduce-overhead")
    if hasattr(blk, "ffn_residual"):
        blk.ffn_residual = torch.compile(blk.ffn_residual, mode="reduce-overhead")
    blk.attn.proj = torch.compile(blk.attn.proj, mode="reduce-overhead")

# Warm-up CUDA graph

torch.compiler.cudagraph_mark_step_begin()
with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.float16):
    model.forward(dummy_input)

# Run streaming inference

with torch.no_grad():
    for frame in stream_of_frames:
        output = model.forward(frame)

Benchmarking the Speed-Up

Measure the performance impact across different sequence lengths:

python scripts/benchmark_gct_memory.py \
    --height 384 \
    --width 518 \
    --frame-counts 64 128 256 \
    --compile

This script reports both peak memory consumption and runtime for each frame count, allowing direct comparison between compiled and eager execution modes.

Key Implementation Files

The torch.compile workflow spans several critical files in the repository:

• demo.py: Top-level entry point that parses --compile, invokes compile_model, and executes _warm_streaming before streaming inference begins.
• scripts/benchmark_gct_memory.py: Benchmarking utility that demonstrates the same compile path and measures latency/memory trade-offs.
• gct_profile.py: Profiling script that explains compilation rationale and reports speed-up statistics when --compile is enabled.
• lingbot_map/models/gct_stream.py: Core streaming model implementation; after compilation, the forward method executes optimized kernels.
• lingbot_map/layers/: Directory containing low-level transformer blocks (block.py, attention.py) that serve as compilation targets.

Summary

• Activation: Pass --compile to demo.py, scripts/benchmark_gct_memory.py, or gct_profile.py to enable optimized inference.
• Strategy: The repository uses selective compilation on frame_blocks, patch_embed.blocks, and attention projections rather than compiling the entire model.
• Mode: Fixed to reduce-overhead to balance JIT compilation costs with kernel fusion benefits for streaming workloads.
• Warm-up: CUDA graph capture via torch.compiler.cudagraph_mark_step_begin() is automatically performed to eliminate first-frame latency.
• Compatibility: The compilation path works safely with other flags like --backend flashinfer and --dtype auto without requiring source code modifications.

Frequently Asked Questions

Does torch.compile work with CPU inference in LingBot-Map?

No, the compilation workflow is CUDA-only. The scripts automatically detect CUDA availability and set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to manage memory during compilation. CPU inference runs in eager mode regardless of the --compile flag.

Which model components are compiled versus left in eager mode?

The compile_model helper compiles the heavy computational blocks including frame_blocks, patch_embed.blocks, attention projections (attn.proj), and optional attn_pre/ffn_residual sub-modules. Dynamic control-flow logic for KV-cache management and frame skipping remains in eager mode to handle variable sequence lengths efficiently.

Is the warm-up step necessary for every inference session?

Yes, the warm-up routine (implemented in _warm_streaming within demo.py) is required to capture the CUDA graph and eliminate first-frame overhead. The demo automatically performs this when --compile is set, executing a few dummy forwards before processing real input data.

Can I use torch.compile with other optimization backends like FlashInfer?

Yes, the compilation path is fully compatible with --backend flashinfer and precision flags like --dtype auto. The compiled kernels operate at the PyTorch level while FlashInfer handles attention computation, allowing both optimizations to work simultaneously without conflicts.

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