# How to Integrate Custom Attention Mechanisms with Robbyant/lingbot-map

> Integrate custom attention mechanisms with lingbot-map by creating a new class and swapping it into the model builder. Learn how to extend Robbyant/lingbot-map without core pipeline modifications.

- Repository: [Robbyant/lingbot-map](https://github.com/Robbyant/lingbot-map)
- Tags: how-to-guide
- Published: 2026-07-07

---

**You can integrate custom attention mechanisms with lingbot-map by creating a class that follows the expected constructor signature and swapping it into the model builder in [`gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_stream_window.py) without modifying the core pipeline.**

The `lingbot-map` repository implements a modular transformer architecture designed for streaming vision-language tasks. Its attention layer is factored into reusable primitives located in `lingbot_map/layers`, making it straightforward to inject custom attention logic while preserving KV-cache management and sliding-window eviction functionality.

## Understanding the Attention Architecture

The repository centralizes all attention implementations in a single module, allowing higher-level blocks to import and instantiate them via a consistent interface.

### Core Attention Classes

The file **[`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/attention.py)** (lines 36–88) defines the generic `Attention` class, while lines 90–130 implement the `CausalAttention` variant optimized for streaming. Additional optimized variants include `FlashInferAttention` and `SDPAAttention`.

Higher-level transformer blocks—such as the GCT streaming blocks—import these classes in **[`lingbot_map/models/gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream_window.py)**:

```python
from lingbot_map.layers.attention import Attention, CausalAttention, FlashInferAttention, SDPAAttention

```

The block constructors accept an *attention class* argument, enabling you to pass any user-defined module that conforms to the expected signature.

### KV-Cache and Streaming Logic

Since KV-cache handling (including sliding-window eviction and special-token bookkeeping) lives inside the existing attention classes—specifically within `CausalAttention._apply_kv_cache_eviction_causal` (lines 993–1010)—you only need to reimplement this logic if your custom attention supports streaming. For feed-forward attention variants, you can inherit from `Attention` and override `forward` while leaving the cache code untouched.

## Step-by-Step Integration Guide

Follow these steps to replace the default attention mechanism with your custom implementation.

### 1. Create Your Custom Attention Module

Define a new class that matches the constructor signature `(dim, num_heads, qkv_bias, …, rope=None)` and returns a tensor of shape `[B, N, dim]`. The `forward` method must accept the following signature:

```python
def forward(self, x, pos=None, num_patches=None, num_special=None,
            num_frames=None, enable_3d_rope=False)

```

### 2. Replace the Class in the Model Builder

Modify the model instantiation in **[`gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/gct_stream_window.py)** to use your custom class. The model builder calls an internal method (typically `_make_attention`) to create each block's attention module, providing a clean injection point.

### 3. Verify with Inference

Run a forward pass with dummy inputs to ensure tensor shapes and gradients flow correctly through your custom layer.

## Custom Attention Implementation Examples

### Example: Minimal Custom Attention with Learnable Bias

The following implementation inherits from the base `Attention` class and adds a per-head learnable bias after the scaled dot-product attention:

```python

# my_attention.py

import torch
import torch.nn.functional as F
from lingbot_map.layers.attention import Attention

class BiasAttention(Attention):
    """Same as the default Attention but adds a per-head bias after SDPA."""
    def __init__(self, *args, bias_dim: int = 0, **kwargs):
        super().__init__(*args, **kwargs)
        if bias_dim:
            self.head_bias = torch.nn.Parameter(torch.zeros(self.num_heads, bias_dim))

    def forward(self, x, pos=None, num_patches=None, num_special=None,
                num_frames=None, enable_3d_rope=False):
        # Run the parent implementation

        out = super().forward(x, pos, num_patches, num_special,
                              num_frames, enable_3d_rope)

        # Optional per-head bias addition (broadcast over batch & seq length)

        if hasattr(self, "head_bias"):
            B, N, _ = out.shape
            bias = self.head_bias.view(1, 1, -1).expand(B, N, -1)
            out = out + bias
        return out

```

### Example: Plugging Custom Attention into a GCT Model

Override the `_make_attention` method in `GCTStreamWindow` to instantiate your custom class:

```python

# build_custom_gct.py

import torch
from lingbot_map.models.gct_stream_window import GCTStreamWindow
from my_attention import BiasAttention

class GCTStreamWindowCustom(GCTStreamWindow):
    def _make_attention(self, dim, num_heads, **kw):
        # Called inside the original __init__ to create each block's attention

        return BiasAttention(dim, num_heads, **kw)

# Instantiate the model

model = GCTStreamWindowCustom(
    dim=256,
    depth=12,
    num_heads=8,
    # other hyper-parameters remain unchanged

)

# Verify a forward pass

dummy = torch.randn(1, 16*16, 256)          # B, N, C

pos = torch.randn(1, 16*16, 3)               # example positional encoding

out = model(dummy, pos=pos)
print(out.shape)  # → torch.Size([1, 16*16, 256])

```

### Example: Using the Custom Model from the CLI

You can load your custom model through the provided [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) entry point:

```bash
python demo.py \
    --model_path /path/to/custom_gct.pt \
    --image_folder example/courthouse \
    --custom_attention my_attention.BiasAttention

```

To expose this functionality, add a thin wrapper in [`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) that dynamically loads the class via `importlib`.

## Key Files for Custom Attention Integration

- **[`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/attention.py)** – Central attention implementations including generic, causal, FlashInfer, and SDPA variants.
- **[`lingbot_map/models/gct_stream_window.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_stream_window.py)** – Model definition that assembles transformer blocks and injects the chosen attention class.
- **[`lingbot_map/layers/block.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/block.py)** – Implements the transformer block; calls the attention module via `self.attn`.
- **[`demo.py`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py)** – End-user entry point demonstrating checkpoint loading and inference.

## Summary

- **Modular Design:** `lingbot-map` isolates attention logic in [`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/attention.py), enabling clean swaps without touching the training or inference pipeline.
- **Inheritance Pattern:** Inherit from `Attention` to reuse KV-cache handling, or implement from scratch for streaming variants that require custom cache logic.
- **Injection Point:** Override `_make_attention` in `GCTStreamWindow` or pass your class directly to block constructors to integrate custom mechanisms.
- **Shape Requirements:** Ensure your `forward` method accepts `(x, pos, num_patches, num_special, num_frames, enable_3d_rope)` and returns `[B, N, dim]`.

## Frequently Asked Questions

### What is the exact signature my custom attention class must implement?

Your `__init__` must accept `(dim, num_heads, qkv_bias, …, rope=None)` and store `self.num_heads`. The `forward` method must accept `(x, pos=None, num_patches=None, num_special=None, num_frames=None, enable_3d_rope=False)` and return a tensor of shape `[B, N, dim]` matching the input batch size `B`, sequence length `N`, and feature dimension `dim`.

### Do I need to reimplement KV-cache logic for streaming models?

Only if your custom attention requires streaming. The base `Attention` class handles standard feed-forward cases. For streaming support, either inherit from `CausalAttention` (which implements `_apply_kv_cache_eviction_causal` at lines 993–1010) or manually implement the sliding-window eviction logic in your subclass.

### Can I use FlashInfer or SDPA optimizations with a custom attention wrapper?

Yes. You can inherit from `FlashInferAttention` or `SDPAAttention` instead of the base `Attention` class, then override specific methods such as `forward`. This preserves the optimized backend while allowing you to inject custom pre- or post-processing.

### Will custom attention break the pretrained checkpoint loading?

Standard checkpoints save weights by module name. If you add new parameters (like `head_bias` in the example), you must either initialize them before loading the base checkpoint (ignoring missing keys) or modify the state dict loading logic to skip incompatible layers. The base attention weights remain compatible as long as you preserve the original `qkv` and `proj` layer names.