Module neuralops.attentions.wgqa_attention
Weighted Grouped-Query Attention (WGQA).
Hardened for production use across diverse transformer architectures: - DistilBERT (Flat MHA) - RoBERTa / BERT (Nested Encoder Self-Attention) - GPT-2 (Causal Decoder with Conv1D) - Qwen / LLaMA (GQA with Rotary Position Embeddings)
Features learnable per-KV-head importance weights, robust GQA pooling, LoRA adapters, RoPE support with fallback, dynamic dimension inference, and causal masking.
Functions
def apply_rotary_pos_emb(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
position_ids: torch.Tensor | None = None) ‑> tuple[torch.Tensor, torch.Tensor]-
Expand source code
def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Applies rotary position embeddings safely handling even or odd dimensions.""" if cos is None or sin is None: return q, k if cos.dim() == 2: cos = cos.unsqueeze(0).unsqueeze(1) sin = sin.unsqueeze(0).unsqueeze(1) elif cos.dim() == 3: cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) while cos.dim() < q.dim(): cos = cos.unsqueeze(0) sin = sin.unsqueeze(0) head_dim = q.shape[-1] rotary_dim = 2 * (min(cos.shape[-1], head_dim) // 2) if rotary_dim == 0: return q, k q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] cos_q = cos[..., : q.shape[-2], :rotary_dim].to(q.dtype) sin_q = sin[..., : q.shape[-2], :rotary_dim].to(q.dtype) cos_k = cos[..., : k.shape[-2], :rotary_dim].to(k.dtype) sin_k = sin[..., : k.shape[-2], :rotary_dim].to(k.dtype) q_rot_embed = (q_rot * cos_q) + (_rotate_half(q_rot) * sin_q) k_rot_embed = (k_rot * cos_k) + (_rotate_half(k_rot) * sin_k) q = torch.cat([q_rot_embed, q_pass], dim=-1) if q_pass.numel() > 0 else q_rot_embed k = torch.cat([k_rot_embed, k_pass], dim=-1) if k_pass.numel() > 0 else k_rot_embed return q, kApplies rotary position embeddings safely handling even or odd dimensions.
def inject_wgqa(layer: Any = None,
model: Any = None,
model_config: Any = None,
comp_config: Any = None,
layer_idx: int = 0,
total_layers: int = 12,
old_attn: nn.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("wgqa_inject") @ATTENTION_REGISTRY.register("inject_wgqa") @ATTENTION_REGISTRY.register("wgqa_attention_inject") def inject_wgqa( layer: Any = None, model: Any = None, model_config: Any = None, comp_config: Any = None, layer_idx: int = 0, total_layers: int = 12, old_attn: nn.Module | None = None, **kwargs: Any, ): target = layer if layer is not None else model if target is None: target = kwargs.get("model", kwargs.get("layer")) if target is None: return None cfg = model_config if model_config is not None else getattr(target, "config", None) # Whole-model traversal layer_list = None if hasattr(target, "distilbert") and hasattr(target.distilbert, "transformer"): layer_list = target.distilbert.transformer.layer elif hasattr(target, "roberta") and hasattr(target.roberta, "encoder"): layer_list = target.roberta.encoder.layer elif hasattr(target, "transformer") and hasattr(target.transformer, "h"): layer_list = target.transformer.h elif hasattr(target, "h") and isinstance(target.h, (torch.nn.ModuleList, list)): layer_list = target.h elif hasattr(target, "model") and hasattr(target.model, "layers"): layer_list = target.model.layers elif hasattr(target, "encoder") and hasattr(target.encoder, "layer"): layer_list = target.encoder.layer elif ( hasattr(target, "bert") and hasattr(target.bert, "encoder") and hasattr(target.bert.encoder, "layer") ): layer_list = target.bert.encoder.layer elif hasattr(target, "layers") and isinstance(target.layers, (torch.nn.ModuleList, list)): layer_list = target.layers if layer_list is not None: num_injected = 0 tot = len(layer_list) for idx, lyr in enumerate(layer_list): res = inject_wgqa( layer=lyr, model_config=cfg, comp_config=( comp_config if "comp_config" in locals() or "comp_config" in kwargs else None ), layer_idx=idx, total_layers=tot, **kwargs, ) if res is not None: num_injected += 1 return target if num_injected > 0 else None layer = target """ Injection hook for WeightedGroupedQueryAttention (WGQA). Replaces self-attention with Weighted Grouped-Query Attention. target_attention_only: Replaces self-attention modules only, preserving all FFN/MLP blocks. """ if old_attn is None: return None if model_config is None and hasattr(layer, "config"): model_config = layer.config if model_config is None and hasattr(old_attn, "config"): model_config = old_attn.config hidden_size = ( _cfg_get(model_config, "hidden_size") or _cfg_get(model_config, "d_model") or _cfg_get(model_config, "dim") or getattr(old_attn, "hidden_size", None) or getattr(old_attn, "embed_dim", None) or ( old_attn.q_proj.in_features if hasattr(old_attn, "q_proj") and hasattr(old_attn.q_proj, "in_features") else None ) or ( old_attn.q_lin.in_features if hasattr(old_attn, "q_lin") and hasattr(old_attn.q_lin, "in_features") else None ) or ( old_attn.query.in_features if hasattr(old_attn, "query") and hasattr(old_attn.query, "in_features") else None ) or 768 ) num_heads = ( _cfg_get(model_config, "num_attention_heads") or _cfg_get(model_config, "num_heads") or _cfg_get(model_config, "n_heads") or getattr(old_attn, "num_heads", None) or getattr(old_attn, "num_attention_heads", None) or 12 ) if model_config is None: import types model_config = types.SimpleNamespace( hidden_size=hidden_size, num_attention_heads=num_heads, num_key_value_heads=getattr(old_attn, "num_key_value_heads", num_heads), is_decoder=getattr(old_attn, "is_decoder", False), ) lora_rank = kwargs.pop("lora_rank", _cfg_get(comp_config, "lora_rank", 16)) rank = kwargs.pop("rank", _cfg_get(comp_config, "rank", 32)) num_kv_heads = kwargs.pop("num_kv_heads", _cfg_get(comp_config, "num_kv_heads", None)) try: new_attn = WeightedGroupedQueryAttention( config=model_config, rank=rank, lora_rank=lora_rank, num_kv_heads=num_kv_heads, pretrained_attn=old_attn, **kwargs, ) 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 return new_attn for attr in ("self_attn", "attn", "self"): if hasattr(layer, attr): setattr(layer, attr, new_attn) return new_attn layer.attention = new_attn return new_attn except Exception as e: logger.error(f"WGQA injection failed: {e}") return None
Classes
class LoRALinear (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = True,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class LoRALinear(nn.Module): """Linear projection with frozen base weights and trainable LoRA adapters.""" def __init__( self, in_features: int, out_features: int, bias: bool = True, lora_rank: int = 16, lora_alpha: float = 32.0, freeze_base: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() factory_kwargs = {"device": device, "dtype": dtype} self.in_features = in_features self.out_features = out_features self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0 self.freeze_base = freeze_base self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs)) if bias: self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs)) else: self.register_parameter("bias", None) if self.lora_rank > 0: self.lora_A = nn.Parameter(torch.empty((lora_rank, in_features), **factory_kwargs)) self.lora_B = nn.Parameter(torch.empty((out_features, lora_rank), **factory_kwargs)) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) self.reset_parameters() if freeze_base: self.freeze_base_weights() def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def freeze_base_weights(self) -> None: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0: if self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_B is not None: self.lora_B.requires_grad = True def forward(self, x: torch.Tensor) -> torch.Tensor: out = F.linear(x, self.weight, self.bias) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling return outLinear projection with frozen base weights and trainable LoRA adapters.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: out = F.linear(x, self.weight, self.bias) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling return outDefine the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them. def freeze_base_weights(self) ‑> None-
Expand source code
def freeze_base_weights(self) -> None: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0: if self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_B is not None: self.lora_B.requires_grad = True def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True
class WeightedGroupedQueryAttention (config: Any | None = None,
rank: int | None = None,
lora_rank: int = 16,
lora_alpha: float = 32.0,
num_kv_heads: int | None = None,
is_decoder: bool | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register( "wgqa", supports_causal=True, supports_bidirectional=True, supports_gqa=True ) @ATTENTION_REGISTRY.register( "wgqa_attention", supports_causal=True, supports_bidirectional=True, supports_gqa=True, ) class WeightedGroupedQueryAttention(BaseAttention): r""" Weighted Grouped-Query Attention (WGQA) with learnable per-KV-head importance weights, robust GQA pooling, LoRA adapters, RoPE support with fallback, and causal masking. """ 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, lora_rank: int = 16, lora_alpha: float = 32.0, num_kv_heads: int | None = None, is_decoder: bool | None = None, pretrained_attn: nn.Module | None = None, old_attn: nn.Module | None = None, **kwargs: Any, ) -> None: super().__init__(config, rank) src_attn = pretrained_attn if pretrained_attn is not None else old_attn if config is not None and not isinstance(config, int): self.adapter = ModelAdapter(config) self.num_heads: int = self.adapter.num_heads self.num_kv_heads: int = self.adapter.num_key_value_heads self.hidden_size: int = self.adapter.hidden_size self.head_dim: int = self.adapter.head_dim self.all_head_size: int = self.num_heads * self.head_dim self.kv_head_size: int = self.num_kv_heads * self.head_dim self.kv_repeat: int = self.adapter.kv_repeat self.is_decoder: bool = getattr(config, "is_decoder", False) else: self.adapter = None self.hidden_size = _cfg_get(config, "hidden_size", kwargs.get("hidden_size", 768)) self.num_heads = _cfg_get(config, "num_attention_heads", kwargs.get("num_heads", 12)) self.num_kv_heads = _cfg_get(config, "num_key_value_heads", self.num_heads) self.head_dim = self.hidden_size // max(1, self.num_heads) self.all_head_size = self.num_heads * self.head_dim self.kv_head_size = self.num_kv_heads * self.head_dim self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads)) self.is_decoder = bool(_cfg_get(config, "is_decoder", kwargs.get("is_decoder", False))) # Prioritize explicit kwargs if "hidden_size" in kwargs: self.hidden_size = kwargs["hidden_size"] self.head_dim = self.hidden_size // max(1, self.num_heads) self.all_head_size = self.num_heads * self.head_dim self.kv_head_size = self.num_kv_heads * self.head_dim if "num_heads" in kwargs: self.num_heads = kwargs["num_heads"] self.head_dim = self.hidden_size // max(1, self.num_heads) self.all_head_size = self.num_heads * self.head_dim self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads)) if num_kv_heads is not None: self.num_kv_heads = int(num_kv_heads) self.kv_head_size = self.num_kv_heads * self.head_dim self.kv_repeat = max(1, self.num_heads // max(1, self.num_kv_heads)) if is_decoder is not None: self.is_decoder = bool(is_decoder) self.is_nested_encoder: bool = False self.lora_rank = kwargs.get("lora_rank", _cfg_get(config, "lora_rank", lora_rank)) self.lora_alpha = kwargs.get("lora_alpha", _cfg_get(config, "lora_alpha", lora_alpha)) self.q_proj = LoRALinear( self.hidden_size, self.all_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) self.k_proj = LoRALinear( self.hidden_size, self.kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) self.v_proj = LoRALinear( self.hidden_size, self.kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) self.out_proj = LoRALinear( self.all_head_size, self.hidden_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=True, ) # Learnable per-KV-head importance weights initialized to 1.0 self.k_weights = nn.Parameter(torch.ones(self.num_kv_heads, 1)) self.v_weights = nn.Parameter(torch.ones(self.num_kv_heads, 1)) dropout_p = _cfg_get( config, "dropout", _cfg_get( config, "hidden_dropout_prob", _cfg_get(config, "attention_probs_dropout_prob", 0.1), ), ) self.dropout = nn.Dropout(dropout_p) self.rotary_emb = getattr(src_attn, "rotary_emb", None) if src_attn is not None: self.copy_weights(src_attn) @property def W_q(self): return self.q_proj @property def W_k(self): return self.k_proj @property def W_v(self): return self.v_proj @property def out(self): return self.out_proj @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 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 dense(self): return self.out_proj @property def o_proj(self): return self.out_proj @property def c_proj(self): return self.out_proj def copy_weights(self, src: Any) -> None: """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.""" extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. UniversalAdapter extraction for proj in ["q", "k", "v", "o"]: try: w, b = UniversalAdapter.get_weight_bias(src, proj) if w is not None: extracted[proj] = (w, b) except Exception: pass # 2. GPT-2 Fused c_attn Conv1D handling if hasattr(src, "c_attn") and ( extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None ): c_mod = src.c_attn base_c = getattr(c_mod, "base_layer", c_mod) weight = getattr(base_c, "weight", getattr(c_mod, "weight", None)) if weight is not None: weight = weight.data.clone() bias = getattr(base_c, "bias", getattr(c_mod, "bias", None)) bias_data = bias.data.clone() if bias is not None else None is_conv1d = ( "Conv1D" in c_mod.__class__.__name__ or "Conv1D" in base_c.__class__.__name__ or hasattr(c_mod, "nf") or hasattr(base_c, "nf") or (weight.dim() == 2 and weight.shape[0] < weight.shape[1]) ) if is_conv1d: weight = weight.t().contiguous() split_size = weight.shape[0] // 3 q_w = weight[:split_size, :] k_w = weight[split_size : 2 * split_size, :] v_w = weight[2 * split_size :, :] q_b = bias_data[:split_size] if bias_data is not None else None k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None v_b = bias_data[2 * split_size :] if bias_data is not None else None extracted["q"] = (q_w, q_b) extracted["k"] = (k_w, k_b) extracted["v"] = (v_w, v_b) if hasattr(src, "c_proj") and extracted["o"][0] is None: p_mod = src.c_proj base_o = getattr(p_mod, "base_layer", p_mod) w_o = getattr(base_o, "weight", getattr(p_mod, "weight", None)) if w_o is not None: w_o = w_o.data.clone() is_conv1d = ( "Conv1D" in p_mod.__class__.__name__ or "Conv1D" in base_o.__class__.__name__ or hasattr(p_mod, "nf") or hasattr(base_o, "nf") or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1]) ) if is_conv1d: w_o = w_o.t().contiguous() b_o = getattr(base_o, "bias", getattr(p_mod, "bias", None)) extracted["o"] = (w_o, b_o.data.clone() if b_o is not None else None) # 3. RoBERTa / BERT Nested Encoder handling if hasattr(src, "self") and ( extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None ): s = src.self if hasattr(s, "query"): extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None)) if hasattr(s, "key"): extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None)) if hasattr(s, "value"): extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None)) if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None: dense = src.output.dense extracted["o"] = (dense.weight.data, getattr(dense, "bias", None)) # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT named_map = { "q": ["q_proj", "q_lin", "query", "W_q"], "k": ["k_proj", "k_lin", "key", "W_k"], "v": ["v_proj", "v_lin", "value", "W_v"], "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"], } for proj, names in named_map.items(): if extracted[proj][0] is None: for name in names: if hasattr(src, name): layer = getattr(src, name) if hasattr(layer, "weight"): extracted[proj] = ( layer.weight.data, getattr(layer, "bias", None), ) break # Detect nested encoder architecture cls_name = getattr(src, "__class__", type(src)).__name__ has_out = ( hasattr(src, "dense") or hasattr(src, "out_proj") or hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") ) is_nested = False if cls_name in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (hasattr(src, "query") and not has_out): is_nested = True if hasattr(src, "self") and hasattr(src, "output"): is_nested = True self.is_nested_encoder = bool(is_nested) missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None] if not self.is_nested_encoder and extracted["o"][0] is None: missing.append("o") if missing: logger.warning( "Attention module copy_weights warning. Proceeding with initialized weights." ) def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor: if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim: return w.t().contiguous() return w with torch.no_grad(): # Copy Q w_q, b_q = extracted["q"] if w_q is not None: w_q_or = _orient(w_q, self.hidden_size) out_d, in_d = self.q_proj.weight.shape s_out, s_in = w_q_or.shape self.q_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_q_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_q is not None and self.q_proj.bias is not None: b_val = b_q.data if isinstance(b_q, nn.Parameter) else b_q self.q_proj.bias.data[: min(len(self.q_proj.bias), len(b_val))].copy_( b_val[: min(len(self.q_proj.bias), len(b_val))] ) # Copy K with GQA pooling if source heads > kv heads w_k, b_k = extracted["k"] if w_k is not None: w_k_or = _orient(w_k, self.hidden_size) if w_k_or.shape[0] == self.kv_head_size: out_d, in_d = self.k_proj.weight.shape s_out, s_in = w_k_or.shape self.k_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_k_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_k is not None and self.k_proj.bias is not None: b_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k self.k_proj.bias.data[: min(len(self.k_proj.bias), len(b_val))].copy_( b_val[: min(len(self.k_proj.bias), len(b_val))] ) else: rep = max(1, w_k_or.shape[0] // max(1, self.kv_head_size)) w_k_pooled = ( w_k_or.view(self.num_kv_heads, rep, self.head_dim, -1) .mean(dim=1) .view(self.kv_head_size, -1) ) out_d, in_d = self.k_proj.weight.shape s_out, s_in = w_k_pooled.shape self.k_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_k_pooled[: min(out_d, s_out), : min(in_d, s_in)] ) if b_k is not None and self.k_proj.bias is not None: b_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k b_k_pooled = ( b_val.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1) ) self.k_proj.bias.data[: min(len(self.k_proj.bias), len(b_k_pooled))].copy_( b_k_pooled[: min(len(self.k_proj.bias), len(b_k_pooled))] ) # Copy V with GQA pooling if source heads > kv heads w_v, b_v = extracted["v"] if w_v is not None: w_v_or = _orient(w_v, self.hidden_size) if w_v_or.shape[0] == self.kv_head_size: out_d, in_d = self.v_proj.weight.shape s_out, s_in = w_v_or.shape self.v_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_v_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_v is not None and self.v_proj.bias is not None: b_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v self.v_proj.bias.data[: min(len(self.v_proj.bias), len(b_val))].copy_( b_val[: min(len(self.v_proj.bias), len(b_val))] ) else: rep = max(1, w_v_or.shape[0] // max(1, self.kv_head_size)) w_v_pooled = ( w_v_or.view(self.num_kv_heads, rep, self.head_dim, -1) .mean(dim=1) .view(self.kv_head_size, -1) ) out_d, in_d = self.v_proj.weight.shape s_out, s_in = w_v_pooled.shape self.v_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_v_pooled[: min(out_d, s_out), : min(in_d, s_in)] ) if b_v is not None and self.v_proj.bias is not None: b_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v b_v_pooled = ( b_val.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1) ) self.v_proj.bias.data[: min(len(self.v_proj.bias), len(b_v_pooled))].copy_( b_v_pooled[: min(len(self.v_proj.bias), len(b_v_pooled))] ) # Copy Out projection w_o, b_o = extracted["o"] if w_o is not None and not self.is_nested_encoder: w_o_or = _orient(w_o, self.all_head_size) out_d, in_d = self.out_proj.weight.shape s_out, s_in = w_o_or.shape self.out_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_o_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_o is not None and self.out_proj.bias is not None: b_val = b_o.data if isinstance(b_o, nn.Parameter) else b_o self.out_proj.bias.data[: min(len(self.out_proj.bias), len(b_val))].copy_( b_val[: min(len(self.out_proj.bias), len(b_val))] ) elif self.is_nested_encoder: min_d = min(self.out_proj.weight.shape[0], self.out_proj.weight.shape[1]) self.out_proj.weight.data.zero_() self.out_proj.weight.data[:min_d, :min_d].fill_diagonal_(1.0) if self.out_proj.bias is not None: self.out_proj.bias.data.zero_() for target in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): target.freeze_base_weights() # Preserve source LoRA adapters if present on source module source_submodules = { "q": ["q_proj", "q_lin", "query", "W_q"], "k": ["k_proj", "k_lin", "key", "W_k"], "v": ["v_proj", "v_lin", "value", "W_v"], "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"], } for proj_key, target in [ ("q", self.q_proj), ("k", self.k_proj), ("v", self.v_proj), ("o", self.out_proj), ]: for name in source_submodules[proj_key]: if hasattr(src, name): sub = getattr(src, name) if ( hasattr(sub, "lora_A") and hasattr(sub, "lora_B") and sub.lora_A is not None and sub.lora_B is not None ): src_lA = getattr(sub.lora_A, "weight", sub.lora_A) src_lB = getattr(sub.lora_B, "weight", sub.lora_B) if hasattr(src_lA, "data") and target.lora_A is not None: if ( src_lA.shape[0] == target.lora_A.shape[0] and src_lA.shape[1] == target.lora_A.shape[1] ): target.lora_A.data.copy_(src_lA.data) target.lora_A.requires_grad = True else: print( f"Warning: lora_A shape mismatch for {proj_key}; keeping zero-init." ) if hasattr(src_lB, "data") and target.lora_B is not None: if ( src_lB.shape[0] == target.lora_B.shape[0] and src_lB.shape[1] == target.lora_B.shape[1] ): target.lora_B.data.copy_(src_lB.data) target.lora_B.requires_grad = True else: print( f"Warning: lora_B shape mismatch for {proj_key}; keeping zero-init." ) if hasattr(sub, "lora_scaling"): target.lora_scaling = sub.lora_scaling break if hasattr(src, "is_decoder"): self.is_decoder = bool(src.is_decoder) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, head_mask: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, past_key_value: Any | None = None, output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: key_value_states = kwargs.get("key_value_states", encoder_hidden_states) is_cross_attention = key_value_states is not None kv_input = key_value_states if is_cross_attention else hidden_states B, N_q, D = hidden_states.shape N_k = kv_input.shape[1] H, d = self.num_heads, self.head_dim device = hidden_states.device dtype = hidden_states.dtype # 1. Projections Q = self.q_proj(hidden_states).view(B, N_q, H, d).transpose(1, 2).contiguous() K = self.k_proj(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous() V = self.v_proj(kv_input).view(B, N_k, self.num_kv_heads, d).transpose(1, 2).contiguous() # Apply learnable per-KV-head importance weights K = K * self.k_weights.view(1, self.num_kv_heads, 1, 1).to(K.dtype) V = V * self.v_weights.view(1, self.num_kv_heads, 1, 1).to(V.dtype) # 2. RoPE support position_embeddings = kwargs.get("position_embeddings") if ( position_embeddings is not None and isinstance(position_embeddings, (tuple, list)) and len(position_embeddings) >= 2 ): cos, sin = position_embeddings[0], position_embeddings[1] Q, K = apply_rotary_pos_emb(Q, K, cos, sin, kwargs.get("position_ids")) elif getattr(self, "rotary_emb", None) is not None: try: rotary_out = self.rotary_emb(hidden_states, position_ids=kwargs.get("position_ids")) if isinstance(rotary_out, tuple) and len(rotary_out) >= 2: cos, sin = rotary_out[0], rotary_out[1] Q, K = apply_rotary_pos_emb(Q, K, cos, sin, kwargs.get("position_ids")) except Exception: try: from neuralops.utils.compat import maybe_apply_rotary Q, K = maybe_apply_rotary( Q, K, getattr(self, "adapter", None), self.rotary_emb, kwargs.get("position_ids"), ) except Exception as e: logger.warning(f"maybe_apply_rotary failed: {e}") # 3. GQA repeat if self.kv_repeat > 1: K = repeat_kv(K, self.kv_repeat) V = repeat_kv(V, self.kv_repeat) scale = d**-0.5 # Check for causal decoders is_decoder = ( self.is_decoder or bool(_cfg_get(self.config, "is_decoder", False)) or str(_cfg_get(self.config, "model_type", "")).lower() in ("gpt2", "qwen", "qwen2", "qwen3_5", "llama") ) is_cross = bool(kwargs.get("is_cross_attention", False)) mask = attention_mask if attention_mask is not None else kwargs.get("mask") if mask is None and head_mask is None and not output_attentions: try: _qs = Q.half() if Q.dtype == torch.float32 else Q _ks = K.half() if K.dtype == torch.float32 else K _vs = V.half() if V.dtype == torch.float32 else V _is_causal = is_decoder and not is_cross and N_q > 1 _sdpa_out = torch.nn.functional.scaled_dot_product_attention( _qs, _ks, _vs, is_causal=_is_causal, dropout_p=self.dropout.p if self.training else 0.0, ) _sdpa_out = ( _sdpa_out.to(hidden_states.dtype) .transpose(1, 2) .contiguous() .view(B, N_q, self.all_head_size) ) _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4) output = ( _sdpa_out if getattr(self, "is_nested_encoder", False) else self.out_proj(_sdpa_out) ) valid_m = get_token_valid_mask(attention_mask, B, N_q, device) if valid_m is not None: output = output * valid_m.unsqueeze(-1).to(output.dtype) return output, None except Exception: pass scores = torch.matmul(Q.float(), K.float().transpose(-2, -1).contiguous()) * scale # Causal mask if is_decoder: causal_tril = ( torch.tril(torch.ones((N_q, N_k), device=device, dtype=torch.bool)) .unsqueeze(0) .unsqueeze(0) ) scores = scores.masked_fill(~causal_tril, -10000.0) # Attention mask mask = attention_mask if attention_mask is not None else kwargs.get("mask") if mask is not None: if mask.dim() == 4: scores = scores + mask.to(scores.dtype) else: valid_m = get_token_valid_mask(mask, B, N_k, device) if valid_m is not None: scores = scores.masked_fill((~valid_m).unsqueeze(1).unsqueeze(2), -10000.0) attn = torch.nan_to_num(F.softmax(scores, dim=-1), nan=0.0).to(dtype) attn = self.dropout(attn) # Head mask broadcasting across 1D/2D/3D/4D if head_mask is not None: if head_mask.dim() == 1: head_mask = head_mask.view(1, -1, 1, 1) elif head_mask.dim() == 2: head_mask = head_mask.view(head_mask.size(0), head_mask.size(1), 1, 1) elif head_mask.dim() == 3: head_mask = head_mask.unsqueeze(-1) elif head_mask.dim() == 4: # Assume [B, H, 1, 1] or [B, 1, 1, 1]; broadcast normally pass attn = attn * head_mask.to(attn.dtype) context = torch.matmul(attn, V) context = context.transpose(1, 2).contiguous().view(B, N_q, self.all_head_size) if self.is_nested_encoder: output = context else: output = self.out_proj(context) valid_m = get_token_valid_mask(attention_mask, B, N_q, device) if valid_m is not None: output = output * valid_m.unsqueeze(-1).to(output.dtype) return output, (attn if output_attentions else None)Weighted Grouped-Query Attention (WGQA) with learnable per-KV-head importance weights, robust GQA pooling, LoRA adapters, RoPE support with fallback, and causal masking.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- BaseAttention
- torch.nn.modules.module.Module
Class variables
var is_custom_attention : bool-
The type of the None singleton.
var supports_bidirectional : bool-
The type of the None singleton.
var supports_causal : bool-
The type of the None singleton.
var supports_gqa : bool-
The type of the None singleton.
var supports_rope : bool-
The type of the None singleton.
Instance variables
prop W_k-
Expand source code
@property def W_k(self): return self.k_proj prop W_q-
Expand source code
@property def W_q(self): return self.q_proj prop W_v-
Expand source code
@property def W_v(self): return self.v_proj 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-
Expand source code
@property def out(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, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: """Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.""" extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. UniversalAdapter extraction for proj in ["q", "k", "v", "o"]: try: w, b = UniversalAdapter.get_weight_bias(src, proj) if w is not None: extracted[proj] = (w, b) except Exception: pass # 2. GPT-2 Fused c_attn Conv1D handling if hasattr(src, "c_attn") and ( extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None ): c_mod = src.c_attn base_c = getattr(c_mod, "base_layer", c_mod) weight = getattr(base_c, "weight", getattr(c_mod, "weight", None)) if weight is not None: weight = weight.data.clone() bias = getattr(base_c, "bias", getattr(c_mod, "bias", None)) bias_data = bias.data.clone() if bias is not None else None is_conv1d = ( "Conv1D" in c_mod.__class__.__name__ or "Conv1D" in base_c.__class__.__name__ or hasattr(c_mod, "nf") or hasattr(base_c, "nf") or (weight.dim() == 2 and weight.shape[0] < weight.shape[1]) ) if is_conv1d: weight = weight.t().contiguous() split_size = weight.shape[0] // 3 q_w = weight[:split_size, :] k_w = weight[split_size : 2 * split_size, :] v_w = weight[2 * split_size :, :] q_b = bias_data[:split_size] if bias_data is not None else None k_b = bias_data[split_size : 2 * split_size] if bias_data is not None else None v_b = bias_data[2 * split_size :] if bias_data is not None else None extracted["q"] = (q_w, q_b) extracted["k"] = (k_w, k_b) extracted["v"] = (v_w, v_b) if hasattr(src, "c_proj") and extracted["o"][0] is None: p_mod = src.c_proj base_o = getattr(p_mod, "base_layer", p_mod) w_o = getattr(base_o, "weight", getattr(p_mod, "weight", None)) if w_o is not None: w_o = w_o.data.clone() is_conv1d = ( "Conv1D" in p_mod.__class__.__name__ or "Conv1D" in base_o.__class__.__name__ or hasattr(p_mod, "nf") or hasattr(base_o, "nf") or (w_o.dim() == 2 and w_o.shape[0] < w_o.shape[1]) ) if is_conv1d: w_o = w_o.t().contiguous() b_o = getattr(base_o, "bias", getattr(p_mod, "bias", None)) extracted["o"] = (w_o, b_o.data.clone() if b_o is not None else None) # 3. RoBERTa / BERT Nested Encoder handling if hasattr(src, "self") and ( extracted["q"][0] is None or extracted["k"][0] is None or extracted["v"][0] is None ): s = src.self if hasattr(s, "query"): extracted["q"] = (s.query.weight.data, getattr(s.query, "bias", None)) if hasattr(s, "key"): extracted["k"] = (s.key.weight.data, getattr(s.key, "bias", None)) if hasattr(s, "value"): extracted["v"] = (s.value.weight.data, getattr(s.value, "bias", None)) if hasattr(src, "output") and hasattr(src.output, "dense") and extracted["o"][0] is None: dense = src.output.dense extracted["o"] = (dense.weight.data, getattr(dense, "bias", None)) # 4. Modern Decoders (Qwen3.5, LLaMA, Mistral) & DistilBERT named_map = { "q": ["q_proj", "q_lin", "query", "W_q"], "k": ["k_proj", "k_lin", "key", "W_k"], "v": ["v_proj", "v_lin", "value", "W_v"], "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"], } for proj, names in named_map.items(): if extracted[proj][0] is None: for name in names: if hasattr(src, name): layer = getattr(src, name) if hasattr(layer, "weight"): extracted[proj] = ( layer.weight.data, getattr(layer, "bias", None), ) break # Detect nested encoder architecture cls_name = getattr(src, "__class__", type(src)).__name__ has_out = ( hasattr(src, "dense") or hasattr(src, "out_proj") or hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") ) is_nested = False if cls_name in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (hasattr(src, "query") and not has_out): is_nested = True if hasattr(src, "self") and hasattr(src, "output"): is_nested = True self.is_nested_encoder = bool(is_nested) missing = [k for k in ["q", "k", "v"] if extracted[k][0] is None] if not self.is_nested_encoder and extracted["o"][0] is None: missing.append("o") if missing: logger.warning( "Attention module copy_weights warning. Proceeding with initialized weights." ) def _orient(w: torch.Tensor, in_dim: int) -> torch.Tensor: if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim: return w.t().contiguous() return w with torch.no_grad(): # Copy Q w_q, b_q = extracted["q"] if w_q is not None: w_q_or = _orient(w_q, self.hidden_size) out_d, in_d = self.q_proj.weight.shape s_out, s_in = w_q_or.shape self.q_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_q_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_q is not None and self.q_proj.bias is not None: b_val = b_q.data if isinstance(b_q, nn.Parameter) else b_q self.q_proj.bias.data[: min(len(self.q_proj.bias), len(b_val))].copy_( b_val[: min(len(self.q_proj.bias), len(b_val))] ) # Copy K with GQA pooling if source heads > kv heads w_k, b_k = extracted["k"] if w_k is not None: w_k_or = _orient(w_k, self.hidden_size) if w_k_or.shape[0] == self.kv_head_size: out_d, in_d = self.k_proj.weight.shape s_out, s_in = w_k_or.shape self.k_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_k_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_k is not None and self.k_proj.bias is not None: b_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k self.k_proj.bias.data[: min(len(self.k_proj.bias), len(b_val))].copy_( b_val[: min(len(self.k_proj.bias), len(b_val))] ) else: rep = max(1, w_k_or.shape[0] // max(1, self.kv_head_size)) w_k_pooled = ( w_k_or.view(self.num_kv_heads, rep, self.head_dim, -1) .mean(dim=1) .view(self.kv_head_size, -1) ) out_d, in_d = self.k_proj.weight.shape s_out, s_in = w_k_pooled.shape self.k_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_k_pooled[: min(out_d, s_out), : min(in_d, s_in)] ) if b_k is not None and self.k_proj.bias is not None: b_val = b_k.data if isinstance(b_k, nn.Parameter) else b_k b_k_pooled = ( b_val.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1) ) self.k_proj.bias.data[: min(len(self.k_proj.bias), len(b_k_pooled))].copy_( b_k_pooled[: min(len(self.k_proj.bias), len(b_k_pooled))] ) # Copy V with GQA pooling if source heads > kv heads w_v, b_v = extracted["v"] if w_v is not None: w_v_or = _orient(w_v, self.hidden_size) if w_v_or.shape[0] == self.kv_head_size: out_d, in_d = self.v_proj.weight.shape s_out, s_in = w_v_or.shape self.v_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_v_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_v is not None and self.v_proj.bias is not None: b_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v self.v_proj.bias.data[: min(len(self.v_proj.bias), len(b_val))].copy_( b_val[: min(len(self.v_proj.bias), len(b_val))] ) else: rep = max(1, w_v_or.shape[0] // max(1, self.kv_head_size)) w_v_pooled = ( w_v_or.view(self.num_kv_heads, rep, self.head_dim, -1) .mean(dim=1) .view(self.kv_head_size, -1) ) out_d, in_d = self.v_proj.weight.shape s_out, s_in = w_v_pooled.shape self.v_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_v_pooled[: min(out_d, s_out), : min(in_d, s_in)] ) if b_v is not None and self.v_proj.bias is not None: b_val = b_v.data if isinstance(b_v, nn.Parameter) else b_v b_v_pooled = ( b_val.view(self.num_kv_heads, rep, self.head_dim).mean(dim=1).view(-1) ) self.v_proj.bias.data[: min(len(self.v_proj.bias), len(b_v_pooled))].copy_( b_v_pooled[: min(len(self.v_proj.bias), len(b_v_pooled))] ) # Copy Out projection w_o, b_o = extracted["o"] if w_o is not None and not self.is_nested_encoder: w_o_or = _orient(w_o, self.all_head_size) out_d, in_d = self.out_proj.weight.shape s_out, s_in = w_o_or.shape self.out_proj.weight.data[: min(out_d, s_out), : min(in_d, s_in)].copy_( w_o_or[: min(out_d, s_out), : min(in_d, s_in)] ) if b_o is not None and self.out_proj.bias is not None: b_val = b_o.data if isinstance(b_o, nn.Parameter) else b_o self.out_proj.bias.data[: min(len(self.out_proj.bias), len(b_val))].copy_( b_val[: min(len(self.out_proj.bias), len(b_val))] ) elif self.is_nested_encoder: min_d = min(self.out_proj.weight.shape[0], self.out_proj.weight.shape[1]) self.out_proj.weight.data.zero_() self.out_proj.weight.data[:min_d, :min_d].fill_diagonal_(1.0) if self.out_proj.bias is not None: self.out_proj.bias.data.zero_() for target in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): target.freeze_base_weights() # Preserve source LoRA adapters if present on source module source_submodules = { "q": ["q_proj", "q_lin", "query", "W_q"], "k": ["k_proj", "k_lin", "key", "W_k"], "v": ["v_proj", "v_lin", "value", "W_v"], "o": ["o_proj", "out_proj", "out_lin", "dense", "out", "c_proj"], } for proj_key, target in [ ("q", self.q_proj), ("k", self.k_proj), ("v", self.v_proj), ("o", self.out_proj), ]: for name in source_submodules[proj_key]: if hasattr(src, name): sub = getattr(src, name) if ( hasattr(sub, "lora_A") and hasattr(sub, "lora_B") and sub.lora_A is not None and sub.lora_B is not None ): src_lA = getattr(sub.lora_A, "weight", sub.lora_A) src_lB = getattr(sub.lora_B, "weight", sub.lora_B) if hasattr(src_lA, "data") and target.lora_A is not None: if ( src_lA.shape[0] == target.lora_A.shape[0] and src_lA.shape[1] == target.lora_A.shape[1] ): target.lora_A.data.copy_(src_lA.data) target.lora_A.requires_grad = True else: print( f"Warning: lora_A shape mismatch for {proj_key}; keeping zero-init." ) if hasattr(src_lB, "data") and target.lora_B is not None: if ( src_lB.shape[0] == target.lora_B.shape[0] and src_lB.shape[1] == target.lora_B.shape[1] ): target.lora_B.data.copy_(src_lB.data) target.lora_B.requires_grad = True else: print( f"Warning: lora_B shape mismatch for {proj_key}; keeping zero-init." ) if hasattr(sub, "lora_scaling"): target.lora_scaling = sub.lora_scaling break if hasattr(src, "is_decoder"): self.is_decoder = bool(src.is_decoder)Robust weight copying from DistilBERT, RoBERTa nested, GPT-2 Conv1D, or Qwen/LLaMA.
Inherited members