# How to Use torch.compile and CUDA Graph Warmup for LingBot-Map Performance Optimization

> Optimize LingBot-Map performance using torch.compile's reduce-overhead mode and CUDA graph warmup. Eliminate cold-start overhead for faster streaming inference.

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

---

**Enable `torch.compile` with `mode="reduce-overhead"` on `GCTStream`'s fixed-shape sub-modules and perform a CUDA graph warmup phase that exercises both keyframe and non-keyframe paths to eliminate cold-start overhead in streaming inference.**

The `Robbyant/lingbot-map` repository implements a high-performance streaming vision model called `GCTStream` that benefits significantly from PyTorch 2.x optimization features. By selectively compiling transformer blocks and capturing CUDA graphs during a dedicated warmup routine, you can eliminate Python-side overhead and kernel launch latency, achieving substantial FPS improvements in real-time inference scenarios.

## Why torch.compile Accelerates LingBot-Map

LingBot-Map’s core architecture consists of **fixed-shape sub-modules**—specifically the per-frame transformer blocks, patch-embedding blocks, and selected global blocks—that are ideal candidates for JIT compilation. When these modules are compiled with **`torch.compile(..., mode="reduce-overhead")`**, PyTorch generates static computation graphs that fuse kernels and eliminate interpreter overhead.

The repository isolates this optimization in two locations:

- **[`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py)** (lines 40-57): Used for pure FPS profiling
- **[`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py)** (lines 70-87): Used for interactive demonstrations

Both implementations follow the same pattern of traversing the model hierarchy:

```python
agg = model.aggregator
for i, b in enumerate(agg.frame_blocks):
    agg.frame_blocks[i] = torch.compile(b, mode="reduce-overhead")

# Repeated for patch_embed.blocks and global block members

```

**Note:** The profiling script specifically omits the point head (`model.point_head = None` on line 57 of [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py)) because it adds approximately 6 ms per frame and is unnecessary for pure throughput measurement.

## Implementing CUDA Graph Warmup for Streaming Inference

Even with compiled modules, LingBot-Map requires **CUDA graph capture** to achieve optimal inference speed. The streaming architecture presents a unique challenge: the first frame processed as a non-keyframe (the "skip-append" path) triggers different KV-cache handling logic than subsequent keyframe captures. Without warmup, this path mismatch causes allocation penalties on the first streamed frame.

The solution is a **warmup phase** implemented in `demo._warm_streaming` (lines 89-122) that:

1. Executes a scale-frame forward pass (phase 1)
2. Runs a short sequence mimicking real keyframe/non-keyframe patterns
3. Wraps each forward call with `torch.compiler.cudagraph_mark_step_begin()` to capture graphs for both execution paths

The critical warmup logic:

```python

# Warmup keyframe path

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

# Warmup non-keyframe path

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(...)
if not is_keyframe:
    model._set_skip_append(False)

```

After warmup, the KV cache is cleared via `model.clean_kv_cache()` to ensure the actual streaming inference starts from a clean state while retaining the captured CUDA graphs.

## End-to-End Optimization Workflow

Follow this sequence to enable full acceleration in your own implementation:

### 1. Load and Initialize the Model

Instantiate `GCTStream` and move it to CUDA:

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

model = GCTStream(
    img_size=480,
    patch_size=14,
    enable_3d_rope=True,
    max_frame_num=256,
    kv_cache_sliding_window=32,
    use_sdpa=True,
).eval().to('cuda')

```

### 2. Compile Hot Modules

Replace fixed-shape sub-modules with compiled versions:

```python
from demo import compile_model  # or gct_profile.compile_model

compile_model(model)  # In-place modification

```

### 3. Warmup CUDA Graphs

Execute the warmup routine to capture both execution paths:

```python
from demo import _warm_streaming

_warm_streaming(
    model=model,
    images=preprocessed_tensor,  # Shape: [S, 3, H, W]

    scale_frames=4,
    warm_stream_n=8,
    dtype=torch.float16,
    passes=1,
    keyframe_interval=3,
)

```

### 4. Execute Streaming Inference

Run your actual inference loop, which now reuses the captured graphs:

```python
from gct_profile import profile_streaming

per_frame_ms, _, _ = profile_streaming(
    model,
    images,
    num_frames=500,
    dtype=torch.float16,
    keyframe_interval=3,
)

```

### Command-Line Interface Options

The repository provides convenience flags for quick benchmarking:

```bash

# Profile with compilation and warmup

python gct_profile.py \
    --backend both \
    --dtype bf16 \
    --compile \
    --num_frames 500 \
    --keyframe_interval 3

# Run interactive demo with full optimizations

python demo.py \
    --input path/to/sequence/ \
    --compile \
    --warmup-stream-n 10 \
    --keyframe_interval 3

```

## Complete Implementation Example

This minimal script demonstrates the entire optimization pipeline:

```python
import torch
from lingbot_map.models.gct_stream import GCTStream
from demo import compile_model, _warm_streaming
from gct_profile import profile_streaming

# 1. Build model

model = GCTStream(
    img_size=480,
    patch_size=14,
    enable_3d_rope=True,
    max_frame_num=256,
    kv_cache_sliding_window=32,
    kv_cache_scale_frames=8,
    use_sdpa=True,
    camera_num_iterations=3,
).eval().to('cuda')

# 2. Compile fixed-shape modules

compile_model(model)

# 3. Load pre-processed images [S, 3, H, W]

images = torch.load("demo_images.pt")

# 4. Warmup CUDA graphs

_warm_streaming(
    model,
    images,
    scale_frames=4,
    warm_stream_n=10,
    dtype=torch.float16,
    passes=2,
    keyframe_interval=3,
)

# 5. Optimized streaming inference

per_frame_ms, scale_frames, phase1_ms = profile_streaming(
    model,
    images,
    num_frames=images.shape[0],
    dtype=torch.float16,
    keyframe_interval=3,
)

avg_latency = sum(per_frame_ms) / len(per_frame_ms)
print(f"Average FPS: {1000.0 / avg_latency:.2f}")

```

## Summary

- **Selective compilation** of `GCTStream`'s fixed-shape blocks in [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py) and [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) eliminates Python overhead while preserving dynamic components like the point head.
- **CUDA graph warmup** via `_warm_streaming` captures static kernel schedules for both keyframe and non-keyframe execution paths, preventing cold-start penalties.
- **Repository structure**: The optimization logic resides in [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py) (profiling) and [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) (interactive use), with the core model implementation in [`lingbot_map/models/gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream.py).
- **Critical API**: Use `torch.compiler.cudagraph_mark_step_begin()` before each forward pass during warmup, and toggle `model._set_skip_append()` to exercise the non-keyframe path.
- **Performance impact**: Omitting the point head during profiling (as done in [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py) line 57) isolates pure transformer throughput, revealing the full benefit of compilation.

## Frequently Asked Questions

### What PyTorch version is required for LingBot-Map compilation?

You need **PyTorch 2.2 or newer** with CUDA 11.8 or later. The `torch.compile` API with `mode="reduce-overhead"` and the `torch.compiler.cudagraph_mark_step_begin()` function were stabilized in these versions and are essential for the optimization workflow described in the repository.

### Can I skip the CUDA graph warmup if I'm only processing single images?

Yes, but you will lose performance benefits. The warmup phase specifically addresses **streaming inference** latency where consistent frame-to-frame timing matters. For single-image inference, you can omit `_warm_streaming`, though you will incur the full kernel launch overhead on every forward pass.

### Why does the profiling script remove the point head?

The point head adds approximately **6 milliseconds per frame** (as noted in [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py) line 57) and is not required for measuring pure transformer throughput. Removing it via `model.point_head = None` isolates the performance impact of `torch.compile` on the `GCTStream` backbone, providing cleaner benchmarking data for the fixed-shape blocks.

### How do I handle different keyframe intervals in production?

The `keyframe_interval` parameter in `_warm_streaming` must match your production streaming configuration. If you warmup with `keyframe_interval=3` but later use `keyframe_interval=5`, the CUDA graphs may not perfectly align with your actual execution pattern, potentially reintroducing overhead on frame transitions. Always warmup with the same interval you intend to use for inference.