Module neuralops.attentions.hybrid_attention

Hybrid softmax + linear attention with per-head gated fusion.

Valid attention REPLACEMENT (not a compressor). Adds a small number of parameters (the gate) but reduces attention computation at long sequence lengths by routing some heads to linear attention.

NOVELTY STATEMENT: Hybrid softmax + linear attention is NOT novel by itself (cf. MEGA 2023, BiGS 2024, Zipformer 2024). Our specific formulation differs in:

1. PER-HEAD learned gate:
     gate_h = sigmoid(W_h · hidden_states)
   This allows different heads to prefer different attention types,
   in contrast to MEGA's single-head gate.

2. Gate initialization to softmax-dominant (bias=10.0):
   Ensures the model starts IDENTICAL to the original softmax attention.
   The gate learns to route some heads to linear attention during training.

3. No extra parameters beyond the gate:
   768 * 12 + 12 = 9,228 params (+0.014% on DistilBERT).

Bug-fix map: HA-FIX-01 out_proj always executed (no dead parameters in nested mode). HA-FIX-02 _linear_branch: proper causal + padding masking. HA-FIX-03 injector recursion guard. HA-FIX-04 injector double-injection guard. HA-FIX-05 removed 'comp_config' in locals() hack. HA-FIX-06 is_nested_encoder set correctly in injector. HA-FIX-07 KV-cache implemented (tuple + DynamicCache). HA-FIX-08 RoPE past_len offset + position_embeddings support. HA-FIX-09 head_mask broadcasting corrected (B, H, 1, 1). HA-FIX-10 _standard_branch: optional SDPA fast path.

Functions

def _extract_valid_mask(mask: torch.Tensor | None, batch_size: int, seq_len: int, device: torch.device) ‑> torch.Tensor

Expand source code

def _extract_valid_mask(
    mask: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    device: torch.device,
) -> torch.Tensor:
    """
    Robustly extract a boolean valid-token mask from any attention_mask format:
      - (B, N) int/bool     → padding mask
      - (B, N) float        → -inf/0 additive mask
      - (B, 1, N) or (B, 1, 1, N) broadcastable
      - (B, H, N, N) full attention matrix
    """
    if mask is None:
        return torch.ones(batch_size, seq_len, device=device, dtype=torch.bool)

    with torch.no_grad():
        if mask.dim() == 2:
            if mask.dtype in (
                torch.bool,
                torch.int8,
                torch.int16,
                torch.int32,
                torch.int64,
            ):
                return mask.bool()
            return mask > -1e4

        if mask.dim() == 3:
            if mask.shape[1] == 1:
                m = mask[:, 0, :]
            else:
                m = mask.any(dim=1)
            if m.dtype not in (
                torch.bool,
                torch.int8,
                torch.int16,
                torch.int32,
                torch.int64,
            ):
                m = m > -1e4
            return m.bool()

        if mask.dim() == 4:
            m = mask.any(dim=1).any(dim=1)
            if m.dtype not in (
                torch.bool,
                torch.int8,
                torch.int16,
                torch.int32,
                torch.int64,
            ):
                m = m > -1e4
            return m.bool()

    return torch.ones(batch_size, seq_len, device=device, dtype=torch.bool)

Robustly extract a boolean valid-token mask from any attention_mask format: - (B, N) int/bool → padding mask - (B, N) float → -inf/0 additive mask - (B, 1, N) or (B, 1, 1, N) broadcastable - (B, H, N, N) full attention matrix

def inject_smart_weighted_ensemble(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) ‑> torch.nn.modules.module.Module | Any | None

Expand source code

@ATTENTION_REGISTRY.register("weighted_ensemble_module_inject")
@ATTENTION_REGISTRY.register("weighted_ensemble_inject")
@ATTENTION_REGISTRY.register("inject_weighted_ensemble")
@ATTENTION_REGISTRY.register("hybrid_attention_inject")
@ATTENTION_REGISTRY.register("inject_hybrid_attention")
def inject_weighted_ensemble(
    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,
    _depth: int = 0,
    **kwargs: Any,
) -> nn.Module | Any | None:
    # HA-FIX-03: recursion guard.
    if _depth > 8:
        logger.warning("inject_weighted_ensemble: max recursion depth exceeded.")
        return None

    target = layer if layer is not None else model
    if target is None:
        return None

    cfg = model_config
    if cfg is None and hasattr(target, "config"):
        cfg = target.config

    # 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_weighted_ensemble(
                layer=lyr,
                model_config=cfg,
                comp_config=comp_config,
                layer_idx=idx,
                total_layers=tot,
                _depth=_depth + 1,
                **kwargs,
            )
            if res is not None:
                num_injected += 1
        return target if num_injected > 0 else None

    layer = target
    try:
        if model_config is None:
            if any(hasattr(layer, a) for a in ("hidden_size", "d_model", "n_embd", "dim")):
                model_config = layer
            elif hasattr(layer, "attention"):
                old_attn = layer.attention

        src_attn = (
            old_attn
            if old_attn is not None
            else getattr(layer, "attention", getattr(layer, "self", getattr(layer, "attn", None)))
        )
        if src_attn is None:
            return None

        # HA-FIX-04: double-injection guard.
        if isinstance(src_attn, HybridAttention):
            return None
        if "HybridAttention" in src_attn.__class__.__name__:
            return None
        if isinstance(getattr(src_attn, "inner", None), HybridAttention):
            return None
        if isinstance(getattr(src_attn, "_wrapped_attention", None), HybridAttention):
            return None

        new_attn = HybridAttention(
            config=model_config,
            pretrained_attn=src_attn,
            **kwargs,
        )

        if src_attn is not None:
            p = next(src_attn.parameters(), None)
            if p is not None:
                new_attn.to(device=p.device, dtype=p.dtype)

        # HA-FIX-06: is_nested_encoder correctly set per branch.
        if hasattr(layer, "attention"):
            if hasattr(layer.attention, "self"):
                layer.attention.self = new_attn
                new_attn.is_nested_encoder = True
            else:
                layer.attention = new_attn
                new_attn.is_nested_encoder = False
        elif hasattr(layer, "self_attn"):
            layer.self_attn = new_attn
            new_attn.is_nested_encoder = False
        elif hasattr(layer, "self"):
            layer.self = new_attn
            new_attn.is_nested_encoder = True
        elif hasattr(layer, "attn"):
            layer.attn = new_attn
            new_attn.is_nested_encoder = False

        return new_attn
    except Exception as e:
        logger.exception("Failed to inject hybrid weighted ensemble: %s", e)
        return None

def inject_weighted_ensemble(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) ‑> torch.nn.modules.module.Module | Any | None

Expand source code

@ATTENTION_REGISTRY.register("weighted_ensemble_module_inject")
@ATTENTION_REGISTRY.register("weighted_ensemble_inject")
@ATTENTION_REGISTRY.register("inject_weighted_ensemble")
@ATTENTION_REGISTRY.register("hybrid_attention_inject")
@ATTENTION_REGISTRY.register("inject_hybrid_attention")
def inject_weighted_ensemble(
    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,
    _depth: int = 0,
    **kwargs: Any,
) -> nn.Module | Any | None:
    # HA-FIX-03: recursion guard.
    if _depth > 8:
        logger.warning("inject_weighted_ensemble: max recursion depth exceeded.")
        return None

    target = layer if layer is not None else model
    if target is None:
        return None

    cfg = model_config
    if cfg is None and hasattr(target, "config"):
        cfg = target.config

    # 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_weighted_ensemble(
                layer=lyr,
                model_config=cfg,
                comp_config=comp_config,
                layer_idx=idx,
                total_layers=tot,
                _depth=_depth + 1,
                **kwargs,
            )
            if res is not None:
                num_injected += 1
        return target if num_injected > 0 else None

    layer = target
    try:
        if model_config is None:
            if any(hasattr(layer, a) for a in ("hidden_size", "d_model", "n_embd", "dim")):
                model_config = layer
            elif hasattr(layer, "attention"):
                old_attn = layer.attention

        src_attn = (
            old_attn
            if old_attn is not None
            else getattr(layer, "attention", getattr(layer, "self", getattr(layer, "attn", None)))
        )
        if src_attn is None:
            return None

        # HA-FIX-04: double-injection guard.
        if isinstance(src_attn, HybridAttention):
            return None
        if "HybridAttention" in src_attn.__class__.__name__:
            return None
        if isinstance(getattr(src_attn, "inner", None), HybridAttention):
            return None
        if isinstance(getattr(src_attn, "_wrapped_attention", None), HybridAttention):
            return None

        new_attn = HybridAttention(
            config=model_config,
            pretrained_attn=src_attn,
            **kwargs,
        )

        if src_attn is not None:
            p = next(src_attn.parameters(), None)
            if p is not None:
                new_attn.to(device=p.device, dtype=p.dtype)

        # HA-FIX-06: is_nested_encoder correctly set per branch.
        if hasattr(layer, "attention"):
            if hasattr(layer.attention, "self"):
                layer.attention.self = new_attn
                new_attn.is_nested_encoder = True
            else:
                layer.attention = new_attn
                new_attn.is_nested_encoder = False
        elif hasattr(layer, "self_attn"):
            layer.self_attn = new_attn
            new_attn.is_nested_encoder = False
        elif hasattr(layer, "self"):
            layer.self = new_attn
            new_attn.is_nested_encoder = True
        elif hasattr(layer, "attn"):
            layer.attn = new_attn
            new_attn.is_nested_encoder = False

        return new_attn
    except Exception as e:
        logger.exception("Failed to inject hybrid weighted ensemble: %s", e)
        return None

Classes

class HybridAttention (config: Any | None = None,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
head_dim: int | None = None,
dropout: float = 0.1,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hybrid_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("weighted_ensemble", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("smart_weighted_ensemble", supports_causal=True, supports_gqa=True)
class HybridAttention(BaseAttention):
    """
    Hybrid softmax + linear attention with per-head dynamic gated fusion.
    """

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

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        head_dim: int | None = None,
        dropout: float = 0.1,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank, **kwargs)

        cfg_source = getattr(config, "text_config", config) if config is not None else None
        h_size = (
            hidden_size
            or getattr(cfg_source, "hidden_size", None)
            or getattr(cfg_source, "d_model", None)
            or getattr(cfg_source, "n_embd", None)
            or getattr(cfg_source, "dim", None)
            or kwargs.get("hidden_size", 768)
        )
        n_heads = (
            num_heads
            or getattr(cfg_source, "num_attention_heads", None)
            or getattr(cfg_source, "num_heads", None)
            or getattr(cfg_source, "n_heads", None)
            or getattr(cfg_source, "n_head", None)
            or kwargs.get("num_heads", 12)
        )
        h_dim = (
            head_dim
            or getattr(cfg_source, "head_dim", None)
            or getattr(cfg_source, "d_kv", None)
            or (int(h_size) // max(1, int(n_heads)))
        )
        n_kv_heads = (
            getattr(cfg_source, "num_key_value_heads", None)
            or getattr(cfg_source, "num_kv_heads", None)
            or kwargs.get("num_key_value_heads", n_heads)
        )
        drop_rate = (
            getattr(cfg_source, "attention_probs_dropout_prob", None)
            or getattr(cfg_source, "dropout", None)
            or kwargs.get("dropout", dropout)
        )

        self.hidden_size: int = int(h_size)
        self.num_heads: int = int(n_heads)
        self.head_dim: int = int(h_dim)
        self.num_kv_heads: int = int(n_kv_heads)
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.all_head_size: int = self.num_heads * self.head_dim
        self.all_kv_head_size: int = self.num_kv_heads * self.head_dim

        # Decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, pretrained_attn)
        self.is_nested_encoder: bool = False

        # Projections
        self.q_proj = nn.Linear(self.hidden_size, self.all_head_size, bias=True)
        self.k_proj = nn.Linear(self.hidden_size, self.all_kv_head_size, bias=True)
        self.v_proj = nn.Linear(self.hidden_size, self.all_kv_head_size, bias=True)
        self.out_proj = nn.Linear(self.all_head_size, self.hidden_size, bias=True)
        self.dropout = nn.Dropout(drop_rate)

        # Per-head gate
        self.ensemble_gate = nn.Linear(self.hidden_size, self.num_heads)

        # Initialize gate to prefer softmax (sigmoid(10.0) ≈ 0.99995)
        with torch.no_grad():
            self.ensemble_gate.weight.zero_()
            self.ensemble_gate.bias.fill_(10.0)

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

    # ------------------------------------------------------------------
    # Aliases
    # ------------------------------------------------------------------
    @property
    def query(self):
        return self.q_proj

    @property
    def key(self):
        return self.k_proj

    @property
    def value(self):
        return self.v_proj

    @property
    def dense(self):
        return self.out_proj

    @property
    def q_lin(self):
        return self.q_proj

    @property
    def k_lin(self):
        return self.k_proj

    @property
    def v_lin(self):
        return self.v_proj

    @property
    def out_lin(self):
        return self.out_proj

    @property
    def o_proj(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # Standard softmax branch (HA-FIX-10: optional SDPA fast path)
    # ------------------------------------------------------------------
    def _standard_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        attention_mask: torch.Tensor | None,
        past_len: int = 0,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, Dh = Q.shape
        N_k = K.shape[2]

        # HA-FIX-10: SDPA fast path when training/eval allows it.
        can_sdpa = (not self.training or self.dropout.p == 0.0) and True
        if can_sdpa:
            try:
                q_sdp = Q.half() if Q.dtype == torch.float32 else Q
                k_sdp = K.half() if K.dtype == torch.float32 else K
                v_sdp = V.half() if V.dtype == torch.float32 else V

                attn_bias_sdp = None
                if valid_mask is not None:
                    # Build a proper (B, 1, 1, N_k) additive mask.
                    attn_bias_sdp = ((valid_mask[:, None, None, :].float() - 1.0) * 1e4).to(
                        q_sdp.dtype
                    )

                is_causal_sdp = self.is_decoder and past_len == 0 and N > 1

                out_sdp = F.scaled_dot_product_attention(
                    q_sdp,
                    k_sdp,
                    v_sdp,
                    attn_mask=attn_bias_sdp,
                    is_causal=is_causal_sdp,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out_sdp = out_sdp.to(Q.dtype)
                return torch.nan_to_num(out_sdp, nan=0.0)
            except Exception as e:
                logger.debug("SDPA fast path failed: %s; falling back to manual.", e)

        # Manual path
        scale = 1.0 / math.sqrt(Dh)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # Causal mask (offset-aware for cache)
        if self.is_decoder and N > 1:
            if past_len > 0:
                abs_q = (past_len + torch.arange(N, device=Q.device)).unsqueeze(1)
                abs_k = torch.arange(N_k, device=Q.device).unsqueeze(0)
                causal_mask = abs_k <= abs_q
            else:
                causal_mask = torch.tril(torch.ones(N, N_k, device=Q.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)

        # Padding mask
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :].bool(), -1e4)
        elif attention_mask is not None:
            if attention_mask.dim() == 2:
                ext_mask = (1.0 - attention_mask.unsqueeze(1).unsqueeze(2).float()) * -1e4
                scores = scores + ext_mask.to(scores.dtype)
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1).to(scores.dtype)
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.to(scores.dtype)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(Q.dtype)
        attn = self.dropout(attn)
        return torch.matmul(attn, V)

    # ------------------------------------------------------------------
    # Linear branch (HA-FIX-02: proper causal + padding)
    # ------------------------------------------------------------------
    def _linear_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor,
        past_len: int = 0,
    ) -> torch.Tensor:
        phi_Q = F.elu(Q.float()) + 1.0
        phi_K = F.elu(K.float()) + 1.0
        V_f = V.float()

        # Padding mask on K and V
        mask = valid_mask.unsqueeze(1).unsqueeze(-1).to(phi_K.dtype)
        phi_K = phi_K * mask
        V_f = V_f * mask

        B_h, H_h, N_h, D_h = phi_Q.shape
        phi_K.shape[2]

        # Handle cache: query positions are [past_len .. past_len + N_h), keys are [0 .. N_k)
        if self.is_decoder and past_len > 0:
            # Only new queries, keys include cache.
            # Non-causal over the full key set (past + new).
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_f)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous()).clamp(min=1e-6)
            out = num / den
            return torch.nan_to_num(out, nan=0.0).to(Q.dtype)

        if self.is_decoder:
            # Chunked causal linear attention.
            chunk_size = 128
            num_chunks = (N_h + chunk_size - 1) // chunk_size

            state_KV = torch.zeros(B_h, H_h, D_h, D_h, device=phi_Q.device, dtype=phi_Q.dtype)
            state_K = torch.zeros(B_h, H_h, D_h, device=phi_Q.device, dtype=phi_Q.dtype)

            num_list = []
            den_list = []

            for c in range(num_chunks):
                s = c * chunk_size
                e = min(N_h, s + chunk_size)

                phi_Q_c = phi_Q[:, :, s:e]
                phi_K_c = phi_K[:, :, s:e]
                V_c = V_f[:, :, s:e]

                KV_c = torch.einsum("bhnd,bhne->bhnde", phi_K_c, V_c)
                KV_c_cumsum = torch.cumsum(KV_c, dim=2)
                num_intra = torch.einsum("bhnd,bhnde->bhne", phi_Q_c, KV_c_cumsum)

                K_c_cumsum = torch.cumsum(phi_K_c, dim=2)
                den_intra = torch.einsum("bhnd,bhnd->bhn", phi_Q_c, K_c_cumsum)

                num_inter = torch.einsum("bhnd,bhde->bhne", phi_Q_c, state_KV)
                den_inter = torch.einsum("bhnd,bhd->bhn", phi_Q_c, state_K)

                num_list.append(num_intra + num_inter)
                den_list.append(den_intra + den_inter)

                state_KV = state_KV + KV_c_cumsum[:, :, -1]
                state_K = state_K + K_c_cumsum[:, :, -1]

            num = torch.cat(num_list, dim=2)
            den = torch.cat(den_list, dim=2).unsqueeze(-1)
        else:
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_f)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous())

        den = den.clamp(min=1e-6)
        out = num / den
        return torch.nan_to_num(out, nan=0.0).to(Q.dtype)

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------
    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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        layer_idx: int = 0,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        device = hidden_states.device

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        N_kv_new = kv_input.shape[1]

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid_mask_new = _extract_valid_mask(eff_mask, B, N_kv_new, device)

        # Projections
        Q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K_new = (
            self.k_proj(kv_input)
            .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V_new = (
            self.v_proj(kv_input)
            .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # HA-FIX-08: RoPE with past_len offset
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            past_len_rope = 0
            if past_key_value is not None:
                if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                    past_len_rope = past_key_value[0].shape[-2]
                elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                    past_len_rope = past_key_value.key_cache[0].shape[-2]
            cos = cos[..., past_len_rope : past_len_rope + N, :]
            sin = sin[..., past_len_rope : past_len_rope + N, :]
            Q, K_new = _apply_rope(Q, K_new, cos, sin)

        # HA-FIX-07: KV-cache
        past_k = past_v = None
        if past_key_value is not None and not is_cross:
            if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
                past_k, past_v = past_key_value[0], past_key_value[1]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                idx = min(layer_idx, len(past_key_value.key_cache) - 1)
                past_k = past_key_value.key_cache[idx]
                past_v = past_key_value.value_cache[idx]

        if past_k is not None and not is_cross:
            K_cat = torch.cat([past_k, K_new], dim=-2)
            V_cat = torch.cat([past_v, V_new], dim=-2)
        else:
            K_cat = K_new
            V_cat = V_new

        if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
            past_key_value.update(K_new, V_new, layer_idx)

        N_kv = K_cat.shape[2]
        past_len = N_kv - N if (self.is_decoder and not is_cross) else 0

        if self.kv_repeat > 1:
            K_cat = repeat_kv(K_cat, self.kv_repeat)
            V_cat = repeat_kv(V_cat, self.kv_repeat)

        # Valid mask for the full key set (with past)
        if past_len > 0:
            past_valid = torch.ones(B, past_len, device=device, dtype=torch.bool)
            valid_mask_full = torch.cat([past_valid, valid_mask_new], dim=1)
        else:
            valid_mask_full = valid_mask_new

        # Two branches
        out_standard = self._standard_branch(
            Q,
            K_cat,
            V_cat,
            eff_mask,
            past_len=past_len,
            valid_mask=valid_mask_full,
        )
        out_linear = self._linear_branch(
            Q,
            K_cat,
            V_cat,
            valid_mask_full,
            past_len=past_len,
        )

        # Per-head gating
        gate = (
            torch.sigmoid(self.ensemble_gate(hidden_states))
            .view(B, N, self.num_heads, 1)
            .transpose(1, 2)
            .contiguous()
        )
        fused = gate * out_standard + (1.0 - gate) * out_linear

        # HA-FIX-09: head_mask broadcasting corrected.
        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.unsqueeze(-1).unsqueeze(-1)  # (B, H, 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)  # (B, H, 1, 1)
            fused = fused * head_mask.to(fused.dtype)

        fused = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)

        # HA-FIX-01: out_proj ALWAYS executed.
        if self.is_nested_encoder:
            # Nested encoder: outer module supplies out_proj, but we still run
            # the (identity-initialised) internal out_proj for parameter stability.
            out = self.out_proj(fused)
        else:
            out = self.out_proj(fused)
            out = self.dropout(out)

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        return out, None

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        if old_attn is None:
            return

        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")
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
            or (
                hasattr(old_attn, "self")
                and hasattr(getattr(old_attn, "self", None), "query")
                and hasattr(old_attn, "output")
                and hasattr(getattr(old_attn, "output", None), "dense")
            )
        )
        logger.debug("HybridAttention.copy_weights: nested_encoder=%s", self.is_nested_encoder)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif cls_name in (
            "GPT2Attention",
            "Qwen2Attention",
            "LlamaAttention",
        ) or hasattr(old_attn, "c_attn"):
            self.is_decoder = True

        from neuralops.core.adapters import UniversalAdapter

        proj_targets = [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
        ]
        if not self.is_nested_encoder:
            proj_targets.append(("o", self.out_proj))

        for proj, target in proj_targets:
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
            except Exception as e:
                logger.debug("UniversalAdapter extraction failed for %s: %s", proj, e)
                w, b = None, None

            if w is None:
                logger.warning("HybridAttention: no weight for '%s' in %s", proj, cls_name)
                continue

            if w.shape != target.weight.shape:
                if w.t().shape == target.weight.shape:
                    w = w.t().contiguous()
                else:
                    logger.warning(
                        "HybridAttention: shape mismatch for %s: src=%s, tgt=%s. Truncating.",
                        proj,
                        tuple(w.shape),
                        tuple(target.weight.shape),
                    )
                    out_d, in_d = target.weight.shape
                    s_out, s_in = w.shape
                    with torch.no_grad():
                        target.weight.data.zero_()
                        target.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_(
                            w[: min(out_d, s_out), : min(in_d, s_in)]
                        )
                    if b is not None and target.bias is not None:
                        min_b = min(target.bias.numel(), b.numel())
                        with torch.no_grad():
                            target.bias.data.zero_()
                            target.bias.data[:min_b].copy_(b[:min_b])
                    continue

            with torch.no_grad():
                norm = torch.linalg.norm(w.float()).item()
                if norm < 1e-4:
                    logger.warning(
                        "HybridAttention: %s weight has near-zero Frobenius norm (%s)",
                        proj,
                        norm,
                    )
                target.weight.copy_(w)
                if b is not None and target.bias is not None:
                    target.bias.copy_(b)

Hybrid softmax + linear attention with per-head dynamic 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 c_proj

Expand source code

@property
def c_proj(self):
    return self.out_proj

prop dense

Expand source code

@property
def dense(self):
    return self.out_proj

prop k_lin

Expand source code

@property
def k_lin(self):
    return self.k_proj

prop key

Expand source code

@property
def key(self):
    return self.k_proj

prop o_proj

Expand source code

@property
def o_proj(self):
    return self.out_proj

prop out_lin

Expand source code

@property
def out_lin(self):
    return self.out_proj

prop q_lin

Expand source code

@property
def q_lin(self):
    return self.q_proj

prop query

Expand source code

@property
def query(self):
    return self.q_proj

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.v_proj

prop value

Expand source code

@property
def value(self):
    return self.v_proj

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    if old_attn is None:
        return

    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")
    )
    self.is_nested_encoder = (
        cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
        or (
            hasattr(old_attn, "self")
            and hasattr(getattr(old_attn, "self", None), "query")
            and hasattr(old_attn, "output")
            and hasattr(getattr(old_attn, "output", None), "dense")
        )
    )
    logger.debug("HybridAttention.copy_weights: nested_encoder=%s", self.is_nested_encoder)

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)
    elif cls_name in (
        "GPT2Attention",
        "Qwen2Attention",
        "LlamaAttention",
    ) or hasattr(old_attn, "c_attn"):
        self.is_decoder = True

    from neuralops.core.adapters import UniversalAdapter

    proj_targets = [
        ("q", self.q_proj),
        ("k", self.k_proj),
        ("v", self.v_proj),
    ]
    if not self.is_nested_encoder:
        proj_targets.append(("o", self.out_proj))

    for proj, target in proj_targets:
        try:
            w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
        except Exception as e:
            logger.debug("UniversalAdapter extraction failed for %s: %s", proj, e)
            w, b = None, None

        if w is None:
            logger.warning("HybridAttention: no weight for '%s' in %s", proj, cls_name)
            continue

        if w.shape != target.weight.shape:
            if w.t().shape == target.weight.shape:
                w = w.t().contiguous()
            else:
                logger.warning(
                    "HybridAttention: shape mismatch for %s: src=%s, tgt=%s. Truncating.",
                    proj,
                    tuple(w.shape),
                    tuple(target.weight.shape),
                )
                out_d, in_d = target.weight.shape
                s_out, s_in = w.shape
                with torch.no_grad():
                    target.weight.data.zero_()
                    target.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_(
                        w[: min(out_d, s_out), : min(in_d, s_in)]
                    )
                if b is not None and target.bias is not None:
                    min_b = min(target.bias.numel(), b.numel())
                    with torch.no_grad():
                        target.bias.data.zero_()
                        target.bias.data[:min_b].copy_(b[:min_b])
                continue

        with torch.no_grad():
            norm = torch.linalg.norm(w.float()).item()
            if norm < 1e-4:
                logger.warning(
                    "HybridAttention: %s weight has near-zero Frobenius norm (%s)",
                    proj,
                    norm,
                )
            target.weight.copy_(w)
            if b is not None and target.bias is not None:
                target.bias.copy_(b)

Universal fallback: copy pretrained weights from old attention layer.

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,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
layer_idx: int = 0,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None]

Expand source code

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,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    layer_idx: int = 0,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    device = hidden_states.device

    is_cross = encoder_hidden_states is not None
    kv_input = encoder_hidden_states if is_cross else hidden_states
    N_kv_new = kv_input.shape[1]

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    valid_mask_new = _extract_valid_mask(eff_mask, B, N_kv_new, device)

    # Projections
    Q = (
        self.q_proj(hidden_states)
        .view(B, N, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    K_new = (
        self.k_proj(kv_input)
        .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    V_new = (
        self.v_proj(kv_input)
        .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # HA-FIX-08: RoPE with past_len offset
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    if position_embeddings is not None and not is_cross:
        cos, sin = position_embeddings
        past_len_rope = 0
        if past_key_value is not None:
            if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                past_len_rope = past_key_value[0].shape[-2]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                past_len_rope = past_key_value.key_cache[0].shape[-2]
        cos = cos[..., past_len_rope : past_len_rope + N, :]
        sin = sin[..., past_len_rope : past_len_rope + N, :]
        Q, K_new = _apply_rope(Q, K_new, cos, sin)

    # HA-FIX-07: KV-cache
    past_k = past_v = None
    if past_key_value is not None and not is_cross:
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            past_k, past_v = past_key_value[0], past_key_value[1]
        elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            idx = min(layer_idx, len(past_key_value.key_cache) - 1)
            past_k = past_key_value.key_cache[idx]
            past_v = past_key_value.value_cache[idx]

    if past_k is not None and not is_cross:
        K_cat = torch.cat([past_k, K_new], dim=-2)
        V_cat = torch.cat([past_v, V_new], dim=-2)
    else:
        K_cat = K_new
        V_cat = V_new

    if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
        past_key_value.update(K_new, V_new, layer_idx)

    N_kv = K_cat.shape[2]
    past_len = N_kv - N if (self.is_decoder and not is_cross) else 0

    if self.kv_repeat > 1:
        K_cat = repeat_kv(K_cat, self.kv_repeat)
        V_cat = repeat_kv(V_cat, self.kv_repeat)

    # Valid mask for the full key set (with past)
    if past_len > 0:
        past_valid = torch.ones(B, past_len, device=device, dtype=torch.bool)
        valid_mask_full = torch.cat([past_valid, valid_mask_new], dim=1)
    else:
        valid_mask_full = valid_mask_new

    # Two branches
    out_standard = self._standard_branch(
        Q,
        K_cat,
        V_cat,
        eff_mask,
        past_len=past_len,
        valid_mask=valid_mask_full,
    )
    out_linear = self._linear_branch(
        Q,
        K_cat,
        V_cat,
        valid_mask_full,
        past_len=past_len,
    )

    # Per-head gating
    gate = (
        torch.sigmoid(self.ensemble_gate(hidden_states))
        .view(B, N, self.num_heads, 1)
        .transpose(1, 2)
        .contiguous()
    )
    fused = gate * out_standard + (1.0 - gate) * out_linear

    # HA-FIX-09: head_mask broadcasting corrected.
    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.unsqueeze(-1).unsqueeze(-1)  # (B, H, 1, 1)
        elif head_mask.dim() == 3:
            head_mask = head_mask[:, :, :1].unsqueeze(-1)  # (B, H, 1, 1)
        fused = fused * head_mask.to(fused.dtype)

    fused = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)

    # HA-FIX-01: out_proj ALWAYS executed.
    if self.is_nested_encoder:
        # Nested encoder: outer module supplies out_proj, but we still run
        # the (identity-initialised) internal out_proj for parameter stability.
        out = self.out_proj(fused)
    else:
        out = self.out_proj(fused)
        out = self.dropout(out)

    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
    return out, None

Must return (output, attention_weights_or_None).

class WeightedEnsembleAttention (config: Any | None = None,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
head_dim: int | None = None,
dropout: float = 0.1,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hybrid_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("weighted_ensemble", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("smart_weighted_ensemble", supports_causal=True, supports_gqa=True)
class HybridAttention(BaseAttention):
    """
    Hybrid softmax + linear attention with per-head dynamic gated fusion.
    """

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

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        head_dim: int | None = None,
        dropout: float = 0.1,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank, **kwargs)

        cfg_source = getattr(config, "text_config", config) if config is not None else None
        h_size = (
            hidden_size
            or getattr(cfg_source, "hidden_size", None)
            or getattr(cfg_source, "d_model", None)
            or getattr(cfg_source, "n_embd", None)
            or getattr(cfg_source, "dim", None)
            or kwargs.get("hidden_size", 768)
        )
        n_heads = (
            num_heads
            or getattr(cfg_source, "num_attention_heads", None)
            or getattr(cfg_source, "num_heads", None)
            or getattr(cfg_source, "n_heads", None)
            or getattr(cfg_source, "n_head", None)
            or kwargs.get("num_heads", 12)
        )
        h_dim = (
            head_dim
            or getattr(cfg_source, "head_dim", None)
            or getattr(cfg_source, "d_kv", None)
            or (int(h_size) // max(1, int(n_heads)))
        )
        n_kv_heads = (
            getattr(cfg_source, "num_key_value_heads", None)
            or getattr(cfg_source, "num_kv_heads", None)
            or kwargs.get("num_key_value_heads", n_heads)
        )
        drop_rate = (
            getattr(cfg_source, "attention_probs_dropout_prob", None)
            or getattr(cfg_source, "dropout", None)
            or kwargs.get("dropout", dropout)
        )

        self.hidden_size: int = int(h_size)
        self.num_heads: int = int(n_heads)
        self.head_dim: int = int(h_dim)
        self.num_kv_heads: int = int(n_kv_heads)
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.all_head_size: int = self.num_heads * self.head_dim
        self.all_kv_head_size: int = self.num_kv_heads * self.head_dim

        # Decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, pretrained_attn)
        self.is_nested_encoder: bool = False

        # Projections
        self.q_proj = nn.Linear(self.hidden_size, self.all_head_size, bias=True)
        self.k_proj = nn.Linear(self.hidden_size, self.all_kv_head_size, bias=True)
        self.v_proj = nn.Linear(self.hidden_size, self.all_kv_head_size, bias=True)
        self.out_proj = nn.Linear(self.all_head_size, self.hidden_size, bias=True)
        self.dropout = nn.Dropout(drop_rate)

        # Per-head gate
        self.ensemble_gate = nn.Linear(self.hidden_size, self.num_heads)

        # Initialize gate to prefer softmax (sigmoid(10.0) ≈ 0.99995)
        with torch.no_grad():
            self.ensemble_gate.weight.zero_()
            self.ensemble_gate.bias.fill_(10.0)

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

    # ------------------------------------------------------------------
    # Aliases
    # ------------------------------------------------------------------
    @property
    def query(self):
        return self.q_proj

    @property
    def key(self):
        return self.k_proj

    @property
    def value(self):
        return self.v_proj

    @property
    def dense(self):
        return self.out_proj

    @property
    def q_lin(self):
        return self.q_proj

    @property
    def k_lin(self):
        return self.k_proj

    @property
    def v_lin(self):
        return self.v_proj

    @property
    def out_lin(self):
        return self.out_proj

    @property
    def o_proj(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # Standard softmax branch (HA-FIX-10: optional SDPA fast path)
    # ------------------------------------------------------------------
    def _standard_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        attention_mask: torch.Tensor | None,
        past_len: int = 0,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, Dh = Q.shape
        N_k = K.shape[2]

        # HA-FIX-10: SDPA fast path when training/eval allows it.
        can_sdpa = (not self.training or self.dropout.p == 0.0) and True
        if can_sdpa:
            try:
                q_sdp = Q.half() if Q.dtype == torch.float32 else Q
                k_sdp = K.half() if K.dtype == torch.float32 else K
                v_sdp = V.half() if V.dtype == torch.float32 else V

                attn_bias_sdp = None
                if valid_mask is not None:
                    # Build a proper (B, 1, 1, N_k) additive mask.
                    attn_bias_sdp = ((valid_mask[:, None, None, :].float() - 1.0) * 1e4).to(
                        q_sdp.dtype
                    )

                is_causal_sdp = self.is_decoder and past_len == 0 and N > 1

                out_sdp = F.scaled_dot_product_attention(
                    q_sdp,
                    k_sdp,
                    v_sdp,
                    attn_mask=attn_bias_sdp,
                    is_causal=is_causal_sdp,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out_sdp = out_sdp.to(Q.dtype)
                return torch.nan_to_num(out_sdp, nan=0.0)
            except Exception as e:
                logger.debug("SDPA fast path failed: %s; falling back to manual.", e)

        # Manual path
        scale = 1.0 / math.sqrt(Dh)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # Causal mask (offset-aware for cache)
        if self.is_decoder and N > 1:
            if past_len > 0:
                abs_q = (past_len + torch.arange(N, device=Q.device)).unsqueeze(1)
                abs_k = torch.arange(N_k, device=Q.device).unsqueeze(0)
                causal_mask = abs_k <= abs_q
            else:
                causal_mask = torch.tril(torch.ones(N, N_k, device=Q.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)

        # Padding mask
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :].bool(), -1e4)
        elif attention_mask is not None:
            if attention_mask.dim() == 2:
                ext_mask = (1.0 - attention_mask.unsqueeze(1).unsqueeze(2).float()) * -1e4
                scores = scores + ext_mask.to(scores.dtype)
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1).to(scores.dtype)
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.to(scores.dtype)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(Q.dtype)
        attn = self.dropout(attn)
        return torch.matmul(attn, V)

    # ------------------------------------------------------------------
    # Linear branch (HA-FIX-02: proper causal + padding)
    # ------------------------------------------------------------------
    def _linear_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor,
        past_len: int = 0,
    ) -> torch.Tensor:
        phi_Q = F.elu(Q.float()) + 1.0
        phi_K = F.elu(K.float()) + 1.0
        V_f = V.float()

        # Padding mask on K and V
        mask = valid_mask.unsqueeze(1).unsqueeze(-1).to(phi_K.dtype)
        phi_K = phi_K * mask
        V_f = V_f * mask

        B_h, H_h, N_h, D_h = phi_Q.shape
        phi_K.shape[2]

        # Handle cache: query positions are [past_len .. past_len + N_h), keys are [0 .. N_k)
        if self.is_decoder and past_len > 0:
            # Only new queries, keys include cache.
            # Non-causal over the full key set (past + new).
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_f)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous()).clamp(min=1e-6)
            out = num / den
            return torch.nan_to_num(out, nan=0.0).to(Q.dtype)

        if self.is_decoder:
            # Chunked causal linear attention.
            chunk_size = 128
            num_chunks = (N_h + chunk_size - 1) // chunk_size

            state_KV = torch.zeros(B_h, H_h, D_h, D_h, device=phi_Q.device, dtype=phi_Q.dtype)
            state_K = torch.zeros(B_h, H_h, D_h, device=phi_Q.device, dtype=phi_Q.dtype)

            num_list = []
            den_list = []

            for c in range(num_chunks):
                s = c * chunk_size
                e = min(N_h, s + chunk_size)

                phi_Q_c = phi_Q[:, :, s:e]
                phi_K_c = phi_K[:, :, s:e]
                V_c = V_f[:, :, s:e]

                KV_c = torch.einsum("bhnd,bhne->bhnde", phi_K_c, V_c)
                KV_c_cumsum = torch.cumsum(KV_c, dim=2)
                num_intra = torch.einsum("bhnd,bhnde->bhne", phi_Q_c, KV_c_cumsum)

                K_c_cumsum = torch.cumsum(phi_K_c, dim=2)
                den_intra = torch.einsum("bhnd,bhnd->bhn", phi_Q_c, K_c_cumsum)

                num_inter = torch.einsum("bhnd,bhde->bhne", phi_Q_c, state_KV)
                den_inter = torch.einsum("bhnd,bhd->bhn", phi_Q_c, state_K)

                num_list.append(num_intra + num_inter)
                den_list.append(den_intra + den_inter)

                state_KV = state_KV + KV_c_cumsum[:, :, -1]
                state_K = state_K + K_c_cumsum[:, :, -1]

            num = torch.cat(num_list, dim=2)
            den = torch.cat(den_list, dim=2).unsqueeze(-1)
        else:
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_f)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous())

        den = den.clamp(min=1e-6)
        out = num / den
        return torch.nan_to_num(out, nan=0.0).to(Q.dtype)

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------
    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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        layer_idx: int = 0,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        device = hidden_states.device

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        N_kv_new = kv_input.shape[1]

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid_mask_new = _extract_valid_mask(eff_mask, B, N_kv_new, device)

        # Projections
        Q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K_new = (
            self.k_proj(kv_input)
            .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V_new = (
            self.v_proj(kv_input)
            .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # HA-FIX-08: RoPE with past_len offset
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            past_len_rope = 0
            if past_key_value is not None:
                if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                    past_len_rope = past_key_value[0].shape[-2]
                elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                    past_len_rope = past_key_value.key_cache[0].shape[-2]
            cos = cos[..., past_len_rope : past_len_rope + N, :]
            sin = sin[..., past_len_rope : past_len_rope + N, :]
            Q, K_new = _apply_rope(Q, K_new, cos, sin)

        # HA-FIX-07: KV-cache
        past_k = past_v = None
        if past_key_value is not None and not is_cross:
            if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
                past_k, past_v = past_key_value[0], past_key_value[1]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                idx = min(layer_idx, len(past_key_value.key_cache) - 1)
                past_k = past_key_value.key_cache[idx]
                past_v = past_key_value.value_cache[idx]

        if past_k is not None and not is_cross:
            K_cat = torch.cat([past_k, K_new], dim=-2)
            V_cat = torch.cat([past_v, V_new], dim=-2)
        else:
            K_cat = K_new
            V_cat = V_new

        if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
            past_key_value.update(K_new, V_new, layer_idx)

        N_kv = K_cat.shape[2]
        past_len = N_kv - N if (self.is_decoder and not is_cross) else 0

        if self.kv_repeat > 1:
            K_cat = repeat_kv(K_cat, self.kv_repeat)
            V_cat = repeat_kv(V_cat, self.kv_repeat)

        # Valid mask for the full key set (with past)
        if past_len > 0:
            past_valid = torch.ones(B, past_len, device=device, dtype=torch.bool)
            valid_mask_full = torch.cat([past_valid, valid_mask_new], dim=1)
        else:
            valid_mask_full = valid_mask_new

        # Two branches
        out_standard = self._standard_branch(
            Q,
            K_cat,
            V_cat,
            eff_mask,
            past_len=past_len,
            valid_mask=valid_mask_full,
        )
        out_linear = self._linear_branch(
            Q,
            K_cat,
            V_cat,
            valid_mask_full,
            past_len=past_len,
        )

        # Per-head gating
        gate = (
            torch.sigmoid(self.ensemble_gate(hidden_states))
            .view(B, N, self.num_heads, 1)
            .transpose(1, 2)
            .contiguous()
        )
        fused = gate * out_standard + (1.0 - gate) * out_linear

        # HA-FIX-09: head_mask broadcasting corrected.
        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.unsqueeze(-1).unsqueeze(-1)  # (B, H, 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)  # (B, H, 1, 1)
            fused = fused * head_mask.to(fused.dtype)

        fused = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)

        # HA-FIX-01: out_proj ALWAYS executed.
        if self.is_nested_encoder:
            # Nested encoder: outer module supplies out_proj, but we still run
            # the (identity-initialised) internal out_proj for parameter stability.
            out = self.out_proj(fused)
        else:
            out = self.out_proj(fused)
            out = self.dropout(out)

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        return out, None

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        if old_attn is None:
            return

        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")
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
            or (
                hasattr(old_attn, "self")
                and hasattr(getattr(old_attn, "self", None), "query")
                and hasattr(old_attn, "output")
                and hasattr(getattr(old_attn, "output", None), "dense")
            )
        )
        logger.debug("HybridAttention.copy_weights: nested_encoder=%s", self.is_nested_encoder)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif cls_name in (
            "GPT2Attention",
            "Qwen2Attention",
            "LlamaAttention",
        ) or hasattr(old_attn, "c_attn"):
            self.is_decoder = True

        from neuralops.core.adapters import UniversalAdapter

        proj_targets = [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
        ]
        if not self.is_nested_encoder:
            proj_targets.append(("o", self.out_proj))

        for proj, target in proj_targets:
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
            except Exception as e:
                logger.debug("UniversalAdapter extraction failed for %s: %s", proj, e)
                w, b = None, None

            if w is None:
                logger.warning("HybridAttention: no weight for '%s' in %s", proj, cls_name)
                continue

            if w.shape != target.weight.shape:
                if w.t().shape == target.weight.shape:
                    w = w.t().contiguous()
                else:
                    logger.warning(
                        "HybridAttention: shape mismatch for %s: src=%s, tgt=%s. Truncating.",
                        proj,
                        tuple(w.shape),
                        tuple(target.weight.shape),
                    )
                    out_d, in_d = target.weight.shape
                    s_out, s_in = w.shape
                    with torch.no_grad():
                        target.weight.data.zero_()
                        target.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_(
                            w[: min(out_d, s_out), : min(in_d, s_in)]
                        )
                    if b is not None and target.bias is not None:
                        min_b = min(target.bias.numel(), b.numel())
                        with torch.no_grad():
                            target.bias.data.zero_()
                            target.bias.data[:min_b].copy_(b[:min_b])
                    continue

            with torch.no_grad():
                norm = torch.linalg.norm(w.float()).item()
                if norm < 1e-4:
                    logger.warning(
                        "HybridAttention: %s weight has near-zero Frobenius norm (%s)",
                        proj,
                        norm,
                    )
                target.weight.copy_(w)
                if b is not None and target.bias is not None:
                    target.bias.copy_(b)

Hybrid softmax + linear attention with per-head dynamic 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 c_proj

Expand source code

@property
def c_proj(self):
    return self.out_proj

prop dense

Expand source code

@property
def dense(self):
    return self.out_proj

prop k_lin

Expand source code

@property
def k_lin(self):
    return self.k_proj

prop key

Expand source code

@property
def key(self):
    return self.k_proj

prop o_proj

Expand source code

@property
def o_proj(self):
    return self.out_proj

prop out_lin

Expand source code

@property
def out_lin(self):
    return self.out_proj

prop q_lin

Expand source code

@property
def q_lin(self):
    return self.q_proj

prop query

Expand source code

@property
def query(self):
    return self.q_proj

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.v_proj

prop value

Expand source code

@property
def value(self):
    return self.v_proj

Methods

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

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    if old_attn is None:
        return

    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")
    )
    self.is_nested_encoder = (
        cls_name in ("BertSelfAttention", "RobertaSelfAttention")
        or (cls_name.endswith("SelfAttention") and not has_out_proj)
        or (
            hasattr(old_attn, "self")
            and hasattr(getattr(old_attn, "self", None), "query")
            and hasattr(old_attn, "output")
            and hasattr(getattr(old_attn, "output", None), "dense")
        )
    )
    logger.debug("HybridAttention.copy_weights: nested_encoder=%s", self.is_nested_encoder)

    if hasattr(old_attn, "is_decoder"):
        self.is_decoder = bool(old_attn.is_decoder)
    elif cls_name in (
        "GPT2Attention",
        "Qwen2Attention",
        "LlamaAttention",
    ) or hasattr(old_attn, "c_attn"):
        self.is_decoder = True

    from neuralops.core.adapters import UniversalAdapter

    proj_targets = [
        ("q", self.q_proj),
        ("k", self.k_proj),
        ("v", self.v_proj),
    ]
    if not self.is_nested_encoder:
        proj_targets.append(("o", self.out_proj))

    for proj, target in proj_targets:
        try:
            w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
        except Exception as e:
            logger.debug("UniversalAdapter extraction failed for %s: %s", proj, e)
            w, b = None, None

        if w is None:
            logger.warning("HybridAttention: no weight for '%s' in %s", proj, cls_name)
            continue

        if w.shape != target.weight.shape:
            if w.t().shape == target.weight.shape:
                w = w.t().contiguous()
            else:
                logger.warning(
                    "HybridAttention: shape mismatch for %s: src=%s, tgt=%s. Truncating.",
                    proj,
                    tuple(w.shape),
                    tuple(target.weight.shape),
                )
                out_d, in_d = target.weight.shape
                s_out, s_in = w.shape
                with torch.no_grad():
                    target.weight.data.zero_()
                    target.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_(
                        w[: min(out_d, s_out), : min(in_d, s_in)]
                    )
                if b is not None and target.bias is not None:
                    min_b = min(target.bias.numel(), b.numel())
                    with torch.no_grad():
                        target.bias.data.zero_()
                        target.bias.data[:min_b].copy_(b[:min_b])
                continue

        with torch.no_grad():
            norm = torch.linalg.norm(w.float()).item()
            if norm < 1e-4:
                logger.warning(
                    "HybridAttention: %s weight has near-zero Frobenius norm (%s)",
                    proj,
                    norm,
                )
            target.weight.copy_(w)
            if b is not None and target.bias is not None:
                target.bias.copy_(b)

Universal fallback: copy pretrained weights from old attention layer.

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,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
layer_idx: int = 0,
*args: Any,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor | None]

Expand source code

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,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    layer_idx: int = 0,
    *args: Any,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    B, N, _ = hidden_states.shape
    device = hidden_states.device

    is_cross = encoder_hidden_states is not None
    kv_input = encoder_hidden_states if is_cross else hidden_states
    N_kv_new = kv_input.shape[1]

    eff_mask = encoder_attention_mask if is_cross else attention_mask
    valid_mask_new = _extract_valid_mask(eff_mask, B, N_kv_new, device)

    # Projections
    Q = (
        self.q_proj(hidden_states)
        .view(B, N, self.num_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    K_new = (
        self.k_proj(kv_input)
        .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )
    V_new = (
        self.v_proj(kv_input)
        .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
        .transpose(1, 2)
        .contiguous()
    )

    # HA-FIX-08: RoPE with past_len offset
    if position_embeddings is None:
        position_embeddings = kwargs.get("position_embeddings")
    if position_embeddings is not None and not is_cross:
        cos, sin = position_embeddings
        past_len_rope = 0
        if past_key_value is not None:
            if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                past_len_rope = past_key_value[0].shape[-2]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                past_len_rope = past_key_value.key_cache[0].shape[-2]
        cos = cos[..., past_len_rope : past_len_rope + N, :]
        sin = sin[..., past_len_rope : past_len_rope + N, :]
        Q, K_new = _apply_rope(Q, K_new, cos, sin)

    # HA-FIX-07: KV-cache
    past_k = past_v = None
    if past_key_value is not None and not is_cross:
        if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
            past_k, past_v = past_key_value[0], past_key_value[1]
        elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
            idx = min(layer_idx, len(past_key_value.key_cache) - 1)
            past_k = past_key_value.key_cache[idx]
            past_v = past_key_value.value_cache[idx]

    if past_k is not None and not is_cross:
        K_cat = torch.cat([past_k, K_new], dim=-2)
        V_cat = torch.cat([past_v, V_new], dim=-2)
    else:
        K_cat = K_new
        V_cat = V_new

    if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
        past_key_value.update(K_new, V_new, layer_idx)

    N_kv = K_cat.shape[2]
    past_len = N_kv - N if (self.is_decoder and not is_cross) else 0

    if self.kv_repeat > 1:
        K_cat = repeat_kv(K_cat, self.kv_repeat)
        V_cat = repeat_kv(V_cat, self.kv_repeat)

    # Valid mask for the full key set (with past)
    if past_len > 0:
        past_valid = torch.ones(B, past_len, device=device, dtype=torch.bool)
        valid_mask_full = torch.cat([past_valid, valid_mask_new], dim=1)
    else:
        valid_mask_full = valid_mask_new

    # Two branches
    out_standard = self._standard_branch(
        Q,
        K_cat,
        V_cat,
        eff_mask,
        past_len=past_len,
        valid_mask=valid_mask_full,
    )
    out_linear = self._linear_branch(
        Q,
        K_cat,
        V_cat,
        valid_mask_full,
        past_len=past_len,
    )

    # Per-head gating
    gate = (
        torch.sigmoid(self.ensemble_gate(hidden_states))
        .view(B, N, self.num_heads, 1)
        .transpose(1, 2)
        .contiguous()
    )
    fused = gate * out_standard + (1.0 - gate) * out_linear

    # HA-FIX-09: head_mask broadcasting corrected.
    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.unsqueeze(-1).unsqueeze(-1)  # (B, H, 1, 1)
        elif head_mask.dim() == 3:
            head_mask = head_mask[:, :, :1].unsqueeze(-1)  # (B, H, 1, 1)
        fused = fused * head_mask.to(fused.dtype)

    fused = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
    fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)

    # HA-FIX-01: out_proj ALWAYS executed.
    if self.is_nested_encoder:
        # Nested encoder: outer module supplies out_proj, but we still run
        # the (identity-initialised) internal out_proj for parameter stability.
        out = self.out_proj(fused)
    else:
        out = self.out_proj(fused)
        out = self.dropout(out)

    out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
    return out, None

Must return (output, attention_weights_or_None).

class SmartWeightedEnsembleAttention (config: Any | None = None,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
hidden_size: int | None = None,
num_heads: int | None = None,
head_dim: int | None = None,
dropout: float = 0.1,
is_decoder: bool | None = None,
**kwargs: Any)

Expand source code

@ATTENTION_REGISTRY.register("hybrid_attention", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("weighted_ensemble", supports_causal=True, supports_gqa=True)
@ATTENTION_REGISTRY.register("smart_weighted_ensemble", supports_causal=True, supports_gqa=True)
class HybridAttention(BaseAttention):
    """
    Hybrid softmax + linear attention with per-head dynamic gated fusion.
    """

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

    def __init__(
        self,
        config: Any | None = None,
        rank: int | None = None,
        pretrained_attn: nn.Module | None = None,
        hidden_size: int | None = None,
        num_heads: int | None = None,
        head_dim: int | None = None,
        dropout: float = 0.1,
        is_decoder: bool | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(config, rank, **kwargs)

        cfg_source = getattr(config, "text_config", config) if config is not None else None
        h_size = (
            hidden_size
            or getattr(cfg_source, "hidden_size", None)
            or getattr(cfg_source, "d_model", None)
            or getattr(cfg_source, "n_embd", None)
            or getattr(cfg_source, "dim", None)
            or kwargs.get("hidden_size", 768)
        )
        n_heads = (
            num_heads
            or getattr(cfg_source, "num_attention_heads", None)
            or getattr(cfg_source, "num_heads", None)
            or getattr(cfg_source, "n_heads", None)
            or getattr(cfg_source, "n_head", None)
            or kwargs.get("num_heads", 12)
        )
        h_dim = (
            head_dim
            or getattr(cfg_source, "head_dim", None)
            or getattr(cfg_source, "d_kv", None)
            or (int(h_size) // max(1, int(n_heads)))
        )
        n_kv_heads = (
            getattr(cfg_source, "num_key_value_heads", None)
            or getattr(cfg_source, "num_kv_heads", None)
            or kwargs.get("num_key_value_heads", n_heads)
        )
        drop_rate = (
            getattr(cfg_source, "attention_probs_dropout_prob", None)
            or getattr(cfg_source, "dropout", None)
            or kwargs.get("dropout", dropout)
        )

        self.hidden_size: int = int(h_size)
        self.num_heads: int = int(n_heads)
        self.head_dim: int = int(h_dim)
        self.num_kv_heads: int = int(n_kv_heads)
        self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads))
        self.all_head_size: int = self.num_heads * self.head_dim
        self.all_kv_head_size: int = self.num_kv_heads * self.head_dim

        # Decoder detection
        self.is_decoder = _detect_decoder(config, is_decoder, pretrained_attn)
        self.is_nested_encoder: bool = False

        # Projections
        self.q_proj = nn.Linear(self.hidden_size, self.all_head_size, bias=True)
        self.k_proj = nn.Linear(self.hidden_size, self.all_kv_head_size, bias=True)
        self.v_proj = nn.Linear(self.hidden_size, self.all_kv_head_size, bias=True)
        self.out_proj = nn.Linear(self.all_head_size, self.hidden_size, bias=True)
        self.dropout = nn.Dropout(drop_rate)

        # Per-head gate
        self.ensemble_gate = nn.Linear(self.hidden_size, self.num_heads)

        # Initialize gate to prefer softmax (sigmoid(10.0) ≈ 0.99995)
        with torch.no_grad():
            self.ensemble_gate.weight.zero_()
            self.ensemble_gate.bias.fill_(10.0)

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

    # ------------------------------------------------------------------
    # Aliases
    # ------------------------------------------------------------------
    @property
    def query(self):
        return self.q_proj

    @property
    def key(self):
        return self.k_proj

    @property
    def value(self):
        return self.v_proj

    @property
    def dense(self):
        return self.out_proj

    @property
    def q_lin(self):
        return self.q_proj

    @property
    def k_lin(self):
        return self.k_proj

    @property
    def v_lin(self):
        return self.v_proj

    @property
    def out_lin(self):
        return self.out_proj

    @property
    def o_proj(self):
        return self.out_proj

    @property
    def c_proj(self):
        return self.out_proj

    # ------------------------------------------------------------------
    # Standard softmax branch (HA-FIX-10: optional SDPA fast path)
    # ------------------------------------------------------------------
    def _standard_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        attention_mask: torch.Tensor | None,
        past_len: int = 0,
        valid_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, H, N, Dh = Q.shape
        N_k = K.shape[2]

        # HA-FIX-10: SDPA fast path when training/eval allows it.
        can_sdpa = (not self.training or self.dropout.p == 0.0) and True
        if can_sdpa:
            try:
                q_sdp = Q.half() if Q.dtype == torch.float32 else Q
                k_sdp = K.half() if K.dtype == torch.float32 else K
                v_sdp = V.half() if V.dtype == torch.float32 else V

                attn_bias_sdp = None
                if valid_mask is not None:
                    # Build a proper (B, 1, 1, N_k) additive mask.
                    attn_bias_sdp = ((valid_mask[:, None, None, :].float() - 1.0) * 1e4).to(
                        q_sdp.dtype
                    )

                is_causal_sdp = self.is_decoder and past_len == 0 and N > 1

                out_sdp = F.scaled_dot_product_attention(
                    q_sdp,
                    k_sdp,
                    v_sdp,
                    attn_mask=attn_bias_sdp,
                    is_causal=is_causal_sdp,
                    dropout_p=self.dropout.p if self.training else 0.0,
                )
                out_sdp = out_sdp.to(Q.dtype)
                return torch.nan_to_num(out_sdp, nan=0.0)
            except Exception as e:
                logger.debug("SDPA fast path failed: %s; falling back to manual.", e)

        # Manual path
        scale = 1.0 / math.sqrt(Dh)
        scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale

        # Causal mask (offset-aware for cache)
        if self.is_decoder and N > 1:
            if past_len > 0:
                abs_q = (past_len + torch.arange(N, device=Q.device)).unsqueeze(1)
                abs_k = torch.arange(N_k, device=Q.device).unsqueeze(0)
                causal_mask = abs_k <= abs_q
            else:
                causal_mask = torch.tril(torch.ones(N, N_k, device=Q.device, dtype=torch.bool))
            scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4)

        # Padding mask
        if valid_mask is not None:
            scores = scores.masked_fill(~valid_mask[:, None, None, :].bool(), -1e4)
        elif attention_mask is not None:
            if attention_mask.dim() == 2:
                ext_mask = (1.0 - attention_mask.unsqueeze(1).unsqueeze(2).float()) * -1e4
                scores = scores + ext_mask.to(scores.dtype)
            elif attention_mask.dim() == 3:
                scores = scores + attention_mask.unsqueeze(1).to(scores.dtype)
            elif attention_mask.dim() == 4:
                scores = scores + attention_mask.to(scores.dtype)

        scores = torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
        attn = F.softmax(scores, dim=-1).to(Q.dtype)
        attn = self.dropout(attn)
        return torch.matmul(attn, V)

    # ------------------------------------------------------------------
    # Linear branch (HA-FIX-02: proper causal + padding)
    # ------------------------------------------------------------------
    def _linear_branch(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        valid_mask: torch.Tensor,
        past_len: int = 0,
    ) -> torch.Tensor:
        phi_Q = F.elu(Q.float()) + 1.0
        phi_K = F.elu(K.float()) + 1.0
        V_f = V.float()

        # Padding mask on K and V
        mask = valid_mask.unsqueeze(1).unsqueeze(-1).to(phi_K.dtype)
        phi_K = phi_K * mask
        V_f = V_f * mask

        B_h, H_h, N_h, D_h = phi_Q.shape
        phi_K.shape[2]

        # Handle cache: query positions are [past_len .. past_len + N_h), keys are [0 .. N_k)
        if self.is_decoder and past_len > 0:
            # Only new queries, keys include cache.
            # Non-causal over the full key set (past + new).
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_f)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous()).clamp(min=1e-6)
            out = num / den
            return torch.nan_to_num(out, nan=0.0).to(Q.dtype)

        if self.is_decoder:
            # Chunked causal linear attention.
            chunk_size = 128
            num_chunks = (N_h + chunk_size - 1) // chunk_size

            state_KV = torch.zeros(B_h, H_h, D_h, D_h, device=phi_Q.device, dtype=phi_Q.dtype)
            state_K = torch.zeros(B_h, H_h, D_h, device=phi_Q.device, dtype=phi_Q.dtype)

            num_list = []
            den_list = []

            for c in range(num_chunks):
                s = c * chunk_size
                e = min(N_h, s + chunk_size)

                phi_Q_c = phi_Q[:, :, s:e]
                phi_K_c = phi_K[:, :, s:e]
                V_c = V_f[:, :, s:e]

                KV_c = torch.einsum("bhnd,bhne->bhnde", phi_K_c, V_c)
                KV_c_cumsum = torch.cumsum(KV_c, dim=2)
                num_intra = torch.einsum("bhnd,bhnde->bhne", phi_Q_c, KV_c_cumsum)

                K_c_cumsum = torch.cumsum(phi_K_c, dim=2)
                den_intra = torch.einsum("bhnd,bhnd->bhn", phi_Q_c, K_c_cumsum)

                num_inter = torch.einsum("bhnd,bhde->bhne", phi_Q_c, state_KV)
                den_inter = torch.einsum("bhnd,bhd->bhn", phi_Q_c, state_K)

                num_list.append(num_intra + num_inter)
                den_list.append(den_intra + den_inter)

                state_KV = state_KV + KV_c_cumsum[:, :, -1]
                state_K = state_K + K_c_cumsum[:, :, -1]

            num = torch.cat(num_list, dim=2)
            den = torch.cat(den_list, dim=2).unsqueeze(-1)
        else:
            KV = torch.matmul(phi_K.transpose(-2, -1).contiguous(), V_f)
            Z = phi_K.sum(dim=-2, keepdim=True)
            num = torch.matmul(phi_Q, KV)
            den = torch.matmul(phi_Q, Z.transpose(-2, -1).contiguous())

        den = den.clamp(min=1e-6)
        out = num / den
        return torch.nan_to_num(out, nan=0.0).to(Q.dtype)

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------
    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,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        layer_idx: int = 0,
        *args: Any,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, N, _ = hidden_states.shape
        device = hidden_states.device

        is_cross = encoder_hidden_states is not None
        kv_input = encoder_hidden_states if is_cross else hidden_states
        N_kv_new = kv_input.shape[1]

        eff_mask = encoder_attention_mask if is_cross else attention_mask
        valid_mask_new = _extract_valid_mask(eff_mask, B, N_kv_new, device)

        # Projections
        Q = (
            self.q_proj(hidden_states)
            .view(B, N, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        K_new = (
            self.k_proj(kv_input)
            .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        V_new = (
            self.v_proj(kv_input)
            .view(B, N_kv_new, self.num_kv_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        # HA-FIX-08: RoPE with past_len offset
        if position_embeddings is None:
            position_embeddings = kwargs.get("position_embeddings")
        if position_embeddings is not None and not is_cross:
            cos, sin = position_embeddings
            past_len_rope = 0
            if past_key_value is not None:
                if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1:
                    past_len_rope = past_key_value[0].shape[-2]
                elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                    past_len_rope = past_key_value.key_cache[0].shape[-2]
            cos = cos[..., past_len_rope : past_len_rope + N, :]
            sin = sin[..., past_len_rope : past_len_rope + N, :]
            Q, K_new = _apply_rope(Q, K_new, cos, sin)

        # HA-FIX-07: KV-cache
        past_k = past_v = None
        if past_key_value is not None and not is_cross:
            if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
                past_k, past_v = past_key_value[0], past_key_value[1]
            elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache:
                idx = min(layer_idx, len(past_key_value.key_cache) - 1)
                past_k = past_key_value.key_cache[idx]
                past_v = past_key_value.value_cache[idx]

        if past_k is not None and not is_cross:
            K_cat = torch.cat([past_k, K_new], dim=-2)
            V_cat = torch.cat([past_v, V_new], dim=-2)
        else:
            K_cat = K_new
            V_cat = V_new

        if past_key_value is not None and hasattr(past_key_value, "update") and not is_cross:
            past_key_value.update(K_new, V_new, layer_idx)

        N_kv = K_cat.shape[2]
        past_len = N_kv - N if (self.is_decoder and not is_cross) else 0

        if self.kv_repeat > 1:
            K_cat = repeat_kv(K_cat, self.kv_repeat)
            V_cat = repeat_kv(V_cat, self.kv_repeat)

        # Valid mask for the full key set (with past)
        if past_len > 0:
            past_valid = torch.ones(B, past_len, device=device, dtype=torch.bool)
            valid_mask_full = torch.cat([past_valid, valid_mask_new], dim=1)
        else:
            valid_mask_full = valid_mask_new

        # Two branches
        out_standard = self._standard_branch(
            Q,
            K_cat,
            V_cat,
            eff_mask,
            past_len=past_len,
            valid_mask=valid_mask_full,
        )
        out_linear = self._linear_branch(
            Q,
            K_cat,
            V_cat,
            valid_mask_full,
            past_len=past_len,
        )

        # Per-head gating
        gate = (
            torch.sigmoid(self.ensemble_gate(hidden_states))
            .view(B, N, self.num_heads, 1)
            .transpose(1, 2)
            .contiguous()
        )
        fused = gate * out_standard + (1.0 - gate) * out_linear

        # HA-FIX-09: head_mask broadcasting corrected.
        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.unsqueeze(-1).unsqueeze(-1)  # (B, H, 1, 1)
            elif head_mask.dim() == 3:
                head_mask = head_mask[:, :, :1].unsqueeze(-1)  # (B, H, 1, 1)
            fused = fused * head_mask.to(fused.dtype)

        fused = fused.transpose(1, 2).contiguous().view(B, N, self.all_head_size)
        fused = torch.nan_to_num(fused, nan=0.0, posinf=1e4, neginf=-1e4)

        # HA-FIX-01: out_proj ALWAYS executed.
        if self.is_nested_encoder:
            # Nested encoder: outer module supplies out_proj, but we still run
            # the (identity-initialised) internal out_proj for parameter stability.
            out = self.out_proj(fused)
        else:
            out = self.out_proj(fused)
            out = self.dropout(out)

        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
        return out, None

    # ------------------------------------------------------------------
    # copy_weights
    # ------------------------------------------------------------------
    def copy_weights(self, old_attn: Any) -> None:
        if old_attn is None:
            return

        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")
        )
        self.is_nested_encoder = (
            cls_name in ("BertSelfAttention", "RobertaSelfAttention")
            or (cls_name.endswith("SelfAttention") and not has_out_proj)
            or (
                hasattr(old_attn, "self")
                and hasattr(getattr(old_attn, "self", None), "query")
                and hasattr(old_attn, "output")
                and hasattr(getattr(old_attn, "output", None), "dense")
            )
        )
        logger.debug("HybridAttention.copy_weights: nested_encoder=%s", self.is_nested_encoder)

        if hasattr(old_attn, "is_decoder"):
            self.is_decoder = bool(old_attn.is_decoder)
        elif cls_name in (
            "GPT2Attention",
            "Qwen2Attention",
            "LlamaAttention",
        ) or hasattr(old_attn, "c_attn"):
            self.is_decoder = True

        from neuralops.core.adapters import UniversalAdapter

        proj_targets = [
            ("q", self.q_proj),
            ("k", self.k_proj),
            ("v", self.v_proj),
        ]
        if not self.is_nested_encoder:
            proj_targets.append(("o", self.out_proj))

        for proj, target in proj_targets:
            try:
                w, b = UniversalAdapter.get_weight_bias(old_attn, proj)
            except Exception as e:
                logger.debug("UniversalAdapter extraction failed for %s: %s", proj, e)
                w, b = None, None

            if w is None:
                logger.warning("HybridAttention: no weight for '%s' in %s", proj, cls_name)
                continue

            if w.shape != target.weight.shape:
                if w.t().shape == target.weight.shape:
                    w = w.t().contiguous()
                else:
                    logger.warning(
                        "HybridAttention: shape mismatch for %s: src=%s, tgt=%s. Truncating.",
                        proj,
                        tuple(w.shape),
                        tuple(target.weight.shape),
                    )
                    out_d, in_d = target.weight.shape
                    s_out, s_in = w.shape
                    with torch.no_grad():
                        target.weight.data.zero_()
                        target.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_(
                            w[: min(out_d, s_out), : min(in_d, s_in)]
                        )
                    if b is not None and target.bias is not None:
                        min_b = min(target.bias.numel(), b.numel())
                        with torch.no_grad():
                            target.bias.data.zero_()
                            target.bias.data[:min_b].copy_(b[:min_b])
                    continue

            with torch.no_grad():
                norm = torch.linalg.norm(w.float()).item()
                if norm < 1e-4:
                    logger.warning(
                        "HybridAttention: %s weight has near-zero Frobenius norm (%s)",
                        proj,
                        norm,
                    )
                target.weight.copy_(w)
                if b is not None and target.bias is not None:
                    target.bias.copy_(b)

Hybrid softmax + linear attention with per-head dynamic 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 c_proj

Expand source code

@property
def c_proj(self):
    return self.out_proj

prop dense

Expand source code

@property
def dense(self):
    return self.out_proj

prop k_lin

Expand source code

@property
def k_lin(self):
    return self.k_proj

prop key

Expand source code

@property
def key(self):
    return self.k_proj

prop o_proj

Expand source code

@property
def o_proj(self):
    return self.out_proj

prop out_lin

Expand source code

@property
def out_lin(self):
    return self.out_proj

prop q_lin

Expand source code

@property
def q_lin(self):
    return self.q_proj

prop query

Expand source code

@property
def query(self):
    return self.q_proj

prop v_lin

Expand source code

@property
def v_lin(self):
    return self.v_proj

prop value

Expand source code

@property
def value(self):
    return self.v_proj

Inherited members

• BaseAttention:
  • copy_weights
  • extra_repr
  • forward