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

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 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 (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:

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:

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

# 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:

# 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 entry point:

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 that dynamically loads the class via importlib.

Key Files for Custom Attention Integration

• lingbot_map/layers/attention.py – Central attention implementations including generic, causal, FlashInfer, and SDPA variants.
• lingbot_map/models/gct_stream_window.py – Model definition that assembles transformer blocks and injects the chosen attention class.
• lingbot_map/layers/block.py – Implements the transformer block; calls the attention module via self.attn.
• 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, 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.

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