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

> Easily modify or extend lingbot-map's attention logic. Subclass existing attention classes and override key methods like forward or prepareqkv for custom functionality.

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

---

**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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/attention.py) without breaking the underlying transformer backbone.

## Understanding the Attention Architecture

The [`lingbot_map/layers/attention.py`](https://github.com/Robbyant/lingbot-map/blob/main/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:

```python

# 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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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.

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

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

```python
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):

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_base.py) instantiates attention blocks via:

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

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

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/layers/rope.py). The `FlashInferAttention` class additionally supports 3‑D RoPE through the `enable_3d_rope` flag in its `prepare_qkv` method.