Module neuralops.attentions.vig_attention

ViG Attention: Vision GNN-inspired 2D Grid Linear Attention with Gated Fusion. Hardened for production use across diverse transformer architectures: - DistilBERT (Flat MHA) - RoBERTa / BERT (Nested Encoder Self-Attention) - GPT-2 (Causal Decoder with Conv1D) - Qwen / LLaMA (GQA with Rotary Position Embeddings)

Functions

def apply_rotary_pos_emb(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
position_ids: torch.Tensor | None = None) ‑> tuple[torch.Tensor, torch.Tensor]

Expand source code

def apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    position_ids: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Applies rotary position embeddings safely handling even or odd dimensions."""
    if cos is None or sin is None:
        return q, k

    if cos.dim() == 2:
        cos = cos.unsqueeze(0).unsqueeze(1)
        sin = sin.unsqueeze(0).unsqueeze(1)
    elif cos.dim() == 3:
        cos = cos.unsqueeze(1)
        sin = sin.unsqueeze(1)

    while cos.dim() < q.dim():
        cos = cos.unsqueeze(0)
        sin = sin.unsqueeze(0)

    head_dim = q.shape[-1]
    rotary_dim = 2 * (min(cos.shape[-1], head_dim) // 2)
    if rotary_dim == 0:
        return q, k

    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]

    cos_q = cos[..., : q.shape[-2], :rotary_dim].to(q.dtype)
    sin_q = sin[..., : q.shape[-2], :rotary_dim].to(q.dtype)

    cos_k = cos[..., : k.shape[-2], :rotary_dim].to(k.dtype)
    sin_k = sin[..., : k.shape[-2], :rotary_dim].to(k.dtype)

    q_rot_embed = (q_rot * cos_q) + (_rotate_half(q_rot) * sin_q)
    k_rot_embed = (k_rot * cos_k) + (_rotate_half(k_rot) * sin_k)

    q = torch.cat([q_rot_embed, q_pass], dim=-1) if q_pass.numel() > 0 else q_rot_embed
    k = torch.cat([k_rot_embed, k_pass], dim=-1) if k_pass.numel() > 0 else k_rot_embed

    return q, k

Applies rotary position embeddings safely handling even or odd dimensions.

def inject_vig(layer: Any = None,
model: Any = None,
model_config: Any = None,
comp_config: Any = None,
layer_idx: int = 0,
total_layers: int = 12,
old_attn: nn.Module | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("vig_inject")
@ATTENTION_REGISTRY.register("vig_attention_inject")
@ATTENTION_REGISTRY.register("inject_vig")
def inject_vig(
    layer: Any = None,
    model: Any = None,
    model_config: Any = None,
    comp_config: Any = None,
    layer_idx: int = 0,
    total_layers: int = 12,
    old_attn: nn.Module | None = None,
    **kwargs: Any,
):
    target = layer if layer is not None else model
    if target is None:
        target = kwargs.get("model", kwargs.get("layer"))
    if target is None:
        return None

    cfg = model_config if model_config is not None else getattr(target, "config", None)

    # Whole-model traversal
    layer_list = None
    if hasattr(target, "distilbert") and hasattr(target.distilbert, "transformer"):
        layer_list = target.distilbert.transformer.layer
    elif hasattr(target, "roberta") and hasattr(target.roberta, "encoder"):
        layer_list = target.roberta.encoder.layer
    elif hasattr(target, "transformer") and hasattr(target.transformer, "h"):
        layer_list = target.transformer.h
    elif hasattr(target, "h") and isinstance(target.h, (torch.nn.ModuleList, list)):
        layer_list = target.h
    elif hasattr(target, "model") and hasattr(target.model, "layers"):
        layer_list = target.model.layers
    elif hasattr(target, "encoder") and hasattr(target.encoder, "layer"):
        layer_list = target.encoder.layer
    elif (
        hasattr(target, "bert")
        and hasattr(target.bert, "encoder")
        and hasattr(target.bert.encoder, "layer")
    ):
        layer_list = target.bert.encoder.layer
    elif hasattr(target, "layers") and isinstance(target.layers, (torch.nn.ModuleList, list)):
        layer_list = target.layers

    if layer_list is not None:
        num_injected = 0
        tot = len(layer_list)
        for idx, lyr in enumerate(layer_list):
            res = inject_vig(
                layer=lyr,
                model_config=cfg,
                comp_config=(
                    comp_config if "comp_config" in locals() or "comp_config" in kwargs else None
                ),
                layer_idx=idx,
                total_layers=tot,
                **kwargs,
            )
            if res is not None:
                num_injected += 1
        return target if num_injected > 0 else None

    layer = target

    """Safe injection helper for ViG attention."""
    if old_attn is None:
        return None

    if model_config is None and hasattr(layer, "config"):
        model_config = layer.config
    if model_config is None and hasattr(old_attn, "config"):
        model_config = old_attn.config

    hidden_size = (
        _cfg_get(model_config, "hidden_size")
        or _cfg_get(model_config, "d_model")
        or _cfg_get(model_config, "dim")
        or getattr(old_attn, "hidden_size", None)
        or getattr(old_attn, "embed_dim", None)
        or (
            old_attn.q_proj.in_features
            if hasattr(old_attn, "q_proj") and hasattr(old_attn.q_proj, "in_features")
            else None
        )
        or (
            old_attn.q_lin.in_features
            if hasattr(old_attn, "q_lin") and hasattr(old_attn.q_lin, "in_features")
            else None
        )
        or (
            old_attn.query.in_features
            if hasattr(old_attn, "query") and hasattr(old_attn.query, "in_features")
            else None
        )
        or 768
    )
    num_heads = (
        _cfg_get(model_config, "num_attention_heads")
        or _cfg_get(model_config, "num_heads")
        or _cfg_get(model_config, "n_heads")
        or getattr(old_attn, "num_heads", None)
        or getattr(old_attn, "num_attention_heads", None)
        or 12
    )

    if model_config is None:
        import types

        model_config = types.SimpleNamespace(
            hidden_size=hidden_size,
            num_attention_heads=num_heads,
            num_key_value_heads=getattr(old_attn, "num_key_value_heads", num_heads),
            is_decoder=getattr(old_attn, "is_decoder", False),
        )

    lora_rank = kwargs.pop("lora_rank", _cfg_get(comp_config, "lora_rank", 16))
    rank = kwargs.pop("rank", _cfg_get(comp_config, "rank", 32))

    try:
        new_attn = ViGAttention(
            config=model_config,
            rank=rank,
            lora_rank=lora_rank,
            pretrained_attn=old_attn,
            **kwargs,
        )

        for attr in ("self_attn", "attn", "self"):
            if hasattr(layer, attr):
                setattr(layer, attr, new_attn)
                return new_attn
        if hasattr(layer, "attention"):
            if hasattr(layer.attention, "self"):
                layer.attention.self = new_attn
            else:
                layer.attention = new_attn
            return new_attn
        layer.attention = new_attn
        return new_attn
    except Exception as e:
        logger.error("inject_vig failed at layer %d: %s", layer_idx, e)
        return None

Classes

class LoRALinear (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = True,
device: torch.device | None = None,
dtype: torch.dtype | None = None)

Expand source code

class LoRALinear(nn.Module):
    """Linear projection with frozen base weights and trainable LoRA adapters."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        freeze_base: bool = True,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        factory_kwargs = {"device": device, "dtype": dtype}
        self.in_features = in_features
        self.out_features = out_features
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0
        self.freeze_base = freeze_base

        self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            self.register_parameter("bias", None)

        if self.lora_rank > 0:
            self.lora_A = nn.Parameter(torch.empty((lora_rank, in_features), **factory_kwargs))
            self.lora_B = nn.Parameter(torch.empty((out_features, lora_rank), **factory_kwargs))
        else:
            self.register_parameter("lora_A", None)
            self.register_parameter("lora_B", None)

        self.reset_parameters()
        if freeze_base:
            self.freeze_base_weights()

    def reset_parameters(self) -> None:
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            nn.init.uniform_(self.bias, -bound, bound)

        if self.lora_rank > 0 and self.lora_A is not None:
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            self.lora_A.requires_grad = True
        if self.lora_rank > 0 and self.lora_B is not None:
            nn.init.zeros_(self.lora_B)
            self.lora_B.requires_grad = True

    def freeze_base_weights(self) -> None:
        self.weight.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False
        if self.lora_rank > 0:
            if self.lora_A is not None:
                self.lora_A.requires_grad = True
            if self.lora_B is not None:
                self.lora_B.requires_grad = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.linear(x, self.weight, self.bias)
        if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
            out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
        return out

Linear projection with frozen base weights and trainable LoRA adapters.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, x: torch.Tensor) -> torch.Tensor:
    out = F.linear(x, self.weight, self.bias)
    if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None:
        out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling
    return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def freeze_base_weights(self) ‑> None

Expand source code

def freeze_base_weights(self) -> None:
    self.weight.requires_grad = False
    if self.bias is not None:
        self.bias.requires_grad = False
    if self.lora_rank > 0:
        if self.lora_A is not None:
            self.lora_A.requires_grad = True
        if self.lora_B is not None:
            self.lora_B.requires_grad = True

def reset_parameters(self) ‑> None

Expand source code

def reset_parameters(self) -> None:
    nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
    if self.bias is not None:
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)

    if self.lora_rank > 0 and self.lora_A is not None:
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.lora_A.requires_grad = True
    if self.lora_rank > 0 and self.lora_B is not None:
        nn.init.zeros_(self.lora_B)
        self.lora_B.requires_grad = True

class ViGAttention (config: Any = None,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register(
    "vig", supports_causal=True, supports_bidirectional=True, supports_gqa=True
)
@ATTENTION_REGISTRY.register(
    "vig_attention",
    supports_causal=True,
    supports_bidirectional=True,
    supports_gqa=True,
)
class ViGAttention(BaseAttention):
    r"""Vision GNN-inspired 2D Grid Linear Attention with Learned Gated Fusion."""

    supports_gqa: bool = True
    supports_causal: bool = True
    supports_bidirectional: bool = True
    supports_rope: bool = True
    is_custom_attention: bool = True

    def __init__(
        self,
        config: Any = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank)
        src_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("old_attn")

        if config is not None and not isinstance(config, int):
            self.adapter = ModelAdapter(config)
            self.num_heads: int = self.adapter.num_heads
            self.num_kv_heads: int = self.adapter.num_key_value_heads
            self.dim: int = self.adapter.hidden_size
            self.head_dim: int = self.adapter.head_dim
            self.all_head_size: int = self.num_heads * self.head_dim
            self.kv_head_size: int = self.num_kv_heads * self.head_dim
            self.kv_repeat: int = self.adapter.kv_repeat
            self.is_decoder: bool = getattr(config, "is_decoder", False)
        else:
            self.adapter = None
            self.dim = _cfg_get(
                config,
                "hidden_size",
                kwargs.get("hidden_size", kwargs.get("d_model", 768)),
            )
            self.num_heads = _cfg_get(config, "num_attention_heads", kwargs.get("num_heads", 12))
            self.num_kv_heads = _cfg_get(config, "num_key_value_heads", self.num_heads)
            self.head_dim = self.dim // max(1, self.num_heads)
            self.all_head_size = self.num_heads * self.head_dim
            self.kv_head_size = self.num_kv_heads * self.head_dim
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))
            self.is_decoder = bool(_cfg_get(config, "is_decoder", kwargs.get("is_decoder", False)))

        # Explicit overrides
        if "hidden_size" in kwargs:
            self.dim = kwargs["hidden_size"]
            self.head_dim = self.dim // max(1, self.num_heads)
            self.all_head_size = self.num_heads * self.head_dim
            self.kv_head_size = self.num_kv_heads * self.head_dim
        if "num_heads" in kwargs:
            self.num_heads = kwargs["num_heads"]
            self.head_dim = self.dim // max(1, self.num_heads)
            self.all_head_size = self.num_heads * self.head_dim
            self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads))

        if is_decoder is not None:
            self.is_decoder = bool(is_decoder)

        self.is_nested_encoder: bool = False
        self.lora_rank = kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank))
        self.lora_alpha = kwargs.get("lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha))

        self.q_proj = LoRALinear(
            self.dim,
            self.all_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.k_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.v_proj = LoRALinear(
            self.dim,
            self.kv_head_size,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )
        self.out_proj = LoRALinear(
            self.all_head_size,
            self.dim,
            bias=True,
            lora_rank=self.lora_rank,
            lora_alpha=self.lora_alpha,
            freeze_base=True,
        )

        self.gate_net = nn.Linear(self.dim, 2 * self.num_heads)
        dropout_p = _cfg_get(
            config,
            "dropout",
            _cfg_get(
                config,
                "hidden_dropout_prob",
                _cfg_get(config, "attention_probs_dropout_prob", 0.1),
            ),
        )
        self.dropout = nn.Dropout(dropout_p)
        self.rotary_emb = getattr(src_attn, "rotary_emb", None)

        if src_attn is not None:
            self.copy_weights(src_attn)

    # State dict compatibility aliases via properties
    @property
    def query(self) -> LoRALinear:
        return self.q_proj

    @property
    def key(self) -> LoRALinear:
        return self.k_proj

    @property
    def value(self) -> LoRALinear:
        return self.v_proj

    @property
    def out(self) -> LoRALinear:
        return self.out_proj

    @property
    def W_q(self) -> LoRALinear:
        return self.q_proj

    @property
    def W_k(self) -> LoRALinear:
        return self.k_proj

    @property
    def W_v(self) -> LoRALinear:
        return self.v_proj

    @property
    def q_lin(self) -> LoRALinear:
        return self.q_proj

    @property
    def k_lin(self) -> LoRALinear:
        return self.k_proj

    @property
    def v_lin(self) -> LoRALinear:
        return self.v_proj

    @property
    def out_lin(self) -> LoRALinear:
        return self.out_proj

    @property
    def dense(self) -> LoRALinear:
        return self.out_proj

    @property
    def o_proj(self) -> LoRALinear:
        return self.out_proj

    @property
    def c_proj(self) -> LoRALinear:
        return self.out_proj

    def copy_weights(self, old_attn: Any) -> None:
        """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
        extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
            "q": (None, None),
            "k": (None, None),
            "v": (None, None),
            "o": (None, None),
        }

        # 1. UniversalAdapter extraction
        for proj in ["q", "k", "v", "o"]:
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
                if w is not None:
                    extracted[proj] = (w, b)
            except Exception as e:
                logger.debug("UniversalAdapter fallback failed: %s", e)

        # 2. GPT-2 Fused c_attn Conv1D handling
        if hasattr(old_attn, "c_attn") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            weight = old_attn.c_attn.weight.data
            bias = getattr(old_attn.c_attn, "bias", None)
            bias_data = bias.data if bias is not None else None
            if (
                ("Conv1D" in old_attn.c_attn.__class__.__name__)
                or hasattr(old_attn.c_attn, "nf")
                or (
                    weight.dim() == 2
                    and (weight.shape[0] < weight.shape[1] or weight.shape[1] == 3 * self.dim)
                )
            ):
                weight = weight.t().contiguous()
            split_size = weight.shape[0] // 3
            q_w = weight[:split_size, :]
            k_w = weight[split_size : 2 * split_size, :]
            v_w = weight[2 * split_size :, :]
            q_b = bias_data[:split_size] if bias_data is not None else None
            k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None
            v_b = bias_data[2 * split_size :] if bias_data is not None else None
            extracted["q"] = (q_w, q_b)
            extracted["k"] = (k_w, k_b)
            extracted["v"] = (v_w, v_b)

        if hasattr(old_attn, "c_proj") and extracted["o"][0] is None:
            w_o = old_attn.c_proj.weight.data
            if (
                ("Conv1D" in old_attn.c_proj.__class__.__name__)
                or hasattr(old_attn.c_proj, "nf")
                or (
                    w_o.dim() == 2
                    and (
                        w_o.shape[0] < w_o.shape[1]
                        or (w_o.shape[0] != self.dim and w_o.shape[1] == self.dim)
                    )
                )
            ):
                w_o = w_o.t().contiguous()
            b_o = getattr(old_attn.c_proj, "bias", None)
            extracted["o"] = (w_o, b_o.data if b_o is not None else None)

        # 3. RoBERTa / BERT Nested Encoder handling
        if hasattr(old_attn, "self") and (
            extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
        ):
            s = old_attn.self
            if hasattr(s, "query"):
                extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
            if hasattr(s, "key"):
                extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
            if hasattr(s, "value"):
                extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

        if (
            hasattr(old_attn, "output")
            and hasattr(old_attn.output, "dense")
            and extracted["o"][0] is None
        ):
            dense = old_attn.output.dense
            extracted["o"] = (dense.weight.data, getattr(dense, "bias", None))

        # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT
        named_map = {
            "q": ["q_proj", "q_lin", "query", "W_q"],
            "k": ["k_proj", "k_lin", "key", "W_k"],
            "v": ["v_proj", "v_lin", "value", "W_v"],
            "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
        }
        for proj, names in named_map.items():
            if extracted[proj][0] is None:
                for name in names:
                    if hasattr(old_attn, name):
                        layer = getattr(old_attn, name)
                        base = getattr(layer, "base_layer", layer)
                        if hasattr(base, "weight"):
                            extracted[proj] = (
                                base.weight.data,
                                getattr(base, "bias", None),
                            )
                            break

        # Detect nested encoder architecture
        cls_name = old_attn.__class__.__name__
        has_out_proj = any(
            hasattr(old_attn, attr)
            for attr in (
                "dense",
                "out_proj",
                "c_proj",
                "out_lin",
                "o_proj",
                "wo",
                "o",
                "proj",
                "out",
            )
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
        ) and not has_out_proj

        missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None]
        if not self.is_nested_encoder and extracted["o"][0] is None:
            missing.append("o")
        if missing:
            logger.warning(
                "Attention module copy_weights warning. Proceeding with initialized weights."
            )

        def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor:
            if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
                return w.t().contiguous()
            return w

        with torch.no_grad():
            for proj_key, target in [
                ("q", self.q_proj),
                ("k", self.k_proj),
                ("v", self.v_proj),
                ("o", self.out_proj),
            ]:
                w, b = extracted[proj_key]
                if w is not None:
                    in_dim = self.all_head_size if proj_key == "o" else self.dim
                    w_oriented = _orient(w, in_dim)
                    target_w = target.weight
                    out_d, in_d = target_w.shape
                    src_out, src_in = w_oriented.shape
                    target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                        w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                    )
                    if b is not None and target.bias is not None:
                        b_val = b.data if isinstance(b, nn.Parameter) else b
                        target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                            b_val[: min(len(target.bias), len(b_val))]
                        )
                elif proj_key == "o" and self.is_nested_encoder:
                    min_d = min(target.weight.shape[0], target.weight.shape[1])
                    target.weight.data.zero_()
                    target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                    if target.bias is not None:
                        target.bias.data.zero_()

                # Frobenius norm verification (>= 1e-4)
                w_norm = float(torch.linalg.norm(target.weight.data.float()).item())
                if w_norm < 1e-4:
                    logger.warning(
                        f"ViGAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                    )

                target.freeze_base_weights()

        # Preserve source LoRA adapters
        source_submodules = {
            "q": ["q_proj", "q_lin", "query", "W_q"],
            "k": ["k_proj", "k_lin", "key", "W_k"],
            "v": ["v_proj", "v_lin", "value", "W_v"],
            "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
        }
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            for name in source_submodules[proj_key]:
                if hasattr(old_attn, name):
                    sub = getattr(old_attn, name)
                    if (
                        hasattr(sub, "lora_A")
                        and hasattr(sub, "lora_B")
                        and sub.lora_A is not None
                        and sub.lora_B is not None
                    ):
                        src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                        src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                        if hasattr(src_lA, "data") and target.lora_A is not None:
                            if (
                                src_lA.shape[0] == target.lora_A.shape[0]
                                and src_lA.shape[1] == target.lora_A.shape[1]
                            ):
                                target.lora_A.data.copy_(src_lA.data)
                                target.lora_A.requires_grad = True
                            else:
                                print(
                                    f"Warning: lora_A shape mismatch for {proj_key}; keeping zero-init."
                                )
                        if hasattr(src_lB, "data") and target.lora_B is not None:
                            if (
                                src_lB.shape[0] == target.lora_B.shape[0]
                                and src_lB.shape[1] == target.lora_B.shape[1]
                            ):
                                target.lora_B.data.copy_(src_lB.data)
                                target.lora_B.requires_grad = True
                            else:
                                print(
                                    f"Warning: lora_B shape mismatch for {proj_key}; keeping zero-init."
                                )
                        if hasattr(sub, "lora_scaling"):
                            target.lora_scaling = sub.lora_scaling
                        break

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_value: Any | None = None,
        output_attentions: bool | None = False,
        output_hidden_states: bool | None = False,
        return_dict: bool | None = True,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        B, N, D = hidden_states.shape
        device = hidden_states.device
        dtype = hidden_states.dtype

        Hh = self.num_heads
        H_kv = self.num_kv_heads
        Dh = self.head_dim

        q = self.q_proj(hidden_states).view(B, N, Hh, Dh).transpose(1, 2).contiguous()
        k = self.k_proj(hidden_states).view(B, N, H_kv, Dh).transpose(1, 2).contiguous()
        v = self.v_proj(hidden_states).view(B, N, H_kv, Dh).transpose(1, 2).contiguous()

        # Rotary position embeddings
        pos_emb = kwargs.get("position_embeddings")
        if pos_emb is not None and isinstance(pos_emb, (tuple, list)) and len(pos_emb) >= 2:
            cos, sin = pos_emb[0], pos_emb[1]
            q, k = apply_rotary_pos_emb(q, k, cos, sin)
        elif getattr(self, "rotary_emb", None) is not None:
            try:
                from neuralops.utils.compat import maybe_apply_rotary

                q, k = maybe_apply_rotary(
                    q,
                    k,
                    getattr(self, "adapter", None),
                    self.rotary_emb,
                    kwargs.get("position_ids"),
                )
            except Exception as e:
                logger.warning("maybe_apply_rotary failed: %s; continuing without RoPE", e)

        # GQA expansion
        if self.kv_repeat > 1:
            k = repeat_kv(k, self.kv_repeat)
            v = repeat_kv(v, self.kv_repeat)

        is_decoder = (
            self.is_decoder
            or getattr(getattr(self, "config", None), "is_decoder", False)
            or getattr(getattr(self, "config", None), "model_type", "")
            in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama")
        )

        # 1. Topological Graph Adjacency for 1D Token Sequences
        # Bounded local window (radius 16) + Global Anchor Tokens (token 0 and token N-1)
        pos = torch.arange(N, device=device)
        diff = pos.unsqueeze(0) - pos.unsqueeze(1)  # diff[i, j] = pos[j] - pos[i]
        local_mask = diff.abs() <= 16

        anchor_mask = torch.zeros((N, N), device=device, dtype=torch.bool)
        anchor_mask[0, :] = True
        anchor_mask[:, 0] = True
        if N > 1:
            anchor_mask[-1, :] = True
            anchor_mask[:, -1] = True

        graph_adj = local_mask | anchor_mask
        if is_decoder:
            causal_mask = diff <= 0  # i >= j
            graph_adj = graph_adj & causal_mask

        # Branch 1: Topological Graph Attention
        scale = 1.0 / math.sqrt(Dh)
        scores_graph = torch.matmul(q, k.transpose(-1, -2).contiguous()) * scale
        scores_graph = scores_graph.masked_fill(~graph_adj.unsqueeze(0).unsqueeze(0), -10000.0)
        if attention_mask is not None:
            scores_graph = apply_attention_mask(scores_graph, attention_mask)
        attn_graph = torch.nan_to_num(F.softmax(scores_graph.float(), dim=-1), nan=0.0).to(dtype)
        out_graph = torch.matmul(attn_graph, v)

        # Branch 2: Global Linear Recurrence / Attention with denom.clamp(min=1.0)
        q_lin = F.elu(q.float()) + 1.0
        k_lin = F.elu(k.float()) + 1.0

        valid_m = get_token_valid_mask(attention_mask, B, N, device)
        if valid_m is not None:
            k_lin = k_lin * valid_m.unsqueeze(1).unsqueeze(-1).to(k_lin.dtype)
            v_f = v.float() * valid_m.unsqueeze(1).unsqueeze(-1).to(v.dtype)
        else:
            v_f = v.float()

        if is_decoder:
            causal_tril = (
                torch.tril(torch.ones(N, N, device=device, dtype=torch.bool))
                .unsqueeze(0)
                .unsqueeze(0)
            )
            scores_lin = torch.matmul(q_lin, k_lin.transpose(-1, -2).contiguous()) * causal_tril
            denom_lin = scores_lin.sum(dim=-1, keepdim=True)
            out_global = (torch.matmul(scores_lin, v_f) / denom_lin.clamp(min=1.0)).to(dtype)
        else:
            kv = torch.matmul(k_lin.transpose(-1, -2).contiguous(), v_f)  # [B, H, Dh, Dh]
            num = torch.matmul(q_lin, kv)  # [B, H, N, Dh]
            denom = torch.matmul(
                q_lin, k_lin.sum(dim=-2, keepdim=True).transpose(-1, -2).contiguous()
            )  # [B, H, N, 1]
            out_global = (num / denom.clamp(min=1.0)).to(dtype)

        # Gated Fusion
        gates = torch.sigmoid(self.gate_net(hidden_states)).view(B, N, self.num_heads, 2)
        gate_graph = gates[..., 0].transpose(1, 2).contiguous().unsqueeze(-1).to(dtype)
        gate_global = gates[..., 1].transpose(1, 2).contiguous().unsqueeze(-1).to(dtype)

        out = gate_graph * out_graph + gate_global * out_global
        out = self.dropout(out)

        # Head mask broadcasting
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.view(1, -1, 1, 1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.view(head_mask.size(0), -1, 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask.unsqueeze(-1)
            out = out * head_mask.to(dtype)

        out = out.transpose(1, 2).contiguous().view(B, N, self.all_head_size)

        if self.is_nested_encoder:
            output = out
        else:
            output = self.out_proj(out)

        return (output, None)

Vision GNN-inspired 2D Grid Linear Attention with Learned Gated Fusion.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• BaseAttention
• torch.nn.modules.module.Module

Class variables

var is_custom_attention : bool

The type of the None singleton.

var supports_bidirectional : bool

The type of the None singleton.

var supports_causal : bool

The type of the None singleton.

var supports_gqa : bool

The type of the None singleton.

var supports_rope : bool

The type of the None singleton.

Instance variables

prop W_k : LoRALinear

Expand source code

@property
def W_k(self) -> LoRALinear:
    return self.k_proj

prop W_q : LoRALinear

Expand source code

@property
def W_q(self) -> LoRALinear:
    return self.q_proj

prop W_v : LoRALinear

Expand source code

@property
def W_v(self) -> LoRALinear:
    return self.v_proj

prop c_proj : LoRALinear

Expand source code

@property
def c_proj(self) -> LoRALinear:
    return self.out_proj

prop dense : LoRALinear

Expand source code

@property
def dense(self) -> LoRALinear:
    return self.out_proj

prop k_lin : LoRALinear

Expand source code

@property
def k_lin(self) -> LoRALinear:
    return self.k_proj

prop key : LoRALinear

Expand source code

@property
def key(self) -> LoRALinear:
    return self.k_proj

prop o_proj : LoRALinear

Expand source code

@property
def o_proj(self) -> LoRALinear:
    return self.out_proj

prop out : LoRALinear

Expand source code

@property
def out(self) -> LoRALinear:
    return self.out_proj

prop out_lin : LoRALinear

Expand source code

@property
def out_lin(self) -> LoRALinear:
    return self.out_proj

prop q_lin : LoRALinear

Expand source code

@property
def q_lin(self) -> LoRALinear:
    return self.q_proj

prop query : LoRALinear

Expand source code

@property
def query(self) -> LoRALinear:
    return self.q_proj

prop v_lin : LoRALinear

Expand source code

@property
def v_lin(self) -> LoRALinear:
    return self.v_proj

prop value : LoRALinear

Expand source code

@property
def value(self) -> LoRALinear:
    return self.v_proj

Methods

def copy_weights(self, old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA."""
    extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = {
        "q": (None, None),
        "k": (None, None),
        "v": (None, None),
        "o": (None, None),
    }

    # 1. UniversalAdapter extraction
    for proj in ["q", "k", "v", "o"]:
        try:
            w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
            if w is not None:
                extracted[proj] = (w, b)
        except Exception as e:
            logger.debug("UniversalAdapter fallback failed: %s", e)

    # 2. GPT-2 Fused c_attn Conv1D handling
    if hasattr(old_attn, "c_attn") and (
        extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
    ):
        weight = old_attn.c_attn.weight.data
        bias = getattr(old_attn.c_attn, "bias", None)
        bias_data = bias.data if bias is not None else None
        if (
            ("Conv1D" in old_attn.c_attn.__class__.__name__)
            or hasattr(old_attn.c_attn, "nf")
            or (
                weight.dim() == 2
                and (weight.shape[0] < weight.shape[1] or weight.shape[1] == 3 * self.dim)
            )
        ):
            weight = weight.t().contiguous()
        split_size = weight.shape[0] // 3
        q_w = weight[:split_size, :]
        k_w = weight[split_size : 2 * split_size, :]
        v_w = weight[2 * split_size :, :]
        q_b = bias_data[:split_size] if bias_data is not None else None
        k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None
        v_b = bias_data[2 * split_size :] if bias_data is not None else None
        extracted["q"] = (q_w, q_b)
        extracted["k"] = (k_w, k_b)
        extracted["v"] = (v_w, v_b)

    if hasattr(old_attn, "c_proj") and extracted["o"][0] is None:
        w_o = old_attn.c_proj.weight.data
        if (
            ("Conv1D" in old_attn.c_proj.__class__.__name__)
            or hasattr(old_attn.c_proj, "nf")
            or (
                w_o.dim() == 2
                and (
                    w_o.shape[0] < w_o.shape[1]
                    or (w_o.shape[0] != self.dim and w_o.shape[1] == self.dim)
                )
            )
        ):
            w_o = w_o.t().contiguous()
        b_o = getattr(old_attn.c_proj, "bias", None)
        extracted["o"] = (w_o, b_o.data if b_o is not None else None)

    # 3. RoBERTa / BERT Nested Encoder handling
    if hasattr(old_attn, "self") and (
        extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None
    ):
        s = old_attn.self
        if hasattr(s, "query"):
            extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None))
        if hasattr(s, "key"):
            extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None))
        if hasattr(s, "value"):
            extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None))

    if (
        hasattr(old_attn, "output")
        and hasattr(old_attn.output, "dense")
        and extracted["o"][0] is None
    ):
        dense = old_attn.output.dense
        extracted["o"] = (dense.weight.data, getattr(dense, "bias", None))

    # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT
    named_map = {
        "q": ["q_proj", "q_lin", "query", "W_q"],
        "k": ["k_proj", "k_lin", "key", "W_k"],
        "v": ["v_proj", "v_lin", "value", "W_v"],
        "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
    }
    for proj, names in named_map.items():
        if extracted[proj][0] is None:
            for name in names:
                if hasattr(old_attn, name):
                    layer = getattr(old_attn, name)
                    base = getattr(layer, "base_layer", layer)
                    if hasattr(base, "weight"):
                        extracted[proj] = (
                            base.weight.data,
                            getattr(base, "bias", None),
                        )
                        break

    # Detect nested encoder architecture
    cls_name = old_attn.__class__.__name__
    has_out_proj = any(
        hasattr(old_attn, attr)
        for attr in (
            "dense",
            "out_proj",
            "c_proj",
            "out_lin",
            "o_proj",
            "wo",
            "o",
            "proj",
            "out",
        )
    )
    self.is_nested_encoder = (
        cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
    ) and not has_out_proj

    missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None]
    if not self.is_nested_encoder and extracted["o"][0] is None:
        missing.append("o")
    if missing:
        logger.warning(
            "Attention module copy_weights warning. Proceeding with initialized weights."
        )

    def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor:
        if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim:
            return w.t().contiguous()
        return w

    with torch.no_grad():
        for proj_key, target in [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
            ("o", self.out_proj),
        ]:
            w, b = extracted[proj_key]
            if w is not None:
                in_dim = self.all_head_size if proj_key == "o" else self.dim
                w_oriented = _orient(w, in_dim)
                target_w = target.weight
                out_d, in_d = target_w.shape
                src_out, src_in = w_oriented.shape
                target_w.data[: min(out_d, src_out), : min(in_d, src_in)].copy_(
                    w_oriented[: min(out_d, src_out), : min(in_d, src_in)]
                )
                if b is not None and target.bias is not None:
                    b_val = b.data if isinstance(b, nn.Parameter) else b
                    target.bias.data[: min(len(target.bias), len(b_val))].copy_(
                        b_val[: min(len(target.bias), len(b_val))]
                    )
            elif proj_key == "o" and self.is_nested_encoder:
                min_d = min(target.weight.shape[0], target.weight.shape[1])
                target.weight.data.zero_()
                target.weight.data[:min_d, :min_d].fill_diagonal_(1.0)
                if target.bias is not None:
                    target.bias.data.zero_()

            # Frobenius norm verification (>= 1e-4)
            w_norm = float(torch.linalg.norm(target.weight.data.float()).item())
            if w_norm < 1e-4:
                logger.warning(
                    f"ViGAttention weight verification failed for '{proj_key}': norm too small ({w_norm:.2e})"
                )

            target.freeze_base_weights()

    # Preserve source LoRA adapters
    source_submodules = {
        "q": ["q_proj", "q_lin", "query", "W_q"],
        "k": ["k_proj", "k_lin", "key", "W_k"],
        "v": ["v_proj", "v_lin", "value", "W_v"],
        "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"],
    }
    for proj_key, target in [
        ("q", self.q_proj),
        ("k", self.k_proj),
        ("v", self.v_proj),
        ("o", self.out_proj),
    ]:
        for name in source_submodules[proj_key]:
            if hasattr(old_attn, name):
                sub = getattr(old_attn, name)
                if (
                    hasattr(sub, "lora_A")
                    and hasattr(sub, "lora_B")
                    and sub.lora_A is not None
                    and sub.lora_B is not None
                ):
                    src_lA = getattr(sub.lora_A, "weight", sub.lora_A)
                    src_lB = getattr(sub.lora_B, "weight", sub.lora_B)
                    if hasattr(src_lA, "data") and target.lora_A is not None:
                        if (
                            src_lA.shape[0] == target.lora_A.shape[0]
                            and src_lA.shape[1] == target.lora_A.shape[1]
                        ):
                            target.lora_A.data.copy_(src_lA.data)
                            target.lora_A.requires_grad = True
                        else:
                            print(
                                f"Warning: lora_A shape mismatch for {proj_key}; keeping zero-init."
                            )
                    if hasattr(src_lB, "data") and target.lora_B is not None:
                        if (
                            src_lB.shape[0] == target.lora_B.shape[0]
                            and src_lB.shape[1] == target.lora_B.shape[1]
                        ):
                            target.lora_B.data.copy_(src_lB.data)
                            target.lora_B.requires_grad = True
                        else:
                            print(
                                f"Warning: lora_B shape mismatch for {proj_key}; keeping zero-init."
                            )
                    if hasattr(sub, "lora_scaling"):
                        target.lora_scaling = sub.lora_scaling
                    break

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)

Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.

Inherited members

• BaseAttention:
  • extra_repr
  • forward