Module neuralops.attentions.leka_v1
LEKAv1: Linear-time Efficient Kernel Attention (Corrected with Positive Features).
Approximates the attention kernel via Positive Feature Mapping (ELU + 1) to guarantee strict positivity and stable gradients, avoiding NaN denominator collapse.
Supports: - Multi-architecture copy_weights (GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, DistilBERT) - Trainable LoRA adapters with frozen base weights (FactorizedLoRALinear) - GQA with repeat_kv - RoPE position embeddings - Causal chunked recurrence for decoders - Nested encoder bypass - Frobenius norm verification (>= 1e-4) - Safe injection returning None on failure
Functions
def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) ‑> 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, ) -> tuple[torch.Tensor, torch.Tensor]: if cos is None or sin is None: return q, k if cos.dim() == 2: cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) 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) if cos.shape[-2] != q.shape[-2]: raise ValueError( f"apply_rotary_pos_emb: cos/sin seq len {cos.shape[-2]} != q seq len {q.shape[-2]}" ) rotary_dim = cos.shape[-1] head_dim = q.shape[-1] if rotary_dim < head_dim: q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] cos_q = cos.to(q_rot.dtype) sin_q = sin.to(q_rot.dtype) cos_k = cos.to(k_rot.dtype) sin_k = sin.to(k_rot.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) k = torch.cat([k_rot_embed, k_pass], dim=-1) else: cos_q = cos[..., :head_dim].to(q.dtype) sin_q = sin[..., :head_dim].to(q.dtype) cos_k = cos[..., :head_dim].to(k.dtype) sin_k = sin[..., :head_dim].to(k.dtype) q = (q * cos_q) + (_rotate_half(q) * sin_q) k = (k * cos_k) + (_rotate_half(k) * sin_k) return q, k def inject_leka_v1(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("leka_v1_inject") @ATTENTION_REGISTRY.register("inject_leka_v1") def inject_leka_v1( 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, ): if _depth > 8: logger.warning("inject_leka_v1: max recursion depth exceeded.") return None 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_leka_v1( 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: src_attn = old_attn if src_attn is None and layer is not None: for attr in ("attention", "self_attn", "attn", "self"): if hasattr(layer, attr): cand = getattr(layer, attr) src_attn = getattr(cand, "self", cand) break if isinstance(src_attn, LEKAV1Attention) or "LEKA" in src_attn.__class__.__name__: return None new_attn = LEKAV1Attention( config=cfg, comp_config=comp_config, layer_idx=layer_idx, total_layers=total_layers, old_attn=src_attn, **kwargs, ) if src_attn is not None: new_attn.copy_weights(src_attn) p = next(src_attn.parameters(), None) if src_attn is not None else None if p is not None: new_attn.to(device=p.device, dtype=p.dtype) if layer is not None: if hasattr(layer, "attention"): if hasattr(layer.attention, "self"): new_attn.is_nested_encoder = True layer.attention.self = new_attn else: new_attn.is_nested_encoder = False layer.attention = new_attn elif hasattr(layer, "self_attn"): new_attn.is_nested_encoder = False layer.self_attn = new_attn elif hasattr(layer, "self"): new_attn.is_nested_encoder = True layer.self = new_attn elif hasattr(layer, "attn"): new_attn.is_nested_encoder = False layer.attn = new_attn return new_attn except Exception as e: logger.exception("inject_leka_v1 failed at layer %d: %s", layer_idx, e) return None
Classes
class FactorizedLoRALinear (in_features: int,
out_features: int,
rank: int = 64,
lora_rank: int = 16,
lora_alpha: float = 32.0,
bias: bool = True,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class FactorizedLoRALinear(nn.Module): """ SVD-factorized linear + LoRA adapter. Forward: y = up(down(x)) + (lora_B @ lora_A)(x) * scaling Base (down, up) frozen; LoRA trainable. """ is_custom_attention = False def __init__( self, in_features: int, out_features: int, rank: int = 64, lora_rank: int = 16, lora_alpha: float = 32.0, bias: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() self.in_features = int(in_features) self.out_features = int(out_features) self.rank = max(1, min(int(rank), self.in_features, self.out_features)) self.svd_rank = self.rank self.lora_rank = int(lora_rank) if lora_rank > 0 else 0 self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0 factory = {"device": device, "dtype": dtype} self.down = nn.Linear(self.in_features, self.rank, bias=False, **factory) self.up = nn.Linear(self.rank, self.out_features, bias=bias, **factory) self.down.weight.requires_grad = False self.up.weight.requires_grad = False if self.up.bias is not None: self.up.bias.requires_grad = False if self.lora_rank > 0: self.lora_A = nn.Parameter(torch.empty(self.lora_rank, self.in_features, **factory)) self.lora_B = nn.Parameter(torch.empty(self.out_features, self.lora_rank, **factory)) nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) @property def weight(self) -> torch.Tensor: w = torch.matmul(self.up.weight, self.down.weight) if self.lora_A is not None and self.lora_B is not None: w = w + (self.lora_B @ self.lora_A) * self.lora_scaling return w @property def bias(self) -> torch.Tensor | None: return self.up.bias def freeze_base_weights(self) -> None: self.down.weight.requires_grad = False self.up.weight.requires_grad = False if self.up.bias is not None: self.up.bias.requires_grad = False 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 = self.up(self.down(x)) if 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 outSVD-factorized linear + LoRA adapter.
Forward: y = up(down(x)) + (lora_B @ lora_A)(x) * scaling Base (down, up) frozen; LoRA trainable.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Class variables
var is_custom_attention-
The type of the None singleton.
Instance variables
prop bias : torch.Tensor | None-
Expand source code
@property def bias(self) -> torch.Tensor | None: return self.up.bias prop weight : torch.Tensor-
Expand source code
@property def weight(self) -> torch.Tensor: w = torch.matmul(self.up.weight, self.down.weight) if self.lora_A is not None and self.lora_B is not None: w = w + (self.lora_B @ self.lora_A) * self.lora_scaling return w
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: out = self.up(self.down(x)) if 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.down.weight.requires_grad = False self.up.weight.requires_grad = False if self.up.bias is not None: self.up.bias.requires_grad = False 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
class LEKAV1Attention (config: Any | None = None,
comp_config: Any | None = None,
layer_idx: int = 0,
total_layers: int = 6,
rank: int | None = None,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
is_decoder: bool | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("leka_v1", supports_causal=True, supports_gqa=True) @ATTENTION_REGISTRY.register("leka_v1_attention", supports_causal=True, supports_gqa=True) class LEKAV1Attention(BaseAttention): supports_causal: bool = True supports_gqa: bool = True is_custom_attention: bool = True def __init__( self, config: Any | None = None, comp_config: Any | None = None, layer_idx: int = 0, total_layers: int = 6, rank: int | None = None, pretrained_attn: nn.Module | None = None, old_attn: nn.Module | None = None, is_decoder: bool | None = None, **kwargs: Any, ) -> None: source_attn = pretrained_attn if pretrained_attn is not None else old_attn try: super().__init__(config, rank) except TypeError: super().__init__(config) hidden_size = self._detect_hidden_size(config, source_attn, kwargs) num_heads = self._detect_num_heads(config, source_attn, kwargs, hidden_size) num_kv_heads = self._detect_num_kv_heads(config, source_attn, kwargs, num_heads) head_dim = self._detect_head_dim(config, source_attn, kwargs, hidden_size, num_heads) self.dim = int(hidden_size) self.hidden_size = self.dim self.num_heads = int(num_heads) self.num_kv_heads = int(num_kv_heads) self.head_dim = int(head_dim) self.all_head_size = self.num_heads * self.head_dim self.all_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.supports_gqa = True self.layer_idx = int(layer_idx) self.feature_dim = int( getattr(comp_config, "feature_dim", 256) if comp_config is not None else 256 ) dropout_p = float( getattr(comp_config, "dropout", getattr(config, "dropout", 0.1)) if config is not None else 0.1 ) self.is_decoder = _detect_decoder(config, is_decoder, source_attn) self.is_causal = self.is_decoder self.is_nested_encoder = False self.lora_rank = int( getattr(config, "lora_rank", None) or (getattr(comp_config, "lora_rank", None) if comp_config else None) or kwargs.get("lora_rank", 16) ) self.lora_alpha = float( getattr(config, "lora_alpha", None) or (getattr(comp_config, "lora_alpha", None) if comp_config else None) or kwargs.get("lora_alpha", 32.0) ) self.proj_rank = int( kwargs.get("proj_rank") or (getattr(comp_config, "proj_rank", None) if comp_config else None) or (rank if rank is not None else 64) ) FF = FactorizedLoRALinear self.query = FF( self.dim, self.all_head_size, rank=self.proj_rank, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, bias=True, ) self.key = FF( self.dim, self.all_kv_head_size, rank=self.proj_rank, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, bias=True, ) self.value = FF( self.dim, self.all_kv_head_size, rank=self.proj_rank, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, bias=True, ) self.out_proj = FF( self.all_head_size, self.dim, rank=self.proj_rank, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, bias=True, ) self.q_proj = self.query self.k_proj = self.key self.v_proj = self.value self.norm = nn.LayerNorm(self.head_dim) self.dropout = nn.Dropout(dropout_p) self.pres_gate = nn.Parameter(torch.tensor(2.0)) proj_data = self._build_projection_buffer() self.register_buffer("proj", proj_data, persistent=True) if source_attn is not None: self.copy_weights(source_attn) # ------------------------------------------------------------------ # Dimension detection # ------------------------------------------------------------------ @staticmethod def _detect_hidden_size(config: Any, src: nn.Module | None, kwargs: dict[str, Any]) -> int: for k in ("hidden_size", "d_model", "n_embd", "dim"): v = kwargs.get(k) if isinstance(v, int) and v > 0: return v if config is not None: for k in ("hidden_size", "d_model", "n_embd", "dim"): v = getattr(config, k, None) if isinstance(v, int) and v > 0: return v if src is not None: for attr in ("out_lin", "out_proj", "o_proj", "dense", "c_proj"): mod = getattr(src, attr, None) if mod is not None and hasattr(mod, "out_features"): return int(mod.out_features) for attr in ("q_proj", "q_lin", "query", "c_attn"): mod = getattr(src, attr, None) if mod is not None and hasattr(mod, "in_features"): return int(mod.in_features) return 768 @staticmethod def _detect_num_heads( config: Any, src: nn.Module | None, kwargs: dict[str, Any], hidden_size: int ) -> int: for k in ("num_heads", "num_attention_heads", "n_heads", "n_head"): v = kwargs.get(k) if isinstance(v, int) and v > 0: return v if config is not None: for k in ("num_attention_heads", "num_heads", "n_heads", "n_head"): v = getattr(config, k, None) if isinstance(v, int) and v > 0: return v if src is not None: for k in ("num_heads", "num_attention_heads", "n_heads"): v = getattr(src, k, None) if isinstance(v, int) and v > 0: return v return max(1, hidden_size // 64) @staticmethod def _detect_num_kv_heads( config: Any, src: nn.Module | None, kwargs: dict[str, Any], num_heads: int ) -> int: for k in ("num_key_value_heads", "num_kv_heads"): v = kwargs.get(k) if isinstance(v, int) and v > 0: return v if config is not None: for k in ("num_key_value_heads", "num_kv_heads"): v = getattr(config, k, None) if isinstance(v, int) and v > 0: return v if src is not None: v = getattr(src, "num_key_value_heads", None) if isinstance(v, int) and v > 0: return v return num_heads @staticmethod def _detect_head_dim( config: Any, src: nn.Module | None, kwargs: dict[str, Any], hidden_size: int, num_heads: int, ) -> int: v = kwargs.get("head_dim") if isinstance(v, int) and v > 0: return v if config is not None: for k in ("head_dim", "kv_channels", "d_kv"): v = getattr(config, k, None) if isinstance(v, int) and v > 0: return v if src is not None: for attr in ("q_proj", "q_lin", "query"): mod = getattr(src, attr, None) if mod is not None and hasattr(mod, "weight") and mod.weight is not None: try: return int(mod.weight.shape[0] // max(1, num_heads)) except Exception: pass return max(1, hidden_size // max(1, num_heads)) # ------------------------------------------------------------------ # LEKA-14: robust QR-orthogonal projection builder # ------------------------------------------------------------------ def _build_projection_buffer(self) -> torch.Tensor: d = self.head_dim F_dim = self.feature_dim proj = torch.empty(self.num_heads, d, F_dim) for h in range(self.num_heads): if F_dim <= d: # (d, F_dim): orthonormal columns q, _ = torch.linalg.qr(torch.randn(d, F_dim)) proj[h] = q else: # (d, F_dim): multi-block orthonormal construction num_blocks = (F_dim + d - 1) // d blocks = [torch.linalg.qr(torch.randn(d, d))[0] for _ in range(num_blocks)] q = torch.cat(blocks, dim=1)[:, :F_dim] proj[h] = q return proj # ------------------------------------------------------------------ # SVD init # ------------------------------------------------------------------ @staticmethod def _svd_init(w: torch.Tensor, b: torch.Tensor | None, proj: FactorizedLoRALinear) -> None: with torch.no_grad(): W = w.detach().float() if W.ndim != 2: W = W.reshape(W.size(0), -1) out_t, in_t = proj.out_features, proj.in_features if W.shape[0] != out_t or W.shape[1] != in_t: padded = torch.zeros(out_t, in_t, dtype=W.dtype, device=W.device) r = min(out_t, W.shape[0]) c = min(in_t, W.shape[1]) padded[:r, :c] = W[:r, :c] W = padded try: U, S, Vh = torch.linalg.svd(W, full_matrices=False) except Exception as e: logger.warning("LEKAv1._svd_init: SVD failed (%s)", e) return r = min(proj.rank, U.size(1), Vh.size(0)) proj.down.weight.data.zero_() proj.up.weight.data.zero_() proj.down.weight.data[:r, :].copy_(Vh[:r, :].to(proj.down.weight.dtype)) US = U[:, :r] @ torch.diag(S[:r]) proj.up.weight.data[:, :r].copy_(US.to(proj.up.weight.dtype)) if b is not None and proj.up.bias is not None: n = min(b.numel(), proj.up.bias.numel()) proj.up.bias.data[:n].copy_(b[:n].to(proj.up.bias.dtype)) elif proj.up.bias is not None: proj.up.bias.data.zero_() # ------------------------------------------------------------------ # Kernel feature map (FIXED: ELU + 1) # ------------------------------------------------------------------ def kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor: """ Positive kernel feature embedding using ELU + 1. Guarantees strict positivity so denominator never collapses or becomes negative. Returns tensor of shape (B, H, N, F). """ x_norm = self.norm(x) proj_matrix = self.proj.to(device=x.device, dtype=x.dtype) # Linear projection to feature space features = torch.einsum("bhnd,hdf->bhnf", x_norm, proj_matrix) # Apply ELU + 1 to ensure outputs are strictly > 0 # Adding 1e-6 for extreme numerical stability return F.elu(features) + 1.0 + 1e-6 # ------------------------------------------------------------------ # LEKA-06: copy_weights raises on missing # ------------------------------------------------------------------ def copy_weights(self, src: Any) -> None: if src is None: return cls_name = src.__class__.__name__ has_out_lin = hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") is_nested = False if ( ( hasattr(src, "self") and hasattr(src.self, "query") and hasattr(src, "output") and hasattr(src.output, "dense") ) or ( cls_name in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (cls_name.endswith("SelfAttention") and "distil" not in cls_name.lower()) ) and not has_out_lin ): is_nested = True self.is_nested_encoder = bool(is_nested) extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. GPT-2 fused c_attn if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() is_conv = ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or (c_w.ndim == 2 and c_w.shape[1] == 3 * self.hidden_size) ) if is_conv: c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None extracted["q"] = ( c_w[:split_size, :], c_b[:split_size] if c_b is not None else None, ) extracted["k"] = ( c_w[split_size : 2 * split_size, :], c_b[split_size : 2 * split_size] if c_b is not None else None, ) extracted["v"] = ( c_w[2 * split_size :, :], c_b[2 * split_size :] if c_b is not None else None, ) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() is_conv_o = ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ) if is_conv_o: o_w = o_w.t().contiguous() bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) extracted["o"] = ( o_w, bias_o.data.clone() if bias_o is not None else None, ) # 2. BERT / RoBERTa nested elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]: mod = getattr(s, sub_name, None) if mod is not None and hasattr(mod, "weight"): extracted[proj_key] = ( mod.weight.data.clone(), mod.bias.data.clone() if mod.bias is not None else None, ) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense extracted["o"] = ( out_mod.weight.data.clone(), out_mod.bias.data.clone() if out_mod.bias is not None else None, ) # 3. Flat modules (DistilBERT) else: proj_candidates = { "q": ["q_lin", "q_proj", "query"], "k": ["k_lin", "k_proj", "key"], "v": ["v_lin", "v_proj", "value"], "o": ["out_lin", "out_proj", "dense", "o_proj"], } for proj_key, names in proj_candidates.items(): for name in names: if hasattr(src, name): mod = getattr(src, name) if mod is not None and hasattr(mod, "weight"): extracted[proj_key] = ( mod.weight.data.clone(), mod.bias.data.clone() if mod.bias is not None else None, ) break missing = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() for k in missing: w, b = adapter.get_weight_bias(src, k) if w is not None: extracted[k] = (w.clone(), b.clone() if b is not None else None) except Exception as e: logger.debug("UniversalAdapter fallback failed: %s", e) missing_after = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing_after: raise RuntimeError( f"LEKAV1Attention.copy_weights: missing mandatory projections " f"{missing_after} from {cls_name}." ) def _orient(w, in_dim): if w.dim() == 2 and w.shape[1] != in_dim and w.shape[0] == in_dim: return w.t().contiguous() return w def _resize_gqa(w, b, out_target): if w is None: return w, b src_out = w.shape[0] if src_out == out_target: return w, b if src_out > out_target and src_out % out_target == 0: rep = src_out // out_target w = w.view(out_target, rep, -1).mean(dim=1) if b is not None: b = b.view(out_target, rep).mean(dim=1) elif src_out < out_target and out_target % src_out == 0: rep = out_target // src_out w = w.repeat(rep, 1) if b is not None: b = b.repeat(rep) return w, b with torch.no_grad(): w_q, _ = extracted["q"] w_k, _ = extracted["k"] w_v, _ = extracted["v"] w_o, b_o = extracted["o"] if w_q is not None: self._svd_init(_orient(w_q, self.hidden_size), None, self.query) if w_k is not None: w_k_o = _orient(w_k, self.hidden_size) w_k_res, _ = _resize_gqa(w_k_o, None, self.key.out_features) self._svd_init(w_k_res, None, self.key) if w_v is not None: w_v_o = _orient(w_v, self.hidden_size) w_v_res, _ = _resize_gqa(w_v_o, None, self.value.out_features) self._svd_init(w_v_res, None, self.value) if w_o is not None: self._svd_init(_orient(w_o, self.all_head_size), b_o, self.out_proj) elif self.is_nested_encoder: eye = torch.eye( self.all_head_size, self.all_head_size, device=self.out_proj.down.weight.device, dtype=self.out_proj.down.weight.dtype, ) self._svd_init(eye, None, self.out_proj) # Norm verification for name, mod in ( ("query", self.query), ("key", self.key), ("value", self.value), ("out_proj", self.out_proj), ): try: w_norm = float(torch.norm(mod.weight).item()) if w_norm < 1e-4: logger.warning( "LEKAV1Attention: weight '%s' has near-zero norm (%.2e).", name, w_norm, ) except Exception: pass for mod in (self.query, self.key, self.value, self.out_proj): mod.freeze_base_weights() try: with torch.no_grad(): d = next(self.parameters()).device dt = self.query.down.weight.dtype dummy = torch.randn(1, 4, self.dim, device=d, dtype=dt) _ = self.forward(dummy) except Exception as e: logger.debug("LEKAv1 forward sanity check warning: %s", e) # ------------------------------------------------------------------ # LEKA-08: chunked causal cumsum # ------------------------------------------------------------------ @staticmethod def _causal_linear_attention( Q_feat: torch.Tensor, # (B, H, N, F) K_feat: torch.Tensor, V: torch.Tensor, # (B, H, N, d) chunk_size: int = 64, ) -> torch.Tensor: B, H, N, F_dim = Q_feat.shape d = V.shape[-1] dtype = Q_feat.dtype state_kv = torch.zeros(B, H, F_dim, d, device=Q_feat.device, dtype=dtype) state_k = torch.zeros(B, H, F_dim, device=Q_feat.device, dtype=dtype) outs = [] for start in range(0, N, chunk_size): end = min(start + chunk_size, N) Qc = Q_feat[:, :, start:end] Kc = K_feat[:, :, start:end] Vc = V[:, :, start:end] KV_local = (Kc.unsqueeze(-1) * Vc.unsqueeze(-2)).cumsum(dim=2) K_local = Kc.cumsum(dim=2) KV_total = KV_local + state_kv.unsqueeze(2) K_total = K_local + state_k.unsqueeze(2) num = torch.einsum("bhcf,bhcfd->bhcd", Qc, KV_total) # The denominator is now perfectly safe because Q and K are strictly positive. den = torch.einsum("bhcf,bhcf->bhc", Qc, K_total).clamp(min=1e-5) outs.append(num / den.unsqueeze(-1)) state_kv = KV_total[:, :, -1] state_k = K_total[:, :, -1] return torch.cat(outs, dim=2) # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, head_mask: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, past_key_value: Any | None = None, layer_idx: int = 0, output_attentions: bool | None = False, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: B, N, D = hidden_states.shape H, d = self.num_heads, self.head_dim dtype = hidden_states.dtype q = self.query(hidden_states).view(B, N, H, d).transpose(1, 2).contiguous() k = self.key(hidden_states).view(B, N, self.num_kv_heads, d).transpose(1, 2).contiguous() v = self.value(hidden_states).view(B, N, self.num_kv_heads, d).transpose(1, 2).contiguous() if position_embeddings is None: position_embeddings = kwargs.get("position_embeddings") if position_embeddings is not None: cos, sin = position_embeddings past_len = 0 if past_key_value is not None: if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 1: past_len = past_key_value[0].shape[-2] elif hasattr(past_key_value, "key_cache") and past_key_value.key_cache: past_len = past_key_value.key_cache[0].shape[-2] if past_len > 0: cos = cos[..., past_len : past_len + N, :] sin = sin[..., past_len : past_len + N, :] q, k = apply_rotary_pos_emb(q, k, cos, sin) if self.kv_repeat > 1: k = repeat_kv(k, self.kv_repeat) v = repeat_kv(v, self.kv_repeat) # Kernel feature map (Now strictly positive) Q_features = self.kernel_feature_map(q) K_features = self.kernel_feature_map(k) if attention_mask is not None: if attention_mask.dim() == 2: valid_mask = (attention_mask > 0).to(dtype) elif attention_mask.dim() == 4: if attention_mask.dtype == torch.bool: valid_mask = attention_mask[:, 0, 0, :].to(dtype) else: valid_mask = (attention_mask[:, 0, 0, :] > -1e4).to(dtype) else: valid_mask = attention_mask.to(dtype) kv_mask = valid_mask[:, None, :, None] # (B, 1, N, 1) K_features = K_features * kv_mask v = v * kv_mask else: kv_mask = None # Compute attention if self.is_decoder: context = self._causal_linear_attention( Q_features.float(), K_features.float(), v.float(), chunk_size=64, ).to(dtype) else: # Bidirectional linear attention KV = torch.einsum("bhnf,bhnd->bhfd", K_features.float(), v.float()) K_sum = K_features.float().sum(dim=2) num = torch.einsum("bhnf,bhfd->bhnd", Q_features.float(), KV) den = ( torch.einsum("bhnf,bhf->bhn", Q_features.float(), K_sum) .unsqueeze(-1) .clamp(min=1e-5) ) context = (num / den).to(dtype) context = torch.nan_to_num(context, nan=0.0) # head_mask if head_mask is not None: if head_mask.dim() == 1: hm = head_mask.view(1, -1, 1, 1) elif head_mask.dim() == 2: hm = head_mask.unsqueeze(-1).unsqueeze(-1) elif head_mask.dim() == 3: hm = head_mask[:, :, :1].unsqueeze(-1) else: hm = head_mask context = context * hm.to(context.dtype) context = context.permute(0, 2, 1, 3).contiguous().view(B, N, self.all_head_size) if self.is_nested_encoder: attn_out = context else: attn_out = self.out_proj(context) # Residual blend via pres_gate alpha = torch.sigmoid(self.pres_gate) out = alpha * attn_out + (1.0 - alpha) * hidden_states out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4) return out, None def get_ortho_loss(self) -> float: return 0.0Every custom attention module must inherit from this.
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_causal : bool-
The type of the None singleton.
var supports_gqa : bool-
The type of the None singleton.
Methods
def get_ortho_loss(self) ‑> float-
Expand source code
def get_ortho_loss(self) -> float: return 0.0 def kernel_feature_map(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def kernel_feature_map(self, x: torch.Tensor) -> torch.Tensor: """ Positive kernel feature embedding using ELU + 1. Guarantees strict positivity so denominator never collapses or becomes negative. Returns tensor of shape (B, H, N, F). """ x_norm = self.norm(x) proj_matrix = self.proj.to(device=x.device, dtype=x.dtype) # Linear projection to feature space features = torch.einsum("bhnd,hdf->bhnf", x_norm, proj_matrix) # Apply ELU + 1 to ensure outputs are strictly > 0 # Adding 1e-6 for extreme numerical stability return F.elu(features) + 1.0 + 1e-6Positive kernel feature embedding using ELU + 1. Guarantees strict positivity so denominator never collapses or becomes negative. Returns tensor of shape (B, H, N, F).
Inherited members