# How to Optimize Inference Speed with torch.compile and CUDA Graph Warmup in Lingvo-Bot-Map

> Boost transformer inference with torch.compile and CUDA graph warmup. Achieve 10-15% throughput gains by optimizing kernel launches and reducing overhead in your Lingvo-Bot-Map projects.

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

---

**You can boost transformer inference throughput by 10–15% using `torch.compile(mode="reduce-overhead")` on static sub-modules and warming up CUDA graphs with `torch.compiler.cudagraph_mark_step_begin()` to eliminate per-frame kernel launch overhead.**

The Lingvo-Bot-Map (Lingo-Bot-Map) repository implements a transformer-based **GCT** (Global-Context-Transformer) that processes video frames sequentially. When deploying this model on CUDA devices, the dominant latency stems from repeated Python interpreter overhead and redundant CUDA kernel launches during per-frame forward passes. By combining ahead-of-time compilation with CUDA graph capture, you can optimize inference speed while maintaining numerical accuracy.

## Compile Static Sub-Modules with torch.compile

The first optimization targets the heavy transformer blocks that execute with fixed input dimensions. In [`scripts/benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py) [lines 62‑74](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py#L62-L74) and [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py) [lines 40‑57](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py#L40-L57), the helper function `compile_model` walks the model hierarchy and applies `torch.compile` with `mode="reduce-overhead"` to specific sub-modules.

**Target these components for compilation:**

- **`frame_blocks`** – The per-frame transformer blocks inside `model.aggregator`
- **`patch_embed.blocks`** – The patch embedding layers
- **`global_blocks`** – Global attention modules including `attn_pre`, `ffn_residual`, and `attn.proj`

This mode generates a fused, ahead-of-time compiled graph that removes Python dispatch overhead and fuses many small operations, yielding a ~5–10 % speed-up for compute-heavy blocks. To save an additional ~5.9 ms per frame during benchmarking, the code drops the optional `point_head` module, which is unnecessary for pure FPS measurement.

## Eliminate Kernel Launch Overhead with CUDA Graph Warmup

Even after compilation, each forward pass normally incurs CPU overhead to record CUDA kernel launches. The second optimization captures these launches into a reusable graph. Before each forward call, the code executes `torch.compiler.cudagraph_mark_step_begin()`.

In [`scripts/benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py) [lines 218‑235](https://github.com/Robbyant/lingbot-map/blob/main/scripts/benchmark_gct_memory.py#L218-L235) and [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py) [lines 92‑102](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py#L92-L102), the `run_streaming_once` and `profile_streaming` functions demonstrate the pattern: perform one warm-up forward pass to capture the graph, then reuse that captured graph for all subsequent frames. This same pattern appears in production code at [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) [lines 221‑233](https://github.com/Robbyant/lingbot-map/blob/main/demo.py#L221-L233) and inside the core model at [`lingbot_map/models/gct_stream.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream.py) [lines 420‑455](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream.py#L420-L455).

**Critical placement:** The marker must appear **inside** the per-frame loop **after** tensors move to the GPU but **before** the forward call. This ensures the captured graph includes the actual inference kernels while excluding CPU-GPU synchronization overhead.

## Step-by-Step Implementation

Below is a minimal implementation that mirrors the logic in [`benchmark_gct_memory.py`](https://github.com/Robbyant/lingbot-map/blob/main/benchmark_gct_memory.py) and [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py). It compiles the hot modules, warms up the CUDA graph, and runs a timed inference loop.

```python
import torch
from lingbot_map.models.gct_stream import GCTStream  # or any model class

def compile_hot_modules(model: torch.nn.Module) -> None:
    """Compile static heavy sub-modules to reduce Python overhead."""
    agg = model.aggregator
    # Frame-level transformer blocks

    for i, blk in enumerate(agg.frame_blocks):
        agg.frame_blocks[i] = torch.compile(blk, mode="reduce-overhead")
    # Patch-embed blocks

    for i, blk in enumerate(agg.patch_embed.blocks):
        agg.patch_embed.blocks[i] = torch.compile(blk, mode="reduce-overhead")
    # Global blocks (may have extra sub-modules)

    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")
    # Point-head not needed for pure FPS measurement

    model.point_head = None

def warmup_and_infer(model, images, dtype=torch.float16, device="cuda"):
    """Run warm-up (captures CUDA graph) then timed streaming pass."""
    model.eval().to(device)
    compile_hot_modules(model)

    # ---- Warm-up (captures the CUDA graph) ----

    torch.cuda.synchronize(device)
    with torch.autocast(device_type="cuda", dtype=dtype):
        torch.compiler.cudagraph_mark_step_begin()
        model.forward(
            images[:1].to(device, dtype),  # single frame or scale batch

            num_frame_for_scale=1,
            num_frame_per_block=1,
            causal_inference=True
        )

    # ---- Timed streaming loop ----

    times = []
    for i in range(images.shape[0]):
        frame = images[i:i+1].to(device, dtype, non_blocking=True)
        torch.cuda.synchronize(device)
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        
        # Mark step immediately before forward

        torch.compiler.cudagraph_mark_step_begin()
        with torch.autocast(device_type="cuda", dtype=dtype):
            model.forward(
                frame,
                num_frame_for_scale=1,
                num_frame_per_block=1,
                causal_inference=True
            )
        
        end.record()
        torch.cuda.synchronize(device)
        times.append(start.elapsed_time(end))  # ms per frame

    return times

```

Note that [`lingbot_map/layers/rope.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/rope.py) contains implementation details ensuring operations remain CUDA-graph-compatible, such as avoiding `torch.view_as_complex` which would break graph capture.

## Performance Gains and Architecture Considerations

When applied together, these techniques deliver a **consistent 10–15 % throughput boost** (approximately 2–3 FPS higher) across benchmark settings. The CUDA graph reuse alone saves roughly **0.3–0.6 ms per frame** by eliminating redundant kernel launch recordings.

**Key architectural constraints:**

- **Fixed-shape compilation only:** Compile only modules with static input dimensions (e.g., frame-level transformers). Leave variable-length logic (such as key-frame handling) in eager mode to preserve correctness.
- **Hot-path isolation:** By compiling only heavy blocks and excluding lightweight utilities, you avoid excessive compilation time and keep startup latency low.
- **Memory efficiency:** Removing the `point_head` module during benchmarking (as seen in `compile_model`) reduces memory pressure and computation without affecting the core transformer path.

## Summary

- Use `torch.compile(mode="reduce-overhead")` on heavy transformer blocks like `frame_blocks` and `patch_embed` to fuse operations and remove Python overhead.
- Call `torch.compiler.cudagraph_mark_step_begin()` immediately before each forward pass to capture and reuse CUDA kernel launch sequences.
- Compile only fixed-shape modules (frame-level transformers) while leaving variable-length logic eager-executed to preserve correctness.
- Expect a 10–15 % throughput increase (approximately 2–3 FPS) when both techniques are applied to the Lingvo-Bot-Map GCT model.

## Frequently Asked Questions

### Where should I place `cudagraph_mark_step_begin` in my inference loop?

Place the call immediately after moving input tensors to the GPU but before the model's forward pass, ensuring it sits inside the per-frame loop. This placement guarantees that the captured CUDA graph contains the actual kernel launches used for inference while excluding CPU-GPU synchronization overhead.

### Why use `mode="reduce-overhead"` instead of other torch.compile modes?

The `reduce-overhead` mode specifically minimizes framework dispatch and Python interpreter costs, which is critical for the small, repeated forward passes in video frame processing. While `max-autotune` might optimize individual kernels further, it introduces additional compilation time that outweighs benefits for this streaming use case.

### Can I compile the entire GCT model at once?

You should not compile the entire model. Instead, selectively compile static sub-modules such as `frame_blocks`, `patch_embed`, and attention projections. Variable-length components—like key-frame logic and the `point_head`—should remain uncompiled to avoid recompilation penalties and maintain flexibility for dynamic sequence lengths.

### How much latency does CUDA graph warmup actually save?

According to benchmarks in [`gct_profile.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_profile.py), reusing captured CUDA graphs eliminates approximately **0.3–0.6 milliseconds of per-frame latency** by removing the need to re-record kernel launches for every forward pass.