How to Modify or Extend the Attention Logic in lingbot‑map's attention.py

You can modify or extend the attention logic in lingbot‑map by subclassing one of the four attention classes in lingbot_map/layers/attention.py and overriding methods like forward, prepare_qkv, or the KV‑cache eviction helpers.

The lingbot_map repository provides a modular Vision‑Transformer implementation with multiple attention mechanisms designed for both batch and streaming inference. To customize behavior for research experiments or production optimizations, developers need to understand how to safely modify or extend the attention logic in attention.py without breaking the underlying transformer backbone.

Understanding the Attention Architecture

The lingbot_map/layers/attention.py module implements four distinct attention classes that share a common structural pattern but serve different computational requirements:

• Attention (lines 36‑87): Standard multi‑head self‑attention supporting both batch and streaming modes
• CausalAttention (lines 90‑196): Autoregressive attention with KV‑cache eviction optimized for streaming video
• FlashInferAttention (lines 252‑370): FlashInfer‑based implementation leveraging paged KV‑cache for accelerated inference
• SDPAAttention (lines 590‑686): Pure‑PyTorch implementation using torch.nn.functional.scaled_dot_product_attention with dict‑based KV caching

All four classes share a common backbone architecture. Each initializes a linear projection self.qkv that creates Q/K/V tensors, optional per‑head normalization layers (q_norm/k_norm), and rotary position embedding support via rope or apply_rotary_emb. The fused_attn flag toggles between the fast scaled_dot_product_attention kernel and a manual q·kᵀ implementation.

Extension Strategies for Custom Attention Logic

Subclassing Core Attention Classes

The most robust method to modify attention logic is to inherit from an existing class and override specific methods. This preserves the initialization logic while allowing custom behavior in the forward pass.

For example, to add a learnable bias term to the Q/K projection:

# my_attention.py (placed next to attention.py)

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

class MyBiasAttention(Attention):
    """Adds a learnable bias term to the Q/K projection."""
    def __init__(self, *args, bias_dim: int = 16, **kwargs):
        super().__init__(*args, **kwargs)
        # Extra bias for each head

        self.head_bias = nn.Parameter(torch.zeros(self.num_heads, bias_dim))

    def forward(self, x, **kw):
        # Call parent to get raw Q/K/V

        qkv = self.qkv(x).reshape(x.shape[0], -1, 3,
                                  self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(2)

        # Inject bias on the query side (broadcast over tokens)

        q = q + self.head_bias[:, None, :, :]  # [B, H, N, D]

        # Continue with standard forward logic

        return super().forward(x, **kw)

Place this file in lingbot_map/layers/ and import it where the model is built in lingbot_map/models/gct_base.py.

Injecting Custom Positional Encoding

To replace or extend the rotary position embedding (RoPE) logic, modify the rope callable passed during initialization or override the position application within forward. The module already supports 3‑D RoPE via apply_rotary_emb from lingbot_map/layers/rope.py.

Substitute the rope parameter with your own implementation that accepts (q, k, pos) tensors and returns modified (q, k) tensors.

Modifying Attention Masks

For selective token attention, customize the mask generation before the scaled_dot_product_attention call. In CausalAttention and SDPAAttention, mask construction occurs at lines 172‑194 and can be intercepted to add custom constraints.

def forward(self, x: torch.Tensor, pos=None, **kwargs) -> torch.Tensor:
    # ... existing code up to attention call ...

    
    # Example: Mask out tokens beyond a fixed distance

    N = x.shape[1]
    distance = torch.arange(N).view(1, -1) - torch.arange(N).view(-1, 1)
    allowed = (distance.abs() <= self.max_distance)   # [N, N] bool mask

    mask = allowed.unsqueeze(0).unsqueeze(0).expand(B, 1, N, N)

    # Pass the mask to the fused kernel

    x = torch.nn.functional.scaled_dot_product_attention(
        q, k, v,
        dropout_p=self.attn_drop.p if self.training else 0.0,
        attn_mask=mask
    )
    # ... remainder unchanged ...

Insert this logic immediately before the F.scaled_dot_product_attention call (see lines 75‑84 in the original Attention class).

Customizing KV‑Cache Eviction Policies

To extend cache eviction logic, override the _apply_kv_cache_eviction_causal method (lines 297‑349) for causal attention or _apply_kv_cache_eviction (lines 688‑730) for SDPA attention. These methods control how the KV‑cache is pruned during streaming inference.

Subclass the relevant attention class and implement alternative retention policies such as exponential decay or attention‑score‑based eviction:

class ExponentialDecayAttention(CausalAttention):
    def _apply_kv_cache_eviction_causal(self, k_cache, v_cache, ...):
        # Custom eviction logic here

        pass

Practical Implementation Examples

Adding Token‑Type Embeddings

To incorporate token‑type information into the attention computation:

class TokenTypeAttention(Attention):
    def __init__(self, token_type_vocab: int, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.token_type_emb = nn.Embedding(token_type_vocab, self.head_dim)

    def forward(self, x, token_type_ids=None, **kw):
        # Usual Q/K/V extraction

        qkv = self.qkv(x).reshape(x.shape[0], -1, 3,
                                  self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(2)

        if token_type_ids is not None:
            # Add token‑type bias to Q and K

            bias = self.token_type_emb(token_type_ids)   # [B, N, D]

            bias = bias.unsqueeze(1)                     # [B, 1, N, D]

            q = q + bias
            k = k + bias

        # Continue with regular forward logic

        return super().forward(x, **kw)

Replicating FlashInfer Logic Without Dependencies

If you require the FlashInferAttention API but cannot install FlashInfer, inherit from Attention and copy the prepare_qkv helper (lines 302‑327):

class MyFastAttention(Attention):
    def prepare_qkv(self, x, pos=None, enable_3d_rope=False):
        # Copy-paste method body from FlashInferAttention

        # ... implementation ...

        return q_nhd, k_nhd, v_nhd

You can then call prepare_qkv manually in a custom forward pass that performs manual attention calculation via torch.einsum.

Integrating Custom Attention into the Model Pipeline

The high‑level model in lingbot_map/models/gct_base.py instantiates attention blocks via:

from lingbot_map.layers.attention import SDPAAttention

# ...

self.attn = SDPAAttention(dim=embed_dim, num_heads=num_heads, rope=self.rope)

To swap in your custom implementation, modify the import and class reference:

from lingbot_map.layers.my_attention import MyBiasAttention

# ...

self.attn = MyBiasAttention(
    dim=embed_dim,
    num_heads=num_heads,
    qkv_bias=True,
    proj_bias=True,
    attn_drop=0.1,
    proj_drop=0.1,
    rope=self.rope,
    bias_dim=self.head_dim // 2,   # Custom argument

)

For dynamic switching between causal and non‑causal modes, implement a factory function:

def get_attention(causal: bool, **kwargs):
    if causal:
        from lingbot_map.layers.attention import CausalAttention
        return CausalAttention(**kwargs)
    else:
        from lingbot_map.layers.attention import SDPAAttention
        return SDPAAttention(**kwargs)

Summary

• Four base classes exist in lingbot_map/layers/attention.py: Attention, CausalAttention, FlashInferAttention, and SDPAAttention, each optimized for different inference scenarios.
• Subclass and override methods like forward, prepare_qkv, or _apply_kv_cache_eviction_causal to inject custom logic without rewriting the entire transformer.
• Hook into the projection layer by modifying self.qkv in __init__ to add convolutions, biases, or token‑type embeddings.
• Customize masking by inserting tensor operations before the scaled_dot_product_attention call in the forward method.
• Integrate seamlessly by importing your custom class in lingbot_map/models/gct_base.py and passing it to the transformer block constructor.

Frequently Asked Questions

Which attention class should I subclass for streaming video applications?

Subclass CausalAttention for streaming video applications. According to the lingbot_map source code, this class implements KV‑cache eviction specifically designed for autoregressive streaming (lines 90‑196), unlike the base Attention class which assumes bidirectional context.

How do I add a learnable bias to the Q/K projections?

Override __init__ to define a nn.Parameter and modify the forward method to add the bias to the query tensor after the qkv projection but before the attention computation. The MyBiasAttention example demonstrates broadcasting a head‑specific bias across the sequence dimension.

Can I use the FlashInfer attention implementation without installing FlashInfer?

Yes. Inherit from the base Attention class and copy the prepare_qkv method implementation from FlashInferAttention (lines 302‑327 in attention.py). This allows you to leverage the same Q/K/V preparation logic and rope handling without the FlashInfer dependency, though you must implement the actual attention computation using standard PyTorch operations.

Where is the rotary position embedding (RoPE) configured?

RoPE is configured in the rope parameter passed to the attention class __init__ and applied via apply_rotary_emb imported from lingbot_map/layers/rope.py. The FlashInferAttention class additionally supports 3‑D RoPE through the enable_3d_rope flag in its prepare_qkv method.

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