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

> Accelerate LingBot-Map inference using torch.compile. Learn how to enable this powerful feature with the --compile flag for faster performance.

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

---

**Add the `--compile` flag when running [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py), [`scripts/benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py), and [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) and [`scripts/benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/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.

```bash
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`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py), the compilation logic explicitly targets:

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

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/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:

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

```

### Minimal Programmatic Usage

Integrate compilation into custom inference scripts:

```python
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:

```bash
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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py)**: Benchmarking utility that demonstrates the same compile path and measures latency/memory trade-offs.
- **[`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py)**: Profiling script that explains compilation rationale and reports speed-up statistics when `--compile` is enabled.
- **[`lingbot_map/models/gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/block.py), [`attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/attention.py)) that serve as compilation targets.

## Summary

- **Activation**: Pass `--compile` to [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py), [`scripts/benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py), or [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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.